02 - Transfer Learning: Riutilizzare Modelli Pre-trained

import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet50_Weights

# ---- API moderna (raccomandata): usa l'enum Weights ----
# Carica ResNet-50 con pesi IMAGENET1K_V2 (80.9% top-1)
weights = ResNet50_Weights.IMAGENET1K_V2
model = models.resnet50(weights=weights)

# I pesi includono le trasformazioni di preprocessing
preprocess = weights.transforms()
# Questo crea automaticamente: Resize(232) -> CenterCrop(224)
#   -> ToTensor() -> Normalize(mean, std)

print(f"Modello caricato: {sum(p.numel() for p in model.parameters()):,} parametri")
print(f"Device: {next(model.parameters()).device}")

# ---- API legacy (deprecata ma ancora funzionante) ----
# model = models.resnet50(pretrained=True)  # NON usare più

# ---- Ispeziona la struttura del modello ----
# ResNet-50 ha: conv1, bn1, relu, maxpool, layer1-4, avgpool, fc
print(model)
    

import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet50_Weights


def create_feature_extractor(num_classes: int, device: str = 'cuda') -> nn.Module:
    """Crea un feature extractor basato su ResNet-50 pre-trained.

    Args:
        num_classes: Numero di classi per il task target
        device: Device su cui caricare il modello

    Returns:
        Modello con backbone congelato e classificatore trainable
    """
    # 1. Carica il modello pre-addestrato
    weights = ResNet50_Weights.IMAGENET1K_V2
    model = models.resnet50(weights=weights)

    # 2. CONGELA tutti i parametri del backbone
    for param in model.parameters():
        param.requires_grad = False

    # 3. Sostituisci il classificatore finale
    # Il layer fc originale: Linear(2048 -> 1000)
    num_features = model.fc.in_features  # 2048

    model.fc = nn.Sequential(
        nn.Linear(num_features, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.3),
        nn.Linear(512, num_classes)
    )
    # I nuovi layer hanno requires_grad=True di default

    # 4. Sposta sul device
    model = model.to(device)

    # Verifica parametri trainable vs congelati
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    frozen = total - trainable
    print(f"Parametri totali:    {total:,}")
    print(f"Parametri congelati: {frozen:,} ({frozen/total*100:.1f}%)")
    print(f"Parametri trainable: {trainable:,} ({trainable/total*100:.1f}%)")

    return model


# Esempio: classificazione binaria (difettoso / non difettoso)
model = create_feature_extractor(num_classes=2)
# Output:
# Parametri totali:    24,607,170
# Parametri congelati: 23,508,032 (95.5%)
# Parametri trainable: 1,099,138  (4.5%)
    

from torchvision import transforms
from torchvision.transforms import v2  # API v2 (PyTorch 2.0+)
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# ---- Normalizzazione ImageNet (OBBLIGATORIA per modelli pre-trained) ----
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

# ---- Trasformazioni per il training (con augmentation) ----
train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.2),
    transforms.RandomRotation(degrees=15),
    transforms.ColorJitter(
        brightness=0.3,
        contrast=0.3,
        saturation=0.2,
        hue=0.1
    ),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.RandomGrayscale(p=0.05),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    transforms.RandomErasing(p=0.2, scale=(0.02, 0.15)),
])

# ---- Trasformazioni per validazione/test (SENZA augmentation) ----
val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

# ---- Caricamento dataset con ImageFolder ----
# Struttura directory richiesta:
#   data/
#     train/
#       classe_A/  (img1.jpg, img2.jpg, ...)
#       classe_B/  (img1.jpg, img2.jpg, ...)
#     val/
#       classe_A/  (...)
#       classe_B/  (...)

train_dataset = ImageFolder(root='data/train', transform=train_transforms)
val_dataset = ImageFolder(root='data/val', transform=val_transforms)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=True
)

