task3.py

import os
import random
from pathlib import Path
from types import MethodType

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix
from skimage.io import imread
from torch.utils.data import Dataset, DataLoader
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
import matplotlib.pyplot as plt
import matplotlib.cm as cm


RANDOM_SEED = 3745356

PROJECT_ROOT = Path(os.getenv("PROJECT_ROOT", ".")).resolve()
DATASET_ROOT_MS = Path(os.getenv("DATASET_ROOT", "EuroSAT_MS")).resolve()
WORK_DIR = Path(os.getenv("SLURM_SUBMIT_DIR", PROJECT_ROOT)).resolve()


def split_ms_by_class_small(train_size=250, val_size=100, test_size=200, seed=RANDOM_SEED):
    """Create small, per-class-balanced splits for the multispectral EuroSAT dataset.

    This mirrors task1.split_by_class_small, but works on .tif files under DATASET_ROOT_MS.
    """
    train_files = []
    val_files = []
    test_files = []

    for class_dir in sorted(DATASET_ROOT_MS.iterdir()):
        if class_dir.is_dir():
            class_files = []
            for img_file in class_dir.glob("*.tif"):
                rel_path = img_file.relative_to(DATASET_ROOT_MS)
                class_files.append(str(rel_path))

            class_files = sorted(class_files)
            if not class_files:
                continue

            total_needed = train_size + val_size + test_size
            if len(class_files) < total_needed:
                subset = class_files
            else:
                subset, _ = train_test_split(class_files, train_size=total_needed, random_state=seed)

            train_temp, test = train_test_split(subset, test_size=test_size, random_state=seed)
            train, val = train_test_split(train_temp, test_size=val_size, random_state=seed)

            train_files.extend(train)
            val_files.extend(val)
            test_files.extend(test)

    return train_files, val_files, test_files


def get_label_mapping(file_paths):
    """Return mapping from class name (directory) to integer label index."""
    labels = [Path(fp).parent.name for fp in file_paths]
    unique_labels = sorted(set(labels))
    label_to_idx = {label: idx for idx, label in enumerate(unique_labels)}
    return label_to_idx


def encode_labels(file_paths, label_to_idx):
    return [label_to_idx[Path(fp).parent.name] for fp in file_paths]


def get_mean_std_multispectral(img_paths, selected_bands, max_images=100):
    """Compute per-channel mean/std for the selected bands on a subset of training images.

    Images are assumed to be uint16 in [0, 65535]; we normalize them to [0, 1] here.
    """
    sum_ = torch.zeros(len(selected_bands), dtype=torch.float64)
    sum_sq = torch.zeros(len(selected_bands), dtype=torch.float64)
    n_pixels = 0

    for i, path in enumerate(img_paths):
        if i >= max_images:
            break

        full_path = DATASET_ROOT_MS / path
        img = imread(full_path)  # shape: H, W, C
        img = img[:, :, selected_bands].astype(np.float32) / 65535.0
        tensor = torch.from_numpy(img).permute(2, 0, 1)  # C, H, W

        num_pixels = tensor.shape[1] * tensor.shape[2]
        sum_ += tensor.sum(dim=[1, 2])
        sum_sq += (tensor ** 2).sum(dim=[1, 2])
        n_pixels += num_pixels

    mean = (sum_ / n_pixels).to(torch.float32)
    std = torch.sqrt(sum_sq / n_pixels - sum_ ** 2 / (n_pixels ** 2)).to(torch.float32)
    return mean, std


class MultispectralDataset(Dataset):
    def __init__(self, file_paths, labels_idx, selected_bands, mean, std):
        self.file_paths = file_paths
        self.labels_idx = labels_idx
        self.selected_bands = selected_bands
        self.mean = mean
        self.std = std

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

    def __getitem__(self, idx):
        img_path = self.file_paths[idx]
        full_path = DATASET_ROOT_MS / img_path
        img = imread(full_path)  # H, W, C (uint16)
        img = img[:, :, self.selected_bands].astype(np.float32) / 65535.0
        tensor = torch.from_numpy(img).permute(2, 0, 1)  # C, H, W

        # channel-wise standardization
        tensor = (tensor - self.mean.view(-1, 1, 1)) / self.std.view(-1, 1, 1)

        label = self.labels_idx[idx]
        return tensor, label


def forward_until_pool(self, x):
    """Forward pass of EfficientNet-B0 truncated right after global average pooling.

    Returns a feature vector of shape (batch, feature_dim).
    """
    x = self.features(x)
    x = self.avgpool(x)
    x = torch.flatten(x, 1)
    return x


class LateFusionNet(nn.Module):
    """Late-fusion classifier for 6-band EuroSAT-MS using a pre-trained EfficientNet-B0.

    The same feature extractor processes two groups of 3 bands; the features are
    concatenated and fed to a single linear layer.
    """

    def __init__(self, num_classes: int):
        super().__init__()
        weights = EfficientNet_B0_Weights.DEFAULT
        self.feature_net = efficientnet_b0(weights=weights)

        # Monkey-patch forward to stop after adaptive pooling
        self.feature_net.forward = MethodType(forward_until_pool, self.feature_net)

        feature_dim = self.feature_net.classifier[1].in_features  # 1280 for EfficientNet-B0
        self.fc = nn.Linear(feature_dim * 2, num_classes)

    def forward(self, x):
        # x: (B, 6, H, W) -> two groups of 3 bands
        x1 = x[:, 0:3, :, :]
        x2 = x[:, 3:6, :, :]

        feat1 = self.feature_net(x1)
        feat2 = self.feature_net(x2)

        fused = torch.cat([feat1, feat2], dim=1)
        output = self.fc(fused)
        return output


