$ show computer-vision-with-marco/yolo-training-template
This repository provides a template for training YOLO models on custom datasets (primarily from Kaggle) and performing inference. It includes command-line scripts for training and inference, a Streamlit app for interactive use, a Jupyter notebook for step-by-step execution, and utility scripts for data preprocessing, dataset exploration, and auto-labeling. The system allows users to download datasets, preprocess them (cleaning and augmentation), train YOLO models of various sizes, run inference on images/videos/webcam, export models to NCNN format, and automatically label images using GroundingDINO. The code is structured into modular components, with notable reusable parts in the data preprocessing and dataset exploration modules.
The LLM identified 2 files worth highlighting for reuse:
This module implements a robust data preprocessing pipeline for YOLO datasets, including cleaning (removing corrupted images, validating annotations) and augmentation (using Albumentations to apply geometric and photometric transformations while updating YOLO labels). It is configured via a YAML file, making it highly adaptable to different projects. The component is isolated, depending only on OpenCV, Albumentations, PyYAML, and standard libraries, and could be extracted for use in any image dataset preparation task.
This module provides utilities to load, visualize, and analyze YOLO-formatted datasets. It includes functions to discover image files in a dataset directory structure, parse YOLO labels into pixel coordinates, draw bounding boxes on images, compute dataset statistics, and extract sample images with annotations. The component handles various image formats and directory structures, making it reusable across different YOLO projects. It demonstrates good practices in error handling and efficient use of OpenCV and NumPy.