154 lines
4.5 KiB
Markdown
154 lines
4.5 KiB
Markdown
# Saw Mill Knot Detection (RF-DETR)
|
|
|
|
This repo contains a minimal training pipeline to fine-tune **RF-DETR** to detect knots in wood.
|
|
|
|
**Dataset Source**: The wood defect images and annotations used in this project come from [Kaggle Wood Surface Defects Dataset](https://www.kaggle.com/datasets/kirs0816/wood-surface-defects?resource=download).
|
|
|
|
## 1) Dataset format (required)
|
|
RF-DETR expects **COCO format**, split into `train/`, `valid/`, `test/`, each with its own `_annotations.coco.json`.
|
|
|
|
Example:
|
|
|
|
```
|
|
dataset/
|
|
├── train/
|
|
│ ├── _annotations.coco.json
|
|
│ ├── 0001.jpg
|
|
│ └── ...
|
|
├── valid/
|
|
│ ├── _annotations.coco.json
|
|
│ ├── 0101.jpg
|
|
│ └── ...
|
|
└── test/
|
|
├── _annotations.coco.json
|
|
├── 0201.jpg
|
|
└── ...
|
|
```
|
|
|
|
Your COCO JSON should include a `categories` entry for your class(es), e.g. `knot`.
|
|
|
|
## 2) Setup
|
|
|
|
Create venv (already created if you used the VS Code prompt) and install deps:
|
|
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python -m pip install -U pip
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python -m pip install -r requirements.txt
|
|
```
|
|
|
|
## 3) Validate dataset
|
|
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python validate_coco_dataset.py --dataset-dir /path/to/dataset
|
|
```
|
|
|
|
## 4) Train
|
|
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python train_rfdetr.py \
|
|
--dataset-dir /path/to/dataset \
|
|
--output-dir runs/knot_rfdetr_medium \
|
|
--model medium \
|
|
--epochs 50 \
|
|
--batch-size 4 \
|
|
--grad-accum-steps 4 \
|
|
--lr 1e-4
|
|
```
|
|
|
|
Notes:
|
|
- Keep **effective batch size** near 16: `batch_size * grad_accum_steps * num_gpus ≈ 16`.
|
|
- Checkpoints are written into `--output-dir` (including `checkpoint_best_total.pth`).
|
|
|
|
## 5) Auto-label new images (automatic)
|
|
|
|
Use your trained model to generate annotations on unlabeled images:
|
|
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python auto_label_images.py \
|
|
--weights runs/knot_rfdetr_medium/checkpoint_best_total.pth \
|
|
--images-dir /path/to/new_images \
|
|
--output-json auto_labeled.json \
|
|
--threshold 0.4
|
|
```
|
|
|
|
This outputs a COCO JSON with predicted bounding boxes. You can then review/correct them manually.
|
|
|
|
## 6) Manual labeling (recommended tools)
|
|
|
|
**Don't build your own GUI** - use these proven open-source tools instead:
|
|
|
|
### Option A: Label Studio (Recommended - Easiest)
|
|
**Best for**: Quick setup, modern UI, ML-assisted labeling
|
|
|
|
```bash
|
|
# Install Label Studio
|
|
pip install label-studio
|
|
|
|
# Start the server
|
|
label-studio start
|
|
```
|
|
|
|
Then open http://localhost:8080 in your browser:
|
|
1. Create a new project for "Object Detection with Bounding Boxes"
|
|
2. Import your images
|
|
3. Start labeling manually OR:
|
|
- Use the auto-label script to generate initial annotations:
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python auto_label_images.py \
|
|
--weights runs/knot_rfdetr_medium/checkpoint_best_total.pth \
|
|
--images-dir /path/to/images \
|
|
--output-json predictions.json \
|
|
--threshold 0.3
|
|
```
|
|
- Import the predictions into Label Studio
|
|
- Review and correct them
|
|
4. Export in COCO format when done
|
|
|
|
### Option B: CVAT (Most Powerful)
|
|
**Best for**: Large-scale projects, team collaboration
|
|
|
|
```bash
|
|
# Using Docker (easiest)
|
|
git clone https://github.com/opencv/cvat
|
|
cd cvat
|
|
docker compose up -d
|
|
```
|
|
|
|
Open http://localhost:8080:
|
|
- Create project → upload images → annotate
|
|
- Supports keyboard shortcuts, interpolation, and advanced features
|
|
- Export directly to COCO JSON
|
|
|
|
[CVAT Documentation](https://opencv.github.io/cvat/docs/)
|
|
|
|
### Option C: labelImg (Simplest Desktop App)
|
|
**Best for**: Offline labeling, no server needed
|
|
|
|
```bash
|
|
pip install labelImg
|
|
labelImg
|
|
```
|
|
|
|
- Simple desktop app with no web server
|
|
- Exports to Pascal VOC (needs conversion to COCO)
|
|
- Good for small datasets
|
|
|
|
### Workflow with Model Assistance:
|
|
1. **Initial batch**: Manually label 50-100 images
|
|
2. **Train RF-DETR**: Use your training script
|
|
3. **Auto-label**: Run `auto_label_images.py` on remaining images
|
|
4. **Review**: Import predictions into Label Studio/CVAT
|
|
5. **Correct**: Fix any mistakes (much faster than labeling from scratch)
|
|
6. **Iterate**: Retrain with corrected labels, repeat
|
|
|
|
This semi-supervised approach is **10-20x faster** than manual labeling alone.
|
|
|
|
## 7) Quick inference sanity check
|
|
|
|
```bash
|
|
/home/dillon/_code/saw_mill_knot_detection/.venv/bin/python predict_rfdetr.py \
|
|
--weights runs/knot_rfdetr_medium/checkpoint_best_total.pth \
|
|
--image /path/to/example.jpg \
|
|
--threshold 0.4
|
|
```
|