RavindranadhM's picture
Refine steel ROI training guide and deployment bundle
26142cd verified
|
Raw
History Blame Contribute Delete
3.79 kB
# Steel ROI Gate Training
This project now supports a **steel-surface detector** in front of defect segmentation.
Why detector and not plain classifier?
1. A classifier can say `steel` or `non_steel`, but it cannot give the bounding box you asked for.
2. A detector can localize the steel region, crop that ROI, and send only that region into the defect segmentation model.
3. This is the right architecture for mixed live-camera scenes with people, tools, and background clutter.
## What the scripts do
`prepare_surface_gate_dataset.py`
- downloads steel-positive sources automatically
- downloads non-steel backgrounds automatically
- creates a YOLO detector dataset
- generates synthetic composites so the detector learns actual steel bounding boxes
- writes `dataset.yaml`
`train_surface_gate_detector.py`
- finetunes a small YOLO detector quickly
- copies the best checkpoint into `models/steel_surface_detector.pt`
- makes the web app use that model automatically after restart
## Automatic data sources
Positive steel sources:
- NEU mirror on Hugging Face
- local `test_images/`
- optional Severstal mirror on Hugging Face
Automatic non-steel/background source:
- `torchvision` datasets downloaded automatically
- default quick path: `CIFAR10`
- optional stronger but heavier path: `STL10`
This avoids making you manually collect a non-steel dataset.
## Recommended quick workflow
1. Install training extras:
```bash
cd "/Users/ravindranadhm/Documents/Projects/steel-surface-inspection copy"
venv/bin/python -m pip install -r training/requirements.txt
```
2. Build the detector dataset:
```bash
cd "/Users/ravindranadhm/Documents/Projects/steel-surface-inspection copy"
venv/bin/python training/prepare_surface_gate_dataset.py
```
This default mode is the fastest path because it avoids the larger Severstal archive and uses `CIFAR10` as an automatic non-steel source.
If you want stronger natural-scene negatives later:
```bash
venv/bin/python training/prepare_surface_gate_dataset.py --negative-source stl10
```
If you want a stronger detector after the first run, rebuild with Severstal included:
```bash
venv/bin/python training/prepare_surface_gate_dataset.py --include-severstal
```
3. Train the lightweight detector:
```bash
cd "/Users/ravindranadhm/Documents/Projects/steel-surface-inspection copy"
venv/bin/python training/train_surface_gate_detector.py
```
4. Restart the app:
```bash
cd "/Users/ravindranadhm/Documents/Projects/steel-surface-inspection copy"
venv/bin/uvicorn api.main:app --host 127.0.0.1 --port 8000
```
## Fast settings
If you want the quickest useful first model:
- `model=yolo11n.pt`
- `epochs=18`
- `imgsz=512`
- `batch=16`
- dataset build without `--include-severstal`
If you have only CPU and want to finish faster:
```bash
venv/bin/python training/train_surface_gate_detector.py --epochs 12 --imgsz 416 --batch 8
```
## If you want a stronger detector later
- increase `--max-severstal`
- increase `--max-negatives`
- raise `--composite-multiplier`
- train `25-40` epochs instead of `12-18`
## Runtime config after training
The app will automatically use:
- `models/steel_surface_detector.pt`
Relevant env keys are already supported:
- `SURFACE_DETECTOR_PATH`
- `SURFACE_DETECTOR_CLASS_NAME`
- `SURFACE_DETECTOR_CONFIDENCE`
- `SURFACE_DETECTOR_MIN_AREA_RATIO`
- `SURFACE_DETECTOR_EXPAND_RATIO`
- `SURFACE_DETECTOR_IMAGE_SIZE`
## Important expectation
The first fast model is meant to be a practical gate, not a perfect production vision model.
The improvement loop should be:
1. train the quick detector
2. deploy it
3. save false accepts / false rejects from live camera
4. add those hard examples back into training
5. retrain
That is the fastest route to a reliable factory-facing steel ROI gate.