administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-04 02:05:21 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Initial commit that adds YOLOv5 and YOLOv8 support for OpenVINO detector

Browse Source
This commit is contained in:
Anil Ozyalcin 2023-02-16 20:30:35 -08:00
parent 27d3676ba5
commit b561f00ff9
2 changed files with 68 additions and 0 deletions
Show Stats 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): class ModelTypeEnum(str, Enum):
ssd = "ssd" ssd = "ssd"
yolox = "yolox" yolox = "yolox"
yolov5 = "yolov5"
yolov8 = "yolov8"
class ModelConfig(BaseModel): class ModelConfig(BaseModel):

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

@ -133,3 +133,69 @@ class OvDetector(DetectionApi):
else: else:
break break
return detections return detections
elif self.ov_model_type == ModelTypeEnum.yolov8:
infer_request = self.interpreter.create_infer_request()
infer_request.infer([tensor_input])
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)
dets = np.concatenate((output_data, scores), axis=1) # add scores to the last column
dets = dets[dets[:,-1] > 0.5,:] # filter out lines with scores below threshold
ordered = dets[dets[:, -1].argsort()[::-1]][:20] # limit to top 20 scores, descending order
detections = np.zeros((20, 6), np.float32)
i = 0
for object_detected in ordered:
if i < 20:
detections[i] = [
np.argmax(object_detected[4:-1]) , # Label ID
object_detected[-1], # 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
return detections
elif self.ov_model_type == ModelTypeEnum.yolov5:
infer_request = self.interpreter.create_infer_request()
infer_request.infer([tensor_input])
out_tensor = infer_request.get_output_tensor()
output_data = out_tensor.data[0]
conf_mask = (output_data[:, 4] >= 0.5).squeeze()
output_data = output_data[conf_mask]
ordered = output_data[output_data[:, 4].argsort()[::-1]][:20] # limit to top 20 scores, descending order
detections = np.zeros((20, 6), np.float32)
i = 0
for object_detected in ordered:
if i < 20:
detections[i] = [
np.argmax(object_detected[5:]) , # Label ID
object_detected[4], # 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
return detections