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