task2.py

import torch
from torchvision import transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
#from torchvision.models import resnet50, ResNet50_Weights
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
import torch.utils
from PIL import Image
import os
from task1 import split_by_class_small
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import re
from sklearn.metrics import confusion_matrix
import numpy as np
from pathlib import Path
PROJECT_ROOT = Path(os.getenv('PROJECT_ROOT', '.')).resolve()
WORK_DIR = Path(os.getenv('SLURM_SUBMIT_DIR', '.')).resolve()
DATASET_ROOT = Path(os.getenv('DATASET_ROOT', '.')).resolve()

def get_mean_std(img_paths):
    """Calculate mean and std efficiently, cache results"""
    cache_dir = WORK_DIR / "output"
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_file = cache_dir / "mean_std_cache.pt"
    
    # Try to load cached values
    if cache_file.exists():
        print("Loading cached mean/std values...")
        cache = torch.load(cache_file)
        print(f"Dataset mean: {cache['mean']}")
        print(f"Dataset std: {cache['std']}")
        return cache['mean'], cache['std']
    
    # Calculate if not cached
    sum_ = torch.zeros(3)
    sum_sq = torch.zeros(3)
    n_pixels = 0
    
    print(f"Calculating mean/std from {len(img_paths)} images...")
    
    for i, path in enumerate(img_paths):
        if i % 500 == 0:
            print(f"  Processing image {i}/{len(img_paths)}...")
        
        full_path = DATASET_ROOT / path
        img = Image.open(full_path).convert("RGB")
        tensor = transforms.ToTensor()(img)
        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
    std = torch.sqrt(sum_sq / n_pixels - mean ** 2)
    
    print(f"Dataset mean: {mean}")
    print(f"Dataset std: {std}")
    
    # Save to cache
    torch.save({'mean': mean, 'std': std}, cache_file)
    print(f"Saved mean/std to {cache_file}")

    return mean, std


class Datasets(Dataset):
    def __init__(self, file_paths, base_transform=None, aug_transforms=None, mode='train'):
        self.file_paths = file_paths
        self.base_transform = base_transform

        # Augmentations in eine Liste von Callables normalisieren
        if aug_transforms is None:
            self.aug_transforms = []
        elif isinstance(aug_transforms, (list, tuple)):
            self.aug_transforms = list(aug_transforms)
        else:
            self.aug_transforms = [aug_transforms]

        self.mode = mode  # 'train', 'val', 'test'

        # Labels automatisch aus Dateinamen extrahieren
        self.labels = [self._extract_label(fp) for fp in self.file_paths]

        # Mappe Labels zu Zahlen
        self.classes = sorted(list(set(self.labels)))
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.labels_idx = [self.class_to_idx[label] for label in self.labels]

    def _extract_label(self, filepath):
        filename = os.path.basename(filepath)
        match = re.match(r"([a-zA-Z]+)_\d+", filename)
        if match:
            return match.group(1)
        else:
            raise ValueError(f"Cannot extract label from {filename}")

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

    def __getitem__(self, idx):
        img_path = self.file_paths[idx]
        full_path = DATASET_ROOT / img_path
        image = Image.open(full_path).convert("RGB")

        label = self.labels_idx[idx]
        description = self.labels[idx]

        if self.mode == 'train' and len(self.aug_transforms) != 0:
            views = []
            img0 = image
            if self.base_transform is not None:
                img0 = self.base_transform(img0)
            views.append(img0)

            for aug in self.aug_transforms:
                img_aug = aug(image)
                if self.base_transform is not None:
                    img_aug = self.base_transform(img_aug)
                views.append(img_aug)

            views_tensor = torch.stack(views, dim=0)  # Shape: [num_views, C, H, W]
            return views_tensor, label, description

        else:
            img = image
            if self.base_transform is not None:
                img = self.base_transform(img)
            return img, label, description

def train_epoch(model, trainloader, criterion, device, optimizer):
    model.train()
    for batch_idx, data in enumerate(trainloader):

        inputs=data[0]
        labels = data[1]
        if len(inputs.shape) == 5:
            b, a = inputs.shape[0], inputs.shape[1]
            inputs = inputs.view(b * a, *inputs.shape[2:]).to(device)
            labels = labels.repeat_interleave(a).to(device)
        else:
            inputs = inputs.to(device)
            labels = labels.to(device)

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

        optimizer.zero_grad() #reset accumulated gradients
        loss.backward() #compute new gradients
        optimizer.step() # apply new gradients to change model parameters

    return

def evaluate(model, dataloader, device, num_classes):
    model.eval()
    all_labels = []
    all_preds = []

    with torch.no_grad():
        for data in dataloader:
            inputs = data[0]
            labels = data[1]
            if len(inputs.shape) == 5:
                b, a = inputs.shape[0], inputs.shape[1]
                inputs = inputs.view(b * a, *inputs.shape[2:]).to(device)
                labels = labels.repeat_interleave(a).to(device)
            else:
                inputs = inputs.to(device)
                labels = labels.to(device)

            outputs = model(inputs)
            preds = torch.argmax(outputs, dim=1)

            all_labels.append(labels.cpu())
            all_preds.append(preds.cpu())

    # Zu 1D-Vektoren zusammenführen
    all_labels = torch.cat(all_labels).numpy()
    all_preds = torch.cat(all_preds).numpy()

    label_lst = list(range(num_classes))

    conf_matrix = confusion_matrix(all_labels, all_preds, labels=label_lst)

    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)  # falls eine Zeile 0 ist
    av_tpr_per_class = tpr_per_class.mean()

    return accuracy, av_tpr_per_class

