Social-Distancing-YOLO/utills.py at master · data-stats/Social-Distancing-YOLO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
# Imports
import cv2
import numpy as np

# Apply perspective transform and get new points.


def get_transformed_points(boxes, mi):

    bottom_points = []
    for box in boxes:
        pnts = np.array(
            [[[int(box[0]+(box[2]*0.5)), int(box[1]+(box[3]*0.5))]]], dtype="float32")
        bd_pnt = cv2.perspectiveTransform(pnts, mi)[0][0]
        pnt = [int(bd_pnt[0]), int(bd_pnt[1])]
        bottom_points.append(pnt)

    return bottom_points

# Function calculates distance between two points of pedestrian bounding boxes.


def cal_dis(p1, p2, distance_w, distance_h):

    h = abs(p2[1]-p1[1])
    w = abs(p2[0]-p1[0])

    dis_w = float((w/distance_w)*180)
    dis_h = float((h/distance_h)*180)

    return int(np.sqrt(((dis_h)**2) + ((dis_w)**2)))

# Function to calculate distance between all pairs and calculate closeness ratio


def get_distances(boxes1, bottom_points, distance_w, distance_h):

    distance_mat = []
    bxs = []

    for i in range(len(bottom_points)):
        for j in range(len(bottom_points)):
            if i != j:
                dist = cal_dis(
                    bottom_points[i], bottom_points[j], distance_w, distance_h)
                if dist <= 150:
                    closeness = 0
                    distance_mat.append(
                        [bottom_points[i], bottom_points[j], closeness])
                    bxs.append([boxes1[i], boxes1[j], closeness])
                elif dist > 150 and dist <= 180:
                    closeness = 1
                    distance_mat.append(
                        [bottom_points[i], bottom_points[j], closeness])
                    bxs.append([boxes1[i], boxes1[j], closeness])
                else:
                    closeness = 2
                    distance_mat.append(
                        [bottom_points[i], bottom_points[j], closeness])
                    bxs.append([boxes1[i], boxes1[j], closeness])

    return distance_mat, bxs

# Function for bird's eye view scale


def get_scale(W, H):

    dis_w = 400
    dis_h = 600

    return float(dis_w/W), float(dis_h/H)


# Get count of different levels of risks i.e. High risk, Low Risk, No risk or Safe
def get_count(distances_mat):

    r = []
    g = []
    y = []

    for i in range(len(distances_mat)):

        if distances_mat[i][2] == 0:
            if (distances_mat[i][0] not in r) and (distances_mat[i][0] not in g) and (distances_mat[i][0] not in y):
                r.append(distances_mat[i][0])
            if (distances_mat[i][1] not in r) and (distances_mat[i][1] not in g) and (distances_mat[i][1] not in y):
                r.append(distances_mat[i][1])

    for i in range(len(distances_mat)):

        if distances_mat[i][2] == 1:
            if (distances_mat[i][0] not in r) and (distances_mat[i][0] not in g) and (distances_mat[i][0] not in y):
                y.append(distances_mat[i][0])
            if (distances_mat[i][1] not in r) and (distances_mat[i][1] not in g) and (distances_mat[i][1] not in y):
                y.append(distances_mat[i][1])

    for i in range(len(distances_mat)):

        if distances_mat[i][2] == 2:
            if (distances_mat[i][0] not in r) and (distances_mat[i][0] not in g) and (distances_mat[i][0] not in y):
                g.append(distances_mat[i][0])
            if (distances_mat[i][1] not in r) and (distances_mat[i][1] not in g) and (distances_mat[i][1] not in y):
                g.append(distances_mat[i][1])

    return (len(r), len(y), len(g))