classifier_train.py

"""
classifier_train.py
"""

import os
import json
import argparse
import warnings
warnings.filterwarnings("ignore")

import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import classification_report, confusion_matrix

# ── Constants ────────────────────────────────────────────────
CLASSES    = ["none", "scratch", "edge_damage", "both"]
CLASS2IDX  = {c: i for i, c in enumerate(CLASSES)}
IMG_SIZE   = 224
BATCH_SIZE = 16
EPOCHS     = 40
LR         = 1e-3
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"


# ── Belt ROI extraction ───────────────────────────────────────
def parse_polygon(label_line, w, h):
    vals = label_line.strip().split()
    coords = [float(v) for v in vals[1:]]
    pts = [(int(coords[i]*w), int(coords[i+1]*h))
           for i in range(0, len(coords)-1, 2)]
    return np.array(pts, dtype=np.int32)


def extract_belt_crop(img_path, label_path, out_size=IMG_SIZE):
    """
    Extract and return the belt region as a square crop.
    Falls back to centre 70% if no label available.
    """
    img = cv2.imread(img_path)
    if img is None:
        return None
    h, w = img.shape[:2]

    # Belt mask from polygon label
    if os.path.exists(label_path):
        with open(label_path) as f:
            lines = [l.strip() for l in f if l.strip()]
        if lines:
            pts = parse_polygon(lines[0], w, h)
            xs  = [p[0] for p in pts]
            ys  = [p[1] for p in pts]
            x1, x2 = max(0, min(xs)), min(w, max(xs))
            y1, y2 = max(0, min(ys)), min(h, max(ys))
            crop = img[y1:y2, x1:x2]
            if crop.size > 0:
                return cv2.resize(crop, (out_size, out_size))

    # Fallback: centre strip
    xl = int(w * 0.15); xr = int(w * 0.85)
    crop = img[:, xl:xr]
    return cv2.resize(crop, (out_size, out_size))


# ── Dataset ───────────────────────────────────────────────────
class BeltDataset(Dataset):
    def __init__(self, records, transform=None):
        """
        records: list of dicts with keys: img_path, label_path, class_idx
        """
        self.records   = records
        self.transform = transform

    def __len__(self):
        return len(self.records)

    def __getitem__(self, idx):
        r    = self.records[idx]
        crop = extract_belt_crop(r["img_path"], r["label_path"])
        if crop is None:
            crop = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)

        # BGR → RGB
        crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
        if self.transform:
            from PIL import Image as PILImage
            crop = self.transform(PILImage.fromarray(crop))
        return crop, r["class_idx"]


# ── Augmentation transforms ───────────────────────────────────
TRAIN_TF = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])

VAL_TF = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])


# ── Model ─────────────────────────────────────────────────────
def build_model(num_classes=4, freeze_backbone=True):
    model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)

    if freeze_backbone:
        # Freeze all layers except last 2 feature blocks + classifier
        for name, param in model.features.named_parameters():
            block_num = int(name.split(".")[0]) if name[0].isdigit() else -1
            if block_num < 6:
                param.requires_grad = False

    # Replace classifier head
    in_features = model.classifier[1].in_features
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.4, inplace=True),
        nn.Linear(in_features, num_classes),
    )
    return model


# ── Training loop ─────────────────────────────────────────────
def train_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0.0
    correct    = 0
    total      = 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        out  = model(imgs)
        loss = criterion(out, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * imgs.size(0)
        preds       = out.argmax(dim=1)
        correct    += (preds == labels).sum().item()
        total      += imgs.size(0)
    return total_loss / total, correct / total


@torch.no_grad()
def eval_epoch(model, loader, criterion):
    model.eval()
    total_loss = 0.0
    correct    = 0
    total      = 0
    all_preds  = []
    all_labels = []
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
        out  = model(imgs)
        loss = criterion(out, labels)
        total_loss += loss.item() * imgs.size(0)
        preds       = out.argmax(dim=1)
        correct    += (preds == labels).sum().item()
        total      += imgs.size(0)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())
    return total_loss / total, correct / total, all_preds, all_labels