val_loader = DataLoader(
    val_dataset,
    batch_size=64,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

print(f"Classi: {train_dataset.classes}")
print(f"Training: {len(train_dataset)} immagini")
print(f"Validation: {len(val_dataset)} immagini")
    

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from typing import Tuple
import time


def train_one_epoch(
    model: nn.Module,
    train_loader,
    criterion: nn.Module,
    optimizer: torch.optim.Optimizer,
    device: str,
    epoch: int
) -> Tuple[float, float]:
    """Addestra il modello per un'epoca."""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (images, labels) in enumerate(train_loader):
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward pass
        optimizer.zero_grad(set_to_none=True)  # Più efficiente di zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        # Metriche
        running_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = 100.0 * correct / total
    return epoch_loss, epoch_acc


@torch.no_grad()
def evaluate(
    model: nn.Module,
    val_loader,
    criterion: nn.Module,
    device: str
) -> Tuple[float, float]:
    """Valuta il modello sul validation set."""
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in val_loader:
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        outputs = model(images)
        loss = criterion(outputs, labels)

        running_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    val_loss = running_loss / total
    val_acc = 100.0 * correct / total
    return val_loss, val_acc


def train_transfer_learning(
    model: nn.Module,
    train_loader,
    val_loader,
    num_epochs: int = 30,
    lr: float = 1e-3,
    device: str = 'cuda'
) -> dict:
    """Pipeline completa di training per Transfer Learning."""

    model = model.to(device)

    # Ottimizzatore: solo parametri trainable (non congelati)
    trainable_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = AdamW(trainable_params, lr=lr, weight_decay=1e-4)

    # Scheduler: cosine annealing con warm restarts
    scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)

    # Loss function con label smoothing
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

    # Tracking
    best_val_acc = 0.0
    history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []}

    for epoch in range(num_epochs):
        start_time = time.time()

        # Training
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device, epoch
        )

        # Validation
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)

        # Scheduler step
        scheduler.step()

        # Salvataggio metriche
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['train_acc'].append(train_acc)
        history['val_acc'].append(val_acc)

        elapsed = time.time() - start_time

        print(
            f"Epoch {epoch+1}/{num_epochs} "
            f"[{elapsed:.1f}s] "
            f"Train: {train_loss:.4f} ({train_acc:.1f}%) | "
            f"Val: {val_loss:.4f} ({val_acc:.1f}%)"
        )

        # Salva il miglior modello
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_acc': val_acc,
            }, 'best_model.pth')
            print(f"  --> Nuovo miglior modello salvato! Val Acc: {val_acc:.2f}%")

    print(f"\nMiglior Val Accuracy: {best_val_acc:.2f}%")
    return history
    

from torchvision import models
from torchvision.models import ResNet50_Weights
import torch.nn as nn
from torch.optim import AdamW


def create_finetuning_model(num_classes: int, device: str = 'cuda'):
    """Crea modello con gruppi di parametri per discriminative LR."""

    weights = ResNet50_Weights.IMAGENET1K_V2
    model = models.resnet50(weights=weights)

    # Sostituisci il classificatore
    num_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(num_features, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.3),
        nn.Linear(512, num_classes)
    )

    model = model.to(device)

    # Definisci gruppi di parametri con LR diversi
    # Gruppo 1: Layer iniziali (conv1, bn1, layer1, layer2) - LR molto basso
    # Gruppo 2: Layer intermedi (layer3) - LR medio
    # Gruppo 3: Layer finali (layer4) - LR più alto
    # Gruppo 4: Classificatore (fc) - LR massimo

    base_lr = 1e-4

    param_groups = [
        {
            'params': list(model.conv1.parameters()) +
                      list(model.bn1.parameters()) +
                      list(model.layer1.parameters()) +
                      list(model.layer2.parameters()),
            'lr': base_lr * 0.01,    # 1e-6
            'name': 'early_layers'
        },
        {
            'params': list(model.layer3.parameters()),
            'lr': base_lr * 0.1,     # 1e-5
            'name': 'mid_layers'
        },
        {
            'params': list(model.layer4.parameters()),
            'lr': base_lr,            # 1e-4
            'name': 'late_layers'
        },
        {
            'params': list(model.fc.parameters()),
            'lr': base_lr * 10,       # 1e-3
            'name': 'classifier'
        },
    ]

    optimizer = AdamW(param_groups, weight_decay=1e-4)

    # Verifica i learning rate
    for group in optimizer.param_groups:
        n_params = sum(p.numel() for p in group['params'])
        print(f"{group.get('name', 'unknown'):<15} LR={group['lr']:.2e}  "
              f"Params={n_params:,}")

    return model, optimizer
    

import torch.nn as nn
from torchvision import models
from torchvision.models import ResNet50_Weights


