Spaces:
Sleeping
Sleeping
| """ | |
| 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") |