voxel51 · jacobmarks · Apr 5, 2024 · Apr 4, 2024 · Apr 4, 2024
diff --git a/fiftyone/utils/ultralytics.py b/fiftyone/utils/ultralytics.py
@@ -42,7 +42,10 @@ def convert_ultralytics_model(model):
     elif isinstance(model.model, ultralytics.nn.tasks.PoseModel):
         return _convert_yolo_pose_model(model)
     elif isinstance(model.model, ultralytics.nn.tasks.DetectionModel):
-        return _convert_yolo_detection_model(model)
+        if isinstance(model.model, ultralytics.nn.tasks.OBBModel):
+            return _convert_yolo_obb_model(model)
+        else:
+            return _convert_yolo_detection_model(model)
     else:
         raise ValueError(
             "Unsupported model type; cannot convert %s to a FiftyOne model"
@@ -166,6 +169,59 @@ def _to_instances(result, confidence_thresh=None):
     return fol.Detections(detections=detections)
 
 
+def obb_to_polylines(results, confidence_thresh=None, filled=False):
+    """Converts ``ultralytics.YOLO`` instance segmentations to FiftyOne format.
+
+    Args:
+        results: a single or list of ``ultralytics.engine.results.Results``
+        confidence_thresh (None): a confidence threshold to filter boxes
+        filled (False): whether the polyline should be filled
+
+    Returns:
+        a single or list of :class:`fiftyone.core.labels.PolyLines`
+    """
+
+    single = not isinstance(results, list)
+    if single:
+        results = [results]
+
+    batch = [
+        _obb_to_polylines(r, filled, confidence_thresh=confidence_thresh)
+        for r in results
+    ]
+
+    if single:
+        return batch[0]
+
+    return batch
+
+
+def _obb_to_polylines(result, filled, confidence_thresh=None):
+    if result.obb is None:
+        return None
+    classes = np.rint(result.obb.cls.detach().cpu().numpy()).astype(int)
+    confs = result.obb.conf.detach().cpu().numpy().astype(float)
+    points = result.obb.xyxyxyxyn.detach().cpu().numpy()
+    polylines = []
+    for cls, _points, conf in zip(classes, points, confs):
+        if confidence_thresh is not None and conf < confidence_thresh:
+            continue
+
+        _points = [_points.astype(float)]
+
+        label = result.names[cls]
+
+        polyline = fol.Polyline(
+            label=label,
+            points=_points,
+            confidence=conf,
+            closed=True,
+            filled=filled,
+        )
+        polylines.append(polyline)
+    return fol.Polylines(polylines=polylines)
+
+
 def to_polylines(results, confidence_thresh=None, tolerance=2, filled=True):
     """Converts ``ultralytics.YOLO`` instance segmentations to FiftyOne format.
 
@@ -381,6 +437,26 @@ def predict_all(self, args):
         return self._format_predictions(predictions)
 
 
+class FiftyOneYOLOOBBConfig(FiftyOneYOLOModelConfig):
+    pass
+
+
+class FiftyOneYOLOOBBModel(FiftyOneYOLOModel):
+    """FiftyOne wrapper around an Ultralytics YOLO OBB detection model.
+
+    Args:
+        config: a :class:`FiftyOneYOLOConfig`
+    """
+
+    def _format_predictions(self, predictions):
+        return obb_to_polylines(predictions)
+
+    def predict_all(self, args):
+        images = [Image.fromarray(arg) for arg in args]
+        predictions = self.model(images, verbose=False)
+        return self._format_predictions(predictions)
+
+
 class FiftyOneYOLOSegmentationModelConfig(FiftyOneYOLOModelConfig):
     pass
 
@@ -426,6 +502,11 @@ def _convert_yolo_detection_model(model):
     return FiftyOneYOLODetectionModel(config)
 
 
+def _convert_yolo_obb_model(model):
+    config = FiftyOneYOLOOBBConfig({"model": model})
+    return FiftyOneYOLOOBBModel(config)
+
+
 def _convert_yolo_segmentation_model(model):
     config = FiftyOneYOLOSegmentationModelConfig({"model": model})
     return FiftyOneYOLOSegmentationModel(config)