Update app.py

app.py

@@ -10,28 +10,28 @@ import torchvision.transforms as transforms
 
 photos_folder = "Photos"
 
-# Pobierz model i config
+# Download model and config
 repo_id = "Kiwinicki/sat2map-generator"
 generator_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
 config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
 model_path = hf_hub_download(repo_id=repo_id, filename="model.py")
 
-# Dodaj ścieżkę do modelu
+# Add path to model
 sys.path.append(os.path.dirname(model_path))
 from model import Generator
 
-# Załaduj konfigurację
+# Load configuration
 with open(config_path, "r") as f:
     config_dict = json.load(f)
 cfg = OmegaConf.create(config_dict)
 
-# Inicjalizacja modelu
+# Initialize model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 generator = Generator(cfg).to(device)
 generator.load_state_dict(torch.load(generator_path, map_location=device))
 generator.eval()
 
-# Transformacje
+# Transformations
 transform = transforms.Compose([
     transforms.Resize((256, 256)),
     transforms.ToTensor(),
@@ -39,86 +39,86 @@ transform = transforms.Compose([
 ])
 
 def process_image(image):
-    # Konwersja do tensora
+    if image is None:
+        return None
+        
+    # Convert to tensor
     image_tensor = transform(image).unsqueeze(0).to(device)
     
-    # Inferencja
+    # Inference
     with torch.no_grad():
         output_tensor = generator(image_tensor)
     
-    # Przygotowanie wyjścia
+    # Prepare output
     output_image = output_tensor.squeeze(0).cpu()
-    output_image = output_image * 0.5 + 0.5  # Denormalizacja
+    output_image = output_image * 0.5 + 0.5  # Denormalization
     output_image = transforms.ToPILImage()(output_image)
     
     return output_image
 
-
 def load_images_from_folder(folder):
     images = []
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+        return images
+        
     for filename in os.listdir(folder):
         if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
             img_path = os.path.join(folder, filename)
-            img = Image.open(img_path)
-            images.append((img, filename))
+            try:
+                img = Image.open(img_path)
+                images.append((img, filename))
+            except Exception as e:
+                print(f"Error loading {filename}: {e}")
     return images
 
-
-
-def load_image_from_gallery(images, index):
-    if images and 0 <= index < len(images):
-        image = images[index]
-        if isinstance(image, tuple):
-            image = image[0]
-        return image
-    return None
-
-
-def gallery_click_event(images, evt: gr.SelectData):
-    index = evt.index
-    selected_img = load_image_from_gallery(images, index)
-    return selected_img
-
-
-def clear_image():
-    return None
-
-
 def app():
     images = load_images_from_folder(photos_folder)
+    gallery_images = [img[0] for img in images] if images else []
 
-    with gr.Blocks(css=".container { background-color: white; }") as demo:
+    with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
-                selected_image = gr.Image(label="Input Image", type="pil")
+                input_image = gr.Image(label="Input Image", type="pil")
                 clear_button = gr.Button("Clear")
 
             with gr.Column():
-                image_gallery = gr.Gallery(label="Image Gallery", elem_id="gallery", type="pil", value=[img for img, _ in images])
+                gallery = gr.Gallery(
+                    label="Image Gallery",
+                    value=gallery_images,
+                    columns=3,
+                    rows=2,
+                    height="auto"
+                ).style(grid=3)
 
             with gr.Column():
-                result_image = gr.Image(label="Result Image", type="pil")
-
-        image_gallery.select(
-            fn=gallery_click_event,
-            inputs=image_gallery,
-            outputs=selected_image
+                output_image = gr.Image(label="Result Image", type="pil")
+        
+        # Handle gallery selection
+        def on_select(evt: gr.SelectData):
+            if 0 <= evt.index < len(images):
+                return images[evt.index][0]
+            return None
+            
+        gallery.select(
+            fn=on_select,
+            outputs=input_image
         )
-
-        selected_image.change(
+        
+        # Process image when input changes
+        input_image.change(
             fn=process_image,
-            inputs=selected_image,
-            outputs=result_image
+            inputs=input_image,
+            outputs=output_image
         )
-
+        
+        # Clear button functionality
         clear_button.click(
-            fn=clear_image,
-            inputs=None,
-            outputs=selected_image
+            fn=lambda: None,
+            outputs=input_image
         )
 
     demo.launch()
 
-
 if __name__ == "__main__":
-    app()
+    app()