def gradual_unfreeze_training(
    num_classes: int,
    train_loader,
    val_loader,
    device: str = 'cuda'
):
    """Training con scongelamento graduale dei layer."""

    # Carica modello e congela tutto
    weights = ResNet50_Weights.IMAGENET1K_V2
    model = models.resnet50(weights=weights)

    # Congela tutti i layer
    for param in model.parameters():
        param.requires_grad = False

    # Sostituisci il classificatore (trainable di default)
    model.fc = nn.Sequential(
        nn.Linear(model.fc.in_features, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.3),
        nn.Linear(512, num_classes)
    )
    model = model.to(device)

    # Definisci le fasi di unfreezing
    unfreeze_schedule = [
        # (epoch_start, layers_to_unfreeze, learning_rate)
        (0,  [],                         1e-3),   # Solo classifier
        (5,  [model.layer4],             5e-4),   # + layer4
        (10, [model.layer3],             2e-4),   # + layer3
        (15, [model.layer2],             1e-4),   # + layer2
        (20, [model.layer1, model.conv1, model.bn1], 5e-5),  # Tutto
    ]

    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    current_phase = 0

    for epoch in range(30):
        # Controlla se dobbiamo scongeleare nuovi layer
        if (current_phase < len(unfreeze_schedule) - 1 and
                epoch >= unfreeze_schedule[current_phase + 1][0]):
            current_phase += 1
            phase = unfreeze_schedule[current_phase]

            # Scongela i nuovi layer
            for layer_group in phase[1]:
                for param in layer_group.parameters():
                    param.requires_grad = True

            # Ricrea optimizer con nuovo LR
            trainable_params = [
                p for p in model.parameters() if p.requires_grad
            ]
            optimizer = torch.optim.AdamW(
                trainable_params, lr=phase[2], weight_decay=1e-4
            )

            n_trainable = sum(
                p.numel() for p in model.parameters() if p.requires_grad
            )
            print(f"\n=== Fase {current_phase}: Epoch {epoch} ===")
            print(f"Parametri trainable: {n_trainable:,}")
            print(f"Learning rate: {phase[2]:.2e}")

        elif epoch == 0:
            # Prima fase: solo classifier
            optimizer = torch.optim.AdamW(
                model.fc.parameters(), lr=1e-3, weight_decay=1e-4
            )

        # Training e validation standard
        # (usa le funzioni train_one_epoch e evaluate definite prima)
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device, epoch
        )
        val_loss, val_acc = evaluate(
            model, val_loader, criterion, device
        )

        print(
            f"Epoch {epoch+1}/30 | "
            f"Train: {train_loss:.4f} ({train_acc:.1f}%) | "
            f"Val: {val_loss:.4f} ({val_acc:.1f}%)"
        )
    

import math
from torch.optim.lr_scheduler import LambdaLR


def get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    min_lr_ratio: float = 0.01
):
    """Cosine annealing con linear warmup.

    Args:
        optimizer: L'ottimizzatore PyTorch
        num_warmup_steps: Numero di step di warmup
        num_training_steps: Numero totale di step di training
        min_lr_ratio: Rapporto LR minimo / LR massimo

    Returns:
        LambdaLR scheduler
    """
    def lr_lambda(current_step: int) -> float:
        # Fase di warmup: LR cresce linearmente da 0 a max
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))

        # Fase di cosine annealing: LR decresce con coseno
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
        return max(min_lr_ratio, cosine_decay)

    return LambdaLR(optimizer, lr_lambda)


# Esempio di utilizzo
num_epochs = 30
steps_per_epoch = len(train_loader)
total_steps = num_epochs * steps_per_epoch
warmup_steps = 3 * steps_per_epoch  # 3 epoche di warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

# Nel training loop, chiama scheduler.step() DOPO ogni batch
# for batch in train_loader:
#     ...
#     optimizer.step()
#     scheduler.step()  # Step per batch, non per epoca
    

from torchvision.transforms import v2


# CutMix e MixUp applicati a livello di batch (non di singola immagine)
cutmix = v2.CutMix(num_classes=10)
mixup = v2.MixUp(num_classes=10, alpha=0.2)

# Combina entrambi con probabilità 50/50
cutmix_or_mixup = v2.RandomChoice([cutmix, mixup])

# Nel training loop:
# for images, labels in train_loader:
#     images, labels = cutmix_or_mixup(images, labels)
#     # Nota: labels ora sono soft (one-hot probabilistico)
#     # La CrossEntropyLoss gestisce automaticamente soft labels
#     outputs = model(images)
#     loss = criterion(outputs, labels)
    

import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import numpy as np
from torch.utils.data import Dataset
from PIL import Image


