Skip to content

Repository files navigation

Segmentation Label Annotator

A local-first, self-improving image-segmentation labelling loop.
Seed labelled images → train YOLOv8-seg → predict → route by confidence → human review in CVAT → retrain → repeat.

Quickstart

pip install -r requirements.txt
# add seed images + COCO JSON to data/raw/labelled_seed/
# add unlabelled images to data/raw/unlabelled_pool/
python train.py --run-name seed-v1   # train
python predict.py run                # predict
python active_learning.py route      # sort by confidence
python review.py start               # launch CVAT for human review

Full walkthrough below.


Features

  • Active-learning loop - confidence-based routing means only uncertain predictions reach a human reviewer.
  • Local-first - everything runs on your own hardware; no images leave the machine.
  • CVAT-integrated review - Docker Compose spins up a self-hosted CVAT instance for in-browser annotation.
  • Class-agnostic - add new categories to your COCO file and retrain; no code changes required.
  • Class-imbalance aware - inverse-frequency loss weighting and stratified train/val splits protect minority classes.
  • Reproducible - every run captures its config, metrics, and weights to models/experiments/.
  • Multiple export formats - COCO JSON, CVAT XML, and per-instance binary masks.

Contents

  1. Prerequisites & Installation
  2. First Run - Adding Your Data
  3. The Loop
  4. Configuration Reference
  5. Folder Structure
  6. Advanced Topics
  7. Troubleshooting
  8. Security Notes
  9. Licence & Third-Party

1. Prerequisites & Installation

Requirement Version
Python 3.11 or newer
CUDA toolkit (GPU only) 12.8 (driver 12.8 - 13.x)
Docker Desktop 4.x (for the CVAT review stack)
GPU (recommended) NVIDIA RTX 3060 or better, 8 GB+ VRAM
# Windows (PowerShell)
py -3 -m venv .venv
.\.venv\Scripts\Activate.ps1

# Linux / macOS
python3 -m venv .venv
source .venv/bin/activate
# Install PyTorch with CUDA 12.8 support (GPU users)
pip install --index-url https://download.pytorch.org/whl/cu128 torch torchvision

# Install everything else
pip install -r requirements.txt

CPU-only / no GPU: skip the first pip install and run pip install -r requirements.txt on its own. The CPU wheels will be pulled from PyPI automatically.

Verify your GPU is detected:

python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU only')"

A conda alternative is also provided:

conda env create -f environment.yml
conda activate seg-label-annotator

Optionally run the test suite to confirm the install:

pytest tests/smoke -q     # 5-second sanity check
pytest                    # full suite

2. First Run - Adding Your Data

Before anything else, populate two directories.

Seed images (hand-labelled)

These bootstrap the first model. Add them once before your first train.py run.

  1. Copy 20–30 JPEG/PNG images into:

    data/raw/labelled_seed/images/
    
  2. Copy your COCO JSON annotation file(s) into:

    data/raw/labelled_seed/annotations/
    

    Multiple COCO files are merged automatically; duplicates are deduplicated by filename.

Unlabelled pool

These are the images the trained model will predict and sort on every cycle.

  1. Copy any number of JPEG/PNG images into:
    data/raw/unlabelled_pool/images/
    

Then validate your seed data before training:

python dataset_manager.py inspect   # show class counts and image stats
python dataset_manager.py validate  # check for broken polygons

3. The Loop

All commands run from the project root.
Steps marked [BROWSER] require opening CVAT - everything else is terminal.


Step 1 - Train initial model (first time only)

python train.py --run-name seed-v1

Trains on your seed images. Early stopping is enabled (patience = 20 epochs). The best model is saved to models/production/ on completion.


Step 2 - Predict unlabelled pool

python predict.py run

Runs the current model over every image in data/raw/unlabelled_pool/images/ and saves a per-image prediction JSON alongside overlay images.


Step 3 - Route predictions

python active_learning.py route --overlays
python active_learning.py status        # see how many went where

Each prediction is scored and routed automatically:

Confidence Destination
≥ 0.90 Auto-approved → data/raw/approved/
0.70 – 0.90 Review queue (configurable - see config)
< 0.70 Review queue → data/raw/review_queue/

If the review queue is empty after routing, skip Steps 4–5 and go straight to Step 6.


Steps 4–5 - Human review via CVAT

CVAT runs locally in Docker. The default administrator credentials are taken from environment variables (ANNOTATOR_CVAT_USERNAME / ANNOTATOR_CVAT_PASSWORD); if you have not set them, the project's defaults (admin / admin) are used. Change these in a .env file before exposing CVAT on anything beyond localhost - see Section 8.

