Skip to content

Request: Add training example to YOLOX model #63

@quietlychris

Description

@quietlychris

For common computer vision annotation tools like cvat (Computer Vision Annotation Tool), there's out-of-the-box support for exporting train/test/validation sets into the COCO format as seen in these docs. I believe that Burn also now has support for this in the burn-dataset crate per this function and potentially some support for image augmentation via tracel-ai/burn#2995.

It would be really wonderful to be able to do low-overhead train the yolox model included in this repository using the output of these training tools, especially without having to necessarily go through the hassle of setting up CUDA via use of the wgpu backend.

I think it would be very cool to have a CLI that looks something like the following (although certainly worth some discussion!)

# The training example has some easily-configurable hyperparameters, probably don't even need to be accessible by CLI
# I like to imagine something like training a detector to find birds in images
$ cargo run --release --features pretrained,webgpu --example training -- --input <cvat_output_as_coco_dir> --output <model_path>
# The model serialized to disk in the previous step can then be imported into the inference bin
# which runs on the CPU backend and then writes the output to file like the current YOLOX inference example does
$ cargo run --release --features=local_model --example inference -- --input samples/has_a_few_birds.jpg --model <model_path_from_training>

Would love any feedback on if this might be something worthwhile as a showcase of some of Burn's capabilities in the computer vision area. Thanks!

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions