Add YOLO OBB models to zoo

jainanisha90 · Apr 5, 2024 · da43822 · da43822
1 parent 4277d83
commit da43822
Show file tree

Hide file tree

Showing 2 changed files with 162 additions and 2 deletions.
diff --git a/fiftyone/utils/ultralytics.py b/fiftyone/utils/ultralytics.py
@@ -437,7 +437,7 @@ def predict_all(self, args):
         return self._format_predictions(predictions)
 
 
-class FiftyOneYOLOOBBConfig(FiftyOneYOLOModelConfig):
+class FiftyOneYOLOOBBModelConfig(FiftyOneYOLOModelConfig):
     pass
 
 
@@ -503,7 +503,7 @@ def _convert_yolo_detection_model(model):
 
 
 def _convert_yolo_obb_model(model):
-    config = FiftyOneYOLOOBBConfig({"model": model})
+    config = FiftyOneYOLOOBBModelConfig({"model": model})
     return FiftyOneYOLOOBBModel(config)
 
 

diff --git a/fiftyone/zoo/models/manifest-torch.json b/fiftyone/zoo/models/manifest-torch.json
@@ -2927,6 +2927,166 @@
             "tags": ["detection", "torch", "yolo", "zero-shot"],
             "date_added": "2024-03-11 19:22:51"
         },
+        {
+            "base_name": "yolov8n-obb-dotav1-torch",
+            "base_filename": "yolov8n-obb.pt",
+            "description": "YOLOv8n Oriented Bounding Box model",
+            "source": "https://docs.ultralytics.com/tasks/obb/",
+            "size_bytes": 6548034,
+            "manager": {
+                "type": "fiftyone.core.models.ModelManager",
+                "config": {
+                    "url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8n-obb.pt"
+                }
+            },
+            "default_deployment_config_dict": {
+                "type": "fiftyone.utils.ultralytics.FiftyOneYOLOOBBModel",
+                "config": {}
+            },
+            "requirements": {
+                "packages": [
+                    "torch>=1.7.0",
+                    "torchvision>=0.8.1",
+                    "ultralytics>=8.1.0"
+                ],
+                "cpu": {
+                    "support": true
+                },
+                "gpu": {
+                    "support": true
+                }
+            },
+            "tags": ["detection", "torch", "yolo", "polylines", "obb"],
+            "date_added": "2024-04-05 19:22:51"
+        },
+        {
+            "base_name": "yolov8s-obb-dotav1-torch",
+            "base_filename": "yolov8s-obb.pt",
+            "description": "YOLOv8s Oriented Bounding Box model",
+            "source": "https://docs.ultralytics.com/tasks/obb/",
+            "size_bytes": 23245186,
+            "manager": {
+                "type": "fiftyone.core.models.ModelManager",
+                "config": {
+                    "url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s-obb.pt"
+                }
+            },
+            "default_deployment_config_dict": {
+                "type": "fiftyone.utils.ultralytics.FiftyOneYOLOOBBModel",
+                "config": {}
+            },
+            "requirements": {
+                "packages": [
+                    "torch>=1.7.0",
+                    "torchvision>=0.8.1",
+                    "ultralytics>=8.1.0"
+                ],
+                "cpu": {
+                    "support": true
+                },
+                "gpu": {
+                    "support": true
+                }
+            },
+            "tags": ["detection", "torch", "yolo", "polylines", "obb"],
+            "date_added": "2024-04-05 19:22:51"
+        },
+        {
+            "base_name": "yolov8m-obb-dotav1-torch",
+            "base_filename": "yolov8m-obb.pt",
+            "description": "YOLOv8m Oriented Bounding Box model",
+            "source": "https://docs.ultralytics.com/tasks/obb/",
+            "size_bytes": 53304682,
+            "manager": {
+                "type": "fiftyone.core.models.ModelManager",
+                "config": {
+                    "url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8m-obb.pt"
+                }
+            },
+            "default_deployment_config_dict": {
+                "type": "fiftyone.utils.ultralytics.FiftyOneYOLOOBBModel",
+                "config": {}
+            },
+            "requirements": {
+                "packages": [
+                    "torch>=1.7.0",
+                    "torchvision>=0.8.1",
+                    "ultralytics>=8.1.0"
+                ],
+                "cpu": {
+                    "support": true
+                },
+                "gpu": {
+                    "support": true
+                }
+            },
+            "tags": ["detection", "torch", "yolo", "polylines", "obb"],
+            "date_added": "2024-04-05 19:22:51"
+        },
+        {
+            "base_name": "yolov8l-obb-dotav1-torch",
+            "base_filename": "yolov8l-obb.pt",
+            "description": "YOLOv8l Oriented Bounding Box model",
+            "source": "https://docs.ultralytics.com/tasks/obb/",
+            "size_bytes": 89504274,
+            "manager": {
+                "type": "fiftyone.core.models.ModelManager",
+                "config": {
+                    "url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8l-obb.pt"
+                }
+            },
+            "default_deployment_config_dict": {
+                "type": "fiftyone.utils.ultralytics.FiftyOneYOLOOBBModel",
+                "config": {}
+            },
+            "requirements": {
+                "packages": [
+                    "torch>=1.7.0",
+                    "torchvision>=0.8.1",
+                    "ultralytics>=8.1.0"
+                ],
+                "cpu": {
+                    "support": true
+                },
+                "gpu": {
+                    "support": true
+                }
+            },
+            "tags": ["detection", "torch", "yolo", "polylines", "obb"],
+            "date_added": "2024-04-05 19:22:51"
+        },
+        {
+            "base_name": "yolov8x-obb-dotav1-torch",
+            "base_filename": "yolov8x-obb.pt",
+            "description": "YOLOv8x Oriented Bounding Box model",
+            "source": "https://docs.ultralytics.com/tasks/obb/",
+            "size_bytes": 139533138,
+            "manager": {
+                "type": "fiftyone.core.models.ModelManager",
+                "config": {
+                    "url": "https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8x-obb.pt"
+                }
+            },
+            "default_deployment_config_dict": {
+                "type": "fiftyone.utils.ultralytics.FiftyOneYOLOOBBModel",
+                "config": {}
+            },
+            "requirements": {
+                "packages": [
+                    "torch>=1.7.0",
+                    "torchvision>=0.8.1",
+                    "ultralytics>=8.1.0"
+                ],
+                "cpu": {
+                    "support": true
+                },
+                "gpu": {
+                    "support": true
+                }
+            },
+            "tags": ["detection", "torch", "yolo", "polylines", "obb"],
+            "date_added": "2024-04-05 19:22:51"
+        },
         {
             "base_name": "yolo-nas-torch",
             "base_filename": null,