Add Ultralytics video detection models - #326
praksharma wants to merge 54 commits into
Conversation
samueljackson92
left a comment
samueljackson92
left a comment
There was a problem hiding this comment.
Broadly looks good. Please also fix the ruff linting warnings.
There was a problem hiding this comment.
Pull request overview
Adds support for Ultralytics-backed object detection models for TokTagger video samples, using TokTagger’s data loader and an in-memory dataset/manifest rather than Ultralytics’ on-disk dataset layout.
Changes:
- Introduces an in-memory Ultralytics detection dataset + custom
DetectionTrainerto train from TokTagger-provided samples/annotations. - Adds YOLO video detection models (including a YOLO26 P2 architecture variant) and frame-by-frame prediction output as
VideoBoundingBoxannotations. - Adds pretrained checkpoint download/caching utilities and pins
ultralyticsas amodelsoptional dependency.
Reviewed changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| toktagger/api/models/ultralytics_detection/video_detection.py | Video frame iteration, training manifest creation, and YOLO-based per-frame prediction emitting VideoBoundingBox annotations. |
| toktagger/api/models/ultralytics_detection/base.py | In-memory dataset, custom Ultralytics trainer adapter, checkpoint discovery, and shared Ultralytics training scaffolding. |
| toktagger/api/models/ultralytics_detection/utils.py | Pretrained checkpoint URL mapping, download/caching, and device/cache-dir helpers. |
| toktagger/api/models/ultralytics_detection/init.py | Package initialization for the new Ultralytics model implementation. |
| toktagger/api/models/init.py | Registers/imports the new Ultralytics YOLO video detection models when model deps are enabled. |
| pyproject.toml | Adds ultralytics==8.4.98 under the models optional dependency group. |
Comments suppressed due to low confidence (1)
toktagger/api/models/ultralytics_detection/video_detection.py:239
decode_frame_image()decodes frames via OpenCV (BGR), but predictions are run on that array without converting to RGB. The training dataset path explicitly converts decoded images to RGB before feeding Ultralytics, so prediction is currently using a different channel order than training/pretrained weights expect, which can significantly degrade detection quality.
image = cv2.imdecode(
encoded_image,
cv2.IMREAD_COLOR,
)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| frame_manifest.append( | ||
| { | ||
| "shot_id": int(sample.shot_id), | ||
| "frame": frame, | ||
| # ImageData stores raw encoded bytes as a JSON-compatible | ||
| # list of integers. Convert it back into bytes here. | ||
| "image": bytes(frame_image.values), | ||
| "boxes": boxes, | ||
| "classes": classes, | ||
| "labels": labels, | ||
| "track_ids": track_ids, | ||
| } | ||
| ) | ||
| sample_record_count += 1 |
There was a problem hiding this comment.
This is acceptable for the initial implementation, but I think we should document an important assumption here. No code changes are required, just think about this and maybe add a comment in the PR description.
The manifest adds every frame from a validated sample. Frames without bounding boxes are passed to YOLO as negative examples. However, validation currently applies to the whole sample, and in practice an annotator may validate a sample without reviewing every frame. In that case, an unreviewed frame containing an object could incorrectly be treated as background.
Are we happy to make the assumption that every frame in a validated sample has been reviewed? If not, a simple short-term alternative would be to include only frames containing validated bounding boxes and discard the remaining frames. The trade-off is that the model would not receive any empty frames as negative examples.
Longer term, we plan to address this through whole-frame labels in #225 as @wk9874 suggested. That would allow annotators to explicitly label frames as, for example, “no UFO”, so the model could distinguish confirmed negative frames from unreviewed frames.
No change is required for this initial PR, but just wanted to make you aware Prakhar, maybe add a small comment at the end of the pr description so it is documented somewhere?
There was a problem hiding this comment.
Good point. I’ve documented this in the PR description.
|
Non-blocking functional feedback: YOLO training currently produces a large amount of terminal output. Would it be worth making Ultralytics quieter by default and keeping the detailed logs behind a verbose/debug setting? The full output is useful for debugging, but most users will probably monitor training through the UI and only need key progress updates, warnings, and errors in the terminal. example of terminal output |
|
To be clear this is a very good PR ! @praksharma Training and prediction works well with no errors. Just some bugs, edge cases and other thoughts I had to get the yolo training to be more robust or user friendly. |
|
@praksharma the weights saving / loading has now changed in the models base class & worker, which should hopefully make your life easier (but will definitely require changes!) Each model is now given a directory named with its model ID, and you can save any number of files in there. There is an optional If See here for more details: #346 |
… for UFO model use case for jet data
|
@wk9874 ready for review and can be merged if your happy with it |
| ) | ||
|
|
||
|
|
||
| def _find_first_useful_frame( |
There was a problem hiding this comment.
Would be good to add tests for this
| assert loaded_paths == [str(expected_path)] | ||
| assert model._trained_weights_path == expected_path | ||
| assert model._prediction_model is sentinel_model | ||
| assert model._trained is True |
There was a problem hiding this comment.
Should add tests for the coarse search stuff here
| return model_path | ||
|
|
||
|
|
||
| def get_torch_device() -> torch.device: |
There was a problem hiding this comment.
I think this may need renaming to not be confusing, since it will ignore cuda GPUs (intentionally)
Alternatively (and probably preferred) - have this function accept a use_cuda: bool flag or something, which basically does the check currently done in base.py L381:
if `use_cuda and torch.cuda.is_available():
return torch.device("cuda")
And then remove that check from base.py L381, instead passing the self.gpu_available into this func as use_cuda.
| return cache_dir | ||
|
|
||
|
|
||
| def resolve_weights_path( |
There was a problem hiding this comment.
The code inside here seems to be duplicated inside the save method in base.py? can we consolidate the two?
| default=0.2, | ||
| ge=0, | ||
| le=1, | ||
| description="Intersection-over-union threshold.", |
There was a problem hiding this comment.
idk what this means - can we rephrase to be clearer?
|
|
||
| return fallback_frame | ||
|
|
||
| return initial_frame |
There was a problem hiding this comment.
I thiiiink the logic in this function is correct, but I find it a bit hard to read. Could we maybe refactor it to be more readable, or add comments which explain the logic?
There was a problem hiding this comment.
I didn't refactor but I did improve variable names and added more comments in commit 0eac05d
… explaining the coarse and fine searches.
Adds the support for Ultralytics-based object detection for video data.
The implementation uses TokTagger’s native data loader and an in-memory dataset, avoiding Ultralytics’ required on-disk dataset structure.
Currently implemented
TODO
RT-DETR x and l modelsUnit testsChanges
optional-dependenciesinpyproject.toml.toktagger/api/models/ultralytics_detection.Model storage
Follow-up
TokTagger tracks whether a model is usable through the Ray actor’s_trainedflag. The generic worker normally restores this state by finding a<model_id>checkpoint and callingwrapped_load().Ultralytics checkpoints use a nested project/model directory instead, so the Ultralytics actor currently finds its checkpoint and restores the trained state itself. Should we think about a more generic checkpoint-discovery hook?Addressed in: #346
Training-data assumption
This initial implementation assumes that every frame in a validated video sample has been reviewed. Frames without validated bounding boxes are included as negative examples.
Whole-frame labels proposed in #225 should eventually allow explicitly reviewed negative frames to be distinguished from unreviewed frames.
More discussion on this topic: #326 (comment)