src/utils.py

import os
import numpy as np
from PIL import Image, ImageOps


def crop_and_pad_image(image_path, threshold=20, target_size=(512, 512)):
    """
    Crop and pad an image to a square with the specified target size.

    Args:
        image_path (str): Path to the input image file.
        threshold (int): Threshold value for binarizing the image.
        target_size (tuple): Target size of the output image (width, height).

    Returns:
        PIL.Image.Image: Cropped and padded image.
    """
    try:
        # Load the image
        image = Image.open(image_path).convert("RGB")
    except Exception as e:
        raise ValueError(f"Error loading image: {str(e)}")

    # Convert the image to a NumPy array
    image_array = np.array(image)

    # Binarize the image
    binary_image_array = np.where(image_array > threshold, 1, 0).astype(np.uint8)

    # Find non-zero elements (non-black pixels)
    non_zero_indices = np.argwhere(binary_image_array)

    # Check if non-zero elements exist
    if non_zero_indices.size == 0:
        raise ValueError(f"No non-zero elements found for the image: {image_path}")

    # Get the bounding box of non-zero elements
    (y1, x1, _), (y2, x2, _) = non_zero_indices.min(0), non_zero_indices.max(0)

    # Crop the Region of Interest (ROI)
    cropped_img = image.crop((x1, y1, x2, y2))

    # Pad the image to make it a square
    squared_img = ImageOps.pad(cropped_img, target_size)

    return squared_img


def track_files(folder_path, extensions=('.jpg', '.jpeg', '.png')):
    """
    Track all the files in a folder and its subfolders.

    Args:
        folder_path (str): The path of the folder to track files in.
        extensions (tuple, optional): Tuple of file extensions to track. Default is ('.jpg', '.jpeg', '.png').

    Returns:
        list: A list containing the paths of all files in the folder and its subfolders.
    """
    # Validate folder_path
    if not os.path.isdir(folder_path):
        raise ValueError("Invalid folder path provided.")

    # Convert extensions to lowercase for case-insensitive comparison
    extensions = tuple(ext.lower() for ext in extensions)

    # Initialize file_list
    file_list = []

    # Walk through the folder and its subfolders
    for root, dirs, files in os.walk(folder_path):
        for filename in files:
            file_path = os.path.join(root, filename)
            _, extension = os.path.splitext(file_path)
            # Check if the file extension is in the list of extensions
            if extension.lower() in extensions:
                file_list.append(file_path)

    return file_list