# ── Main ──────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image_dir",   required=True)
    parser.add_argument("--label_dir",   default=None)
    parser.add_argument("--review_xlsx", required=True)
    parser.add_argument("--output_dir",  default="classifier")
    parser.add_argument("--epochs",      type=int, default=EPOCHS)
    parser.add_argument("--no_freeze",   action="store_true",
                        help="Unfreeze entire backbone (slower, needs more data)")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # Auto-detect label dir
    if args.label_dir:
        label_dir = args.label_dir
    else:
        label_dir = os.path.join(
            os.path.dirname(args.image_dir.rstrip("/\\")), "labels")

    print(f"\n{'='*60}")
    print(f"  Belt Damage Classifier — Training")
    print(f"{'='*60}")
    print(f"  Device     : {DEVICE}")
    print(f"  Image dir  : {args.image_dir}")
    print(f"  Label dir  : {label_dir}")
    print(f"  Review xlsx: {args.review_xlsx}")
    print(f"  Output dir : {args.output_dir}")
    print(f"  Epochs     : {args.epochs}")
    print(f"{'='*60}\n")

    # ── Load reviewed data ──────────────────────────────────
    df = pd.read_excel(args.review_xlsx, sheet_name="Manual Review", header=2)
    df.columns = ["num","image","daynight","n_scratch","n_edge","n_total",
                  "scratch_ok","edge_ok","fp","miss","severity","dtype","notes","bboxes"]
    reviewed = df[df["scratch_ok"].notna()].copy()
    reviewed = reviewed[reviewed["dtype"].notna()].copy()

    # Normalise dtype labels
    reviewed["dtype"] = reviewed["dtype"].str.strip().str.lower()
    # Map any variants
    dtype_map = {
        "scratch only": "scratch",
        "edge": "edge_damage",
        "edge damage": "edge_damage",
    }
    reviewed["dtype"] = reviewed["dtype"].replace(dtype_map)
    reviewed = reviewed[reviewed["dtype"].isin(CLASSES)].copy()

    print(f"Labelled images found : {len(reviewed)}")
    print(f"Class distribution    :")
    for cls, cnt in reviewed["dtype"].value_counts().items():
        print(f"  {cls:15s}: {cnt}")

    # Build records list
    records = []
    missing = 0
    for _, row in reviewed.iterrows():
        img_name   = str(row["image"])
        img_path   = os.path.join(args.image_dir, img_name)
        stem       = os.path.splitext(img_name)[0]
        label_path = os.path.join(label_dir, stem + ".txt")

        if not os.path.exists(img_path):
            missing += 1
            continue

        records.append({
            "img_path"  : img_path,
            "label_path": label_path,
            "class_idx" : CLASS2IDX[row["dtype"]],
            "dtype"     : row["dtype"],
        })

    if missing:
        print(f"\n  [WARN] {missing} images not found in image_dir — skipped")

    print(f"\n  Total training samples: {len(records)}")

    if len(records) < 20:
        print("\n[ERROR] Too few images to train. Check --image_dir path.")
        return

    # ── Class weights for imbalance ──────────────────────────
    class_counts = np.zeros(len(CLASSES))
    for r in records:
        class_counts[r["class_idx"]] += 1
    class_weights = torch.FloatTensor(
        [len(records) / (len(CLASSES) * max(c, 1)) for c in class_counts]
    ).to(DEVICE)
    print(f"\n  Class weights: {dict(zip(CLASSES, class_weights.cpu().numpy().round(2)))}")

    # ── Train/Val split (stratified 80/20) ───────────────────
    labels_arr = [r["class_idx"] for r in records]
    skf        = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    train_idx, val_idx = next(skf.split(records, labels_arr))

    train_records = [records[i] for i in train_idx]
    val_records   = [records[i] for i in val_idx]
    print(f"  Train: {len(train_records)}  |  Val: {len(val_records)}\n")

    train_ds = BeltDataset(train_records, transform=TRAIN_TF)
    val_ds   = BeltDataset(val_records,   transform=VAL_TF)

    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE,
                              shuffle=True,  num_workers=0, drop_last=False)
    val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE,
                              shuffle=False, num_workers=0)

    # ── Build model ──────────────────────────────────────────
    model     = build_model(num_classes=len(CLASSES),
                            freeze_backbone=not args.no_freeze).to(DEVICE)
    criterion = nn.CrossEntropyLoss(weight=class_weights)
    optimizer = optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=LR, weight_decay=1e-4
    )
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs, eta_min=1e-5
    )

    # ── Training ─────────────────────────────────────────────
    best_val_acc = 0.0
    log_rows     = []
    print(f"{'Epoch':>6} {'TrainLoss':>10} {'TrainAcc':>9} {'ValLoss':>8} {'ValAcc':>7}")
    print("-" * 48)

    for epoch in range(1, args.epochs + 1):
        tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, criterion)
        vl_loss, vl_acc, preds, true_labels = eval_epoch(model, val_loader, criterion)
        scheduler.step()

        print(f"{epoch:>6}   {tr_loss:>9.4f}   {tr_acc*100:>7.2f}%  "
              f"{vl_loss:>7.4f}  {vl_acc*100:>6.2f}%")

        log_rows.append({
            "epoch": epoch,
            "train_loss": round(tr_loss, 4),
            "train_acc":  round(tr_acc, 4),
            "val_loss":   round(vl_loss, 4),
            "val_acc":    round(vl_acc, 4),
        })

        if vl_acc > best_val_acc:
            best_val_acc = vl_acc
            torch.save(model.state_dict(),
                       os.path.join(args.output_dir, "best_model.pth"))

    # ── Final evaluation on val set ───────────────────────────
    model.load_state_dict(
        torch.load(os.path.join(args.output_dir, "best_model.pth"),
                   map_location=DEVICE)
    )
    _, final_acc, final_preds, final_true = eval_epoch(
        model, val_loader, criterion)

    print(f"\n{'='*60}")
    print(f"  Best Val Accuracy : {best_val_acc*100:.2f}%")
    print(f"{'='*60}")
    print("\nClassification Report (Validation Set):")
    print(classification_report(
        final_true, final_preds,
        target_names=CLASSES, zero_division=0
    ))
    print("Confusion Matrix (rows=true, cols=pred):")
    print(pd.DataFrame(
        confusion_matrix(final_true, final_preds),
        index=[f"true_{c}" for c in CLASSES],
        columns=[f"pred_{c}" for c in CLASSES]
    ).to_string())

    # ── Save artefacts ────────────────────────────────────────
    with open(os.path.join(args.output_dir, "class_map.json"), "w") as f:
        json.dump({"classes": CLASSES, "class2idx": CLASS2IDX}, f, indent=2)

    pd.DataFrame(log_rows).to_csv(
        os.path.join(args.output_dir, "train_log.csv"), index=False)

    print(f"\n  Saved → {args.output_dir}/best_model.pth")
    print(f"  Saved → {args.output_dir}/class_map.json")
    print(f"  Saved → {args.output_dir}/train_log.csv")
    print(f"\nNext step:")
    print(f"  python classifier_predict.py \\")
    print(f'      --image_dir   "{args.image_dir}" \\')
    print(f'      --label_dir   "{label_dir}" \\')
    print(f'      --review_xlsx "{args.review_xlsx}" \\')
    print(f'      --model_path  "{args.output_dir}/best_model.pth" \\')
    print(f'      --output_xlsx "manual_review_completed.xlsx"')


if __name__ == "__main__":
    main()