Full Project Guide

This README explains how to use the trained classifier (ResNet) and segmentation model (U-Net) in your project.

---

Project Structure

scratch_detector/
├── data/
│   ├── raw/
│   │   ├── good/
│   │   ├──bad/
│   │   ├── masks/
│   ├── splits/
│
├── src/
│   ├── preprocess_dataset.py
│   ├── train_classifier.py
│   ├── train_segmentation.py
│   ├── inference.py
│
├── models/
│   ├── trained_weights
│
├── results/
│   ├── metrics/
│   └── plots/
│
├── outputs/
│   ├── test_results/
│
├── README.md
├── requirements.txt

---

Installation

1. Clone the repository

git clone <your_repo_link>
cd scratch_detector

2. Install dependencies

pip install -r requirements.txt

Make sure you have Python 3.10+.

---

Dataset Preparation

Place your images in:

data/raw/good/
data/raw/bad/
data/raw/masks/   # segmentation masks for bad images

Then generate train/val/test splits:

python src/preprocess_dataset.py

This creates:

data/splits/train.csv
data/splits/val.csv
data/splits/test.csv

---

Train the Classifier (ResNet)

To train the ResNet-based scratch vs no-scratch classifier:

python src/train_classifier.py

After training, the model will be saved in:

models/trained_weights/resnet.pth

---

Train the Segmentation Model (U-Net)

To train the segmentation model:

python src/train_segmentation.py

After training, the model will be saved in:

models/trained_weights/unet.pth

---

Visualizing Metrics and Plots

• All Performance metrics and plots will appear inside the models or results folders.
• Check them after training.

---

Inference (Run Trained Models on Test Data)

Use inference.py for classification and segmentation.

1. Classify a single image

python src/inference.py --img path/to/image.jpg --task classify --model models/resnet_classifier.pth

2. Segment a scratch

python src/inference.py --img path/to/image.jpg --task segment --model models/unet_segmenter.pth

3. Run on a folder of images

python src/inference.py --folder path/to/images --task classify

The predictions will be saved in the ouput folder.

---

Tips for Best Results

• Ensure masks are binary
• Maintain consistent image sizes
• Use CUDA for faster training
• Keep folder names exactly good/, bad/, masks/

---

If you use the project, consider giving your repo a star!