Skip to content

Commit 5c12977

Browse files
makseqmicaelakaplancaitlinwheeless
authored
feat: RND-118: YOLO for TimelineLabels (#626)
Co-authored-by: micaelakaplan <[email protected]> Co-authored-by: caitlinwheeless <[email protected]>
1 parent c910e30 commit 5c12977

31 files changed

+1833
-220
lines changed

.github/workflows/tests.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ jobs:
9191
run: docker compose -f label_studio_ml/examples/${{ matrix.backend_dir_name }}/docker-compose.yml up -d --build
9292

9393
- name: Wait for stack
94-
timeout-minutes: 10
94+
timeout-minutes: 20
9595
run: |
9696
while [ "$(curl -s -o /dev/null -L -w ''%{http_code}'' "http://localhost:9090/health")" != "200" ]; do
9797
echo "=> Waiting for service to become available" && sleep 2s

label_studio_ml/api.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def _predict():
5656
data = request.json
5757
tasks = data.get('tasks')
5858
label_config = data.get('label_config')
59-
project = data.get('project')
59+
project = str(data.get('project'))
6060
project_id = project.split('.', 1)[0] if project else None
6161
params = data.get('params', {})
6262
context = params.pop('context', {})
@@ -123,8 +123,14 @@ def webhook():
123123
project_id = str(data['project']['id'])
124124
label_config = data['project']['label_config']
125125
model = MODEL_CLASS(project_id, label_config=label_config)
126-
model.fit(event, data)
127-
return jsonify({}), 201
126+
result = model.fit(event, data)
127+
128+
try:
129+
response = jsonify({'result': result, 'status': 'ok'})
130+
except Exception as e:
131+
response = jsonify({'error': str(e), 'status': 'error'})
132+
133+
return response, 201
128134

129135

130136
@_server.route('/health', methods=['GET'])

label_studio_ml/default_configs/_wsgi.py.tmpl

Lines changed: 28 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4,30 +4,33 @@ import json
44
import logging
55
import logging.config
66

7-
logging.config.dictConfig({{
8-
"version": 1,
9-
"disable_existing_loggers": False,
10-
"formatters": {{
11-
"standard": {{
12-
"format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
13-
}}
14-
}},
15-
"handlers": {{
16-
"console": {{
17-
"class": "logging.StreamHandler",
18-
"level": os.getenv('LOG_LEVEL'),
19-
"stream": "ext://sys.stdout",
20-
"formatter": "standard"
21-
}}
22-
}},
23-
"root": {{
24-
"level": os.getenv('LOG_LEVEL'),
25-
"handlers": [
26-
"console"
27-
],
28-
"propagate": True
29-
}}
30-
}})
7+
# Set a default log level if LOG_LEVEL is not defined
8+
log_level = os.getenv("LOG_LEVEL", "INFO")
9+
10+
logging.config.dictConfig(
11+
{
12+
"version": 1,
13+
"disable_existing_loggers": False, # Prevent overriding existing loggers
14+
"formatters": {
15+
"standard": {
16+
"format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
17+
}
18+
},
19+
"handlers": {
20+
"console": {
21+
"class": "logging.StreamHandler",
22+
"level": log_level,
23+
"stream": "ext://sys.stdout",
24+
"formatter": "standard",
25+
}
26+
},
27+
"root": {
28+
"level": log_level,
29+
"handlers": ["console"],
30+
"propagate": True,
31+
},
32+
}
33+
)
3134

3235
from label_studio_ml.api import init_app
3336
from {script} import {model_class}
@@ -60,7 +63,7 @@ if __name__ == "__main__":
6063
'-d', '--debug', dest='debug', action='store_true',
6164
help='Switch debug mode')
6265
parser.add_argument(
63-
'--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=None,
66+
'--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=log_level,
6467
help='Logging level')
6568
parser.add_argument(
6669
'--model-dir', dest='model_dir', default=os.path.dirname(__file__),

label_studio_ml/examples/yolo/.dockerignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
!tests/*
1212
!control_models/*
1313
!models/*
14+
!utils/*
1415

1516
# Include requirements files
1617
!requirements*.txt

label_studio_ml/examples/yolo/Dockerfile

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,14 @@ WORKDIR /app
4949

5050
COPY . ./
5151

52-
WORKDIR /app/models
53-
5452
# Download the YOLO models
55-
RUN yolo predict model=yolov8m.pt source=/app/tests/car.jpg \
56-
&& yolo predict model=yolov8n.pt source=/app/tests/car.jpg \
57-
&& yolo predict model=yolov8n-cls.pt source=/app/tests/car.jpg \
58-
&& yolo predict model=yolov8n-seg.pt source=/app/tests/car.jpg
59-
60-
WORKDIR /app
53+
RUN /bin/sh -c 'if [ ! -f /app/models/yolov8m.pt ]; then \
54+
yolo predict model=/app/models/yolov8m.pt source=/app/tests/car.jpg \
55+
&& yolo predict model=/app/models/yolov8n.pt source=/app/tests/car.jpg \
56+
&& yolo predict model=/app/models/yolov8n-cls.pt source=/app/tests/car.jpg \
57+
&& yolo predict model=/app/models/yolov8n-seg.pt source=/app/tests/car.jpg; \
58+
fi'
59+
60+
ENV PYTHONPATH=/app
6161

6262
CMD ["/app/start.sh"]

label_studio_ml/examples/yolo/README.md

Lines changed: 65 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -36,16 +36,17 @@ making it easier to annotate large datasets and ensure high-quality predictions.
3636

3737
**Supported Features**
3838

39-
| YOLO Task Name | LS Control Tag | Prediction Supported | LS Import Supported | LS Export Supported |
40-
|---------------------------------------|--------------------------------------|----------------------|---------------------|---------------------|
41-
| Object Detection | `<RectangleLabels>` || YOLO, COCO | YOLO, COCO |
42-
| Oriented Bounding Boxes (OBB) | `<RectangleLabels model_obb="true">` || YOLO | YOLO |
43-
| Image Instance Segmentation: Polygons | `<PolygonLabels>` || COCO | YOLO, COCO |
44-
| Image Semantic Segmentation: Masks | `<BrushLabels>` || Native | Native |
45-
| Image Classification | `<Choices>` || Native | Native |
46-
| Pose Detection | `<KeyPoints>` || Native | Native |
47-
| Video Object Tracking | `<VideoRectangle>` || Native | Native |
48-
| Video Temporal Classification | `<TimelineLabels>` | Coming soon | Native | Native |
39+
| YOLO Task Name | LS Control Tag | Prediction Supported | LS Import Supported | LS Export Supported |
40+
|--------------------------------------------------------------|--------------------------------------|----------------------|---------------------|---------------------|
41+
| Object Detection | `<RectangleLabels>` || YOLO, COCO | YOLO, COCO |
42+
| Oriented Bounding Boxes (OBB) | `<RectangleLabels model_obb="true">` || YOLO | YOLO |
43+
| Image Instance Segmentation: Polygons | `<PolygonLabels>` || COCO | YOLO, COCO |
44+
| Image Semantic Segmentation: Masks | `<BrushLabels>` || Native | Native |
45+
| Image Classification | `<Choices>` || Native | Native |
46+
| Pose Detection | `<KeyPoints>` || Native | Native |
47+
| Video Object Tracking | `<VideoRectangle>` || Native | Native |
48+
| [Video Temporal Classification](./README_TIMELINE_LABELS.md) | `<TimelineLabels>` || Native | Native |
49+
4950

5051
* **LS Control Tag**: Label Studio [control tag](https://labelstud.io/tags/) from the labeling configuration.
5152
* **LS Import Supported**: Indicates whether Label Studio supports Import from YOLO format to Label Studio (using the LS converter).
@@ -82,7 +83,7 @@ This tutorial uses the [YOLO example](https://github.com/HumanSignal/label-studi
8283

8384
4. Then from the **Model** page in the project settings, [connect the model](https://labelstud.io/guide/ml#Connect-the-model-to-Label-Studio). The default URL is `http://localhost:9090`.
8485

85-
5. Add images to Label Studio.
86+
5. Add images or video (depending on tasks you are going to solve) to Label Studio.
8687

8788
6. Open any task in the Data Manager and see the predictions from the YOLO model.
8889

@@ -97,11 +98,13 @@ This tutorial uses the [YOLO example](https://github.com/HumanSignal/label-studi
9798

9899
**Control tags**
99100

101+
- `<Choices>` - [Classification](https://labelstud.io/tags/choices); image classification task
100102
- `<RectangleLabels>` - [Bounding boxes](https://labelstud.io/tags/rectanglelabels); object detection task
101103
- `<PolygonLabels>` - [Polygons](https://labelstud.io/tags/polygonlables); segmentation task
102-
- `<VideoRectangle>` - [Video bounding boxes](https://labelstud.io/tags/videorectangle); video object tracking task
104+
- `<VideoRectangle>` - [Video bounding boxes](https://labelstud.io/tags/videorectangle); object tracking task for videos
103105
- `<KeyPointLabels>` - [Key points](https://labelstud.io/tags/keypointlabels); pose detection task
104-
- `<Choices>` - [Classification](https://labelstud.io/tags/choices)
106+
- `<TimelineLabels>` - [Temporal labels for videos](https://labelstud.io/tags/timelinelabels); multi-label temporal classification task for videos
107+
105108

106109
**How to skip the control tag?**
107110

@@ -681,6 +684,55 @@ Small models like `yolov8n.pt` are recommended for real-time tracking, however,
681684

682685
<br>
683686

687+
688+
## Video temporal classification using `TimelineLabels`
689+
690+
This ML backend supports temporal multi-label video classification for the [`<TimelineLabels>` control tag](https://labelstud.io/tags/timelinelabels) in Label Studio.
691+
There are two modes available:
692+
- **Simple:** In the simple mode, the model uses pre-trained YOLO classes to generate predictions without additional training.
693+
- **Trainable:** In the [trainable mode](README_TIMELINE_LABELS.md), the model can be trained on custom labels and annotations submitted in Label Studio using few-shot learning as training is performed on a small number of annotations.
694+
695+
<div align="left">
696+
<a href="https://www.youtube.com/watch?v=tfMn5q1tqKI" title="Video Frame Classification with YOLOv8 and Label Studio">
697+
<img src="http://img.youtube.com/vi/tfMn5q1tqKI/0.jpg" alt="Video Temporal Classification video" style="width:50%;"/>
698+
<br>
699+
Check the video tutorial
700+
</a>
701+
</div>
702+
<br/>
703+
704+
### Labeling config
705+
706+
```xml
707+
<View>
708+
<Video name="video" value="$video"/>
709+
<TimelineLabels
710+
name="label" toName="video"
711+
model_trainable="false" model_score_threshold="0.25">
712+
<Label value="Ball" predicted_values="soccer_ball" />
713+
<Label value="hamster" />
714+
</TimelineLabels>
715+
</View>
716+
```
717+
718+
### Model training
719+
720+
For more details on using the `TimelineLabels` ML backend, including training the model
721+
and adjusting neural network classifier parameters, please refer to
722+
**[README_TIMELINE_LABELS.md](README_TIMELINE_LABELS.md)**.
723+
724+
### Default model
725+
726+
`yolov8n-cls.pt` is the default classification model for simple mode.
727+
728+
729+
<br>
730+
731+
-------------------
732+
733+
<br>
734+
735+
684736
## Run the YOLO ML backend
685737

686738

label_studio_ml/examples/yolo/README_DEVELOP.md

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,13 +47,19 @@ classDiagram
4747
+create_video_rectangles(results, path) List[Dict]
4848
+update_tracker_params(yaml_path: str, prefix: str) str | None
4949
}
50+
51+
class TimelineLabelsModel {
52+
+predict_regions(path: str) List[Dict]
53+
+fit(event, data)
54+
}
5055
5156
ControlModel <|-- RectangleLabelsModel
5257
ControlModel <|-- RectangleLabelsObbModel
5358
ControlModel <|-- PolygonLabelsModel
5459
ControlModel <|-- ChoicesModel
5560
ControlModel <|-- KeyPointLabelsModel
5661
ControlModel <|-- VideoRectangleModel
62+
ControlModel <|-- TimelineLabelsModel
5763
5864
```
5965

@@ -122,6 +128,12 @@ The architecture of the project is modular and is primarily centered around inte
122128
- `create_video_rectangles()`: Processes the output of the tracking model to create a sequence of bounding boxes across video frames.
123129
- `update_tracker_params()`: Customizes the tracking parameters based on settings in Label Studio’s configuration.
124130

131+
8. **`control_models/timelinelabels.py` (TimelineLabelsModel)**:
132+
- **Purpose**: Supports the training of a YOLO model on video data incrementally by updating the model with new annotations as they are submitted.
133+
- **Key Functions**:
134+
- `predict_regions()`: Runs YOLO on video frames and returns the predictions.
135+
- `fit()`: Placeholder method for updating the model with new annotations.
136+
125137
### **Module Interaction**
126138

127139
- **Workflow**: The main workflow begins with `model.py`, which reads tasks and the Label Studio configuration to detect and instantiate the appropriate control models. These control models are responsible for making predictions using the YOLO model and converting the results into a format that Label Studio can use for annotations.

0 commit comments

Comments
 (0)