src/utils.py

import numpy as np
from PIL import Image, ImageOps
import os
from tqdm import tqdm
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import pickle
from collections import Counter
import random


class utils:

    def __init__(self):
        pass

    @staticmethod
    def calculate_padding(
        image_height: int, image_width: int, window_size: tuple, stride: tuple
    ):
        """
        Calculate padding needed for height and width

        :param image_height:
        :param image_width:
        :param window_size:
        :param stride:
        :return: pad_height, pad_width
        """

        # Calculate padding needed for height and width
        if (image_height - window_size[0]) // stride[0] != (
            image_height - window_size[0]
        ) / stride[0]:
            if window_size[0] == stride[0]:
                pad_height = window_size[0] - (image_height % stride[0])
            else:
                pad_height = image_height - (
                    ((image_height - window_size[0]) // stride[0]) * stride[0]
                    + window_size[0]
                )
        else:
            pad_height = 0
        if (image_width - window_size[1]) // stride[1] != (
            image_width - window_size[1]
        ) / stride[1]:
            if window_size[1] == stride[1]:
                pad_width = window_size[1] - (image_width % stride[1])
            else:
                pad_width = image_width - (
                    ((image_width - window_size[1]) // stride[1]) * stride[1]
                    + window_size[1]
                )
        else:
            pad_width = 0

        return pad_height, pad_width

    @staticmethod
    def sliding_window(
        image_dir: str, window_size: tuple, stride: tuple, padding_value: int = 0
    ):
        """
        Slide a window across the image and extract patches.
        :param image_dir: Path to the image file.
        :param window_size: Tuple (height, width) of the window.
        :param stride: Tuple (height, width) of the stride.
        :param padding_value: Value to use for padding.

        :return: List of tuples (window, (y, x)) and the size of the padded image.
        """
        # Open the image
        if isinstance(image_dir, str):
            image = Image.open(image_dir)
        else:
            image = image_dir
        # Convert image to a numpy array
        image_np = np.array(image)

        # Get image dimensions
        image_height, image_width = image_np.shape[:2]

        # if window_size[0] > image_height or window_size[1] > image_width:
        #     raise ValueError("Window size should be smaller than the image size.")
        # if stride[0] > window_size[0] or stride[1] > window_size[1]:
        #     raise ValueError("Stride should be smaller than the window size.")

        # Calculate padding needed for height and width
        pad_height, pad_width = utils.calculate_padding(
            image_height, image_width, window_size, stride
        )

        # Pad the image using Pillow (symmetric padding)
        padded_image = ImageOps.expand(
            image, border=(0, 0, pad_width, pad_height), fill=padding_value
        )

        # Convert padded image back to a numpy array
        padded_image_np = np.array(padded_image)

        # Get padded image dimensions
        padded_height, padded_width = padded_image_np.shape[:2]

        # List to store cropped windows and their positions
        windows = []

        # Slide the window across the padded image
        for y in range(0, padded_height - window_size[0] + 1, stride[0]):
            for x in range(0, padded_width - window_size[1] + 1, stride[1]):
                # Crop the window from the image
                window = padded_image_np[y : y + window_size[0], x : x + window_size[1]]
                # Append the window and its position to the list
                windows.append((Image.fromarray(window), (y, x)))

        return windows, padded_image.size  # Return windows and padded image size

    @staticmethod
    def save_windows(windows: list, out_dir: str):
        """
        Save the cropped windows to the output directory.

        Args:
            windows: List of tuples (window, (y, x)) from the sliding_window function.
            out_dir: Path to the output directory.
        """
        # Create the output directory if it doesn't exist
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        # Save each window to the output directory
        for i, (window, _) in enumerate(windows):
            window.save(os.path.join(out_dir, f"window_{i}.png"))

    @staticmethod
    def reconstruct_image(windows: list, padded_size: tuple, window_size: tuple):
        """
        Reconstruct the original image from the sliding window patches. Strides are averaged.

        :param windows: List of tuples (window, (y, x)) from the sliding_window function.
        :param padded_size: Tuple (height, width) of the padded image.
        :param window_size: Tuple (height, width) of the window.

        :return: Reconstructed image and count map.
        """
        # Initialize an empty numpy array for the reconstructed image
        num_channels = 4 if windows[0][0].mode == "RGBA" else 3
        reconstructed_image = np.zeros(
            (padded_size[1], padded_size[0], num_channels), dtype=np.float32
        )

        # Initialize an array to count the number of overlapping windows
        count_map = np.zeros(
            (padded_size[1], padded_size[0], num_channels), dtype=np.float32
        )

        # Place each window back into the reconstructed image
        for window, (y, x) in windows:
            window_np = np.array(window, dtype=np.float32)

            # Add the window to the corresponding position in the reconstructed image
            reconstructed_image[
                y : y + window_size[1], x : x + window_size[0], :
            ] += window_np

            # Increment the count map to handle overlaps
            count_map[y : y + window_size[1], x : x + window_size[0], :] += 1

        # Divide by the count map to average overlapping areas
        reconstructed_image = np.divide(
            reconstructed_image, count_map, where=count_map != 0
        )

        # Clip and convert to 8-bit image
        reconstructed_image = np.clip(reconstructed_image, 0, 255).astype(np.uint8)
        mode = "RGBA" if num_channels == 4 else "RGB"
        reconstructed_image = Image.fromarray(reconstructed_image, mode=mode)

        return reconstructed_image, count_map

    @staticmethod
    def zip_images(directory: str, label: int):
        """
        This function reads images from a directory and returns them as a list of tuples (image, label).
        :param directory: The directory containing images
        :param label: The label to assign to all images in the directory
        :return: List of tuples where each tuple is (image, label)
        """
        data = []

        range_dir = os.listdir(directory)
        for file_name in tqdm(range_dir):
            img_path = os.path.join(directory, file_name)
            img = Image.open(img_path).convert("RGB")  # Open image and convert to RGB
            img_array = np.array(img)  # Convert image to a numpy array
            data.append((img_array, label))  # Append tuple (image, label) to the list

        return data

    @staticmethod
    def unzip_images(data: list, base_directory: str):
        """
        Unzips a list of tuples (image, label) and saves the images into directories named after their labels.
        :param data: List of tuples (image, label)
        :param base_directory: Base directory where the images will be saved
        """
        for i, (img_array, label) in enumerate(data):
            label_directory = os.path.join(base_directory, str(label))

            # Create the label directory if it doesn't exist
            if not os.path.exists(label_directory):
                os.makedirs(label_directory)

            # Define the image path (e.g., "label_directory/image_0.png")
            img_path = os.path.join(label_directory, f"image_{i}.png")

            # Convert the numpy array back to an image and save it
            img = Image.fromarray(img_array)
            img.save(img_path)

    @staticmethod
    def plot_image(image: np.array, title: str = None):
        """
        This function plots an image using matplotlib.
        :param image: Numpy array representing the image.
        :param title: Title of the plot.
        :return: None
        """
        plt.imshow(image)
        plt.axis("off")
        if title:
            plt.title(title)
        plt.show()

    @staticmethod
    def plot_images(images: list, titles: list):
        """
        This function plots multiple images side by side.
        :param images: List of numpy arrays representing the images.
        :param titles: List of titles for each image.
        :return: None
        """
        fig, axes = plt.subplots(1, len(images), figsize=(20, 20))
        for i, (image, title) in enumerate(zip(images, titles)):
            axes[i].imshow(image)
            axes[i].axis("off")
            axes[i].set_title(title)
        plt.show()

    @staticmethod
    def split_data(
        data,
        train_size: float,
        val_size: float,
        test_size: float,
        random_state: int = None,
    ):
        """
        This function splits the dataset into training, validation, and test sets.

        :param data: List of tuples where each tuple is (image, label)
        :param train_size: Proportion of the dataset to include in the training set
        :param val_size: Proportion of the dataset to include in the validation set
        :param test_size: Proportion of the dataset to include in the test set
        :param random_state: Controls the shuffling applied to the data before applying the split
        :return: Tuple of (train_data, val_data, test_data) where each is a list of (image, label)
        """
        # Ensure the split sizes add up to 1
        assert np.isclose(
            train_size + val_size + test_size, 1.0
        ), "Split sizes must add up to 1"

        # First split: Train + (Val + Test)
        train_data, temp_data = train_test_split(
            data, train_size=train_size, random_state=random_state
        )

        # Second split: Val + Test
        val_ratio = val_size / (
            val_size + test_size
        )  # Adjust val_size to be relative to the size of temp_data
        val_data, test_data = train_test_split(
            temp_data, train_size=val_ratio, random_state=random_state
        )

        return train_data, val_data, test_data

    @staticmethod
    def save_to_pickle(data: list, file_path: str):
        """
        This function saves data to a pickle file.
        :param data: Data to save.
        :param file_path: Path to save the pickle file.
        :return: None
        """
        with open(file_path, "wb") as f:
            pickle.dump(data, f)

    @staticmethod
    def load_from_pickle(file_path: str):
        """
        This function loads data from a pickle file.
        :param file_path: Path to the pickle file.
        :return: Loaded data.
        """
        with open(file_path, "rb") as f:
            return pickle.load(f)

    @staticmethod
    def plot_distribution(
        train_data,
        val_data,
        test_data,
        class_names,
        title: str = "Split Data Distribution",
    ):
        """
        Plots the distribution of data after splitting into training, validation, and test sets.
        :param train_data: List of tuples (image, label) for the training set
        :param val_data: List of tuples (image, label) for the validation set
        :param test_data: List of tuples (image, label) for the test set
        :param class_names: List of class names corresponding to the label
        :param title: Title of the plot
        :return: None
        """
        # Count the number of samples for each label in each dataset
        train_labels = [label for _, label in train_data]
        val_labels = [label for _, label in val_data]
        test_labels = [label for _, label in test_data]

        train_counter = Counter(train_labels)
        val_counter = Counter(val_labels)
        test_counter = Counter(test_labels)

        # Prepare the data for plotting
        labels = sorted(set(train_labels + val_labels + test_labels))
        train_counts = [train_counter[label] for label in labels]
        val_counts = [val_counter[label] for label in labels]
        test_counts = [test_counter[label] for label in labels]

        # Plot the distribution
        x = range(len(labels))
        width = 0.25  # Width of the bars

        plt.figure(figsize=(10, 6))

        plt.bar(
            x, train_counts, width=width, label="Train", color="blue", align="center"
        )
        plt.bar(
            [p + width for p in x],
            val_counts,
            width=width,
            label="Validation",
            color="orange",
            align="center",
        )
        plt.bar(
            [p + width * 2 for p in x],
            test_counts,
            width=width,
            label="Test",
            color="green",
            align="center",
        )

        plt.xlabel("Classes")
        plt.ylabel("Number of Samples")
        plt.title(title)
        plt.xticks([p + width for p in x], [class_names[label] for label in labels])
        plt.legend()
        plt.show()

    @staticmethod
    def balance_data(train_data: list, random_state=None):
        """
        This function balances the data by randomly selecting samples from the majority class to match the number of
        samples in the minority class.
        :param train_data: List of tuples (image, label)
        :param random_state: int, random seed for reproducibility
        :return: List of tuples (image, label)
        with balanced classes and a list of tuples (image, label) with the excess data
        """

        # get all labels from the train data
        labels = [sample[1] for sample in train_data]

        # count the number of samples for each label
        counter = Counter(labels)

        # find the label with the fewest samples
        min_samples = min(counter.values())

        # create a list to store the balanced data
        balanced_data = []

        # create a list to store the excess data
        excess_data = []

        # get how many classes we have
        num_classes = len(set(labels))

        # for each class
        for i in range(num_classes):
            # get all the samples for that class
            samples = [sample for sample in train_data if sample[1] == i]

            # shuffle the samples
            random.seed(random_state)

            random.shuffle(samples)

            # add the first min_samples samples to the balanced data
            balanced_data += samples[:min_samples]

            # add the remaining samples to the excess data
            excess_data += samples[min_samples:]

        return balanced_data, excess_data

    @staticmethod
    def shuffle_data(data: list, random_state=None):
        """
        This function shuffles the data.
        :param data: List of tuples (image, label)
        :param random_state: int, random seed for reproducibility
        :return: List of tuples (image, label) with shuffled data
        """
        random.seed(random_state)
        random.shuffle(data)
        return data