Merge pull request voxel51#4238 from voxel51/yolo-obb-inference

[Enhancement] Add YOLOv8 OBB Models Directly to Zoo and Document
jainanisha90 · Apr 6, 2024 · aff97ce · aff97ce
2 parents 5f271ca + 04344ae
commit aff97ce
Show file tree

Hide file tree

Showing 3 changed files with 212 additions and 2 deletions.
diff --git a/docs/source/integrations/ultralytics.rst b/docs/source/integrations/ultralytics.rst
@@ -241,6 +241,56 @@ manually convert Ultralytics predictions to :ref:`FiftyOne format <keypoints>`:
    :alt: ultralytics-keypoints
    :align: center
 
+
+.. _ultralytics-oriented-bounding-boxes:
+
+Oriented bounding boxes
+-----------------------
+
+You can directly pass Ultralytics YOLO oriented bounding box models to
+:meth:`apply_model() <fiftyone.core.collections.SampleCollection.apply_model>`:
+
+.. code-block:: python
+    :linenos:
+
+    model = YOLO("yolov8n-obb.pt")
+    # model = YOLO("yolov8s-obb.pt")
+    # model = YOLO("yolov8m-obb.pt")
+    # model = YOLO("yolov8l-obb.pt")
+    # model = YOLO("yolov8x-obb.pt")
+
+    dataset.apply_model(model, label_field="oriented_boxes")
+
+    session = fo.launch_app(dataset)
+
+
+You can also load YOLOv8 oriented bounding box models from the
+:ref:`FiftyOne Model Zoo <model-zoo>`:
+
+.. code-block:: python
+    :linenos:
+
+    model_name = "yolov8n-obb-dotav1-torch"
+    # model_name = "yolov8s-obb-dotav1-torch"
+    # model_name = "yolov8m-obb-dotav1-torch"
+    # model_name = "yolov8l-obb-dotav1-torch"
+    # model_name = "yolov8x-obb-dotav1-torch"
+
+    model = foz.load_zoo_model(model_name)
+
+    dataset.apply_model(model, label_field="oriented_boxes")
+
+    session = fo.launch_app(dataset)
+
+
+.. note::
+
+    The oriented bounding box models are trained on the `DOTA dataset
+    <https://captain-whu.github.io/DOTA/index.html>`_, which consists of
+    drone images with oriented bounding boxes. The models are trained to
+    predict on bird's eye view images, so applying them to regular images
+    may not yield good results.
+
 .. _ultralytics-open-vocabulary-object-detection:
 
 Open vocabulary detection

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,