def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0.0
    for inputs, labels in dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    return total_loss / max(len(dataloader), 1)


def evaluate(model, dataloader, device, num_classes):
    """Evaluate model and return accuracy and average TPR."""
    model.eval()
    all_labels = []
    all_preds = []
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            preds = outputs.argmax(dim=1)
            
            all_labels.append(labels.cpu())
            all_preds.append(preds.cpu())
    
    all_labels = torch.cat(all_labels).numpy()
    all_preds = torch.cat(all_preds).numpy()
    
    conf_matrix = confusion_matrix(all_labels, all_preds, labels=list(range(num_classes)))
    accuracy = np.trace(conf_matrix) / np.sum(conf_matrix)
    tpr_per_class = conf_matrix.diagonal() / conf_matrix.sum(axis=1)
    tpr_per_class = np.nan_to_num(tpr_per_class)
    avg_tpr = tpr_per_class.mean()
    
    return accuracy, avg_tpr


def plot_training_curves(accuracy_history, tpr_history, num_epochs, output_dir="plots"):
    """Plot training curves for accuracy and TPR."""
    out_dir = WORK_DIR / "output" / output_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    
    epochs = list(range(1, num_epochs + 1))
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    ax1.plot(epochs, accuracy_history, 'o-', color='blue', label='Validation Accuracy')
    ax1.set_xlabel('Epochs')
    ax1.set_ylabel('Accuracy')
    ax1.set_title('Validation Accuracy')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    ax2.plot(epochs, tpr_history, 's-', color='green', label='Avg TPR')
    ax2.set_xlabel('Epochs')
    ax2.set_ylabel('Avg TPR')
    ax2.set_title('Average TPR per Class')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    fig.savefig(out_dir / "task3_training_curves.png", dpi=300, bbox_inches='tight')
    plt.close(fig)
    print(f"Saved training curves to {out_dir / 'task3_training_curves.png'}")


def run():
    # Reproducibility
    random.seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)
    torch.manual_seed(RANDOM_SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(RANDOM_SEED)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Use 6 informative Sentinel-2 bands: B02,B03,B04,B08,B11,B12
    # Assuming band order in the TIFF is [B01, B02, B03, B04, B05, B06, B07, B08, B8A, B09, B11, B12]
    # These correspond to Blue, Green, Red, NIR, SWIR-2, SWIR-3 (informative bands per EuroSAT paper)
    selected_bands = [1, 2, 3, 7, 10, 11]

    print("Creating train/val/test splits for EuroSAT-MS...")
    train_files, val_files, test_files = split_ms_by_class_small()

    label_to_idx = get_label_mapping(train_files + val_files + test_files)
    num_classes = len(label_to_idx)

    train_labels_idx = encode_labels(train_files, label_to_idx)
    val_labels_idx = encode_labels(val_files, label_to_idx)
    test_labels_idx = encode_labels(test_files, label_to_idx)

    print("Computing mean/std for selected bands on training set...")
    mean, std = get_mean_std_multispectral(train_files, selected_bands)
    print(f"Mean: {mean}")
    print(f"Std : {std}")

    train_dataset = MultispectralDataset(train_files, train_labels_idx, selected_bands, mean, std)
    val_dataset = MultispectralDataset(val_files, val_labels_idx, selected_bands, mean, std)
    test_dataset = MultispectralDataset(test_files, test_labels_idx, selected_bands, mean, std)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    model = LateFusionNet(num_classes=num_classes).to(device)
    criterion = nn.CrossEntropyLoss()
    # Use SGD with momentum (same as Task 2)
    optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

    num_epochs = 10
    best_val_acc = 0.0
    best_state_dict = None
    
    # Track metrics for plotting
    accuracy_history = []
    tpr_history = []

    for epoch in range(num_epochs):
        train_loss = train_epoch(model, train_loader, criterion, optimizer, device)
        val_acc, val_tpr = evaluate(model, val_loader, device, num_classes)
        
        accuracy_history.append(val_acc)
        tpr_history.append(val_tpr)

        print(f"Epoch {epoch + 1}/{num_epochs} - loss: {train_loss:.4f}, val_acc: {val_acc:.4f}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_state_dict = model.state_dict()

    if best_state_dict is not None:
        output_dir = WORK_DIR / "output"
        output_dir.mkdir(parents=True, exist_ok=True)
        ckpt_path = output_dir / "best_model_task3.pth"
        torch.save(best_state_dict, ckpt_path)
        print(f"Saved best model weights to {ckpt_path}")

        model.load_state_dict(best_state_dict)
    
    # Plot training curves
    plot_training_curves(accuracy_history, tpr_history, num_epochs)

    test_acc, test_tpr = evaluate(model, test_loader, device, num_classes)
    print(f"Test accuracy: {test_acc:.4f}")


if __name__ == "__main__":
    run()