From 7514ccf804da48d667d5e488ccc82f334bbc3058 Mon Sep 17 00:00:00 2001 From: Koushik Dutta Date: Mon, 8 Apr 2024 14:57:47 -0700 Subject: [PATCH] detect: remove old models --- plugins/coreml/package-lock.json | 4 +- plugins/coreml/package.json | 2 +- plugins/coreml/src/coreml/__init__.py | 79 +++--- plugins/openvino/package-lock.json | 4 +- plugins/openvino/package.json | 2 +- plugins/openvino/src/ov/__init__.py | 237 ++++++++++-------- plugins/openvino/src/yolo/__init__.py | 2 +- plugins/tensorflow-lite/package-lock.json | 4 +- plugins/tensorflow-lite/package.json | 2 +- .../tensorflow-lite/src/tflite/__init__.py | 60 +++-- 10 files changed, 215 insertions(+), 181 deletions(-) diff --git a/plugins/coreml/package-lock.json b/plugins/coreml/package-lock.json index bedf0affc..5472799a8 100644 --- a/plugins/coreml/package-lock.json +++ b/plugins/coreml/package-lock.json @@ -1,12 +1,12 @@ { "name": "@scrypted/coreml", - "version": "0.1.33", + "version": "0.1.34", "lockfileVersion": 2, "requires": true, "packages": { "": { "name": "@scrypted/coreml", - "version": "0.1.33", + "version": "0.1.34", "devDependencies": { "@scrypted/sdk": "file:../../sdk" } diff --git a/plugins/coreml/package.json b/plugins/coreml/package.json index b63596e35..0823a710b 100644 --- a/plugins/coreml/package.json +++ b/plugins/coreml/package.json @@ -42,5 +42,5 @@ "devDependencies": { "@scrypted/sdk": "file:../../sdk" }, - "version": "0.1.33" + "version": "0.1.34" } diff --git a/plugins/coreml/src/coreml/__init__.py b/plugins/coreml/src/coreml/__init__.py index ef31eafd8..23e3f525b 100644 --- a/plugins/coreml/src/coreml/__init__.py +++ b/plugins/coreml/src/coreml/__init__.py @@ -18,9 +18,18 @@ from predict import Prediction, PredictPlugin, Rectangle predictExecutor = concurrent.futures.ThreadPoolExecutor(8, "CoreML-Predict") +availableModels = [ + "Default", + "scrypted_yolov9c_320", + "scrypted_yolov9c", + "yolov9c_320", + "yolov9c", + "ssdlite_mobilenet_v2", + "yolov4-tiny", +] def parse_label_contents(contents: str): - lines = contents.split(',') + lines = contents.split(",") ret = {} for row_number, content in enumerate(lines): pair = re.split(r"[:\s]+", content.strip(), maxsplit=1) @@ -45,23 +54,21 @@ def parse_labels(userDefined): raise Exception("no classes found in model metadata") return parse_label_contents(classes) + class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProvider): def __init__(self, nativeId: str | None = None): super().__init__(nativeId=nativeId) model = self.storage.getItem("model") or "Default" - if model == "Default": - # if arch is arm64, use yolov9c, otherwise use yolov8n - if platform.machine() == 'arm64': - model = "scrypted_yolov9c_320" - else: - model = "scrypted_yolov8n_320" + if model == "Default" or model not in availableModels: + if model != "Default": + self.storage.setItem("model", "Default") + model = "scrypted_yolov9c_320" self.yolo = "yolo" in model - self.yolov8 = "yolov8" in model self.yolov9 = "yolov9" in model self.scrypted_model = "scrypted" in model - model_version = "v4" - mlmodel = 'model' if self.yolov8 or self.yolov9 else model + model_version = "v5" + mlmodel = "model" if self.yolov9 else model print(f"model: {model}") @@ -72,7 +79,7 @@ class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProv f"{model}.mlmodel", ) else: - if self.yolov8 or self.yolov9: + if self.yolov9: files = [ f"{model}/{model}.mlpackage/Data/com.apple.CoreML/weights/weight.bin", f"{model}/{model}.mlpackage/Data/com.apple.CoreML/{mlmodel}.mlmodel", @@ -112,30 +119,34 @@ class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProv self.loop = asyncio.get_event_loop() self.minThreshold = 0.2 - asyncio.ensure_future(self.prepareVisionFramework(), loop = self.loop) + asyncio.ensure_future(self.prepareVisionFramework(), loop=self.loop) async def prepareVisionFramework(self): try: from vision import VisionPlugin + if not VisionPlugin: raise Exception("no vision plugin") - await scrypted_sdk.deviceManager.onDevicesChanged({ - "devices": [ - { - "nativeId": "vision", - "type": scrypted_sdk.ScryptedDeviceType.Builtin.value, - "interfaces": [ - scrypted_sdk.ScryptedInterface.ObjectDetection.value, - ], - "name": "Vision Framework", - } - ] - }) + await scrypted_sdk.deviceManager.onDevicesChanged( + { + "devices": [ + { + "nativeId": "vision", + "type": scrypted_sdk.ScryptedDeviceType.Builtin.value, + "interfaces": [ + scrypted_sdk.ScryptedInterface.ObjectDetection.value, + ], + "name": "Vision Framework", + } + ] + } + ) except: pass async def getDevice(self, nativeId: str) -> Any: from vision import VisionPlugin + return VisionPlugin(nativeId) async def getSettings(self) -> list[Setting]: @@ -145,19 +156,7 @@ class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProv "key": "model", "title": "Model", "description": "The detection model used to find objects.", - "choices": [ - "Default", - "scrypted_yolov9c_320", - "scrypted_yolov8n_320", - "scrypted_yolov9c", - "scrypted_yolov8n", - "yolov8n_320", - "yolov9c_320", - "yolov8n", - "yolov9c", - "ssdlite_mobilenet_v2", - "yolov4-tiny", - ], + "choices": availableModels, "value": model, }, ] @@ -179,7 +178,7 @@ class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProv # run in executor if this is the plugin loop if self.yolo: - input_name = "image" if self.yolov8 or self.yolov9 else "input_1" + input_name = "image" if self.yolov9 else "input_1" if asyncio.get_event_loop() is self.loop: out_dict = await asyncio.get_event_loop().run_in_executor( predictExecutor, lambda: self.model.predict({input_name: input}) @@ -187,9 +186,9 @@ class CoreMLPlugin(PredictPlugin, scrypted_sdk.Settings, scrypted_sdk.DeviceProv else: out_dict = self.model.predict({input_name: input}) - if self.yolov8 or self.yolov9: + if self.yolov9: results = list(out_dict.values())[0][0] - objs = yolo.parse_yolov8(results) + objs = yolo.parse_yolov9(results) ret = self.create_detection_result(objs, src_size, cvss) return ret diff --git a/plugins/openvino/package-lock.json b/plugins/openvino/package-lock.json index 3845b66d4..48f7ff0cb 100644 --- a/plugins/openvino/package-lock.json +++ b/plugins/openvino/package-lock.json @@ -1,12 +1,12 @@ { "name": "@scrypted/openvino", - "version": "0.1.57", + "version": "0.1.58", "lockfileVersion": 2, "requires": true, "packages": { "": { "name": "@scrypted/openvino", - "version": "0.1.57", + "version": "0.1.58", "devDependencies": { "@scrypted/sdk": "file:../../sdk" } diff --git a/plugins/openvino/package.json b/plugins/openvino/package.json index 8d74a2fc1..5bd7d8e64 100644 --- a/plugins/openvino/package.json +++ b/plugins/openvino/package.json @@ -41,5 +41,5 @@ "devDependencies": { "@scrypted/sdk": "file:../../sdk" }, - "version": "0.1.57" + "version": "0.1.58" } diff --git a/plugins/openvino/src/ov/__init__.py b/plugins/openvino/src/ov/__init__.py index a018da05c..6111c5bd4 100644 --- a/plugins/openvino/src/ov/__init__.py +++ b/plugins/openvino/src/ov/__init__.py @@ -14,39 +14,58 @@ from predict import PredictPlugin, Prediction, Rectangle import numpy as np import yolo +availableModels = [ + "Default", + "scrypted_yolov9c_320", + "scrypted_yolov9c", + "ssd_mobilenet_v1_coco", + "ssdlite_mobilenet_v2", + "yolo-v3-tiny-tf", + "yolo-v4-tiny-tf", +] + def parse_label_contents(contents: str): lines = contents.splitlines() ret = {} for row_number, content in enumerate(lines): - pair = re.split(r'[:\s]+', content.strip(), maxsplit=1) + pair = re.split(r"[:\s]+", content.strip(), maxsplit=1) if len(pair) == 2 and pair[0].strip().isdigit(): ret[int(pair[0])] = pair[1].strip() else: ret[row_number] = content.strip() return ret + def param_to_string(parameters) -> str: """Convert a list / tuple of parameters returned from IE to a string.""" if isinstance(parameters, (list, tuple)): - return ', '.join([str(x) for x in parameters]) + return ", ".join([str(x) for x in parameters]) else: return str(parameters) + def dump_device_properties(core): - print('Available devices:') + print("Available devices:") for device in core.available_devices: - print(f'{device} :') - print('\tSUPPORTED_PROPERTIES:') - for property_key in core.get_property(device, 'SUPPORTED_PROPERTIES'): - if property_key not in ('SUPPORTED_METRICS', 'SUPPORTED_CONFIG_KEYS', 'SUPPORTED_PROPERTIES'): + print(f"{device} :") + print("\tSUPPORTED_PROPERTIES:") + for property_key in core.get_property(device, "SUPPORTED_PROPERTIES"): + if property_key not in ( + "SUPPORTED_METRICS", + "SUPPORTED_CONFIG_KEYS", + "SUPPORTED_PROPERTIES", + ): try: property_val = core.get_property(device, property_key) except TypeError: - property_val = 'UNSUPPORTED TYPE' - print(f'\t\t{property_key}: {param_to_string(property_val)}') - print('') + property_val = "UNSUPPORTED TYPE" + print(f"\t\t{property_key}: {param_to_string(property_val)}") + print("") -class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Settings): + +class OpenVINOPlugin( + PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Settings +): def __init__(self, nativeId: str | None = None): super().__init__(nativeId=nativeId) @@ -54,75 +73,94 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S dump_device_properties(self.core) available_devices = self.core.available_devices self.available_devices = available_devices - print('available devices: %s' % available_devices) + print("available devices: %s" % available_devices) - mode = self.storage.getItem('mode') - if mode == 'Default': - mode = 'AUTO' - mode = mode or 'AUTO' + mode = self.storage.getItem("mode") + if mode == "Default": + mode = "AUTO" + mode = mode or "AUTO" - precision = self.storage.getItem('precision') or 'Default' - if precision == 'Default': + precision = self.storage.getItem("precision") or "Default" + if precision == "Default": using_mode = mode - if using_mode == 'AUTO': - if 'GPU' in available_devices: - using_mode = 'GPU' - if using_mode == 'GPU': - precision = 'FP16' + if using_mode == "AUTO": + if "GPU" in available_devices: + using_mode = "GPU" + if using_mode == "GPU": + precision = "FP16" else: - precision = 'FP32' + precision = "FP32" - - model = self.storage.getItem('model') or 'Default' - if model == 'Default': - model = 'scrypted_yolov9c_320' - self.yolo = 'yolo' in model - self.yolov8 = "yolov8" in model + model = self.storage.getItem("model") or "Default" + if model == "Default" or model not in availableModels: + if model != "Default": + self.storage.setItem("model", "Default") + model = "scrypted_yolov9c_320" + self.yolo = "yolo" in model self.yolov9 = "yolov9" in model self.scrypted_model = "scrypted" in model - self.sigmoid = model == 'yolo-v4-tiny-tf' + self.sigmoid = model == "yolo-v4-tiny-tf" - print(f'model/mode/precision: {model}/{mode}/{precision}') + print(f"model/mode/precision: {model}/{mode}/{precision}") - ovmodel = 'best' if self.scrypted_model else model + ovmodel = "best" if self.scrypted_model else model - model_version = 'v4' - xmlFile = self.downloadFile(f'https://raw.githubusercontent.com/koush/openvino-models/main/{model}/{precision}/{ovmodel}.xml', f'{model_version}/{precision}/{ovmodel}.xml') - binFile = self.downloadFile(f'https://raw.githubusercontent.com/koush/openvino-models/main/{model}/{precision}/{ovmodel}.bin', f'{model_version}/{precision}/{ovmodel}.bin') + model_version = "v4" + xmlFile = self.downloadFile( + f"https://raw.githubusercontent.com/koush/openvino-models/main/{model}/{precision}/{ovmodel}.xml", + f"{model_version}/{precision}/{ovmodel}.xml", + ) + binFile = self.downloadFile( + f"https://raw.githubusercontent.com/koush/openvino-models/main/{model}/{precision}/{ovmodel}.bin", + f"{model_version}/{precision}/{ovmodel}.bin", + ) if self.scrypted_model: - labelsFile = self.downloadFile('https://raw.githubusercontent.com/koush/openvino-models/main/scrypted_labels.txt', 'scrypted_labels.txt') + labelsFile = self.downloadFile( + "https://raw.githubusercontent.com/koush/openvino-models/main/scrypted_labels.txt", + "scrypted_labels.txt", + ) elif self.yolo: - labelsFile = self.downloadFile('https://raw.githubusercontent.com/koush/openvino-models/main/coco_80cl.txt', 'coco_80cl.txt') + labelsFile = self.downloadFile( + "https://raw.githubusercontent.com/koush/openvino-models/main/coco_80cl.txt", + "coco_80cl.txt", + ) else: - labelsFile = self.downloadFile('https://raw.githubusercontent.com/koush/openvino-models/main/coco_labels.txt', 'coco_labels.txt') + labelsFile = self.downloadFile( + "https://raw.githubusercontent.com/koush/openvino-models/main/coco_labels.txt", + "coco_labels.txt", + ) print(xmlFile, binFile, labelsFile) try: self.compiled_model = self.core.compile_model(xmlFile, mode) - print("EXECUTION_DEVICES", self.compiled_model.get_property("EXECUTION_DEVICES")) + print( + "EXECUTION_DEVICES", + self.compiled_model.get_property("EXECUTION_DEVICES"), + ) except: import traceback + traceback.print_exc() print("Reverting all settings.") - self.storage.removeItem('mode') - self.storage.removeItem('model') - self.storage.removeItem('precision') + self.storage.removeItem("mode") + self.storage.removeItem("model") + self.storage.removeItem("precision") self.requestRestart() # mobilenet 1,300,300,3 # yolov3/4 1,416,416,3 - # yolov8 1,3,320,320 + # yolov9 1,3,320,320 # second dim is always good. self.model_dim = self.compiled_model.inputs[0].shape[2] - labels_contents = open(labelsFile, 'r').read() + labels_contents = open(labelsFile, "r").read() self.labels = parse_label_contents(labels_contents) async def getSettings(self) -> list[Setting]: - mode = self.storage.getItem('mode') or 'Default' - model = self.storage.getItem('model') or 'Default' - precision = self.storage.getItem('precision') or 'Default' + mode = self.storage.getItem("mode") or "Default" + model = self.storage.getItem("model") or "Default" + precision = self.storage.getItem("precision") or "Default" return [ { "title": "Available Devices", @@ -132,52 +170,38 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S "key": "available_devices", }, { - 'key': 'model', - 'title': 'Model', - 'description': 'The detection model used to find objects.', - 'choices': [ - 'Default', - 'scrypted_yolov9c_320', - 'scrypted_yolov8n_320', - 'scrypted_yolov9c', - 'scrypted_yolov8n', - 'yolov9c_320', - 'yolov8n_320', - 'yolov9c', - 'yolov8n', - 'ssd_mobilenet_v1_coco', - 'ssdlite_mobilenet_v2', - 'yolo-v3-tiny-tf', - 'yolo-v4-tiny-tf', - ], - 'value': model, + "key": "model", + "title": "Model", + "description": "The detection model used to find objects.", + "choices": availableModels, + "value": model, }, { - 'key': 'mode', - 'title': 'Mode', - 'description': 'AUTO, CPU, or GPU mode to use for detections. Requires plugin reload. Use CPU if the system has unreliable GPU drivers.', - 'choices': [ - 'Default', - 'AUTO', - 'CPU', - 'GPU', + "key": "mode", + "title": "Mode", + "description": "AUTO, CPU, or GPU mode to use for detections. Requires plugin reload. Use CPU if the system has unreliable GPU drivers.", + "choices": [ + "Default", + "AUTO", + "CPU", + "GPU", ], - 'value': mode, - 'combobox': True, + "value": mode, + "combobox": True, }, { - 'key': 'precision', - 'title': 'Precision', - 'description': 'The model floating point precision. FP16 is recommended for GPU. FP32 is recommended for CPU.', - 'choices': [ - 'Default', - 'FP16', - 'FP32', + "key": "precision", + "title": "Precision", + "description": "The model floating point precision. FP16 is recommended for GPU. FP32 is recommended for CPU.", + "choices": [ + "Default", + "FP16", + "FP32", ], - 'value': precision, - } + "value": precision, + }, ] - + async def putSetting(self, key: str, value: SettingValue): self.storage.setItem(key, value) await self.onDeviceEvent(scrypted_sdk.ScryptedInterface.Settings.value, None) @@ -195,7 +219,7 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S infer_request = self.compiled_model.create_infer_request() # the input_tensor can be created with the shared_memory=True parameter, # but that seems to cause issues on some platforms. - if self.yolov8 or self.yolov9: + if self.yolov9: im = np.stack([input]) im = im.transpose((0, 3, 1, 2)) # BHWC to BCHW, (n, 3, h, w) im = im.astype(np.float32) / 255.0 @@ -203,7 +227,9 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S im = ov.Tensor(array=im) input_tensor = im elif self.yolo: - input_tensor = ov.Tensor(array=np.expand_dims(np.array(input), axis=0).astype(np.float32)) + input_tensor = ov.Tensor( + array=np.expand_dims(np.array(input), axis=0).astype(np.float32) + ) else: input_tensor = ov.Tensor(array=np.expand_dims(np.array(input), axis=0)) # Set input tensor for model with one input @@ -213,8 +239,8 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S objs = [] - if self.yolov8 or self.yolov9: - objs = yolo.parse_yolov8(infer_request.output_tensors[0].data[0]) + if self.yolov9: + objs = yolo.parse_yolov9(infer_request.output_tensors[0].data[0]) return objs if self.yolo: @@ -224,17 +250,21 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S out_blob = infer_request.outputs[0] else: out_blob = infer_request.outputs[1] - + # 13 13 - objects = yolo.parse_yolo_region(out_blob.data, (input.width, input.height),(81,82, 135,169, 344,319), self.sigmoid) + objects = yolo.parse_yolo_region( + out_blob.data, + (input.width, input.height), + (81, 82, 135, 169, 344, 319), + self.sigmoid, + ) for r in objects: - obj = Prediction(r['classId'], r['confidence'], Rectangle( - r['xmin'], - r['ymin'], - r['xmax'], - r['ymax'] - )) + obj = Prediction( + r["classId"], + r["confidence"], + Rectangle(r["xmin"], r["ymin"], r["xmax"], r["ymax"]), + ) objs.append(obj) # what about output[1]? @@ -243,7 +273,6 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S return objs - output = infer_request.get_output_tensor(0) for values in output.data[0][0].astype(float): valid, index, confidence, l, t, r, b = values @@ -258,12 +287,7 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S r = torelative(r) b = torelative(b) - obj = Prediction(index - 1, confidence, Rectangle( - l, - t, - r, - b - )) + obj = Prediction(index - 1, confidence, Rectangle(l, t, r, b)) objs.append(obj) return objs @@ -272,6 +296,7 @@ class OpenVINOPlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.S objs = await predict() except: import traceback + traceback.print_exc() raise diff --git a/plugins/openvino/src/yolo/__init__.py b/plugins/openvino/src/yolo/__init__.py index 79f025509..28fa12e0b 100644 --- a/plugins/openvino/src/yolo/__init__.py +++ b/plugins/openvino/src/yolo/__init__.py @@ -6,7 +6,7 @@ from predict import Prediction, Rectangle defaultThreshold = .2 -def parse_yolov8(results, threshold = defaultThreshold, scale = None, confidence_scale = None): +def parse_yolov9(results, threshold = defaultThreshold, scale = None, confidence_scale = None): objs = [] keep = np.argwhere(results[4:] > threshold) for indices in keep: diff --git a/plugins/tensorflow-lite/package-lock.json b/plugins/tensorflow-lite/package-lock.json index f4cc0f353..f1bb6eba1 100644 --- a/plugins/tensorflow-lite/package-lock.json +++ b/plugins/tensorflow-lite/package-lock.json @@ -1,12 +1,12 @@ { "name": "@scrypted/tensorflow-lite", - "version": "0.1.52", + "version": "0.1.54", "lockfileVersion": 2, "requires": true, "packages": { "": { "name": "@scrypted/tensorflow-lite", - "version": "0.1.52", + "version": "0.1.54", "devDependencies": { "@scrypted/sdk": "file:../../sdk" } diff --git a/plugins/tensorflow-lite/package.json b/plugins/tensorflow-lite/package.json index 3ae74c14a..db4f36ac2 100644 --- a/plugins/tensorflow-lite/package.json +++ b/plugins/tensorflow-lite/package.json @@ -53,5 +53,5 @@ "devDependencies": { "@scrypted/sdk": "file:../../sdk" }, - "version": "0.1.52" + "version": "0.1.54" } diff --git a/plugins/tensorflow-lite/src/tflite/__init__.py b/plugins/tensorflow-lite/src/tflite/__init__.py index b3d3cafbc..6255bd8a7 100644 --- a/plugins/tensorflow-lite/src/tflite/__init__.py +++ b/plugins/tensorflow-lite/src/tflite/__init__.py @@ -29,6 +29,18 @@ from scrypted_sdk.types import Setting, SettingValue import yolo from predict import PredictPlugin +availableModels = [ + "Default", + "ssd_mobilenet_v2_coco_quant_postprocess", + "tf2_ssd_mobilenet_v2_coco17_ptq", + "ssdlite_mobiledet_coco_qat_postprocess", + "scrypted_yolov9c_320", + "efficientdet_lite0_320_ptq", + "efficientdet_lite1_384_ptq", + "efficientdet_lite2_448_ptq", + "efficientdet_lite3_512_ptq", + "efficientdet_lite3x_640_ptq", +] def parse_label_contents(contents: str): lines = contents.splitlines() @@ -59,26 +71,32 @@ class TensorFlowLitePlugin( edge_tpus = None pass - model_version = "v12" + model_version = "v13" model = self.storage.getItem("model") or "Default" + if model not in availableModels: + self.storage.setItem("model", "Default") + model = "Default" defaultModel = model == "Default" branch = "main" labelsFile = None + def configureModel(): nonlocal labelsFile nonlocal model if defaultModel: - if edge_tpus and next((obj for obj in edge_tpus if obj['type'] == 'usb'), None): + if edge_tpus and next( + (obj for obj in edge_tpus if obj["type"] == "usb"), None + ): model = "ssdlite_mobiledet_coco_qat_postprocess" else: model = "efficientdet_lite0_320_ptq" self.yolo = "yolo" in model - self.yolov8 = "yolov8" in model + self.yolov9 = "yolov9" in model self.scrypted_model = "scrypted" in model - print(f'model: {model}') + print(f"model: {model}") if self.scrypted_model: labelsFile = self.downloadFile( @@ -172,20 +190,7 @@ class TensorFlowLitePlugin( "key": "model", "title": "Model", "description": "The detection model used to find objects.", - "choices": [ - "Default", - "ssd_mobilenet_v2_coco_quant_postprocess", - "tf2_ssd_mobilenet_v2_coco17_ptq", - "ssdlite_mobiledet_coco_qat_postprocess", - "yolov8n_full_integer_quant", - "yolov8n_full_integer_quant_320", - "scrypted_yolov8n_320", - "efficientdet_lite0_320_ptq", - "efficientdet_lite1_384_ptq", - "efficientdet_lite2_448_ptq", - "efficientdet_lite3_512_ptq", - "efficientdet_lite3x_640_ptq", - ], + "choices": availableModels, "value": model, }, ] @@ -202,12 +207,12 @@ class TensorFlowLitePlugin( interpreter = self.interpreters.get() try: if self.yolo: - tensor_index = input_details(interpreter, 'index') + tensor_index = input_details(interpreter, "index") im = np.stack([input]) i = interpreter.get_input_details()[0] - if i['dtype'] == np.int8: - scale, zero_point = i['quantization'] + if i["dtype"] == np.int8: + scale, zero_point = i["quantization"] if scale == 0.003986024297773838 and zero_point == -128: # fast path for quantization 1/255 = 0.003986024297773838 im = im.view(np.int8) @@ -222,16 +227,21 @@ class TensorFlowLitePlugin( interpreter.invoke() output_details = interpreter.get_output_details() output = output_details[0] - x = interpreter.get_tensor(output['index']) + x = interpreter.get_tensor(output["index"]) input_scale = self.get_input_details()[0] if x.dtype == np.int8: - scale, zero_point = output['quantization'] + scale, zero_point = output["quantization"] threshold = yolo.defaultThreshold / scale + zero_point combined_scale = scale * input_scale - objs = yolo.parse_yolov8(x[0], threshold, scale=lambda v: (v - zero_point) * combined_scale, confidence_scale=lambda v: (v - zero_point) * scale) + objs = yolo.parse_yolov9( + x[0], + threshold, + scale=lambda v: (v - zero_point) * combined_scale, + confidence_scale=lambda v: (v - zero_point) * scale, + ) else: # this code path is unused. - objs = yolo.parse_yolov8(x[0], scale=lambda v: v * input_scale) + objs = yolo.parse_yolov9(x[0], scale=lambda v: v * input_scale) else: common.set_input(interpreter, input) interpreter.invoke()