Update app.py

app.py

@@ -161,25 +161,23 @@
 #     demo.launch()
 
 
-
 import gradio as gr
 import subprocess
 import os
 import random
 from PIL import Image
+import shutil
 import requests
-import io
-import warnings
-
-warnings.filterwarnings("ignore")
 
 # === Setup Paths ===
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 GEN_SCRIPT = os.path.join(BASE_DIR, "stylegan3", "gen_images.py")
 OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")
 MODEL_PATH = os.path.join(BASE_DIR, "dress.pkl")
+SAVE_DIR = os.path.join(BASE_DIR, "saved_images")
 
 os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.makedirs(SAVE_DIR, exist_ok=True)
 
 # === Image Generation Function ===
 def generate_images():
@@ -206,21 +204,21 @@ def get_random_images():
     image_paths = [os.path.join(OUTPUT_DIR, img) for img in random_images]
     return image_paths
 
-# === Upload to NestJS Backend ===
+# === Send Image to Backend ===
 def send_to_backend(img_path, user_id):
     if not user_id:
-        return "❌ user_id not found."
-    if not os.path.exists(img_path):
-        return "⚠️ Image path invalid."
+        return "❌ user_id not found in URL."
 
-    try:
-        with open(img_path, "rb") as f:
-            file_bytes = f.read()
+    if not img_path or not os.path.exists(img_path):
+        return "⚠️ No image selected or image not found."
 
-        files = {'file': ('generated_image.png', file_bytes, 'image/png')}
+    try:
+        with open(img_path, 'rb') as f:
+            files = {'file': ('generated_image.png', f, 'image/png')}
 
-        url = f"https://5a4d-103-40-74-78.ngrok-free.app/customisation/upload/{user_id}"
-        response = requests.post(url, files=files)
+            # Your backend endpoint here
+            url = f"https://5a4d-103-40-74-78.ngrok-free.app/customisation/upload/{user_id}"
+            response = requests.post(url, files=files)
 
         if response.status_code == 201:
             return "✅ Image uploaded and saved to database!"
@@ -232,46 +230,50 @@ def send_to_backend(img_path, user_id):
 
 # === Gradio Interface ===
 with gr.Blocks() as demo:
-    user_id_state = gr.State()
-
-    @demo.load(inputs=None, outputs=[user_id_state])
-    def on_load(request: gr.Request):
-        user_id = request.query_params.get('user_id', '')
-        return user_id
-
     gr.Markdown("# 🎨 AI-Generated Clothing Designs - Dresses")
 
     generate_button = gr.Button("Generate New Designs")
+    user_id_state = gr.State()
+
+    @demo.load(inputs=None, outputs=[user_id_state])
+    def get_user_id(request: gr.Request):
+        return request.query_params.get("user_id", "")
 
     image_components = []
+    file_paths = []
     save_buttons = []
     outputs = []
-    file_paths = []
-
-    with gr.Row():
-        for i in range(10):  # 10 images
-            with gr.Column(scale=1):
-                img = gr.Image(width=200, height=200)
-                image_components.append(img)
-
-                file_path = gr.Textbox(visible=False)
-                file_paths.append(file_path)
-
-                save_btn = gr.Button("💾 Save")
-                output = gr.Textbox(visible=False)
-                outputs.append(output)
-
-                # Save button uses the backend-saving version
-                save_btn.click(
-                    fn=send_to_backend,
-                    inputs=[file_path, user_id_state],
-                    outputs=output
-                )
-                save_buttons.append(save_btn)
 
+    # Use 3 columns layout
+    for row_idx in range(4):  # 4 rows (to cover 10 images)
+        with gr.Row():
+            for col_idx in range(3):  # 3 columns
+                i = row_idx * 3 + col_idx
+                if i >= 10:
+                    break
+                with gr.Column():
+                    img = gr.Image(width=180, height=180, label=f"Design {i+1}")
+                    image_components.append(img)
+
+                    file_path = gr.Textbox(visible=False)
+                    file_paths.append(file_path)
+
+                    save_btn = gr.Button("💾 Save to DB")
+                    save_buttons.append(save_btn)
+
+                    output = gr.Textbox(label="Status", interactive=False)
+                    outputs.append(output)
+
+                    save_btn.click(
+                        fn=send_to_backend,
+                        inputs=[file_path, user_id_state],
+                        outputs=output
+                    )
+
+    # Generate button logic
     def generate_and_display_images():
         image_paths = get_random_images()
-        return image_paths + image_paths  # one set for images, one for file paths
+        return image_paths + image_paths  # One for display, one for hidden path tracking
 
     generate_button.click(
         fn=generate_and_display_images,