File size: 3,788 Bytes
a3abb2d 1925f28 a3abb2d 26142cd a3abb2d 1925f28 a3abb2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | # 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.
|