Spaces:
Sleeping
Sleeping
File size: 3,801 Bytes
0f049b3 | 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 | # -------------------------------------------------------------
# HOW TO RUN THIS SCRIPT:
# -------------------------------------------------------------
# 1️⃣ Create a virtual environment:
# python -m venv .venv && .\.venv\Scripts\Activate
#
# 2️⃣ Install dependencies:
# pip install -r requirements.txt
#
# 3️⃣ Train and evaluate model:
# python src/main.py --data data/combined --epochs 15 --lr 3e-4 --bs 32
# -------------------------------------------------------------
import argparse, random
from pathlib import Path
from typing import Dict
import numpy as np
import torch
from data import load_data
from model import build_model
from train import train_model
from eval import eval_on_test
# -------------------------------------------------------------
# FUNCTION: set_seed
# -------------------------------------------------------------
def set_seed(seed: int = 42):
random.seed(seed); np.random.seed(seed)
torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# -------------------------------------------------------------
# FUNCTION: make_dirs
# -------------------------------------------------------------
def make_dirs():
Path("models").mkdir(parents=True, exist_ok=True)
Path("reports").mkdir(parents=True, exist_ok=True)
# -------------------------------------------------------------
# FUNCTION: get_device
# -------------------------------------------------------------
def get_device():
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
# -------------------------------------------------------------
# FUNCTION: parse_args
# -------------------------------------------------------------
def parse_args():
p = argparse.ArgumentParser(description="Recycle Material Classifier (ResNet-18)")
p.add_argument("--data", type=str, default="data/images", help="root with train/val/test")
p.add_argument("--epochs", type=int, default=10)
p.add_argument("--bs", type=int, default=32)
p.add_argument("--lr", type=float, default=1e-3)
p.add_argument("--freeze", action="store_true", help="freeze ResNet backbone")
p.add_argument("--seed", type=int, default=42)
return p.parse_args()
# -------------------------------------------------------------
# FUNCTION: main
# -------------------------------------------------------------
def main():
# 1️⃣ Read input arguments
args = parse_args()
# 2️⃣ Ensure results are reproducible
set_seed(args.seed)
# 3️⃣ Create folders for saving models/reports
make_dirs()
# 4️⃣ Detect whether to use GPU or CPU
device = get_device()
print("Device:", device)
# 5️⃣ Load dataset and dataloaders
out = load_data(args.data, args.bs)
dsets, loaders = out["dsets"], out["loaders"]
# Extract class names in correct order
class_names = [k for k, _ in sorted(dsets["train"].class_to_idx.items(), key=lambda x: x[1])]
print("Classes:", class_names)
# 6️⃣ Build model (ResNet18 with custom classification head)
model = build_model(num_classes=len(class_names), freeze_backbone=args.freeze, device=device)
# 7️⃣ Train model using training and validation sets
model, best_val = train_model(model, loaders, epochs=args.epochs, lr=args.lr, device=device)
print(f"Best val acc: {best_val:.4f}")
# 8️⃣ Evaluate final model on the test set and save reports
eval_on_test(model, loaders["test"], class_names, device)
# -------------------------------------------------------------
# ENTRY POINT (runs when executed directly)
# -------------------------------------------------------------
if __name__ == "__main__":
main()
|