Skip to content

Commit

Permalink
Merge pull request #4544 from voxel51/add-zoo-detection-models
Browse files Browse the repository at this point in the history 
Add YOLOv10 and RT-DETR Models to Zoo
  • Loading branch information
@jacobmarks
jacobmarks authored Jul 1, 2024
2 parents 1004229 + e72a63f commit ca249a0
Show file tree
Hide file tree
Showing 3 changed files with 294 additions and 4 deletions.
10 changes: 8 additions & 2 deletions docs/source/integrations/ultralytics.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ the following sample dataset:
Object detection
----------------

You can directly pass Ultralytics YOLO detection models to
You can directly pass Ultralytics `YOLO` or `RTDETR` detection models to
:meth:`apply_model() <fiftyone.core.collections.SampleCollection.apply_model>`:

.. code-block:: python
Expand Down Expand Up @@ -81,6 +81,10 @@ You can directly pass Ultralytics YOLO detection models to
# model = YOLO("yolov10l.pt)
# model = YOLO("yolov10x.pt)
# RTDETR
# model = YOLO("rtdetr-l.pt")
# model = YOLO("rtdetr-x.pt")
dataset.apply_model(model, label_field="boxes")
session = fo.launch_app(dataset)
Expand Down Expand Up @@ -111,6 +115,8 @@ You can also load any of these models directly from the
model_name = "yolov5l-coco-torch"
# model_name = "yolov8m-coco-torch"
# model_name = "yolov9e-coco-torch"
# model_name = "yolov10s-coco-torch"
# model_name = "rtdetr-l-coco-torch"
model = foz.load_zoo_model(
model_name,
Expand All @@ -131,7 +137,7 @@ available YOLO models that are compatible with Ultralytics or SuperGradients:
print(foz.list_zoo_models(tags="yolo"))
In general, model names will contain "yolov", followed by the version number,
In general, YOLO model names will contain "yolov", followed by the version number,
then the model size ("n", "s", "m", "l", or "x"), and an indicator of the
label classes ("coco" for MS COCO or "world" for open-world), followed by
"torch".
Expand Down
64 changes: 62 additions & 2 deletions fiftyone/utils/ultralytics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,6 @@
import numpy as np
from PIL import Image

import eta.core.utils as etau

from fiftyone.core.config import Config
import fiftyone.core.labels as fol
from fiftyone.core.models import Model
Expand Down Expand Up @@ -437,6 +435,68 @@ def predict_all(self, args):
return self._format_predictions(predictions)


class FiftyOneRTDETRModelConfig(FiftyOneYOLOModelConfig):
"""Configuration for a :class:`FiftyOneRTDETRModel`.
Args:
model (None): an ``ultralytics.RTDETR`` model to use
model_name (None): the name of an ``ultralytics.RTDETR`` model to load
model_path (None): the path to an ``ultralytics.RTDETR`` model checkpoint
"""

def __init__(self, d):
pass


class FiftyOneRTDETRModel(Model):
"""FiftyOne wrapper around an ``ultralytics.RTDETR`` model.
Args:
config: a :class:`FiftyOneRTDETRModelConfig`
"""

def __init__(self, config):
self.config = config
self.model = self._load_model(config)

def _load_model(self, config):
if config.model is not None:
return config.model

if config.model_path is not None:
model = ultralytics.RTDETR(config.model_path)
elif config.model_name is not None:
model = ultralytics.RTDETR(config.model_name)
else:
model = ultralytics.RTDETR()

return model

@property
def media_type(self):
return "image"

@property
def ragged_batches(self):
return False

@property
def transforms(self):
return None

@property
def preprocess(self):
return False

def _format_predictions(self, predictions):
return to_detections(predictions)

def predict(self, arg):
image = Image.fromarray(arg)
predictions = self.model(image, verbose=False)
return self._format_predictions(predictions[0])


class FiftyOneYOLOOBBModelConfig(FiftyOneYOLOModelConfig):
pass

Expand Down
224 changes: 224 additions & 0 deletions fiftyone/zoo/models/manifest-torch.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2799,6 +2799,230 @@
"tags": ["segmentation", "coco", "torch", "yolo"],
"date_added": "2024-04-02 19:22:51"
},
{
"base_name": "yolov10n-coco-torch",
"base_filename": "yolov10n-coco.pt",
"description": "YOLOv10-N model trained on COCO",
"source": "https://docs.ultralytics.com/models/yolov10/",
"size_bytes": 5860383,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov10n.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneYOLODetectionModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "yolo"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "yolov10s-coco-torch",
"base_filename": "yolov10s-coco.pt",
"description": "YOLOv10-S model trained on COCO",
"source": "https://docs.ultralytics.com/models/yolov10/",
"size_bytes": 16623111,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov10s.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneYOLODetectionModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "yolo"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "yolov10m-coco-torch",
"base_filename": "yolov10m-coco.pt",
"description": "YOLOv10-M model trained on COCO",
"source": "https://docs.ultralytics.com/models/yolov10/",
"size_bytes": 33643667,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov10m.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneYOLODetectionModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "yolo"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "yolov10l-coco-torch",
"base_filename": "yolov10l-coco.pt",
"description": "YOLOv10-L model trained on COCO",
"source": "https://docs.ultralytics.com/models/yolov10/",
"size_bytes": 52425230,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov10l.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneYOLODetectionModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "yolo"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "yolov10x-coco-torch",
"base_filename": "yolov10x-coco.pt",
"description": "YOLOv10-X model trained on COCO",
"source": "https://docs.ultralytics.com/models/yolov10/",
"size_bytes": 64395854,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov10x.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneYOLODetectionModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "yolo"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "rtdetr-l-coco-torch",
"base_filename": "rtdetr-l-coco.pt",
"description": "RT-DETR-l model trained on COCO",
"source": "https://docs.ultralytics.com/models/rtdetr/",
"size_bytes": 66511432,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/rtdetr-l.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneRTDETRModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "transformer"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "rtdetr-x-coco-torch",
"base_filename": "rtdetr-x-coco.pt",
"description": "RT-DETR-x model trained on COCO",
"source": "https://docs.ultralytics.com/models/rtdetr/",
"size_bytes": 135755662,
"manager": {
"type": "fiftyone.core.models.ModelManager",
"config": {
"url": "https://github.com/ultralytics/assets/releases/download/v8.2.0/rtdetr-x.pt"
}
},
"default_deployment_config_dict": {
"type": "fiftyone.utils.ultralytics.FiftyOneRTDETRModel",
"config": {}
},
"requirements": {
"packages": [
"torch>=1.7.0",
"torchvision>=0.8.1",
"ultralytics>=8.2.0"
],
"cpu": {
"support": true
},
"gpu": {
"support": true
}
},
"tags": ["detection", "coco", "torch", "transformer"],
"date_added": "2024-07-01 19:22:51"
},
{
"base_name": "yolov8s-world-torch",
"base_filename": "yolov8s-world.pt",
Expand Down

0 comments on commit ca249a0

Please sign in to comment.