def get_albumentations_transforms(is_training: bool = True):
    """Pipeline di augmentation con Albumentations."""

    if is_training:
        return A.Compose([
            A.RandomResizedCrop(height=224, width=224, scale=(0.7, 1.0)),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.2),

            # Augmentation geometriche
            A.ShiftScaleRotate(
                shift_limit=0.1,
                scale_limit=0.15,
                rotate_limit=15,
                border_mode=cv2.BORDER_REFLECT_101,
                p=0.5
            ),

            # Augmentation di colore/illuminazione
            A.OneOf([
                A.RandomBrightnessContrast(
                    brightness_limit=0.3, contrast_limit=0.3, p=1.0
                ),
                A.HueSaturationValue(
                    hue_shift_limit=10, sat_shift_limit=30,
                    val_shift_limit=20, p=1.0
                ),
                A.CLAHE(clip_limit=4.0, p=1.0),
            ], p=0.7),

            # Rumore e blur
            A.OneOf([
                A.GaussNoise(var_limit=(10, 50), p=1.0),
                A.GaussianBlur(blur_limit=(3, 5), p=1.0),
                A.MotionBlur(blur_limit=5, p=1.0),
            ], p=0.3),

            # Dropout spaziale
            A.CoarseDropout(
                max_holes=8, max_height=16, max_width=16,
                min_holes=1, min_height=8, min_width=8,
                fill_value=0, p=0.3
            ),

            # Normalizzazione ImageNet
            A.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
            ToTensorV2(),
        ])
    else:
        return A.Compose([
            A.Resize(256, 256),
            A.CenterCrop(224, 224),
            A.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
            ToTensorV2(),
        ])


class AlbumentationsDataset(Dataset):
    """Dataset wrapper che usa Albumentations per le trasformazioni."""

    def __init__(self, image_paths: list, labels: list, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int):
        # Leggi immagine con OpenCV (BGR -> RGB)
        image = cv2.imread(self.image_paths[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        label = self.labels[idx]

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        return image, label
    

"""
Pipeline completa per classificazione difetti industriali.
Transfer Learning con ResNet-50 e gradual unfreezing.
"""

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.models import ResNet50_Weights
from torchvision.datasets import ImageFolder
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
    f1_score
)
import numpy as np
from pathlib import Path


# ============================================================
# 1. CONFIGURAZIONE
# ============================================================
class Config:
    """Configurazione centralizzata del training."""
    data_dir = Path('data/pcb_defects')
    num_classes = 4
    class_names = ['OK', 'missing_component', 'scratch', 'solder_bridge']

    # Training
    batch_size = 32
    num_workers = 4
    phase1_epochs = 10   # Feature extraction
    phase2_epochs = 15   # Fine-tuning layer4
    phase3_epochs = 10   # Full fine-tuning
    phase1_lr = 1e-3
    phase2_lr = 1e-4
    phase3_lr = 1e-5

    # Device
    device = 'cuda' if torch.cuda.is_available() else 'cpu'


# ============================================================
# 2. PREPARAZIONE DATI
# ============================================================
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomVerticalFlip(p=0.5),
    transforms.RandomRotation(degrees=20),
    transforms.ColorJitter(brightness=0.3, contrast=0.3),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    transforms.RandomErasing(p=0.2),
])

val_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

cfg = Config()
train_dataset = ImageFolder(cfg.data_dir / 'train', transform=train_transforms)
val_dataset = ImageFolder(cfg.data_dir / 'val', transform=val_transforms)
test_dataset = ImageFolder(cfg.data_dir / 'test', transform=val_transforms)

train_loader = DataLoader(
    train_dataset, batch_size=cfg.batch_size,
    shuffle=True, num_workers=cfg.num_workers, pin_memory=True
)
val_loader = DataLoader(
    val_dataset, batch_size=cfg.batch_size * 2,
    shuffle=False, num_workers=cfg.num_workers, pin_memory=True
)
test_loader = DataLoader(
    test_dataset, batch_size=cfg.batch_size * 2,
    shuffle=False, num_workers=cfg.num_workers, pin_memory=True
)


# ============================================================
# 3. MODELLO
# ============================================================
def build_model(num_classes: int, device: str) -> nn.Module:
    """Crea il modello ResNet-50 per il Transfer Learning."""
    weights = ResNet50_Weights.IMAGENET1K_V2
    model = models.resnet50(weights=weights)

    # Congela tutto
    for param in model.parameters():
        param.requires_grad = False

    # Nuovo classificatore
    model.fc = nn.Sequential(
        nn.Linear(2048, 512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.4),
        nn.Linear(512, 128),
        nn.ReLU(inplace=True),
        nn.Dropout(0.2),
        nn.Linear(128, num_classes)
    )

    return model.to(device)


