precommit format

dnth · Nov 1, 2024 · 7c7eb5d · 7c7eb5d
1 parent ebb47df
commit 7c7eb5d
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 33 deletions.
diff --git a/pre-commit-config.yaml → .pre-commit-config.yaml b/pre-commit-config.yaml → .pre-commit-config.yaml
diff --git a/nbs/serve_openai.py b/nbs/serve_openai.py
@@ -3,7 +3,6 @@
 from typing import Dict, List
 
 import numpy as np
-import ray
 import requests
 import timm
 import torch

diff --git a/requirements_dev.txt b/requirements_dev.txt
@@ -33,4 +33,5 @@ transformers
 timm
 ultralytics
 pytest
-pytest-mock
+pytest-mock
+pre-commit
diff --git a/xinfer/ultralytics/ultralytics_model.py b/xinfer/ultralytics/ultralytics_model.py
@@ -13,59 +13,68 @@ def __init__(
         self, model_id: str, device: str = "cpu", dtype: str = "float32", **kwargs
     ):
         super().__init__(model_id, device, dtype)
-        self.model_type = 'classification' if 'cls' in model_id else 'detection'
+        self.model_type = "classification" if "cls" in model_id else "detection"
         self.load_model(**kwargs)
 
     def load_model(self, **kwargs):
-        if self.model_type == 'classification':
-            self.model = YOLO(self.model_id.replace("ultralytics/", ""), task='classification', **kwargs)
+        if self.model_type == "classification":
+            self.model = YOLO(
+                self.model_id.replace("ultralytics/", ""),
+                task="classification",
+                **kwargs,
+            )
         else:
             self.model = YOLO(self.model_id.replace("ultralytics/", ""), **kwargs)
 
     @track_inference
     def infer_batch(self, images: str | List[str], **kwargs) -> List[List[Dict]]:
         half = self.dtype == torch.float16
-        self.results = self.model.predict(images, device=self.device, half=half, **kwargs)
+        self.results = self.model.predict(
+            images, device=self.device, half=half, **kwargs
+        )
         batch_results = []
         for result in self.results:
-
-            if self.model_type == 'classification':
+            if self.model_type == "classification":
                 classification_results = []
                 probs = result.probs
-                classification_results.append({
-                    "class_id": int(probs.top1),
-                    "score": float(probs.top1conf.cpu().numpy()),
-                    "class_name": result.names[int(probs.top1)],
-                })
+                classification_results.append(
+                    {
+                        "class_id": int(probs.top1),
+                        "score": float(probs.top1conf.cpu().numpy()),
+                        "class_name": result.names[int(probs.top1)],
+                    }
+                )
                 batch_results.append(classification_results)
-                
+
             else:
                 detection_results = []
                 boxes = result.boxes
                 for box in boxes:
                     x1, y1, x2, y2 = box.xyxy[0].tolist()
                     width = x2 - x1
                     height = y2 - y1
-                    detection_results.append({
-                        "bbox": [x1, y1, width, height],
-                        "category_id": int(box.cls),
-                        "score": float(box.conf),
-                        "class_name": result.names[int(box.cls)],
-                    })
+                    detection_results.append(
+                        {
+                            "bbox": [x1, y1, width, height],
+                            "category_id": int(box.cls),
+                            "score": float(box.conf),
+                            "class_name": result.names[int(box.cls)],
+                        }
+                    )
                 batch_results.append(detection_results)
-                
+
         return batch_results
 
     @track_inference
     def infer(self, image: str, **kwargs) -> List[List[Dict]]:
         results = self.infer_batch([image], **kwargs)
         return results[0]
-    
-    def render(self, save_path: str = './', **kwargs):
+
+    def render(self, save_path: str = "./", **kwargs):
         for _, r in enumerate(self.results):
             # im_bgr = r.plot()
             # im_rgb = Image.fromarray(im_bgr[..., ::-1])
-            
+
             # plot results (such as bounding boxes, masks, keypoints, and probabilities)
             file_name = os.path.basename(r.path)
             file_name = os.path.join(save_path, file_name)

diff --git a/xinfer/ultralytics/yolov11.py b/xinfer/ultralytics/yolov11.py
@@ -6,10 +6,18 @@
 @register_model("ultralytics/yolov11s", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
 @register_model("ultralytics/yolov11m", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
 @register_model("ultralytics/yolov11l", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
-@register_model("ultralytics/yolov11n-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov11s-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov11m-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov11l-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
+@register_model(
+    "ultralytics/yolov11n-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov11s-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov11m-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov11l-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
 class YOLOv11(UltralyticsModel):
     def __init__(self, model_id: str, **kwargs):
         model_id = model_id.replace("v", "")

diff --git a/xinfer/ultralytics/yolov8.py b/xinfer/ultralytics/yolov8.py
@@ -7,11 +7,21 @@
 @register_model("ultralytics/yolov8l", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
 @register_model("ultralytics/yolov8m", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
 @register_model("ultralytics/yolov8x", "ultralytics", ModelInputOutput.IMAGE_TO_BOXES)
-@register_model("ultralytics/yolov8n-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov8s-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov8l-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov8m-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
-@register_model("ultralytics/yolov8x-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES)
+@register_model(
+    "ultralytics/yolov8n-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov8s-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov8l-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov8m-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
+@register_model(
+    "ultralytics/yolov8x-cls", "ultralytics", ModelInputOutput.IMAGE_TO_CATEGORIES
+)
 class YOLOv8(UltralyticsModel):
     def __init__(self, model_id: str, **kwargs):
         super().__init__(model_id, **kwargs)