chris/frigate
1
0
Fork 0
forked from Github/frigate
Code Pull Requests Packages Activity

YOLOv5 & YOLOv8 support for the OpenVINO Detector (#5523)

Browse Source 
* Initial commit that adds YOLOv5 and YOLOv8 support for OpenVINO detector

* Fixed double inference bug with YOLOv5 and YOLOv8

* Modified documentation to mention YOLOv5 and YOLOv8

* Changes to pass lint checks

* Change minimum threshold to improve model performance

* Fix link

* Clean up YOLO post-processing

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
This commit is contained in:
Anil Ozyalcin
2023-02-19 05:39:47 -08:00
committed by GitHub
parent 2b685ac343
commit 0592c8b0e2
4 changed files with 61 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
frigate/detectors/detector_config.py
View File

@@ -26,6 +26,8 @@ class InputTensorEnum(str, Enum):
class ModelTypeEnum(str, Enum):
ssd = "ssd"
yolox = "yolox"
yolov5 = "yolov5"
yolov8 = "yolov8"
class ModelConfig(BaseModel):

75
frigate/detectors/plugins/openvino.py
View File

@@ -67,6 +67,18 @@ class OvDetector(DetectionApi):
self.grids = np.concatenate(grids, 1)
self.expanded_strides = np.concatenate(expanded_strides, 1)
## Takes in class ID, confidence score, and array of [x, y, w, h] that describes detection position,
## returns an array that's easily passable back to Frigate.
def process_yolo(self, class_id, conf, pos):
return [
class_id, # class ID
conf, # confidence score
(pos[1] - (pos[3] / 2)) / self.h, # y_min
(pos[0] - (pos[2] / 2)) / self.w, # x_min
(pos[1] + (pos[3] / 2)) / self.h, # y_max
(pos[0] + (pos[2] / 2)) / self.w, # x_max
]
def detect_raw(self, tensor_input):
infer_request = self.interpreter.create_infer_request()
infer_request.infer([tensor_input])
@@ -113,23 +125,50 @@ class OvDetector(DetectionApi):
ordered = dets[dets[:, 5].argsort()[::-1]][:20]
detections = np.zeros((20, 6), np.float32)
i = 0
for object_detected in ordered:
if i < 20:
detections[i] = [
object_detected[6], # Label ID
object_detected[5], # Confidence
(object_detected[1] - (object_detected[3] / 2))
/ self.h, # y_min
(object_detected[0] - (object_detected[2] / 2))
/ self.w, # x_min
(object_detected[1] + (object_detected[3] / 2))
/ self.h, # y_max
(object_detected[0] + (object_detected[2] / 2))
/ self.w, # x_max
]
i += 1
else:
break
for i, object_detected in enumerate(ordered):
detections[i] = self.process_yolo(
object_detected[6], object_detected[5], object_detected[:4]
)
return detections
elif self.ov_model_type == ModelTypeEnum.yolov8:
out_tensor = infer_request.get_output_tensor()
results = out_tensor.data[0]
output_data = np.transpose(results)
scores = np.max(output_data[:, 4:], axis=1)
if len(scores) == 0:
return np.zeros((20, 6), np.float32)
scores = np.expand_dims(scores, axis=1)
# add scores to the last column
dets = np.concatenate((output_data, scores), axis=1)
# filter out lines with scores below threshold
dets = dets[dets[:, -1] > 0.5, :]
# limit to top 20 scores, descending order
ordered = dets[dets[:, -1].argsort()[::-1]][:20]
detections = np.zeros((20, 6), np.float32)
for i, object_detected in enumerate(ordered):
detections[i] = self.process_yolo(
np.argmax(object_detected[4:-1]),
object_detected[-1],
object_detected[:4],
)
return detections
elif self.ov_model_type == ModelTypeEnum.yolov5:
out_tensor = infer_request.get_output_tensor()
output_data = out_tensor.data[0]
# filter out lines with scores below threshold
conf_mask = (output_data[:, 4] >= 0.5).squeeze()
output_data = output_data[conf_mask]
# limit to top 20 scores, descending order
ordered = output_data[output_data[:, 4].argsort()[::-1]][:20]
detections = np.zeros((20, 6), np.float32)
for i, object_detected in enumerate(ordered):
detections[i] = self.process_yolo(
np.argmax(object_detected[5:]),
object_detected[4],
object_detected[:4],
)
return detections