def train_model(data_train, data_test, model, criterion, optimizer, num_epochs, num_classes, device):

    best_measure = 0
    best_epoch =-1
    bestweights = None
    accuracy_per_epoch = []
    av_tpr_per_epoch = []

    for epoch in range(num_epochs):
        train_epoch(model,  data_train,  criterion,  device , optimizer)
        measure, av_tpr_per_class = evaluate(model, data_test, device, num_classes)
        accuracy_per_epoch.append(measure)
        av_tpr_per_epoch.append(av_tpr_per_class)

        if measure > best_measure:
          bestweights = model.state_dict()
          best_measure = measure
          best_epoch = epoch

    return best_epoch, best_measure, bestweights, accuracy_per_epoch, av_tpr_per_epoch


def plot_all(histories, num_epochs, out_dir):
    out_dir = WORK_DIR / "output" / out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    epochs = list(range(1, num_epochs + 1))
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 6))
    n = len(histories)
    colors = cm.get_cmap('tab10')(np.linspace(0, 1, n))

    # Plots zeichnen
    for idx, (name, hist) in enumerate(histories.items()):
        ax1.plot(epochs, hist["accuracy"], 'o-', color=colors[idx], label=name)
        ax2.plot(epochs, hist["avg_tpr"], 's-', color=colors[idx], label=name)

    # Achsen beschriften
    ax1.legend(loc='upper left')
    ax1.set_xlabel('Epochs')
    ax1.set_ylabel('Accuracy')
    ax1.set_title('Accuracy')

    ax2.legend(loc='upper left')
    ax2.set_xlabel('Epochs')
    ax2.set_ylabel('Avg TPR')
    ax2.set_title('Avg TPR')

    # Speichern + zeigen
    plt.tight_layout()
    fig.savefig(out_dir / "results.png", dpi=300, bbox_inches='tight')
    plt.close(fig)

def run():
    random_seed = 3745356
    
    # seed festlegen
    import random
    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")

    num_epochs = 10
    batchsize = 60
    training_data, validation_data, test_data = split_by_class_small()
    mean, std = get_mean_std(training_data)
    basic_transform = transforms.Compose([
        transforms.Resize(128),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std),
    ])
    aug_colorjitters = transforms.Compose([
        transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.3),
    ])
    aug_blur = transforms.Compose([
        transforms.GaussianBlur(kernel_size=5),
    ])

    loss = torch.nn.CrossEntropyLoss()
    histories = {}
    lr_rates = [0.01, 0.001]
    best_measure = 0
    best_epoch = -1
    bestweights = None
    bestlr = None
    bestaug = None
    augmentation_setting = {
        "none": None,
        "color": aug_colorjitters,
        "blur": aug_blur,
    }

    for lr in lr_rates:
        for aug_name, aug in augmentation_setting.items():
            dataloaders = {
                'train': DataLoader(
                    Datasets(training_data, base_transform=basic_transform, aug_transforms=aug, mode='train'),
                    batch_size=batchsize, shuffle=True, ),

                'val': torch.utils.data.DataLoader(Datasets(validation_data, base_transform=basic_transform),
                                                   batch_size=batchsize, shuffle=False),
            }
            #model = resnet50(weights=ResNet50_Weights.DEFAULT)
            #model.fc = nn.Linear(2048, 10)
            #model = model.to(device)

            model = efficientnet_b0(weights=EfficientNet_B0_Weights.DEFAULT).to(device)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, 10).to(device)
            model.to(device)

            optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
            epoch_best_epoch, epoch_best_measure, epoch_bestweights, accuracy_per_epoch, av_tpr_per_epoch \
                = train_model(
                dataloaders['train'],
                dataloaders['val'],
                model=model,
                criterion=loss,
                optimizer=optimizer,
                num_epochs=num_epochs,
                num_classes=10,
                device=device)
            histories[f"LR={lr}, aug={aug_name}"] = {
                "accuracy": accuracy_per_epoch,
                "avg_tpr": av_tpr_per_epoch
            }
            if epoch_best_measure > best_measure:
                bestlr = lr
                bestaug = aug_name
                bestweights = epoch_bestweights
                best_measure = epoch_best_measure
                best_epoch = epoch_best_epoch

    print(f'Best Learning Rate: {bestlr}', f'Best Augmentation: {bestaug}',f'Best measure: {best_measure}', f'Best epoch: {best_epoch}')

    plot_all(histories, num_epochs=num_epochs, out_dir='plots')
    
    model_path = WORK_DIR / "output" / "best_model_weights.pth"
    torch.save(bestweights, model_path)
    print(f"best weights saved: {model_path}")

    return

if __name__ == "__main__":
    run()