Add Ultralytics video detection models - #326
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.
| def build_video_frame_manifest( | ||
| samples: list[Sample], | ||
| annotations: list[list[Annotation]], | ||
| class_map: dict[str, int], | ||
| data_loader: TokTaggerDataLoader, | ||
| ) -> list[DetectionRecord]: | ||
| """ | ||
| Convert validated video samples into frame-level training records. | ||
| Hash table has been used to speed up things. | ||
| """ |
| if isinstance(frame_image.values, str): | ||
| raise TypeError( | ||
| "Expected raw image bytes but received a base64 string." | ||
| ) |
There was a problem hiding this comment.
Some video loaders, including ArrayDataLoader and UDACameraDataLoader, return a base64 string even when return_raw=True. This check therefore makes YOLO training and prediction fail for those loaders.
Could we support both raw bytes and base64 images here? Otherwise, we should clearly restrict this model to loaders that support raw image bytes.
Here is an idea to support both image formats in this model below.
The model could decode base64 values as well as accepting raw bytes:
if isinstance(frame_image.values, str):
image_bytes = base64.b64decode(frame_image.values)
else:
image_bytes = bytes(frame_image.values)This would allow images returned by ArrayDataLoader and UDACameraDataLoader to reach YOLO.
There was a problem hiding this comment.
Instead of that, I would just make sure all image-based data loaders have the return_raw option implemented...
Potentially we should make another image dataloader base class which implements that logic, which each of ImageDataLoader, ArrayDataLoader, UDACameraDataLoader can inherit from
| ), | ||
| ) | ||
| except FileNotFoundError: | ||
| break |
There was a problem hiding this comment.
This frame loop currently reaches the end of a video cleanly only when using ImageDataLoader, because it expects the loader to raise FileNotFoundError.
I was looking into the other valid video loaders and found they behave differently:
ArrayDataLoaderraisesDataLoaderErrorfor an out-of-range frameUDACameraDataLoaderwraps its loading failures inDataLoaderError.
As a result, YOLO training or prediction can fail for other data loaders instead of stopping after the final frame.
Below is an idea to give all video loaders the same end-of-video signal.
Simply catching every DataLoaderError would not be safe:
except DataLoaderError:
breakDataLoaderError can also mean that a NumPy file is invalid, the data has the wrong shape, or an external data source failed. Catching it here would hide those real errors and incorrectly treat them as the end of the video.
One option would be to introduce a more specific exception in the data-loader layer:
class FrameNotFoundError(DataLoaderError):
"""The requested video frame does not exist."""Each video loader would raise FrameNotFoundError only when the requested frame is unavailable. The YOLO iterator could then stop safely without hiding other failures:
except FrameNotFoundError:
breakThis approach would require changes in:
toktagger/api/core/data_loaders.pytoktagger/api/models/ultralytics_detection/video_detection.pytoktagger/api/models/temp.py
That shared loader change may be outside the scope of this PR to be fair. If this PR intentionally supports only ImageDataLoader, could we validate that explicitly, return a clear error for unsupported loaders and put this in the PR description? A follow up PR could then introduce the common frame @wk9874, what do you think? or should we do it here.
There was a problem hiding this comment.
I agree with implementing FrameNotFoundError, probably worth doing it either as part of this PR, or as a separate small PR which is merged before this one
| 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.
| overrides = { | ||
| "model": model_path, | ||
| "epochs": epochs, | ||
| "batch": self.batch, | ||
| "imgsz": self.imgsz, | ||
| "workers": self.workers, | ||
| "device": self.get_device().type, | ||
| "project": str(training_output_root), | ||
| "name": self.id, | ||
| "exist_ok": True, | ||
| "save": True, | ||
| "plots": False, | ||
| "val": has_validation_data, | ||
| "close_mosaic": 0, | ||
| } | ||
| if learning_rate > 0: | ||
| overrides["lr0"] = learning_rate | ||
| return overrides |
There was a problem hiding this comment.
Could you do me a favour and try GPU training and prediction, then check git status? After GPU training, I consistently get an untracked yolo26n.pt in the TokTagger repository root:
Untracked files:
yolo26n.pt
It looks like Ultralytics’ AMP check downloads a second copy instead of using the checkpoint already in TokTagger’s cache. Do you see the same behaviour?
If so, could we configure Ultralytics to reuse the pretrained checkpoint under MODEL_STORAGE so weight files are not created in the source directory?
Although I comment on this override block the main code block responsible is below on line 461 I think.
trainer = ToktaggerDetectionTrainer(
overrides=overrides,
train_dataset=train_dataset,
val_dataset=validation_dataset,
class_names=self.class_names,
progress_callback=self.log_progress,
)
trainer.train()There was a problem hiding this comment.
Thanks to copilot, I was able to find the culprit in thousands of lines of code. For GPU training there is a list of checks being carried out. And this specific line is the culprit:
INterestingly, Ultralytics has two different settings. overrides are per-training-run arguments and settings are global Ultralytics configuration, stored in a temp file settings.json.
All the settings are defined here:
More specifically, I think we just change the weights_dir in the __init__ method of that class and I hope it works. Unfortunately, I don't have a laptop with GPU to test this. So, I am not sure if the nano model will be created in every model_id folder for AMP checks after this commit 7036a28.
|
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 |
|
The PR is ready for a review. I have created a new branch (will create a PR once I commit something) to standardise the image-based dataloaders |
|
I understand now that this selection is valid for P2 in this training form. It controls the P2 model size and selects the regular checkpoint used to initialize compatible layers. There are two ways we could make this clearer:
Could we also update the PR description? It currently says the P2 model trains from random initialization, but the latest implementation transfers compatible weights from the selected checkpoint. 
|
There was a problem hiding this comment.
This is very close to being ready. Training and prediction both worked with no errors during my nvidia GPU testing, and most of my original comments have been addressed.
Before merging, could you complete the follow up data-loader PR? YOLO is offered to all video projects but currently relies on ImageDataLoader behaviour, so I think that work should be merged first.
I’m happy to update uv.lock. I’ve also left a few comments about test coverage, rare edge cases, and separating the normal YOLO and P2 schemas. Otherwise, this looks nearly ready to merge.
| @ModelRegistry.register( | ||
| "yolo_ufo_p2", | ||
| ["video"], | ||
| YoloTrainParams, | ||
| YoloPredictParams, | ||
| ) |
There was a problem hiding this comment.
Was YOLO11n intentionally removed? I suspect it was left out because there is no matching P2 architecture, but that should not prevent users from selecting it for the regular YOLO model.
The normal YOLO and P2 models currently share YoloTrainParams, so they must expose the same checkpoint options. It looks like yolo11n.pt was removed because it has no matching P2 architecture, but this also removes YOLO11n from the normal yolo_ufo form even though the PR description says it is supported.
Could we give P2 its own training parameter schema? The normal schema could continue to offer YOLO11n, while the P2 schema would offer only checkpoints with matching P2 architectures. This would also let us give the P2 checkpoint field the clearer description discussed above. If you disagree and there is a good reason to omit yolo11n and keep the current schema that is fine, at least update the pr description.
other related code:
YoloModelName = Literal[
"yolov8n.pt",
"yolo26n.pt",
"yolo26m.pt",
"yolo26l.pt",
"yolo26x.pt",
]..
..
@ModelRegistry.register(
"yolo_ufo",
["video"],
YoloTrainParams,
YoloPredictParams,
)| def iter_sample_frames( | ||
| data_loader: TokTaggerDataLoader, | ||
| sample: Sample, | ||
| ) -> Iterator[ImageData]: |
There was a problem hiding this comment.
Add a basic unit test for this? see this comment for more detail #326 (comment)
| def build_video_frame_manifest( | ||
| samples: list[Sample], | ||
| annotations: list[list[Annotation]], | ||
| class_map: dict[str, int], | ||
| data_loader: TokTaggerDataLoader, | ||
| ) -> list[DetectionRecord]: |
There was a problem hiding this comment.
Add a basic unit test for this? see this comment for more detail #326 (comment)
| def decode_frame_image(frame_image: ImageData) -> np.ndarray: | ||
| """Decode raw TokTagger image bytes for Ultralytics prediction.""" | ||
| if isinstance(frame_image.values, str): |
There was a problem hiding this comment.
Add a basic unit test for this? see this comment for more detail #326 (comment)
| def predict( | ||
| self, | ||
| samples: list[Sample], | ||
| params: YoloPredictParams, | ||
| data_params=None, | ||
| ) -> list[list[AnnotationBase]]: |
There was a problem hiding this comment.
Add a basic unit test for this? see this comment for more detail #326 (comment)
There was a problem hiding this comment.
The functions below needs some test coverage:
- iter_sample_frames()
- build_video_frame_manifest()
- decode_frame_image
- predict()
Could we add some focused unit tests for this new functionality? The PR currently has no tests for frame iteration, manifest creation, image decoding or converting YOLO results into bounding-box annotations.
These tests should not need to train a real model. Small test images and mocked data loaders/YOLO predictions could cover:
- Reaching the end of a video
- Frames with no annotations
- Multiple boxes on one frame
- Bounding-box coordinate conversion
- Predictions containing no detections
- Loading best.pt, with fallback to last.pt
- Restoring a model after its Ray actor is recreated
| # Download to a temporary file so an interrupted transfer cannot leave a | ||
| # partial checkpoint at the final cache path. | ||
| temporary_path = model_path.with_name(f"{model_path.name}.tmp") | ||
| temporary_path.unlink(missing_ok=True) | ||
|
|
||
| try: | ||
| urlretrieve(MODEL_URLS[model_name], temporary_path) | ||
| temporary_path.replace(model_path) | ||
| except Exception: | ||
| temporary_path.unlink(missing_ok=True) | ||
| raise |
There was a problem hiding this comment.
Could two training jobs try to download the same checkpoint at the same time?
Every Ray actor uses the same .tmp path, so one actor could delete or replace another actor’s download. Also, urlretrieve() has no explicit timeout, so a failed connection could leave training waiting for a long time.
Could we use the existing filelock dependency to allow only one download per checkpoint at a time, and add a network timeout? After obtaining the lock, we should check again whether another actor has already completed the download.
Here is a example code snippet of a potential solution:
from filelock import FileLock
from urllib.request import urlopen
lock_path =
model_path.with_suffix(".lock")
with FileLock(lock_path):
# Another actor may have downloaded it
while we waited.
if model_path.exists() and not
force_download:
return model_path
with urlopen(MODEL_URLS[model_name],
timeout=60) as response:
temporary_path.write_bytes(respons
e.read())
temporary_path.replace(model_path)| REMOVE BELOW COMMENT AFTER SOMEONE HAS TESTED IT ON A NVIDIA GPU. | ||
| https://github.com/ukaea/toktagger/pull/326#discussion_r3754183103 |
There was a problem hiding this comment.
Yep you can remove this I tested your branch on a nvidia gpu.
|
|
||
| # TYPE_CHECKING is always false while the application runs.Type checkers treat it as true | ||
| # Othewise this will become a circular import. | ||
| if TYPE_CHECKING: |
There was a problem hiding this comment.
This is a bit of a code smell - can we just define YoloTrainParams in this file, and import it inside video_detection instead?
| shot_id: int | ||
| frame: int | ||
| image: bytes | ||
| boxes: list[list[float]] |
There was a problem hiding this comment.
If this is always x1, x2, y1, y2, could narrow the typing to list[tuple[float, float, float, float]]
| "https://github.com/ultralytics/assets/releases/download/v8.4.0" | ||
| ) | ||
|
|
||
| MODEL_URLS = { |
There was a problem hiding this comment.
This seems somewhat redundant - isn't it always just f"{BASE_URL/{file_name}"?
| The model is downloaded into TokTagger's cache when it is not already | ||
| available. Setting ``force_download`` replaces an existing cached model. | ||
| """ | ||
| if model_name not in MODEL_URLS: |
There was a problem hiding this comment.
Feels like you can get rid of MODEL_URLs and just use MODEL_FAMILIES here
There was a problem hiding this comment.
Otherwise we risk a key being added to one but not the other, leading to a KeyError on line 54
| ) | ||
|
|
||
|
|
||
| def iter_sample_frames( |
There was a problem hiding this comment.
At some point you had that logic which did a coarse search until there was a non black frame, and then iterated from there - should we add that in here?
| self, | ||
| samples: list[Sample], | ||
| params: YoloPredictParams, | ||
| data_params=None, |
There was a problem hiding this comment.
Here, you could add a this_frame_only parameter to YoloPredictParams. If this is enabled (and data_params is available, which means the prediction was requested from an individual sample page), then take the frame number from data_params and return predictions only for that single frame.
Add in the description of the parameter that 'this is ignored for multi-sample predictions' or something
|
|
||
| result = results[0] | ||
|
|
||
| if result.boxes is None or len(result.boxes) == 0: |
There was a problem hiding this comment.
can this just be if not result.boxes?
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
Changes
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)