# ============================================================
# 4. TRAINING MULTI-FASE
# ============================================================
def run_training():
    """Esegue il training multi-fase con gradual unfreezing."""

    model = build_model(cfg.num_classes, cfg.device)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    best_val_acc = 0.0

    # ---- Fase 1: Feature Extraction ----
    print("=" * 60)
    print("FASE 1: Feature Extraction (solo classifier)")
    print("=" * 60)

    optimizer = AdamW(model.fc.parameters(), lr=cfg.phase1_lr, weight_decay=1e-4)

    for epoch in range(cfg.phase1_epochs):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, cfg.device, epoch
        )
        val_loss, val_acc = evaluate(
            model, val_loader, criterion, cfg.device
        )
        print(f"  Epoch {epoch+1}/{cfg.phase1_epochs} | "
              f"Train: {train_acc:.1f}% | Val: {val_acc:.1f}%")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')

    # ---- Fase 2: Fine-tuning layer4 ----
    print("\n" + "=" * 60)
    print("FASE 2: Fine-tuning layer4 + classifier")
    print("=" * 60)

    for param in model.layer4.parameters():
        param.requires_grad = True

    optimizer = AdamW([
        {'params': model.layer4.parameters(), 'lr': cfg.phase2_lr * 0.1},
        {'params': model.fc.parameters(), 'lr': cfg.phase2_lr},
    ], weight_decay=1e-4)

    for epoch in range(cfg.phase2_epochs):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, cfg.device, epoch
        )
        val_loss, val_acc = evaluate(
            model, val_loader, criterion, cfg.device
        )
        print(f"  Epoch {epoch+1}/{cfg.phase2_epochs} | "
              f"Train: {train_acc:.1f}% | Val: {val_acc:.1f}%")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')

    # ---- Fase 3: Full Fine-tuning ----
    print("\n" + "=" * 60)
    print("FASE 3: Full Fine-tuning (tutti i layer)")
    print("=" * 60)

    for param in model.parameters():
        param.requires_grad = True

    optimizer = AdamW([
        {'params': model.conv1.parameters(), 'lr': cfg.phase3_lr * 0.01},
        {'params': model.layer1.parameters(), 'lr': cfg.phase3_lr * 0.1},
        {'params': model.layer2.parameters(), 'lr': cfg.phase3_lr * 0.1},
        {'params': model.layer3.parameters(), 'lr': cfg.phase3_lr},
        {'params': model.layer4.parameters(), 'lr': cfg.phase3_lr * 2},
        {'params': model.fc.parameters(), 'lr': cfg.phase3_lr * 10},
    ], weight_decay=1e-4)

    for epoch in range(cfg.phase3_epochs):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, cfg.device, epoch
        )
        val_loss, val_acc = evaluate(
            model, val_loader, criterion, cfg.device
        )
        print(f"  Epoch {epoch+1}/{cfg.phase3_epochs} | "
              f"Train: {train_acc:.1f}% | Val: {val_acc:.1f}%")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')

    print(f"\nMiglior Val Accuracy: {best_val_acc:.2f}%")
    return model
    

@torch.no_grad()
def full_evaluation(
    model: nn.Module,
    test_loader,
    class_names: list,
    device: str = 'cuda'
) -> dict:
    """Valutazione completa con metriche dettagliate."""

    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []

    for images, labels in test_loader:
        images = images.to(device)
        outputs = model(images)
        probs = torch.softmax(outputs, dim=1)

        _, predicted = outputs.max(1)
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())
        all_probs.extend(probs.cpu().numpy())

    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)

    # Classification Report
    report = classification_report(
        all_labels, all_preds,
        target_names=class_names,
        digits=4
    )
    print("Classification Report:")
    print(report)

    # Confusion Matrix
    cm = confusion_matrix(all_labels, all_preds)
    print("\nConfusion Matrix:")
    print(f"{'':.<15} " + " ".join(f"{name:>12}" for name in class_names))
    for i, row in enumerate(cm):
        print(f"{class_names[i]:<15} " + " ".join(f"{val:>12d}" for val in row))

    # F1-Score (macro e weighted)
    f1_macro = f1_score(all_labels, all_preds, average='macro')
    f1_weighted = f1_score(all_labels, all_preds, average='weighted')
    print(f"\nF1-Score Macro:    {f1_macro:.4f}")
    print(f"F1-Score Weighted: {f1_weighted:.4f}")

    return {
        'report': report,
        'confusion_matrix': cm,
        'f1_macro': f1_macro,
        'f1_weighted': f1_weighted,
        'predictions': all_preds,
        'labels': all_labels,
        'probabilities': np.array(all_probs)
    }


