helper_functions.py

#imports and setup
import os #enable changing directory to run notebooks
from PIL import Image
import matplotlib.pyplot as plt
import torch

classes = ["basket", "eye", "binoculars", "rabbit", "hand"]
image_folder = './Data_train_test/own/'
DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")

#Helper functions
def generate_classification_report(model, model_path, image_transformer, softmax=False):
    model.eval()
    model.load_state_dict(torch.load(f'{model_path}/latest_weights.pth', map_location=DEVICE))
    checkpoint = torch.load(f'{model_path}/checkpoint.pth', map_location=DEVICE)  #Load the checkpoint to get the accuracy on the test set
    train_loss_records = checkpoint['train_loss_records']
    test_loss_records = checkpoint['test_loss_records']
    test_accuracy_records = checkpoint['test_accuracy_records']

    print(f"Accuracy on test set: {checkpoint['best_accuracy']}")
    print(f"Trained epochs: {checkpoint['epoch']+1}")
    print()
    fig, ax = plt.subplots(1, 2, figsize=(8, 4))
    ax[0].plot(train_loss_records, '-b', label='train')
    ax[0].plot(test_loss_records, 'r', label='test')
    ax[0].legend(loc='best')
    ax[0].set_title("Loss function")
    ax[1].plot(test_accuracy_records)
    ax[1].set_title("Test accuracy")
    plt.tight_layout()

    for filename in os.listdir(image_folder):
        img_path = os.path.join(image_folder, filename)
        image = Image.open(img_path)
        image_tensor = image_transformer(image).to(DEVICE).unsqueeze(0)

        with torch.no_grad():
            probs = model(image_tensor)
            if softmax:
                probs = torch.nn.functional.softmax(probs, dim=1)
            confidence, predicted_class = torch.max(probs, 1)

        # Create a figure with 1 row and 2 columns
        fig, axs = plt.subplots(1, 2, figsize=(12, 5))

        # Left plot: show image
        axs[0].imshow(image,cmap="gray")
        axs[0].set_title(f"Predicted: {classes[predicted_class.item()]} ({confidence.item() * 100:.2f}%)")
        axs[0].axis('off')

        # Right plot: confidence bar chart
        axs[1].bar(range(len(classes)), probs[0].cpu().numpy() * 100, tick_label=classes)
        axs[1].set_xlabel('Class')
        axs[1].set_ylabel('Confidence (%)')
        axs[1].set_title(f'Confidence Scores for {filename}')
        axs[1].set_ylim(0, 100)
        plt.tight_layout()
        plt.show()

def generate_generation_report(model, model_path, multiModel, labelEmbedding = False):
    model.eval()
    if multiModel:
        for class_name in classes:
            checkpoint = torch.load(f'{model_path}/checkpoints/checkpoint_{class_name}.pth', map_location=DEVICE)
            print(f"Trained epochs: {checkpoint['epoch']+1}")

            model.load_state_dict(checkpoint['model_state_dict'])
            with torch.no_grad():
                num_generated_images = 3
                generated = model.generate_images(num_generated_images).detach().cpu()

            fig, axes = plt.subplots(1, 3, figsize=(8, 8))
            for idx, ax in enumerate(axes.flat):
                ax.imshow(generated[idx], cmap='gray')
                ax.set_title(f'{class_name}')
                ax.axis('off')
            plt.tight_layout()
            plt.show()
    elif not labelEmbedding:
        checkpoint = torch.load(f'{model_path}/checkpoints/checkpoint.pth', map_location=DEVICE)
        print(f"Trained epochs: {checkpoint['epoch'] + 1}")

        model.load_state_dict(checkpoint['model_state_dict'])
        with torch.no_grad():
            num_generated_images = 6
            generated = model.generate_images(num_generated_images).detach().cpu()

        fig, axes = plt.subplots(2, 3, figsize=(8, 8))
        for idx, ax in enumerate(axes.flat):
            ax.imshow(generated[idx], cmap='gray')
            ax.set_title(f'Generated')
            ax.axis('off')
        plt.tight_layout()
        plt.show()
    else:
        for i in range(len(classes)):
            checkpoint = torch.load(f'{model_path}/checkpoints/checkpoint.pth', map_location=DEVICE)
            print(f"Trained epochs: {checkpoint['epoch']+1}")

            model.load_state_dict(checkpoint['model_state_dict'])
            with torch.no_grad():
                num_generated_images = 3
                generated = model.generate_images(num_generated_images, i).detach().cpu()

            fig, axes = plt.subplots(1, 3, figsize=(8, 8))
            for idx, ax in enumerate(axes.flat):
                ax.imshow(generated[idx], cmap='gray')
                ax.set_title(f'{classes[i]}')
                ax.axis('off')
            plt.tight_layout()
            plt.show()

def generate_gan_report(generator, model_path, multiMode):
    generator.eval()
    if multiMode:
        for class_name in classes:
            checkpoint = torch.load(f'{model_path}/checkpoints/checkpoint_{class_name}.pth', map_location=DEVICE)
            print(f"Trained epochs: {checkpoint['epoch']+1}")

            generator.load_state_dict(checkpoint['generator_state_dict'])
            with torch.no_grad():
                noise = torch.randn(3, 100, 1, 1, device=DEVICE)
                generated = generator(noise).detach().cpu()
            # Display the first 3 grayscale images
            num_images_to_show = 3
            fig, axs = plt.subplots(1, num_images_to_show, figsize=(10, 4))
            for j in range(num_images_to_show):
                img = generated[j].squeeze(0)  # remove channel dimension (1, H, W) → (H, W)
                axs[j].imshow(img, cmap='gray')
                axs[j].axis('off')
                axs[j].set_title(f"{class_name} {j + 1}")
            plt.tight_layout()
            plt.show()
    else:
        checkpoint = torch.load(f'{model_path}/checkpoints/checkpoint.pth', map_location=DEVICE)
        print(f"Trained epochs: {checkpoint['epoch']+1}")

        generator.load_state_dict(checkpoint['generator_state_dict'])
        with torch.no_grad():
            noise = torch.randn(3, 100, 1, 1, device=DEVICE)
            generated = generator(noise).detach().cpu()
        # Display the first 3 grayscale images
        num_images_to_show = 3
        fig, axs = plt.subplots(1, num_images_to_show, figsize=(10, 4))
        for j in range(num_images_to_show):
            img = generated[j].squeeze(0)  # remove channel dimension (1, H, W) → (H, W)
            axs[j].imshow(img, cmap='gray')
            axs[j].axis('off')
            axs[j].set_title(f"Image {j + 1}")
        plt.tight_layout()
        plt.show()