diff --git a/plugins/coreml/package-lock.json b/plugins/coreml/package-lock.json index 4984580951..ad666fe16c 100644 --- a/plugins/coreml/package-lock.json +++ b/plugins/coreml/package-lock.json @@ -1,12 +1,12 @@ { "name": "@scrypted/coreml", - "version": "0.1.17", + "version": "0.1.18", "lockfileVersion": 2, "requires": true, "packages": { "": { "name": "@scrypted/coreml", - "version": "0.1.17", + "version": "0.1.18", "devDependencies": { "@scrypted/sdk": "file:../../sdk" } diff --git a/plugins/coreml/package.json b/plugins/coreml/package.json index d723ea8722..2da20be627 100644 --- a/plugins/coreml/package.json +++ b/plugins/coreml/package.json @@ -40,5 +40,5 @@ "devDependencies": { "@scrypted/sdk": "file:../../sdk" }, - "version": "0.1.17" + "version": "0.1.18" } diff --git a/plugins/coreml/src/coreml/__init__.py b/plugins/coreml/src/coreml/__init__.py index fe97143680..90360bf2ac 100644 --- a/plugins/coreml/src/coreml/__init__.py +++ b/plugins/coreml/src/coreml/__init__.py @@ -10,6 +10,7 @@ import asyncio import concurrent.futures import yolo +import numpy as np predictExecutor = concurrent.futures.ThreadPoolExecutor(8, "CoreML-Predict") @@ -33,9 +34,12 @@ def __init__(self, nativeId: str | None = None): model = self.storage.getItem("model") or "Default" if model == "Default": model = "ssdlite_mobilenet_v2" - self.yolo = model == "yolov4-tiny" + self.yolo = 'yolo' in model + self.yolov8 = 'yolov8' in model model_version = "v1" + print(f'model: {model}') + if not self.yolo: # todo convert these to mlpackage labelsFile = self.downloadFile( @@ -44,23 +48,29 @@ def __init__(self, nativeId: str | None = None): ) modelFile = self.downloadFile( f"https://github.com/koush/coreml-models/raw/main/{model}/{model}.mlmodel", - "{model}.mlmodel", + f"{model}.mlmodel", ) else: - files = [ - f"{model}/{model}.mlpackage/Data/com.apple.CoreML/FeatureDescriptions.json", - f"{model}/{model}.mlpackage/Data/com.apple.CoreML/Metadata.json", - f"{model}/{model}.mlpackage/Data/com.apple.CoreML/weights/weight.bin", - f"{model}/{model}.mlpackage/Data/com.apple.CoreML/{model}.mlmodel", - f"{model}/{model}.mlpackage/Manifest.json", - ] - - for f in files: - p = self.downloadFile( - f"https://github.com/koush/coreml-models/raw/main/{f}", - f"{model_version}/{f}", - ) - modelFile = os.path.dirname(p) + if self.yolov8: + modelFile = self.downloadFile( + f"https://github.com/koush/coreml-models/raw/main/{model}/{model}.mlmodel", + f"{model}.mlmodel", + ) + else: + files = [ + f"{model}/{model}.mlpackage/Data/com.apple.CoreML/FeatureDescriptions.json", + f"{model}/{model}.mlpackage/Data/com.apple.CoreML/Metadata.json", + f"{model}/{model}.mlpackage/Data/com.apple.CoreML/weights/weight.bin", + f"{model}/{model}.mlpackage/Data/com.apple.CoreML/{model}.mlmodel", + f"{model}/{model}.mlpackage/Manifest.json", + ] + + for f in files: + p = self.downloadFile( + f"https://github.com/koush/coreml-models/raw/main/{f}", + f"{model_version}/{f}", + ) + modelFile = os.path.dirname(p) labelsFile = self.downloadFile( f"https://github.com/koush/coreml-models/raw/main/{model}/coco_80cl.txt", @@ -92,6 +102,7 @@ async def getSettings(self) -> list[Setting]: "Default", "ssd_mobilenet_v1_coco", "yolov4-tiny", + "yolov8n", ], "value": model, }, @@ -114,14 +125,43 @@ async def detect_once(self, input: Image.Image, settings: Any, src_size, cvss): # run in executor if this is the plugin loop if self.yolo: + input_name = 'image' if self.yolov8 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_1": input}) + predictExecutor, lambda: self.model.predict({input_name: input}) ) else: - out_dict = self.model.predict({"input_1": input}) + out_dict = self.model.predict({input_name: input}) + + if self.yolov8: + out_blob = out_dict["var_914"] + var_914 = out_dict['var_914'] + results = var_914[0] + keep = np.argwhere(results[4:] > 0.2) + for indices in keep: + class_id = indices[0] + index = indices[1] + confidence = results[class_id + 4, index] + x = results[0][index].astype(float) + y = results[1][index].astype(float) + w = results[2][index].astype(float) + h = results[3][index].astype(float) + obj = Prediction( + class_id, + confidence.astype(float), + Rectangle( + x - w / 2, + y - h / 2, + x + w / 2, + y + h / 2, + ), + ) + objs.append(obj) + + ret = self.create_detection_result(objs, src_size, cvss) + return ret + out_blob = out_dict["Identity"] - # out_blob = out_dict["Identity_1"] objects = yolo.parse_yolo_region( out_blob,