yeah

app.py

@@ -2,7 +2,7 @@ import os, torch
 import gradio as gr
 from PIL import Image
 
-from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing
+from utils.imageHandling import hfImageToTensor, preprocessing, postprocessing, loadPreloadedImages
 from model.modelLoading import loadModel
 
 
@@ -71,22 +71,17 @@ with gr.Blocks(title="Semantic Segmentation Predictors") as demo:
     with gr.Row():
         gr.Markdown("## Preloaded GTA V images to be used for testing the model")
     with gr.Row():
-        gta_gallery = gr.Gallery(
-            value=sorted([Image.open(os.path.join(gta_image_dir, f)).convert("RGB") for f in os.listdir(gta_image_dir) if f.endswith(".png")]),
-            label="GTA V Examples",
-            show_label=False,
-            columns=5,
-            rows=1,
-            height=200,
-            type="pil"
+        gta_gallery = gr.Gallery(value=loadPreloadedImages(gta_image_dir),
+            label="GTA V Examples", show_label=False, columns=5,
+            rows=1, height=200, type="pil"
         )
 
     with gr.Row():
         gr.Markdown("## Preloaded Cityscapes images to be used for testing the model")
     with gr.Row():
-        city_gallery = gr.Gallery(value=sorted([Image.open(os.path.join(city_image_dir, f)).convert("RGB") for f in os.listdir(city_image_dir) if f.endswith(".png")]),
-            label="Cityscapes Examples", show_label=False, columns=5, rows=1,
-            height=256, type="pil"
+        city_gallery = gr.Gallery(value=loadPreloadedImages(city_image_dir),
+            label="Cityscapes Examples", show_label=False, columns=5,
+            rows=1, height=256, type="pil"
     )
     
     submit_btn.click(

utils/imageHandling.py

@@ -1,4 +1,5 @@
-import torch, torchvision
+import torch, torchvision, os
+from PIL import Image
 
 # %% image loading
 def hfImageToTensor(image, width:int=1024, height:int=512)->torch.Tensor:
@@ -81,3 +82,17 @@ def postprocessing(pred: torch.Tensor) -> torch.Tensor:
         torch.Tensor: Processed tensor of shape (3, H, W) for visualization.
     """ 
     return torchvision.transforms.functional.to_pil_image(print_mask(pred.squeeze(0).cpu().to(torch.uint8)))
+
+
+# %% preloaded images
+def loadPreloadedImages(image_dir: str) -> list[Image.Image]:
+    """
+    Load preloaded images from a directory.
+
+    Args:
+        image_dir (str): Path to the directory containing images.
+
+    Returns:
+        list[Image.Image]: List of loaded images.
+    """
+    return list(map(lambda f: Image.open(os.path.join(image_dir, f)).convert("RGB"), sorted([f for f in os.listdir(image_dir) if f.endswith(".png")])))