""" Fish Freshness Classifier - Training Script ========================================== Dataset struktur (folder-based classification): classification/ train/ eye-fresh/ -> label 0 eye-non-fresh/ -> label 1 gill-fresh/ -> label 2 gill-non-fresh/ -> label 3 valid/ test/ Model: MobileNetV3-Small (ringan, cocok untuk deployment) Grade logic (dari kombinasi prediksi): eye-fresh + gill-fresh = Grade A (Sangat Segar) eye-fresh + gill-non-fresh = Grade B (Cukup Segar) eye-non-fresh + gill-fresh = Grade B (Cukup Segar) eye-non-fresh + gill-non-fresh = Grade C (Tidak Segar) Fix v2: - num_workers=0 (Windows tidak support num_workers > 0 tanpa if __name__) - Early stopping (patience=7) untuk cegah overfitting - Dropout tambahan di classifier head - evaluate() pakai labels= agar tidak crash jika test set kurang kelas - Augmentasi lebih agresif untuk dataset kecil """ import os import json import copy import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms, models from sklearn.metrics import classification_report, confusion_matrix import matplotlib.pyplot as plt import seaborn as sns import numpy as np # ── Config ─────────────────────────────────────────────────────────────────── DATA_DIR = "../data/classification" SAVE_DIR = "models" NUM_EPOCHS = 50 BATCH_SIZE = 16 # lebih kecil → update lebih sering, generalisasi lebih baik LR = 5e-4 # lebih rendah → lebih stabil IMG_SIZE = 224 NUM_CLASSES = 4 PATIENCE = 10 # early stopping: berhenti jika val_acc tidak naik N epoch DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {DEVICE}") # ── Transforms ─────────────────────────────────────────────────────────────── # Augmentasi lebih agresif karena dataset kecil (522 train) train_transforms = transforms.Compose([ transforms.Resize((IMG_SIZE + 32, IMG_SIZE + 32)), transforms.RandomCrop(IMG_SIZE), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(20), transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.3, hue=0.1), transforms.RandomGrayscale(p=0.05), transforms.RandomAffine(degrees=0, translate=(0.15, 0.15), shear=10), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), transforms.RandomErasing(p=0.2, scale=(0.02, 0.15)), ]) val_transforms = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) # ── Dataset ────────────────────────────────────────────────────────────────── def get_dataloaders(data_dir): train_ds = datasets.ImageFolder(os.path.join(data_dir, "train"), train_transforms) valid_ds = datasets.ImageFolder(os.path.join(data_dir, "valid"), val_transforms) test_ds = datasets.ImageFolder(os.path.join(data_dir, "test"), val_transforms) print(f"Train: {len(train_ds)} | Valid: {len(valid_ds)} | Test: {len(test_ds)}") print(f"Classes: {train_ds.classes}") # Handle imbalance dengan WeightedRandomSampler class_counts = [0] * NUM_CLASSES for _, label in train_ds.samples: class_counts[label] += 1 print(f"Class counts (train): {dict(zip(train_ds.classes, class_counts))}") weights = [1.0 / class_counts[label] for _, label in train_ds.samples] sampler = torch.utils.data.WeightedRandomSampler(weights, len(weights)) # num_workers=0 wajib di Windows agar tidak error loaders = { "train": DataLoader(train_ds, batch_size=BATCH_SIZE, sampler=sampler, num_workers=0), "valid": DataLoader(valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0), "test": DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0), } return loaders, train_ds.classes # ── Model ───────────────────────────────────────────────────────────────────── def build_model(): model = models.mobilenet_v3_small(weights=None) # Tambah Dropout sebelum layer final untuk kurangi overfitting in_features = model.classifier[3].in_features model.classifier[3] = nn.Sequential( nn.Dropout(p=0.4), nn.Linear(in_features, NUM_CLASSES), ) return model.to(DEVICE) # ── Training ────────────────────────────────────────────────────────────────── def train_model(model, loaders, num_epochs): criterion = nn.CrossEntropyLoss(label_smoothing=0.1) optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-3) # ReduceLROnPlateau: turunkan LR jika val_acc tidak naik 5 epoch scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="max", factor=0.5, patience=5, min_lr=1e-6 ) best_acc = 0.0 best_wts = copy.deepcopy(model.state_dict()) no_improve = 0 history = {"train_loss": [], "val_loss": [], "train_acc": [], "val_acc": []} for epoch in range(num_epochs): current_lr = optimizer.param_groups[0]["lr"] print(f"\nEpoch {epoch+1}/{num_epochs} lr={current_lr:.6f}") print("-" * 50) for phase in ["train", "valid"]: model.train() if phase == "train" else model.eval() running_loss, running_correct = 0.0, 0 for inputs, labels in loaders[phase]: inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() with torch.set_grad_enabled(phase == "train"): outputs = model(inputs) loss = criterion(outputs, labels) preds = outputs.argmax(dim=1) if phase == "train": loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_correct += (preds == labels).sum().item() n = len(loaders[phase].dataset) loss = running_loss / n acc = running_correct / n if phase == "train": history["train_loss"].append(loss) history["train_acc"].append(acc) else: history["val_loss"].append(loss) history["val_acc"].append(acc) scheduler.step(acc) if acc > best_acc: best_acc = acc no_improve = 0 best_wts = copy.deepcopy(model.state_dict()) torch.save(model.state_dict(), os.path.join(SAVE_DIR, "best_model.pth")) print(f" ✓ Best model saved val_acc={acc:.4f}") else: no_improve += 1 print(f" No improvement ({no_improve}/{PATIENCE})") print(f" {phase:5s} loss={loss:.4f} acc={acc:.4f}") # Early stopping if no_improve >= PATIENCE: print(f"\nEarly stopping triggered at epoch {epoch+1}.") break print(f"\nBest Val Accuracy: {best_acc:.4f}") model.load_state_dict(best_wts) return model, history # ── Evaluation ──────────────────────────────────────────────────────────────── def evaluate(model, loader, class_names): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for inputs, labels in loader: inputs = inputs.to(DEVICE) preds = model(inputs).argmax(dim=1).cpu() all_preds.extend(preds.numpy()) all_labels.extend(labels.numpy()) # Gunakan labels= agar tidak crash jika test set tidak punya semua kelas present_labels = sorted(set(all_labels)) present_names = [class_names[i] for i in present_labels] print("\nClassification Report:") print(classification_report( all_labels, all_preds, labels=present_labels, target_names=present_names, zero_division=0, )) # Confusion matrix hanya untuk kelas yang ada di test set cm = confusion_matrix(all_labels, all_preds, labels=present_labels) fig, ax = plt.subplots(figsize=(7, 5)) sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=present_names, yticklabels=present_names, ax=ax) ax.set_xlabel("Predicted"); ax.set_ylabel("Actual") ax.set_title("Confusion Matrix - Test Set") plt.tight_layout() plt.savefig(os.path.join(SAVE_DIR, "confusion_matrix.png"), dpi=150) print("Confusion matrix saved.") # ── Plot History ────────────────────────────────────────────────────────────── def plot_history(history): fig, axes = plt.subplots(1, 2, figsize=(12, 4)) axes[0].plot(history["train_loss"], label="Train") axes[0].plot(history["val_loss"], label="Val") axes[0].set_title("Loss"); axes[0].legend(); axes[0].set_xlabel("Epoch") axes[1].plot(history["train_acc"], label="Train") axes[1].plot(history["val_acc"], label="Val") axes[1].set_title("Accuracy"); axes[1].legend(); axes[1].set_xlabel("Epoch") plt.tight_layout() plt.savefig(os.path.join(SAVE_DIR, "training_history.png"), dpi=150) print("Training history saved.") # ── Main ────────────────────────────────────────────────────────────────────── if __name__ == "__main__": os.makedirs(SAVE_DIR, exist_ok=True) loaders, classes = get_dataloaders(DATA_DIR) with open(os.path.join(SAVE_DIR, "class_names.json"), "w") as f: json.dump(classes, f) print(f"Class order saved: {classes}") model = build_model() model, history = train_model(model, loaders, NUM_EPOCHS) print("\n=== Test Set Evaluation ===") evaluate(model, loaders["test"], classes) plot_history(history) with open(os.path.join(SAVE_DIR, "history.json"), "w") as f: json.dump(history, f) print("\n✓ Training complete! Model saved to models/best_model.pth")