From efe97aef06d756a7967cc8e9efe6f4c87999345c Mon Sep 17 00:00:00 2001
From: Luigi311 <git@luigi311.com>
Date: Mon, 9 Oct 2023 23:34:06 -0600
Subject: [PATCH] Support png images

Signed-off-by: Luigi311 <git@luigi311.com>
---
 images/images.py | 55 +++++++++++++++++++++++++++++++++++++-----------
 utils/utils.py   | 24 ++++++++++++---------
 2 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/images/images.py b/images/images.py
index 283e4e0..d0ae1b9 100644
--- a/images/images.py
+++ b/images/images.py
@@ -31,6 +31,44 @@ def process_raw(dng_file, half_size=False, auto_white_balance=False):
 
     return image
 
+def process_image(image_path, half_size=False, auto_white_balance=False):
+    if image_path.endswith("dng"):
+        image = process_raw(image_path, half_size, auto_white_balance)
+    else:
+        image = cv2.imread(image_path)
+
+    return image
+
+
+def resize_images(images):
+    # Get the minimum width and height of all the images
+    min_x = min([image.shape[1] for image in images])
+    min_y = min([image.shape[0] for image in images])
+
+    # Resize all the images to the minimum width and height if they are larger and aspect ratio is compatible
+    resized_images = []
+    for image in images:
+        if image.shape[1] > min_x or image.shape[0] > min_y:
+            aspect_ratio = image.shape[1] / image.shape[0]
+            if aspect_ratio > 1.0:  # Landscape orientation
+                new_width = min_x
+                new_height = int(min_x / aspect_ratio)
+            else:  # Portrait or square orientation
+                new_height = min_y
+                new_width = int(min_y * aspect_ratio)
+
+            # Check if resizing is possible based on new dimensions
+            if new_width <= min_x and new_height <= min_y:
+                resized_image = cv2.resize(image, (new_width, new_height))
+                resized_images.append(resized_image)
+            else:
+                # Image can't be resized while maintaining aspect ratio, so skip it
+                print(f"Skipping image due to aspect ratio: {image.shape}")
+        else:
+            resized_images.append(image)
+    
+    return resized_images
+
 
 def loadImages(path, threads=None, half_size=False, auto_white_balance=False):
     """
@@ -44,7 +82,7 @@ def loadImages(path, threads=None, half_size=False, auto_white_balance=False):
 
     """
 
-    dng_files, tiff_files = files(path)
+    file_list = files(path)
 
     images = []
 
@@ -52,18 +90,11 @@ def loadImages(path, threads=None, half_size=False, auto_white_balance=False):
     workers = threads if threads else max(floor(os.cpu_count() / 2), 1)
     # Instead of appending, concatenate the result images to numpy_array
     with ProcessPoolExecutor(max_workers=workers) as executor:
-        if dng_files:
-            for result_image in executor.map(
-                process_raw,
-                dng_files,
-                [half_size] * len(dng_files),
-                [auto_white_balance] * len(dng_files),
-            ):
-                images.append(result_image)
-
-        if tiff_files:
-            for result_image in executor.map(cv2.imread, tiff_files):
+        if file_list:
+            for result_image in executor.map(process_image, file_list, [half_size] * len(file_list), [auto_white_balance] * len(file_list)):
                 images.append(result_image)
+    
+    images = resize_images(images)
 
     return np.array(images)
 
diff --git a/utils/utils.py b/utils/utils.py
index a388415..81ea39a 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -4,7 +4,7 @@
 from concurrent.futures import ThreadPoolExecutor
 
 
-def files(path):
+def files(path: str) -> list:
     # Check if the path exists and is a directory
     if not os.path.exists(path) or not os.path.isdir(path):
         raise ValueError(f"ERROR {path} is not a valid directory.")
@@ -17,15 +17,19 @@ def files(path):
     file_list = os.listdir(path)
 
     # Filter the list to only include dng, tiff, and hdf5 files
-    process_file_list = [
-        os.path.join(path, x) for x in file_list if x.endswith(("dng", "tiff"))
-    ]
-
-    # Check if there are any dng or tiff files in the directory
-    dng_files = [x for x in process_file_list if x.endswith("dng")]
-    tiff_files = [x for x in process_file_list if x.endswith("tiff")]
-
-    return dng_files, tiff_files
+    process_file_list = []
+
+    for file in file_list:
+        if file.endswith("dng") or file.endswith("tiff"):
+            process_file_list.append(os.path.join(path, file))
+        # Allow other unncompressed formats to be processed
+        elif file.endswith("png"):
+            # exclude all that include main in the name to 
+            # avoid processing processed images
+            if "main" not in file:
+                process_file_list.append(os.path.join(path, file))
+
+    return process_file_list
 
 
 def downloader(url, file_name, directory):