diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index b401d18d5..54b32d92f 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -17,9 +17,11 @@ logger = logging.getLogger(__name__) DETECTOR_KEY = "onnx" + class ONNXDetectorConfig(BaseDetectorConfig): type: Literal[DETECTOR_KEY] + class ONNXDetector(DetectionApi): type_key = DETECTOR_KEY @@ -34,12 +36,23 @@ class ONNXDetector(DetectionApi): ) raise - assert detector_config.model.model_type == 'yolov8', "ONNX: detector_config.model.model_type: only yolov8 supported" - assert detector_config.model.input_tensor == 'nhwc', "ONNX: detector_config.model.input_tensor: only nhwc supported" - if detector_config.model.input_pixel_format != 'rgb': - logger.warn("ONNX: detector_config.model.input_pixel_format: should be 'rgb' for yolov8, but '{detector_config.model.input_pixel_format}' specified!") + assert ( + detector_config.model.model_type == "yolov8" + ), "ONNX: detector_config.model.model_type: only yolov8 supported" + assert ( + detector_config.model.input_tensor == "nhwc" + ), "ONNX: detector_config.model.input_tensor: only nhwc supported" + if detector_config.model.input_pixel_format != "rgb": + logger.warn( + "ONNX: detector_config.model.input_pixel_format: should be 'rgb' for yolov8, but '{detector_config.model.input_pixel_format}' specified!" + ) - assert detector_config.model.path is not None, "ONNX: No model.path configured, please configure model.path and model.labelmap_path; some suggestions: " + ', '.join(glob.glob("/config/model_cache/yolov8/*.onnx")) + " and " + ', '.join(glob.glob("/config/model_cache/yolov8/*_labels.txt")) + assert detector_config.model.path is not None, ( + "ONNX: No model.path configured, please configure model.path and model.labelmap_path; some suggestions: " + + ", ".join(glob.glob("/config/model_cache/yolov8/*.onnx")) + + " and " + + ", ".join(glob.glob("/config/model_cache/yolov8/*_labels.txt")) + ) path = detector_config.model.path logger.info(f"ONNX: loading {detector_config.model.path}") @@ -55,4 +68,3 @@ class ONNXDetector(DetectionApi): tensor_output = self.model.run(None, {model_input_name: tensor_input})[0] return yolov8_postprocess(model_input_shape, tensor_output) - diff --git a/frigate/detectors/plugins/rocm.py b/frigate/detectors/plugins/rocm.py index fa88e2e5c..ff5376d62 100644 --- a/frigate/detectors/plugins/rocm.py +++ b/frigate/detectors/plugins/rocm.py @@ -18,33 +18,47 @@ logger = logging.getLogger(__name__) DETECTOR_KEY = "rocm" + def detect_gfx_version(): - return subprocess.getoutput("unset HSA_OVERRIDE_GFX_VERSION && /opt/rocm/bin/rocminfo | grep gfx |head -1|awk '{print $2}'") + return subprocess.getoutput( + "unset HSA_OVERRIDE_GFX_VERSION && /opt/rocm/bin/rocminfo | grep gfx |head -1|awk '{print $2}'" + ) + def auto_override_gfx_version(): # If environment varialbe already in place, do not override gfx_version = detect_gfx_version() - old_override = os.getenv('HSA_OVERRIDE_GFX_VERSION') - if old_override not in (None, ''): - logger.warning(f"AMD/ROCm: detected {gfx_version} but HSA_OVERRIDE_GFX_VERSION already present ({old_override}), not overriding!") + old_override = os.getenv("HSA_OVERRIDE_GFX_VERSION") + if old_override not in (None, ""): + logger.warning( + f"AMD/ROCm: detected {gfx_version} but HSA_OVERRIDE_GFX_VERSION already present ({old_override}), not overriding!" + ) return old_override mapping = { - 'gfx90c': '9.0.0', - 'gfx1031': '10.3.0', - 'gfx1103': '11.0.0', + "gfx90c": "9.0.0", + "gfx1031": "10.3.0", + "gfx1103": "11.0.0", } override = mapping.get(gfx_version) if override is not None: - logger.warning(f"AMD/ROCm: detected {gfx_version}, overriding HSA_OVERRIDE_GFX_VERSION={override}") - os.putenv('HSA_OVERRIDE_GFX_VERSION', override) + logger.warning( + f"AMD/ROCm: detected {gfx_version}, overriding HSA_OVERRIDE_GFX_VERSION={override}" + ) + os.putenv("HSA_OVERRIDE_GFX_VERSION", override) return override - return '' + return "" class ROCmDetectorConfig(BaseDetectorConfig): type: Literal[DETECTOR_KEY] - conserve_cpu: bool = Field(default=True, title="Conserve CPU at the expense of latency (and reduced max throughput)") - auto_override_gfx: bool = Field(default=True, title="Automatically detect and override gfx version") + conserve_cpu: bool = Field( + default=True, + title="Conserve CPU at the expense of latency (and reduced max throughput)", + ) + auto_override_gfx: bool = Field( + default=True, title="Automatically detect and override gfx version" + ) + class ROCmDetector(DetectionApi): type_key = DETECTOR_KEY @@ -59,24 +73,33 @@ class ROCmDetector(DetectionApi): logger.info(f"AMD/ROCm: loaded migraphx module") except ModuleNotFoundError: - logger.error( - "AMD/ROCm: module loading failed, missing ROCm environment?" - ) + logger.error("AMD/ROCm: module loading failed, missing ROCm environment?") raise if detector_config.conserve_cpu: logger.info(f"AMD/ROCm: switching HIP to blocking mode to conserve CPU") - ctypes.CDLL('/opt/rocm/lib/libamdhip64.so').hipSetDeviceFlags(4) - assert detector_config.model.model_type == 'yolov8', "AMD/ROCm: detector_config.model.model_type: only yolov8 supported" - assert detector_config.model.input_tensor == 'nhwc', "AMD/ROCm: detector_config.model.input_tensor: only nhwc supported" - if detector_config.model.input_pixel_format != 'rgb': - logger.warn("AMD/ROCm: detector_config.model.input_pixel_format: should be 'rgb' for yolov8, but '{detector_config.model.input_pixel_format}' specified!") + ctypes.CDLL("/opt/rocm/lib/libamdhip64.so").hipSetDeviceFlags(4) + assert ( + detector_config.model.model_type == "yolov8" + ), "AMD/ROCm: detector_config.model.model_type: only yolov8 supported" + assert ( + detector_config.model.input_tensor == "nhwc" + ), "AMD/ROCm: detector_config.model.input_tensor: only nhwc supported" + if detector_config.model.input_pixel_format != "rgb": + logger.warn( + "AMD/ROCm: detector_config.model.input_pixel_format: should be 'rgb' for yolov8, but '{detector_config.model.input_pixel_format}' specified!" + ) - assert detector_config.model.path is not None, "No model.path configured, please configure model.path and model.labelmap_path; some suggestions: " + ', '.join(glob.glob("/config/model_cache/yolov8/*.onnx")) + " and " + ', '.join(glob.glob("/config/model_cache/yolov8/*_labels.txt")) + assert detector_config.model.path is not None, ( + "No model.path configured, please configure model.path and model.labelmap_path; some suggestions: " + + ", ".join(glob.glob("/config/model_cache/yolov8/*.onnx")) + + " and " + + ", ".join(glob.glob("/config/model_cache/yolov8/*_labels.txt")) + ) path = detector_config.model.path - mxr_path = os.path.splitext(path)[0] + '.mxr' - if path.endswith('.mxr'): + mxr_path = os.path.splitext(path)[0] + ".mxr" + if path.endswith(".mxr"): logger.info(f"AMD/ROCm: loading parsed model from {mxr_path}") self.model = migraphx.load(mxr_path) elif os.path.exists(mxr_path): @@ -84,30 +107,38 @@ class ROCmDetector(DetectionApi): self.model = migraphx.load(mxr_path) else: logger.info(f"AMD/ROCm: loading model from {path}") - if path.endswith('.onnx'): + if path.endswith(".onnx"): self.model = migraphx.parse_onnx(path) - elif path.endswith('.tf') or path.endswith('.tf2') or path.endswith('.tflite'): + elif ( + path.endswith(".tf") + or path.endswith(".tf2") + or path.endswith(".tflite") + ): # untested self.model = migraphx.parse_tf(path) else: raise Exception(f"AMD/ROCm: unkown model format {path}") logger.info(f"AMD/ROCm: compiling the model") - self.model.compile(migraphx.get_target('gpu'), offload_copy=True, fast_math=True) + self.model.compile( + migraphx.get_target("gpu"), offload_copy=True, fast_math=True + ) logger.info(f"AMD/ROCm: saving parsed model into {mxr_path}") os.makedirs("/config/model_cache/rocm", exist_ok=True) migraphx.save(self.model, mxr_path) logger.info(f"AMD/ROCm: model loaded") def detect_raw(self, tensor_input): - model_input_name = self.model.get_parameter_names()[0]; - model_input_shape = tuple(self.model.get_parameter_shapes()[model_input_name].lens()); - + model_input_name = self.model.get_parameter_names()[0] + model_input_shape = tuple( + self.model.get_parameter_shapes()[model_input_name].lens() + ) tensor_input = preprocess(tensor_input, model_input_shape, np.float32) detector_result = self.model.run({model_input_name: tensor_input})[0] addr = ctypes.cast(detector_result.data_ptr(), ctypes.POINTER(ctypes.c_float)) - tensor_output = np.ctypeslib.as_array(addr, shape=detector_result.get_shape().lens()) + tensor_output = np.ctypeslib.as_array( + addr, shape=detector_result.get_shape().lens() + ) return yolov8_postprocess(model_input_shape, tensor_output) - diff --git a/frigate/detectors/util.py b/frigate/detectors/util.py index 7a8f48e4b..4133fe9c6 100644 --- a/frigate/detectors/util.py +++ b/frigate/detectors/util.py @@ -5,30 +5,50 @@ import cv2 logger = logging.getLogger(__name__) + def preprocess(tensor_input, model_input_shape, model_input_element_type): model_input_shape = tuple(model_input_shape) - assert tensor_input.dtype == np.uint8, f'tensor_input.dtype: {tensor_input.dtype}' + assert tensor_input.dtype == np.uint8, f"tensor_input.dtype: {tensor_input.dtype}" if len(tensor_input.shape) == 3: tensor_input = tensor_input[np.newaxis, :] if model_input_element_type == np.uint8: # nothing to do for uint8 model input - assert model_input_shape == tensor_input.shape, f'model_input_shape: {model_input_shape}, tensor_input.shape: {tensor_input.shape}' + assert ( + model_input_shape == tensor_input.shape + ), f"model_input_shape: {model_input_shape}, tensor_input.shape: {tensor_input.shape}" return tensor_input - assert model_input_element_type == np.float32, f'model_input_element_type: {model_input_element_type}' + assert ( + model_input_element_type == np.float32 + ), f"model_input_element_type: {model_input_element_type}" # tensor_input must be nhwc - assert tensor_input.shape[3] == 3, f'tensor_input.shape: {tensor_input.shape}' + assert tensor_input.shape[3] == 3, f"tensor_input.shape: {tensor_input.shape}" if tensor_input.shape[1:3] != model_input_shape[2:4]: - logger.warn(f"preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!") + logger.warn( + f"preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!" + ) # cv2.dnn.blobFromImage is faster than numpying it - return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False) + return cv2.dnn.blobFromImage( + tensor_input[0], + 1.0 / 255, + (model_input_shape[3], model_input_shape[2]), + None, + swapRB=False, + ) -def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20, score_threshold = 0.5, nms_threshold = 0.5): + +def yolov8_postprocess( + model_input_shape, + tensor_output, + box_count=20, + score_threshold=0.5, + nms_threshold=0.5, +): model_box_count = tensor_output.shape[2] probs = tensor_output[0, 4:, :] all_ids = np.argmax(probs, axis=0) all_confidences = probs.T[np.arange(model_box_count), all_ids] all_boxes = tensor_output[0, 0:4, :].T - mask = (all_confidences > score_threshold) + mask = all_confidences > score_threshold class_ids = all_ids[mask] confidences = all_confidences[mask] cx, cy, w, h = all_boxes[mask].T @@ -37,7 +57,17 @@ def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20, score_t scale_y, scale_x = 1 / model_input_shape[1], 1 / model_input_shape[2] else: scale_y, scale_x = 1 / model_input_shape[2], 1 / model_input_shape[3] - detections = np.stack((class_ids, confidences, scale_y * (cy - h / 2), scale_x * (cx - w / 2), scale_y * (cy + h / 2), scale_x * (cx + w / 2)), axis=1) + detections = np.stack( + ( + class_ids, + confidences, + scale_y * (cy - h / 2), + scale_x * (cx - w / 2), + scale_y * (cy + h / 2), + scale_x * (cx + w / 2), + ), + axis=1, + ) if detections.shape[0] > box_count: # if too many detections, do nms filtering to suppress overlapping boxes boxes = np.stack((cx - w / 2, cy - h / 2, w, h), axis=1) @@ -45,8 +75,9 @@ def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20, score_t detections = detections[indexes] # if still too many, trim the rest by confidence if detections.shape[0] > box_count: - detections = detections[np.argpartition(detections[:,1], -box_count)[-box_count:]] + detections = detections[ + np.argpartition(detections[:, 1], -box_count)[-box_count:] + ] detections = detections.copy() detections.resize((box_count, 6)) return detections -