Metadata-Version: 2.4
Name: traincv
Version: 0.2.0
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: 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

# TrainCV

<p align="center"><img src="logo.svg" alt="traincv logo" width="120"></p>

![PyPI](https://img.shields.io/pypi/v/traincv)
![Python](https://img.shields.io/pypi/pyversions/traincv)
![License](https://img.shields.io/pypi/l/traincv)

A no-code labeling and training toolkit for computer vision. Train object detection, classification, and segmentation models with minimal code via CLI or Python API.

**Supports local HuggingFace models.** Pre-trained models can be downloaded from HuggingFace Hub and cached locally for offline training and fine-tuning.

**Built by [NRL.ai](https://www.nrl.ai)**

---

## Installation

```bash
# Core (no training backends)
pip install traincv

# With PyTorch backend
pip install traincv[torch]

# With YOLO backend
pip install traincv[yolo]

# Everything
pip install traincv[full]
```

## Quick Start

### CLI

```bash
# Train an object detection model
traincv train --task detection --data ./dataset --model yolov8n --epochs 50

# Train an image classifier
traincv train --task classification --data ./images --model mobilenetv2

# Export a trained model to ONNX
traincv export --model ./best.pt --format onnx

# Validate dataset format
traincv validate --data ./dataset --format coco
```

### Python API

```python
import traincv

# Train a detector
traincv.train(task="detection", data="./dataset", model="yolov8n", epochs=50)

# Export to ONNX
traincv.export(model="./best.pt", format="onnx")

# Load and inspect a dataset
dataset = traincv.Dataset.load("./dataset", format="coco")
print(dataset.summary())
```

## Supported Tasks

| Task | Models | Backend |
|------|--------|---------|
| Object Detection | YOLOv8n/s/m/l/x, YOLOv5 | ultralytics |
| Image Classification | MobileNetV2, ResNet18/34/50, EfficientNet-B0 | torchvision |
| Segmentation | YOLOv8n/s/m-seg, DeepLabV3 | ultralytics / torchvision |

## Dataset Formats

TrainCV supports the following annotation formats:

| Format | Description | File Type |
|--------|-------------|-----------|
| COCO | COCO JSON format | `.json` |
| YOLO | YOLO darknet format | `.txt` per image |
| VOC | Pascal VOC XML format | `.xml` per image |

### COCO Format

```
dataset/
  images/
    img001.jpg
    img002.jpg
  annotations/
    instances.json
```

### YOLO Format

```
dataset/
  images/
    img001.jpg
    img002.jpg
  labels/
    img001.txt
    img002.txt
```

### Pascal VOC Format

```
dataset/
  JPEGImages/
    img001.jpg
  Annotations/
    img001.xml
```

## Model Zoo

| Model | Task | Size | Speed |
|-------|------|------|-------|
| `yolov8n` | Detection | 6.2M | Fast |
| `yolov8s` | Detection | 11.2M | Fast |
| `yolov8m` | Detection | 25.9M | Medium |
| `yolov8l` | Detection | 43.7M | Slow |
| `yolov8x` | Detection | 68.2M | Slow |
| `mobilenetv2` | Classification | 3.5M | Fast |
| `resnet18` | Classification | 11.7M | Fast |
| `resnet50` | Classification | 25.6M | Medium |
| `efficientnet_b0` | Classification | 5.3M | Fast |

## Export Formats

| Format | Extension | Use Case |
|--------|-----------|----------|
| ONNX | `.onnx` | Cross-platform inference |
| TorchScript | `.torchscript` | PyTorch production |

## Configuration

Training parameters can be set via CLI flags, Python kwargs, or a YAML config file:

```yaml
# config.yaml
task: detection
model: yolov8n
data: ./dataset
epochs: 50
batch_size: 16
lr: 0.001
augmentation: medium
input_size: [640, 640]
device: auto
```

```bash
traincv train --config config.yaml
```

## Data Augmentation

Built-in augmentation presets:

| Preset | Transforms |
|--------|------------|
| `light` | Horizontal flip, slight rotation |
| `medium` | Flip, rotation, color jitter, scale |
| `heavy` | Flip, rotation, color jitter, scale, mosaic, mixup |

## Local-First / Edge AI

This package supports downloading pre-trained models from HuggingFace Hub
for offline use. After initial download, no internet connection is required
for training or fine-tuning.

```bash
# Pre-download a model for offline use
python -m traincv download yolov8n

# Download a classification model
python -m traincv download mobilenetv2 --task classification
```

```python
import traincv

# Download and cache a model locally
model_path = traincv.download_model("yolov8n")

# Train using the cached model (works offline)
traincv.train(task="detection", data="./dataset", model=str(model_path))
```

## Links

- **Website:** [nrl.ai](https://www.nrl.ai)
- **GitHub:** [github.com/vietanhdev/traincv](https://github.com/vietanhdev/traincv)
- **Author:** [Viet-Anh Nguyen](https://github.com/vietanhdev)

## License

MIT License - see [LICENSE](LICENSE) for details.