# Esegui la valutazione
model.load_state_dict(torch.load('best_model.pth', weights_only=True))
results = full_evaluation(model, test_loader, Config.class_names)
    

import torch


def export_model(model: torch.nn.Module, save_dir: str = 'exported_models'):
    """Esporta il modello in formati ottimizzati per la produzione."""
    from pathlib import Path

    save_path = Path(save_dir)
    save_path.mkdir(exist_ok=True)

    model.eval()
    model = model.cpu()
    dummy_input = torch.randn(1, 3, 224, 224)

    # ---- TorchScript (tracing) ----
    scripted = torch.jit.trace(model, dummy_input)
    scripted_path = save_path / 'model_scripted.pt'
    scripted.save(str(scripted_path))
    print(f"TorchScript salvato: {scripted_path} "
          f"({scripted_path.stat().st_size / 1024 / 1024:.1f} MB)")

    # ---- ONNX ----
    onnx_path = save_path / 'model.onnx'
    torch.onnx.export(
        model,
        dummy_input,
        str(onnx_path),
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'},
        },
        opset_version=17
    )
    print(f"ONNX salvato: {onnx_path} "
          f"({onnx_path.stat().st_size / 1024 / 1024:.1f} MB)")

    # ---- Verifica ----
    loaded = torch.jit.load(str(scripted_path))
    with torch.no_grad():
        orig_out = model(dummy_input)
        loaded_out = loaded(dummy_input)
        diff = (orig_out - loaded_out).abs().max().item()
        print(f"Differenza max originale vs TorchScript: {diff:.2e}")

    return scripted_path, onnx_path


export_model(model)
    

import torchvision
from torchvision.models.detection import (
    fasterrcnn_resnet50_fpn_v2,
    FasterRCNN_ResNet50_FPN_V2_Weights
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


def get_detection_model(num_classes: int) -> torchvision.models.detection.FasterRCNN:
    """Crea un Faster R-CNN pre-trained per il fine-tuning.

    Args:
        num_classes: Numero di classi target + 1 (background)

    Returns:
        Modello Faster R-CNN pronto per il fine-tuning
    """
    # Carica il modello pre-trained su COCO (91 classi)
    weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT
    model = fasterrcnn_resnet50_fpn_v2(weights=weights)

    # Sostituisci la testa del box predictor
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(
        in_features, num_classes
    )

    return model


# Esempio: rilevamento difetti (3 tipi + background = 4 classi)
det_model = get_detection_model(num_classes=4)

# Training: Faster R-CNN usa una loss composita
# (classification_loss + box_regression_loss + RPN losses)
# Non serve definire la loss manualmente: model(images, targets)
# restituisce direttamente le losses in training mode
    

from ultralytics import YOLO

# Carica modello pre-trained su COCO
model = YOLO('yolo11n.pt')  # nano (più veloce) o yolo11s.pt, yolo11m.pt

# Fine-tuning sul tuo dataset
# Il dataset deve essere in formato YOLO (images/ + labels/ con .txt)
results = model.train(
    data='pcb_defects.yaml',    # Config del dataset
    epochs=100,
    imgsz=640,
    batch=16,
    lr0=0.001,                  # Learning rate iniziale
    lrf=0.01,                   # Learning rate finale (ratio)
    warmup_epochs=3,
    freeze=10,                  # Congela i primi 10 layer del backbone
    augment=True,               # Augmentation automatica
    patience=20,                # Early stopping
    device='0'                  # GPU 0
)

# Valutazione sul test set
metrics = model.val(data='pcb_defects.yaml', split='test')
print(f"mAP50:    {metrics.box.map50:.4f}")
print(f"mAP50-95: {metrics.box.map:.4f}")
    

def freeze_batchnorm(model: nn.Module):
    """Congela i layer BatchNorm anche quando il modello e in train mode.

    Importante quando si fa fine-tuning con batch size piccoli:
    le statistiche di batch sarebbero instabili e rovinerebbero
    le running_mean e running_var accumulate durante il pre-training.
    """
    for module in model.modules():
        if isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d)):
            module.eval()        # Usa running stats, non batch stats
            module.weight.requires_grad = False
            module.bias.requires_grad = False


# Utilizzo nel training loop:
model.train()
freeze_batchnorm(model)  # BN resta in eval mode