Add utility functions for image processing and tracking files

src/utils.py

@@ -0,0 +1,79 @@
+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