python review.py start
python review.py import

[BROWSER] Open http://localhost:8080 and log in with the credentials configured above.
Open the task, correct any wrong polygons or labels, and mark all images Done.
Note the Task ID in the CVAT URL (e.g. /tasks/2/).

python review.py export 2          # replace 2 with your task ID
python review.py stop

Corrected annotations and images are written to data/raw/approved/ automatically.


Step 6 - Retrain

python train.py --run-name reviewed-v1

Retrains on all approved data (seed + everything reviewed so far). Increment the run name each cycle (reviewed-v2, reviewed-v3, …) to preserve a history.

→ Go back to Step 2 and repeat.

Each cycle the model improves: more images are auto-approved and fewer need human review.


4. Configuration Reference

All config lives in configs/*.yaml. Any value can be overridden with an environment variable using the ANNOTATOR_ prefix and __ as the nested delimiter.

configs/model.yaml

Key Default Description
architecture yolov8n-seg Ultralytics model variant
epochs 100 Maximum training epochs
patience 20 Early-stopping patience (stops if val loss stalls for N epochs)
batch_size 8 Images per batch - reduce to 4 if GPU OOM
image_size 640 Input resolution
val_split 0.2 Fraction of data held out for validation
workers 4 DataLoader workers - reduce to 2 if CPU RAM OOM

configs/active_learning.yaml

Key Default Description
auto_approve_threshold 0.90 Confidence floor for auto-approval
review_threshold 0.70 Below this, image goes to human review
medium_confidence_action review review | auto_approve | reject
retrain_triggers [20, 50, 100, 500, …] Approved-image milestones that fire auto-retrain

configs/preprocessing.yaml

Key Default Description
image_size [640, 640] Output resolution
enable_circle_crop false Auto-crop to the dominant circular region (Hough Circle Transform)
preserve_originals true Never overwrite source files

Environment variable examples

ANNOTATOR_MODEL__EPOCHS=50
ANNOTATOR_MODEL__BATCH_SIZE=4
ANNOTATOR_ACTIVE_LEARNING__AUTO_APPROVE_THRESHOLD=0.85
ANNOTATOR_CVAT_PASSWORD=secret

5. Folder Structure

seg-label-annotator/
│
├── configs/
│   ├── model.yaml              # YOLOv8 hyperparameters
│   ├── active_learning.yaml    # Routing thresholds and retrain triggers
│   └── preprocessing.yaml      # Image size and optional circular-crop settings
│
├── data/raw/
│   ├── labelled_seed/          ← PUT YOUR SEED IMAGES + COCO JSON HERE
│   │   ├── images/
│   │   └── annotations/
│   ├── unlabelled_pool/        ← PUT IMAGES TO PREDICT HERE
│   │   └── images/
│   ├── review_queue/           # Auto-populated by active_learning.py route
│   │   ├── images/
│   │   └── predictions/
│   ├── approved/               # Auto-populated by review.py export
│   │   ├── images/
│   │   └── annotations/
│   └── exports/                # Output from export.py
│       ├── coco/
│       ├── cvat/
│       └── masks/
│
├── models/
│   ├── checkpoints/            # Per-run training checkpoints
│   ├── experiments/            # Metrics history JSON + class distribution plots
│   └── production/             # best.pt + class_names.json (live model)
│
├── src/                        # Library code (not called directly)
│   ├── active_learning/        # Confidence-based routing logic
│   ├── cvat/                   # Docker lifecycle + CVAT REST client
│   ├── export/                 # COCO / CVAT XML / mask exporters
│   ├── inference/              # YOLOv8 predictor wrapper
│   ├── preprocessing/          # Letterbox + optional circular-crop pipeline
│   ├── training/               # DatasetManager, COCO → YOLO conversion
│   └── utils/                  # Logging, image IO, annotation helpers
│
├── tests/                      # pytest suite (unit / integration / smoke)
│
# CLI entry points  (python <script>.py --help)
├── train.py            # Train or resume a YOLOv8-seg model
├── predict.py          # Run inference over image directory or single file
├── active_learning.py  # Route predictions to approved / review queues
├── review.py           # Manage CVAT: start, import, export, stop
├── export.py           # Export approved data to COCO / CVAT XML / masks
├── dataset_manager.py  # Inspect and validate the COCO dataset
├── metrics.py          # Log, show, and compare training metrics
└── settings.py         # Pydantic v2 settings (YAML + env-var overrides)

6. Advanced Topics

Adding new label classes

  1. Add the new class to your COCO annotation files as a new categories entry - no code changes needed.
  2. Run python dataset_manager.py inspect to confirm it was discovered.
  3. Retrain: python train.py --run-name new-class-v1.

train.py options

Flag Default Description
--run-name run-<timestamp> Experiment name
--epochs from config Override epoch count
--batch from config Override batch size
--device 0 (first GPU) cpu, 0, 0,1
--resume false Resume from last checkpoint

Export annotated data

python export.py coco    # COCO JSON of all approved predictions
python export.py cvat    # CVAT XML
python export.py masks   # Per-instance binary PNG masks
python export.py all     # All three formats

Automated pipeline

python main.py pipeline             # one-shot: predict → route → retrain if triggered
python main.py loop --interval 300  # continuous: run pipeline every 5 minutes

Starting fresh

To reset to a clean slate (keeping all code and configs), delete the contents of these directories - the directories themselves should remain:

Directory What it holds
data/raw/labelled_seed/images/ Seed images
data/raw/labelled_seed/annotations/ Seed COCO JSON files
data/raw/unlabelled_pool/images/ Images queued for prediction
data/raw/review_queue/images/ Images queued for human review
data/raw/review_queue/predictions/ Pre-annotation JSON files
data/raw/approved/images/ Reviewed / auto-approved images
data/raw/approved/annotations/ Reviewed / auto-approved annotations
data/raw/exports/ COCO / CVAT / mask exports
data/predictions/ Per-image prediction JSONs from predict.py
data/yolo_dataset/ Built YOLO dataset (auto-rebuilt on next train)
models/checkpoints/ All training checkpoints
models/experiments/ Metrics history and plots
models/production/ Live best.pt and class_names.json

After clearing, return to Section 2 - First Run and add your new seed data.


7. Troubleshooting

Symptom Fix
ModuleNotFoundError: No module named 'torch' Activate venv: .venv\Scripts\activate, then pip install -r requirements.txt
CUDA out of memory python train.py --batch 4 or set ANNOTATOR_MODEL__BATCH_SIZE=4
No images found in unlabelled pool Check data/raw/unlabelled_pool/images/ contains .jpg / .png files
CVAT not reachable at http://localhost:8080 Ensure Docker Desktop is running and port 8080 is free; check docker compose logs cvat_server
No module named 'cv2' pip install opencv-python-headless
Predictions show 0 detections Model not trained yet - complete Step 1 first
models/production/best.pt missing after training Check models/experiments/<run-name>/weights/ - the script copies it on completion
CPU RAM OOM during training Set workers: 2 in configs/model.yaml
Training dataset has far more images than expected Multiple COCO exports of the same images in labelled_seed/annotations/ - deduplication is automatic; ensure files are genuine separate batches

8. Security Notes

This project is designed to run on a single trusted machine. The defaults are tuned for a developer workstation, not a production deployment.

  • CVAT credentials and database passwords in docker-compose.yml use placeholder defaults (admin / admin, cvat_local_password, local-dev-secret-change-in-production). These are convenient for local use; always override them before exposing the stack on a network. Copy .env.example to .env and set:

    ANNOTATOR_CVAT_USERNAME=your_user
    ANNOTATOR_CVAT_PASSWORD=a_strong_password
    CVAT_DB_PASSWORD=a_strong_db_password
    CVAT_SECRET_KEY=a_long_random_string
    
  • CVAT is bound to localhost:8080 by default. Do not publish that port directly to the internet; put it behind a reverse proxy with TLS and authentication.

  • The project does not include any telemetry: no data, metrics, or images leave your machine.

  • Trust your training data: anything you place in data/raw/labelled_seed/ is loaded and parsed at training time. Treat COCO JSON files from third parties as untrusted input.


9. Licence & Third-Party

The code in this repository is released under the MIT Licence - see LICENSE.

The project depends on third-party components with their own licences. In particular, Ultralytics YOLOv8 is distributed under the AGPL-3.0 licence, which has network-distribution obligations that the MIT licence does not override. See NOTICE.md for the full third-party licence summary before you publish a derivative work or host the model over a network.

About

A semi-automated pipeline for generating label files for segmentation training using a handful of manually labelled seeds. The pipeline will auto approve predicted labels above a confidence threshold, and require manual intervention for anything below a threshold.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages