Prompt2YOLO

From a text prompt or an image folder to a trained object detector.

Prompt2YOLO collects images, removes bad and duplicate data, generates bounding boxes with Grounding DINO, and trains a YOLO model. It also includes web and desktop tools for quickly reviewing generated labels.

See the complete Setup and Usage Guide for installation, configuration, and commands.

Why Prompt2YOLO?

Building a custom detector usually means finding images, drawing boxes, converting dataset formats, and writing training code. Prompt2YOLO connects those steps into one simple workflow:

Prompt / local images
        ↓
Auto Collect → Curate → Grounding DINO labels → Optional review → YOLO training
                                                               ↓
                                                     Onnx/TensorRT inference

• Search images with Openverse and Wikimedia Commons, or use a local folder.
• Reject invalid images and remove exact or visually similar duplicates.
• Generate YOLO bounding boxes from text prompts with Grounding DINO Base.
• Review labels using either a web app or a Tkinter desktop app.
• Build train, validation, and test splits automatically.
• Fine-tune YOLO26 through a small YAML config.
• Export the best model to TensorRT FP32 or INT8 and evaluate it.

Label review

The review tool lets you draw missing boxes, delete incorrect ones, reject images, and skip directly to the next machine-labeled sample. Accepted edits are marked as human-reviewed.

Web

Desktop

Both interfaces edit the same YOLO labels and labeling.jsonl records.

Demo result

As a quick end-to-end test, I trained YOLO26s to detect safety helmets:

• 922 images downloaded from Wikimedia Commons
• 891 images kept after curation and deduplication
• 810 images automatically labeled by Grounding DINO
• 0 images manually reviewed
• 50 training epochs at 640 px on an NVIDIA L4

Precision	Recall	mAP@50	mAP@50:95
88.05%	83.65%	88.62%	70.66%

The validation labels were also generated by Grounding DINO, so these metrics show that the automatic pipeline trains a useful YOLO model; a manually labeled test set would be needed for a strict accuracy benchmark.

The exported TensorRT FP32 engine took 2.47–3.23 ms for model inference in two notebook examples on the L4. These are single-image observations rather than a full speed benchmark.

Quick start

Install the environment:

conda env create -f environment.yaml
conda activate prompt2yolo

Edit config.yaml with your classes, labeling prompts, search queries, and image limits. Then run:

# Collect, curate, and auto-label
python main.py run --config config.yaml

# Optional label review
python main.py review --ui web --config config.yaml
# Or: python main.py review --ui tkinter --config config.yaml

# Build the dataset and train YOLO
python main.py train --config config.yaml

# Export the best model to TensorRT
python main.py postprocess --config config.yaml

# Evaluate on the test split
python main.py evaluate --config config.yaml

Use python main.py postprocess --int8 --config config.yaml for INT8 export.

Open inference_demo.ipynb to load the exported engine, run inference, display detections, and save the result.

For installation alternatives, every configuration option, command details, and workspace structure, see docs/USAGE.md.

Main outputs

All generated data stays under workspace/<project>/, which is ignored by Git:

workspace/<project>/
├── raw/                 # downloaded images
├── curated/             # validated and deduplicated images
├── labeled/             # images with bounding boxes
├── labels/              # YOLO TXT annotations
├── dataset/             # train / val / test dataset
├── weights/             # best.pt, last.pt, TensorRT engine
└── reports/             # pipeline metadata and metrics

Built with PyTorch, Hugging Face Transformers, Grounding DINO, Ultralytics YOLO, FastAPI, Tkinter, and TensorRT.

_{Demo source images: Wikimedia Commons (inference, desktop reviewer, web reviewer). Original licenses apply.}