-
Notifications
You must be signed in to change notification settings - Fork 0
/
YOLO.py
102 lines (88 loc) · 3.83 KB
/
YOLO.py
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
import tensorflow as tf
import numpy as np
class Yolo:
def __init__(self, model_path, CLASSES, int8 = False):
self.int8 = int8
self.CLASSES = CLASSES
Interpreter = tf.lite.Interpreter
self.interpreter = Interpreter(model_path=model_path)
self.interpreter.allocate_tensors()
self.input_details = self.interpreter.get_input_details()
self.output_details = self.interpreter.get_output_details()
def pred(self, im):
'''
im: image in the form of normalized numpy array
'''
input, output = self.input_details[0], self.output_details[0]
if self.int8:
scale, zero_point = input['quantization']
im = (im / scale + zero_point).astype(np.uint8)
self.interpreter.set_tensor(self.input_details[0]['index'], im)
self.interpreter.invoke()
self.output_data = self.interpreter.get_tensor(self.output_details[0]['index'])
if self.int8:
scale, zero_point = output['quantization']
self.output_data = (self.output_data.astype(np.float32) - zero_point) * scale # re-scale
def classFilter(self, classdata):
'''
returns the index of the class with highest probability score
'''
classes = []
for i in range(classdata.shape[0]):
classes.append(classdata[i].argmax())
return classes
def extract_results(self):
'''
extract scores, class probabilities and bounding boxes from the output tensor
'''
output_data = self.output_data[0]
boxes = np.squeeze(output_data[..., :4]) # the first 4 elements are contain box coordinates
self.scores = np.squeeze(output_data[..., 4:5]) # the 5th element is the confidence score
self.classes = self.classFilter(output_data[..., 5:]) # the remaining elements are class probabilities
x, y, w, h = boxes[..., 0], boxes[..., 1], boxes[..., 2], boxes[..., 3] # extract the coordinates of the bounding boxes
self.xyxy = [x - w / 2, y - h / 2, x + w / 2, y + h / 2] # convert xywh to xyxy
return self.scores
def return_results(self, H, W):
''''
return the results in the form of json
'''
class_names = []
class_scores = []
coordinates = []
final = []
for i in self.bbox_index:
xmin = int(self.xyxy[0][i] * W)
ymin = int(self.xyxy[1][i] * H)
xmax = int(self.xyxy[2][i] * W)
ymax = int(self.xyxy[3][i] * H)
if xmin < 0:
xmin = 0
if ymin > H - 1:
ymin = H - 1
if xmax < 0:
xmax = 0
if ymax > H - 1:
ymax = H - 1
coordinates.append([xmin, ymin, xmax, ymax])
class_names.append(self.CLASSES[self.classes[i]])
class_scores.append(int(self.scores[i] * 100))
for coordinate, class_name, class_score in zip(coordinates, class_names, class_scores):
final.append(
[
coordinate,
class_score,
class_name,
]
)
return {"class_names": class_names,
"class_scores": class_scores,
"coordinates": coordinates,
"final": final}
def return_bbox(self, score_threshold = 0.65, max_size = 10, iou_threshold = 0.5):
'''
filters bounding boxes based on the score threshold and the maximum size to prevent overlapping boxes.
'''
self.bbox_index = tf.image.non_max_suppression(np.squeeze(self.output_data[..., :4]), self.scores, max_output_size=max_size, \
iou_threshold=iou_threshold, score_threshold = score_threshold)
self.bbox_index = np.array(self.bbox_index)
return self.bbox_index