Update app.py

app.py

@@ -56,65 +56,79 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-
 import gradio as gr
-import requests
+import subprocess
+import os
+import random
 from PIL import Image
-import numpy as np
-import io
-
-# Dummy function to simulate image generation
-def generate_images():
-    # Replace with your actual image generation logic
-    img = Image.new("RGB", (256, 256), color="blue")
-    return [img] * 6
-
-# Function to send image to backend
-def save_image_to_backend(image):
-    if image is None:
-        return "No image selected."
 
-    # Convert image to bytes
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    buffered.seek(0)
+# === 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")
 
-    files = {"file": ("image.png", buffered, "image/png")}
-    response = requests.post("http://localhost:5000/save-image", files=files)
+# Ensure output directory exists
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-    if response.status_code == 200:
-        return "Image saved successfully!"
-    else:
-        return "Failed to save image."
+# === Image Generation Function ===
+def generate_images():
+    command = [
+        "python",
+        GEN_SCRIPT,
+        f"--outdir={OUTPUT_DIR}",
+        "--trunc=1",
+        "--seeds=3-5,7,9,12-14,16-26,29,31,32,34,40,41",
+        f"--network={MODEL_PATH}"
+    ]
+    
+    try:
+        result = subprocess.run(command, check=True, capture_output=True, text=True)
+        print("stdout:\n", result.stdout)
+        print("stderr:\n", result.stderr)
+    except subprocess.CalledProcessError as e:
+        print("Image generation failed with error:")
+        print(e.stderr)
+        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")]
+    random_images = random.sample(image_files, min(10, len(image_files)))
+    return [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images], random_images
+
+# === Save Image Function ===
+def save_image(image_name):
+    image_path = os.path.join(OUTPUT_DIR, image_name)
+    return image_path
+
+# === Gradio Interface ===
+def generate_and_display():
+    images, image_names = get_random_images()
+    return images, image_names
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Generated Images Gallery")
-
-    gallery = gr.Gallery(label="Generated Images", columns=3, height=300)
-    selected_image = gr.Image(label="Selected Image", visible=False)
-    save_button = gr.Button("Save Image", visible=False)
-    save_status = gr.Textbox(label="Status", interactive=False, visible=False)
-
-    def handle_select(evt: gr.SelectData):
-        # Get selected image index
-        index = evt.index
-        image = gallery.value[index]
-        return image, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
-
-    gallery.select(
-        handle_select,
-        outputs=[selected_image, selected_image, save_button, save_status]
-    )
-
-    save_button.click(
-        fn=save_image_to_backend,
-        inputs=selected_image,
-        outputs=save_status
-    )
-
-    save_status.change(fn=None, outputs=None, show_progress=False)
+    gr.Markdown("# 🎨 AI-Generated Clothing Designs - Dresses")
+    generate_button = gr.Button("Generate New Designs")
+    output_gallery = gr.Gallery(label="Generated Designs", columns=5).style(height=300)
+    
+    # Output for image names
+    image_names = gr.Variable()
+    
+    generate_button.click(fn=generate_and_display, outputs=[output_gallery, image_names])
+    
+    # Add a button to save the selected image
+    save_button = gr.Button("Save Selected Image")
+    
+    def save_selected_image(selected_image_name):
+        if selected_image_name:
+            return save_image(selected_image_name)
+        return None
 
-    # Generate initial images
-    gallery.value = generate_images()
+    output_gallery.select(fn=save_selected_image, inputs=image_names, outputs=save_button)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()