Update app.py

app.py

@@ -69,6 +69,7 @@
 
 
 
+
 # import gradio as gr
 # import subprocess
 # import os
@@ -113,38 +114,50 @@
 #     return f"✅ Saved {filename}"
 
 # with gr.Blocks() as demo:
-#     gr.Markdown("# 👗 AI-Generated Clothing Designs")
+#     gr.Markdown("## 👗 AI-Generated Clothing Designs")
 
 #     generate_btn = gr.Button("✨ Generate New Designs")
 
 #     image_components = []
+#     file_paths = []
 #     save_buttons = []
 #     outputs = []
 
-#     with gr.Column():
-#         for _ in range(10):
-#             with gr.Row():
-#                 img = gr.Image()
-#                 btn = gr.Button("💾 Save")
-#                 out = gr.Textbox(visible=False)
+#     with gr.Row():
+#         for row in range(2):  # 2 rows
+#             with gr.Column():
+#                 for col in range(5):  # 5 images per row
+#                     index = row * 5 + col
+#                     with gr.Column():
+#                         img = gr.Image()
+#                         file_paths.append(gr.Textbox(visible=False))  # to pass file path
+#                         btn = gr.Button("💾 Save")
+#                         out = gr.Textbox(visible=False)
 
-#                 btn.click(fn=save_image, inputs=img, outputs=out)
+#                         btn.click(fn=save_image, inputs=file_paths[index], outputs=out)
 
-#                 image_components.append(img)
-#                 save_buttons.append(btn)
-#                 outputs.append(out)
+#                         image_components.append(img)
+#                         save_buttons.append(btn)
+#                         outputs.append(out)
 
 #     def update_gallery():
 #         files = get_random_images()
-#         return files + [""] * (10 - len(files))
-
-#     generate_btn.click(fn=update_gallery, outputs=image_components)
+#         updates = []
+#         for i, path in enumerate(files):
+#             updates.append(path)  # for img.update(value=...)
+#             updates.append(path)  # for file_path.update(value=...)
+#         # Fill remaining slots if less than 10
+#         for i in range(10 - len(files)):
+#             updates.append(None)
+#             updates.append("")
+#         return updates
+
+#     generate_btn.click(fn=update_gallery, outputs=sum(zip(image_components, file_paths), ()))
 
 # demo.launch()
 
 
 
-
 import gradio as gr
 import subprocess
 import os
@@ -152,14 +165,17 @@ import random
 from PIL import Image
 import shutil
 
+# === 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(os.path.join(BASE_DIR, "saved_images"), exist_ok=True)
+os.makedirs(SAVE_DIR, exist_ok=True)
 
+# === Image Generation Function ===
 def generate_images():
     command = [
         "python",
@@ -169,64 +185,71 @@ def generate_images():
         "--seeds=3-5,7,9,12-14,16-26,29,31,32,34,40,41",
         f"--network={MODEL_PATH}"
     ]
+    
     try:
         subprocess.run(command, check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
         return f"Error generating images:\n{e.stderr}"
 
+# === Select Random Images from Output Folder ===
 def get_random_images():
     image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
     if len(image_files) < 10:
         generate_images()
         image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
-    selected_files = random.sample(image_files, min(10, len(image_files)))
-    return [os.path.join(OUTPUT_DIR, f) for f in selected_files]
-
-def save_image(image_path):
-    filename = os.path.basename(image_path)
-    dest = os.path.join(BASE_DIR, "saved_images", filename)
-    shutil.copy(image_path, dest)
-    return f"✅ Saved {filename}"
+    random_images = random.sample(image_files, min(10, len(image_files)))
+    image_paths = [os.path.join(OUTPUT_DIR, img) for img in random_images]
+    return image_paths
 
+# === Save Image Function ===
+def save_image(img_path):
+    try:
+        filename = os.path.basename(img_path)
+        dest_path = os.path.join(SAVE_DIR, filename)
+        shutil.copy(img_path, dest_path)
+        return f"Saved to: {dest_path}"
+    except Exception as e:
+        return f"Error saving image: {str(e)}"
+
+# === Gradio Interface ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 👗 AI-Generated Clothing Designs")
+    gr.Markdown("# 🎨 AI-Generated Clothing Designs - Dresses")
 
-    generate_btn = gr.Button("✨ Generate New Designs")
+    generate_button = gr.Button("Generate New Designs")
 
     image_components = []
-    file_paths = []
     save_buttons = []
     outputs = []
+    file_paths = []
 
     with gr.Row():
         for row in range(2):  # 2 rows
             with gr.Column():
-                for col in range(5):  # 5 images per row
+                for col in range(5):  # 5 columns
                     index = row * 5 + col
                     with gr.Column():
-                        img = gr.Image()
-                        file_paths.append(gr.Textbox(visible=False))  # to pass file path
-                        btn = gr.Button("💾 Save")
-                        out = gr.Textbox(visible=False)
+                        img = gr.Image(width=200, height=200)
+                        image_components.append(img)
 
-                        btn.click(fn=save_image, inputs=file_paths[index], outputs=out)
+                        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_btn.click(fn=save_image, inputs=file_path, outputs=output)
+                        save_buttons.append(save_btn)
+
+    def generate_and_display_images():
+        image_paths = get_random_images()
+        return image_paths + image_paths  # for setting file_paths and images
+
+    generate_button.click(
+        fn=generate_and_display_images,
+        outputs=image_components + file_paths
+    )
+
+if __name__ == "__main__":
+    demo.launch()
 
-                        image_components.append(img)
-                        save_buttons.append(btn)
-                        outputs.append(out)
-
-    def update_gallery():
-        files = get_random_images()
-        updates = []
-        for i, path in enumerate(files):
-            updates.append(path)  # for img.update(value=...)
-            updates.append(path)  # for file_path.update(value=...)
-        # Fill remaining slots if less than 10
-        for i in range(10 - len(files)):
-            updates.append(None)
-            updates.append("")
-        return updates
-
-    generate_btn.click(fn=update_gallery, outputs=sum(zip(image_components, file_paths), ()))
-
-demo.launch()