Metadata-Version: 2.4
Name: traincv
Version: 0.2.4
Summary: No-code labeling and training toolkit for computer vision
Project-URL: Homepage, https://www.nrl.ai
Project-URL: Repository, https://github.com/vietanhdev/traincv
Project-URL: Documentation, https://github.com/vietanhdev/traincv#readme
Project-URL: Issues, https://github.com/vietanhdev/traincv/issues
Author-email: Viet-Anh Nguyen <vietanh.dev@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: computer-vision,deep-learning,image-classification,no-code,object-detection,segmentation,training
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Image Recognition
Requires-Python: >=3.8
Requires-Dist: click
Requires-Dist: numpy
Requires-Dist: pillow
Requires-Dist: pyyaml
Provides-Extra: clip
Requires-Dist: transformers>=4.20.0; extra == 'clip'
Provides-Extra: dev
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Provides-Extra: full
Requires-Dist: torch>=1.9.0; extra == 'full'
Requires-Dist: torchvision>=0.10.0; extra == 'full'
Requires-Dist: transformers>=4.20.0; extra == 'full'
Requires-Dist: ultralytics>=8.0.0; extra == 'full'
Provides-Extra: progress
Requires-Dist: tqdm>=4.60.0; extra == 'progress'
Provides-Extra: torch
Requires-Dist: torch>=1.9.0; extra == 'torch'
Requires-Dist: torchvision>=0.10.0; extra == 'torch'
Provides-Extra: yolo
Requires-Dist: ultralytics>=8.0.0; extra == 'yolo'
Description-Content-Type: text/markdown

<h1 align="center">traincv</h1>
<p align="center"><em>No-code computer vision training — from dataset to deployable model in one command.</em></p>

<p align="center">
<img src="https://img.shields.io/pypi/v/traincv.svg" alt="PyPI">
<img src="https://img.shields.io/pypi/pyversions/traincv.svg" alt="Python">
<img src="https://img.shields.io/pypi/l/traincv.svg" alt="License">
</p>

**traincv** lets you train production-grade computer vision models (detection, classification, segmentation) with a single command or a 3-line Python API. It wraps Ultralytics YOLO and torchvision behind a unified interface, auto-converts between common dataset formats (COCO, YOLO, Pascal VOC), offers zero-shot auto-labeling via SAM2 / YOLO-World / CLIP, and exports trained models to ONNX, TorchScript, or TFLite ready for deployment.

