utils.py

import argparse
import random
import cv2
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader

from data import create_image_scale_dataset
from main import get_args

def display_class_samples(dataset, num_classes, n_samples, filename):
    class_images = {i: [] for i in range(num_classes)}
    random.shuffle(dataset)
    for image, label in dataset:
        if len(class_images[label]) < n_samples:  # Collect only n images per class
            class_images[label].append(image)
    
    # Create figure and set up subplots
    fig, axes = plt.subplots(num_classes, n_samples, figsize=(n_samples * 2, num_classes * 2))
    fig.suptitle(f"Sample Images - {filename.split('/')[-1]}", fontsize=16)

    for class_label in range(num_classes):
        # Display the class label as text at the beginning of each row
        for i in range(n_samples):
            ax = axes[class_label, i]  # Access subplot for each class and sample

            # If we have an image for this class and sample
            if i < len(class_images[class_label]):
                image = class_images[class_label][i]
                ax.imshow(image, cmap="gray")  # Display the image (assuming grayscale)
            else:
                # Hide empty subplots if there aren't enough images for the class
                ax.axis('off')
            
            ax.set_xticks([])
            ax.set_yticks([])

            # Write class label text only on the first image in the row
            if i == 0:
                # ax.set_ylabel(f"Class {class_label}", fontsize=12, labelpad=20)
                ax.set_ylabel(f"{2**(num_classes - 1 - class_label)}x", fontsize=12, labelpad=20)
    
    # Adjust layout and save the figure
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Leaves space for the title
    plt.savefig(filename)
    plt.show()


if __name__ == "__main__":
    args = get_args()

    image = cv2.imread(args.image_file, cv2.IMREAD_GRAYSCALE)
    train_dataset, test_dataset = create_image_scale_dataset(
            image,
            args.num_scales,
            args.block_size,
            args.stride,
            norm_before=False
        )

    display_class_samples(train_dataset.image_class_pairs, args.num_scales, 10, f"figs/samples/{args.block_size[0]}x{args.block_size[1]}_{args.image_file.split('/')[-1]}")