🚀 YOLOE Data Engine

The YOLOE Data Engine builds and refines training data for YOLOE, the real-time open-vocabulary model in the Ultralytics family. It closes the loop between a dataset and a model: existing labels are loaded, a trained YOLOE model predicts over the same images, and any prediction not already covered by a ground-truth box is merged back into the labels. The result is a denser, better-covered dataset for the next training run.

The pipeline handles both grounding data (free-form text phrases tied to image regions, such as Flickr30k Entities or mixed-grounding GQA) and detection data (a fixed class vocabulary, such as Objects365), and it includes tooling to refine the text prompts attached to grounding labels and to visually inspect everything it produces.

This is a working research repository: a collection of scripts rather than an installable package, with dataset, model, and buffer paths configured inside each script rather than passed on the command line.

🔧 Requirements

A Python environment with the ultralytics package and its dependencies, plus PyTorch, NumPy, Pillow, OpenCV, and Matplotlib for the visualizers:

pip install ultralytics matplotlib

There is no requirements.txt — the scripts import from ultralytics and are meant to run alongside a local Ultralytics checkout. A CUDA GPU is required for prediction; DataEngineAgent takes a list of devices and defaults to a single cuda:0, while the example at the bottom of data_engine_agent.py spreads the work across four.

⚙️ How It Works

The engine has two layers over the same logic:

• DataEngine (data_engine.py) — the single-process core. It loads an Ultralytics .cache label file, loads a YOLOE model, sets the class vocabulary, predicts, merges the predictions into the in-memory labels, and writes a new cache file.
• DataEngineAgent (data_engine_agent.py) — the multi-GPU wrapper. It runs the same stages as parallel worker processes that communicate through JSON files in a buffer directory, one file per image, so a large dataset can be processed in stages and inspected in between.

A typical run moves through four stages:

1. Load labels. DataEngine.load_cached_label() reads an Ultralytics .cache file in either grounding or detection style. For raw grounding annotations, DataEngineAgent.multi_process_load_grounding_data() first explodes the source JSON into one sample JSON per image under the buffer directory.
2. Predict. DataEngine.load_yoloe() loads the YOLOE segmentation checkpoint, and set_classes() sets the vocabulary — from a dataset YAML, an explicit list of names, or the keys of a MobileCLIP text-embedding .pt file (that file supplies the class names only; the prompt embeddings are recomputed by the model). DataEngineAgent.multi_process_batch_model_predict() then writes one prediction JSON per image.
3. Merge. Predictions are matched against the existing boxes, and any that overlap a ground-truth box above the IoU threshold are dropped; the rest are appended to the label. multi_process_merge_prediction() folds the results back into the sample files, and save_cached_label() writes the updated cache.
4. Refine and inspect. refine_text.py re-scores the text phrases on grounding labels using YOLOE visual prompt embeddings and writes an updated cache. The data_visual_*.py and grounding_dataset_visualizer.py scripts render labels and predictions so the output can be checked by eye.

🖥️ Usage

Both end-to-end flows are shell scripts. They invoke python3 yoloe_data_engine/<script>.py, so run them from the parent directory of this checkout:

# Flickr grounding: refine text prompts, then visualize
bash yoloe_data_engine/do_flickr.sh

# Mixed-grounding (GQA): refine text prompts across two GPUs
bash yoloe_data_engine/do_mixed.sh

Individual scripts run from the repository root. refine_text.py is the only one that takes command-line arguments:

# Refine the text prompts on a grounding dataset and write an updated .cache
python3 refine_text.py \
  --img_path ../datasets/flickr/full_images/ \
  --json_file ../datasets/flickr/annotations/final_flickr_separateGT_train_segm.json

# Render generated labels and predictions
python3 data_visual_flickr.py

Every other entry point is configured by editing the paths at the top of the file — dataset roots, the YOLOE checkpoint, text-embedding .pt files, and buffer directories are hard-coded, and the shell scripts also activate a clipenv conda environment. Set these for your own machine before running anything.

📁 Repository Structure

Path	Purpose
data_engine.py	DataEngine — cache loading, YOLOE prediction, label merging, cache saving
data_engine_agent.py	DataEngineAgent — multi-GPU multiprocess versions of the same stages
refine_text.py	Refines grounding text prompts with YOLOE visual prompt embeddings
remove_segment.py	Strips masks and segments from a cached dataset
visual_json.py, grounding_dataset_visualizer.py	Render labels and predictions for inspection
data_visual*.py	Per-dataset visualization entry points (Flickr, mixed-grounding, Objects365)
utils.py	Small dataset counting helpers
do_flickr.sh, do_mixed.sh	End-to-end example flows
log.md	Scratch record of past runs — label and box counts, not documentation

Buffer directories and generated .cache files are local artifacts and are not committed.

💡 Contribute

Ultralytics thrives on community collaboration, and we deeply value your contributions! Whether it's reporting bugs, suggesting features, or submitting code changes, your involvement is crucial.

• Reporting Issues: Encounter a bug? Please report it on GitHub Issues.
• Feature Requests: Have an idea for improvement? Share it via GitHub Issues.
• Pull Requests: Want to contribute code? Please read our Contributing Guide first, then submit a Pull Request.
• Feedback: Share your thoughts and experiences by participating in our official Survey.

A heartfelt thank you 🙏 goes out to all our contributors! Your efforts help make Ultralytics tools better for everyone.

📄 License

Ultralytics offers two licensing options to accommodate diverse needs:

• AGPL-3.0 License: Ideal for students, researchers, and enthusiasts passionate about open collaboration and knowledge sharing. This OSI-approved open-source license promotes transparency and community involvement. See the LICENSE file for details.
• Enterprise License: Designed for commercial applications, this license permits the seamless integration of Ultralytics software and AI models into commercial products and services, bypassing the copyleft requirements of AGPL-3.0. For commercial use cases, please inquire about an Ultralytics Enterprise License.

📮 Contact

For bug reports or feature suggestions, please use GitHub Issues. For general questions, discussions, and community support, join our Discord server!