Built by [Viet-Anh Nguyen](https://github.com/vietanhdev) at [NRL.ai](https://www.nrl.ai).

## Why traincv?

- **One-liner API** — `traincv.train("dataset/", task="detect")` is a complete pipeline
- **Plugin architecture** — Register custom backbones, losses, and exporters
- **Local-first** — Trains on your GPU, no cloud credits or upload required
- **Minimal core deps** — Training stacks (YOLO / torchvision) are optional extras
- **Production-ready** — Auto-splitting, logging, checkpointing, resumable runs

## Installation

```bash
pip install traincv
```

For training backends:

```bash
pip install traincv[yolo]          # Ultralytics YOLOv8/v11 for detect + segment
pip install traincv[torch]         # torchvision for classification + DeepLabV3
pip install traincv[autolabel]     # SAM2 + YOLO-World + CLIP for auto-annotation
pip install traincv[export]        # ONNX + TorchScript + TFLite exporters
pip install traincv[all]           # everything
```

**Python 3.8+ supported** (tested on 3.8, 3.9, 3.10, 3.11, 3.12, 3.13)

## Quick Start

```python
import traincv

# 1. Train an object detector (auto-detects YOLO-format dataset and uses YOLOv8n)
run = traincv.train(
    "datasets/pets/",      # expects images/ + labels/ or data.yaml
    task="detect",
    model="yolov8n",
    epochs=50,
    imgsz=640,
)
print(run.best_map50)       # best mAP50 on validation
print(run.weights_path)     # path to best checkpoint

# 2. Export to ONNX for deployment (compatible with anycv / anydeploy)
traincv.export(run.weights_path, format="onnx", out="pets.onnx")

# 3. Use the trained model immediately
preds = traincv.predict("pets.onnx", "new_image.jpg")
```

## Models & Methods

### Training backends

| Task | Backend | Models | Notes |
|---|---|---|---|
| **Object Detection** | Ultralytics | `yolov8n/s/m/l/x`, `yolov11n/s/m/l/x` | COCO-pretrained, transfer-learning by default |
| **Instance Segmentation** | Ultralytics | `yolov8n-seg`, `yolov11n-seg` | Mask + box outputs |
| **Classification** | torchvision | `mobilenetv2`, `mobilenetv3`, `resnet18/50`, `efficientnet_b0` | ImageNet-pretrained |
| **Semantic Segmentation** | torchvision | `deeplabv3_resnet50`, `deeplabv3_mobilenetv3` | Pascal VOC-pretrained |

### Dataset formats (auto-detected)

- **YOLO** — `images/` + `labels/` (one `.txt` per image) + `data.yaml`
- **COCO** — `annotations.json` with `images` / `annotations` / `categories`
- **Pascal VOC** — `JPEGImages/` + `Annotations/*.xml`
- **ImageFolder** — One sub-folder per class (classification only)

traincv converts between formats on the fly:

```python
traincv.convert("coco_dataset/", to="yolo", out="yolo_dataset/")
```

### Auto-labeling (zero-shot annotation)

Install `traincv[autolabel]` to bootstrap labels without manual work:

| Method | Model | Best for |
|---|---|---|
| **SAM2** | Meta Segment Anything 2 | Instance masks from point/box prompts |
| **YOLO-World** | Open-vocabulary YOLO | Detection from text prompts (`"a red car"`) |
| **CLIP** | OpenAI CLIP | Zero-shot image classification |

```python
# Generate YOLO-format bounding boxes using only text prompts
traincv.autolabel(
    "unlabeled_images/",
    method="yolo-world",
    classes=["person", "dog", "bicycle"],
    out="auto_labels/",
)
```

### Export formats

| Format | Via | Use case |
|---|---|---|
| **ONNX** | `torch.onnx.export` / ultralytics exporter | Cross-platform inference (anycv, anydeploy) |
| **TorchScript** | `torch.jit.trace` | Python-free PyTorch inference |
| **TFLite** | TensorFlow converter | Mobile / embedded |

## API Reference

| Function | Purpose |
|---|---|
| `traincv.train(dataset, task, model, epochs, ...)` | Train a model, returns `TrainRun` |
| `traincv.predict(weights, image)` | Inference with a trained model |
| `traincv.evaluate(weights, dataset)` | Compute metrics on a held-out split |
| `traincv.export(weights, format, out)` | Export to ONNX / TorchScript / TFLite |
| `traincv.convert(src, to, out)` | Dataset format conversion |
| `traincv.autolabel(images, method, classes)` | Zero-shot annotation |
| `traincv.split(dataset, ratios=(0.8, 0.1, 0.1))` | Auto train/val/test split |

## CLI Usage

```bash
# Train
traincv train datasets/pets/ --task detect --model yolov8n --epochs 50

# Evaluate
traincv evaluate runs/detect/best.pt --data datasets/pets/

# Export
traincv export runs/detect/best.pt --format onnx --out pets.onnx

# Auto-label with text prompts
traincv autolabel unlabeled/ --method yolo-world --classes "person,dog,cat"

# Convert datasets
traincv convert coco_ds/ --to yolo --out yolo_ds/
```

## Examples

### Full end-to-end: label, train, export

```python
import traincv

# 1. Auto-label raw images with SAM2 + YOLO-World
traincv.autolabel("raw/", method="yolo-world",
                  classes=["product", "price_tag"], out="labeled/")

# 2. Split into train/val/test (80/10/10)
traincv.split("labeled/", ratios=(0.8, 0.1, 0.1))

# 3. Train YOLOv8s
run = traincv.train("labeled/", task="detect", model="yolov8s", epochs=100)

# 4. Export for mobile deployment
traincv.export(run.weights_path, format="tflite", out="products.tflite")
```

### Fine-tune a classifier on ImageFolder data

```python
import traincv

run = traincv.train(
    "flowers/",            # flowers/daisy/*.jpg, flowers/rose/*.jpg, ...
    task="classify",
    model="mobilenetv3",
    epochs=30,
    lr=1e-3,
)
print(run.best_accuracy)
```

### Resume an interrupted run

```python
traincv.train("datasets/pets/", task="detect", resume="runs/detect/exp5/last.pt")
```

## License

MIT (c) Viet-Anh Nguyen
