update

ultralytics/models/yolov10/__init__.py

@@ -0,0 +1,5 @@
+from .model import YOLOv10
+from .predict import YOLOv10DetectionPredictor
+from .val import YOLOv10DetectionValidator
+
+__all__ = "YOLOv10DetectionPredictor", "YOLOv10DetectionValidator", "YOLOv10"

ultralytics/models/yolov10/card.py

@@ -0,0 +1,64 @@
+card_template_text = """
+---
+license: agpl-3.0
+library_name: ultralytics
+repo_url: https://github.com/THU-MIG/yolov10
+tags:
+- object-detection
+- computer-vision
+- yolov10
+datasets:
+- detection-datasets/coco
+inference: false
+---
+
+### Model Description
+[YOLOv10: Real-Time End-to-End Object Detection](https://arxiv.org/abs/2405.14458v1)
+
+- arXiv: https://arxiv.org/abs/2405.14458v1
+- github: https://github.com/THU-MIG/yolov10
+
+### Installation
+```
+pip install git+https://github.com/THU-MIG/yolov10.git
+```
+
+### Training and validation
+```python
+from ultralytics import YOLOv10
+
+model = YOLOv10.from_pretrained('jameslahm/yolov10n')
+# Training
+model.train(...)
+# after training, one can push to the hub
+model.push_to_hub("your-hf-username/yolov10-finetuned")
+
+# Validation
+model.val(...)
+```
+
+### Inference
+
+Here's an end-to-end example showcasing inference on a cats image:
+
+```python
+from ultralytics import YOLOv10
+
+model = YOLOv10.from_pretrained('jameslahm/yolov10n')
+source = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+model.predict(source=source, save=True)
+```
+which shows:
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/628ece6054698ce61d1e7be3/tBwAsKcQA_96HCYQp7BRr.png)
+
+### BibTeX Entry and Citation Info
+```
+@article{wang2024yolov10,
+  title={YOLOv10: Real-Time End-to-End Object Detection},
+  author={Wang, Ao and Chen, Hui and Liu, Lihao and Chen, Kai and Lin, Zijia and Han, Jungong and Ding, Guiguang},
+  journal={arXiv preprint arXiv:2405.14458},
+  year={2024}
+}
+```
+""".strip()

ultralytics/models/yolov10/model.py

@@ -0,0 +1,36 @@
+from ultralytics.engine.model import Model
+from ultralytics.nn.tasks import YOLOv10DetectionModel
+from .val import YOLOv10DetectionValidator
+from .predict import YOLOv10DetectionPredictor
+from .train import YOLOv10DetectionTrainer
+
+from huggingface_hub import PyTorchModelHubMixin
+from .card import card_template_text
+
+class YOLOv10(Model, PyTorchModelHubMixin, model_card_template=card_template_text):
+
+    def __init__(self, model="yolov10n.pt", task=None, verbose=False, 
+                 names=None):
+        super().__init__(model=model, task=task, verbose=verbose)
+        if names is not None:
+            setattr(self.model, 'names', names)
+
+    def push_to_hub(self, repo_name, **kwargs):
+        config = kwargs.get('config', {})
+        config['names'] = self.names
+        config['model'] = self.model.yaml['yaml_file']
+        config['task'] = self.task
+        kwargs['config'] = config
+        super().push_to_hub(repo_name, **kwargs)
+
+    @property
+    def task_map(self):
+        """Map head to model, trainer, validator, and predictor classes."""
+        return {
+            "detect": {
+                "model": YOLOv10DetectionModel,
+                "trainer": YOLOv10DetectionTrainer,
+                "validator": YOLOv10DetectionValidator,
+                "predictor": YOLOv10DetectionPredictor,
+            },
+        }

ultralytics/models/yolov10/predict.py

@@ -0,0 +1,38 @@
+from ultralytics.models.yolo.detect import DetectionPredictor
+import torch
+from ultralytics.utils import ops
+from ultralytics.engine.results import Results
+
+
+class YOLOv10DetectionPredictor(DetectionPredictor):
+    def postprocess(self, preds, img, orig_imgs):
+        if isinstance(preds, dict):
+            preds = preds["one2one"]
+
+        if isinstance(preds, (list, tuple)):
+            preds = preds[0]
+
+        if preds.shape[-1] == 6:
+            pass
+        else:
+            preds = preds.transpose(-1, -2)
+            bboxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, preds.shape[-1]-4)
+            bboxes = ops.xywh2xyxy(bboxes)
+            preds = torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
+
+        mask = preds[..., 4] > self.args.conf
+        if self.args.classes is not None:
+            mask = mask & (preds[..., 5:6] == torch.tensor(self.args.classes, device=preds.device).unsqueeze(0)).any(2)
+        
+        preds = [p[mask[idx]] for idx, p in enumerate(preds)]
+
+        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
+            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
+
+        results = []
+        for i, pred in enumerate(preds):
+            orig_img = orig_imgs[i]
+            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
+            img_path = self.batch[0][i]
+            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
+        return results

ultralytics/models/yolov10/train.py

@@ -0,0 +1,20 @@
+from ultralytics.models.yolo.detect import DetectionTrainer
+from .val import YOLOv10DetectionValidator
+from .model import YOLOv10DetectionModel
+from copy import copy
+from ultralytics.utils import RANK
+
+class YOLOv10DetectionTrainer(DetectionTrainer):
+    def get_validator(self):
+        """Returns a DetectionValidator for YOLO model validation."""
+        self.loss_names = "box_om", "cls_om", "dfl_om", "box_oo", "cls_oo", "dfl_oo", 
+        return YOLOv10DetectionValidator(
+            self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
+        )
+
+    def get_model(self, cfg=None, weights=None, verbose=True):
+        """Return a YOLO detection model."""
+        model = YOLOv10DetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
+        if weights:
+            model.load(weights)
+        return model

ultralytics/models/yolov10/val.py

@@ -0,0 +1,24 @@
+from ultralytics.models.yolo.detect import DetectionValidator
+from ultralytics.utils import ops
+import torch
+
+class YOLOv10DetectionValidator(DetectionValidator):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.args.save_json |= self.is_coco
+
+    def postprocess(self, preds):
+        if isinstance(preds, dict):
+            preds = preds["one2one"]
+
+        if isinstance(preds, (list, tuple)):
+            preds = preds[0]
+        
+        # Acknowledgement: Thanks to sanha9999 in #190 and #181!
+        if preds.shape[-1] == 6:
+            return preds
+        else:
+            preds = preds.transpose(-1, -2)
+            boxes, scores, labels = ops.v10postprocess(preds, self.args.max_det, self.nc)
+            bboxes = ops.xywh2xyxy(boxes)
+            return torch.cat([bboxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)