From 9fb7f4aa186612806af2becfb621f6ed8d9fdbaf Mon Sep 17 00:00:00 2001
From: Max Tkachenko <makseq@gmail.com>
Date: Mon, 30 Sep 2024 18:21:16 +0200
Subject: [PATCH] feat: Add YOLO11 docs (#641)

---
 label_studio_ml/examples/yolo/README.md       | 36 ++++++++++++++++---
 .../yolo/control_models/timeline_labels.py    |  7 ++++
 .../examples/yolo/utils/neural_nets.py        |  8 ++++-
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/label_studio_ml/examples/yolo/README.md b/label_studio_ml/examples/yolo/README.md
index ab5736c5..7f35c244 100644
--- a/label_studio_ml/examples/yolo/README.md
+++ b/label_studio_ml/examples/yolo/README.md
@@ -223,13 +223,23 @@ Here is an example of a prompt for this. It includes 1000 labels from YOLOv8 cla
 
 </details>
 
-## YOLOv5 and other YOLO versions
+## Supported YOLO Versions: YOLOv5, YOLO11, and others
 
-YOLOv8 models have been successfully tested with this ML backend.
+- **YOLOv5**: This model is supported for object detection tasks. Make sure to specify `model_path="yolov5nu.pt"` (don't forget the **`u`**) to use the YOLOv5 model.
+
+- **YOLOv8**: These models have been successfully tested with the current ML backend. Check the full list of [v8 models here](https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes).
+
+- **YOLO11**: YOLO11 models have also been successfully tested with this ML backend. Check the full list of [v11 models here](https://docs.ultralytics.com/models/yolo11/#supported-tasks-and-modes).
+
+    **Warning 1**: You must upgrade the `ultralytics` package to the latest version (`pip install -U ultralytics`) or rebuild the ML backend Docker from scratch (`docker-compose build --no-cache`) if you used it before the latest Ultralytics update on **Monday, September 30, 2024**.
+
+    **Warning 2**: YOLO11 models do not use the `v` in their naming convention. For example, use **`yolo11n.pt`** instead of `yolov11n.pt`, unlike the naming convention in YOLOv8.
+
+- For a full list of supported YOLO versions and models, refer to the Ultralytics documentation:
+  [Ultralytics Supported YOLO Models](https://docs.ultralytics.com/models/)
+
+    **Note**: Some YOLO models listed in the Ultralytics documentation have not been tested with this ML backend, but they might still work.
 
-Attempts to run YOLOv5 were unsuccessful without modifications. 
-It may be possible to run it by applying some changes, such as installing additional dependencies. 
-The same applies to other YOLO models.
 
 ## Your own custom YOLO models
 
@@ -248,6 +258,22 @@ You can load your own YOLO labels using the following steps:
 You can integrate your own custom-trained YOLOv8 models with the YOLO ML backend for Label Studio. 
 Follow these detailed steps to set up your custom model in the ML backend Docker:
 
+### Step 0: Install Label Studio and clone this Github repository
+
+1. Install and run Label Studio (or see [more ways here](https://labelstud.io/guide/install)). 
+
+```
+pip install label-studio
+label-studio
+```
+
+2. Clone the Label Studio ML backend repository and go to the `yolo` example folder:
+
+```
+git clone https://github.com/HumanSignal/label-studio-ml-backend.git
+cd label-studio-ml-backend/examples/yolo
+```
+
 ### Step 1: Prepare your custom YOLOv8 model
 
 Ensure that your custom YOLOv8 model is saved as a `.pt` file, which is the standard format for PyTorch models. 
diff --git a/label_studio_ml/examples/yolo/control_models/timeline_labels.py b/label_studio_ml/examples/yolo/control_models/timeline_labels.py
index e0e63160..ebf11721 100644
--- a/label_studio_ml/examples/yolo/control_models/timeline_labels.py
+++ b/label_studio_ml/examples/yolo/control_models/timeline_labels.py
@@ -112,6 +112,13 @@ def create_timelines_trainable(self, video_path):
         return regions
 
     def fit(self, event, data, **kwargs):
+        if not self.trainable:
+            logger.debug(
+                'TimelineLabels model is in not trainable mode. '
+                'Use model_trainable="true" to enable training.'
+            )
+            return
+
         """Fit the model."""
         if event == "START_TRAINING":
             # TODO: the full training makes a lot of sense here, but it's not implemented yet
diff --git a/label_studio_ml/examples/yolo/utils/neural_nets.py b/label_studio_ml/examples/yolo/utils/neural_nets.py
index 9c5ba068..3caf3ef4 100644
--- a/label_studio_ml/examples/yolo/utils/neural_nets.py
+++ b/label_studio_ml/examples/yolo/utils/neural_nets.py
@@ -83,7 +83,13 @@ def get_label_map(self):
         return self.label_map
 
     def save(self, path):
-        torch.save(self, path)
+        # ultralytics yolo11 patches torch.save to use dill,
+        # however it leads to serialization errors,
+        # so let's check for use_dill and disable it
+        if 'use_dill' in torch.save.__code__.co_varnames:
+            torch.save(self, path, use_dill=False)
+        else:
+            torch.save(self, path)
         logger.info(f"Model saved to {path}")
 
     @classmethod