| # 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. |
|
|