Update app.py

app.py

@@ -23,16 +23,16 @@ MODEL_MAP = {
 # Sample clothing items from your clothes folder
 SAMPLE_CLOTHES = {
     "top": [
-        "clothes/shirt1.jpg",
-        "clothes/shirt2.jpg", 
-        "clothes/dress1.jpg",
-        "clothes/jacket1.jpg"
+        "clothes/shirt.jpg",
+        "clothes/cardigan.jpg", 
+        "clothes/jacket.jpg",
+        "clothes/dress.jpg"
     ],
     "bottom": [
-        "clothes/pants1.jpg",
-        "clothes/skirt1.jpg",
-        "clothes/jeans1.jpg", 
-        "clothes/shorts1.jpg"
+        "clothes/pants.jpg",
+        "clothes/jeans.jpg",
+        "clothes/skirt.jpg", 
+        "clothes/shorts.jpg"
     ]
 }
 
@@ -55,6 +55,18 @@ def load_sample_image(image_path):
         print(f"Error loading image {image_path}: {e}")
         return np.zeros((400, 400, 3), dtype=np.uint8)
 
+def load_model_image(model_path):
+    """Load a model image and return the file path"""
+    try:
+        if os.path.exists(model_path):
+            return model_path
+        else:
+            print(f"Model image not found: {model_path}")
+            return model  # fallback to default
+    except Exception as e:
+        print(f"Error loading model {model_path}: {e}")
+        return model  # fallback to default
+
 def get_tryon_result(model_name, garment1, garment2, seed=1234):
     # model_name = "AI Model " + model_name.split("\\")[-1].split(".")[0] # windows
     model_name = "AI Model " + model_name.split("/")[-1].split(".")[0] # linux
@@ -133,7 +145,38 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
         """)
     with gr.Row():
         with gr.Column():
-            init_image = gr.Image(sources='clipboard', type="filepath", label="model", value=model)
+            # Model selection buttons with preview images
+            gr.HTML("<h3>Select AI Model:</h3>")
+            with gr.Row():
+                with gr.Column():
+                    model_preview1 = gr.Image(value=MODEL_MAP["AI Model female_model"], 
+                                            label="Female Model", 
+                                            height=150, width=100, 
+                                            interactive=False, show_label=True)
+                    model_btn1 = gr.Button("Select Female Model", variant="secondary")
+                
+                with gr.Column():
+                    model_preview2 = gr.Image(value=MODEL_MAP["AI Model female_model1"], 
+                                            label="Female Model 1", 
+                                            height=150, width=100, 
+                                            interactive=False, show_label=True)
+                    model_btn2 = gr.Button("Select Female Model 1", variant="secondary")
+                
+                with gr.Column():
+                    model_preview3 = gr.Image(value=MODEL_MAP["AI Model male_model"], 
+                                            label="Male Model", 
+                                            height=150, width=100, 
+                                            interactive=False, show_label=True)
+                    model_btn3 = gr.Button("Select Male Model", variant="secondary")
+                
+                with gr.Column():
+                    model_preview4 = gr.Image(value=MODEL_MAP["AI Model male_model2"], 
+                                            label="Male Model 2", 
+                                            height=150, width=100, 
+                                            interactive=False, show_label=True)
+                    model_btn4 = gr.Button("Select Male Model 2", variant="secondary")
+            
+            init_image = gr.Image(sources='clipboard', type="filepath", label="Selected Model", value=model)
             
         with gr.Column():
             gr.HTML(
@@ -146,24 +189,68 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
                 </div>
                 """)
             
-            # Sample clothing buttons section
+            # Sample clothing buttons section with preview images
             gr.HTML("<h3>Quick Select Sample Clothes:</h3>")
             
+            gr.HTML("<h4>Top Garments:</h4>")
             with gr.Row():
-                gr.HTML("<h4>Top Garments:</h4>")
-            with gr.Row():
-                top_btn1 = gr.Button("Shirt 1", variant="secondary")
-                top_btn2 = gr.Button("Shirt 2", variant="secondary") 
-                top_btn3 = gr.Button("Dress", variant="secondary")
-                top_btn4 = gr.Button("Jacket", variant="secondary")
+                with gr.Column():
+                    top_preview1 = gr.Image(value=SAMPLE_CLOTHES["top"][0], 
+                                          label="Shirt", 
+                                          height=150, width=100, 
+                                          interactive=False, show_label=True)
+                    top_btn1 = gr.Button("Select Shirt", variant="secondary")
+                
+                with gr.Column():
+                    top_preview2 = gr.Image(value=SAMPLE_CLOTHES["top"][1], 
+                                          label="Cardigan", 
+                                          height=150, width=100, 
+                                          interactive=False, show_label=True)
+                    top_btn2 = gr.Button("Select Cardigan", variant="secondary")
+                
+                with gr.Column():
+                    top_preview3 = gr.Image(value=SAMPLE_CLOTHES["top"][2], 
+                                          label="Jacket", 
+                                          height=150, width=100, 
+                                          interactive=False, show_label=True)
+                    top_btn3 = gr.Button("Select Jacket", variant="secondary")
+                
+                with gr.Column():
+                    top_preview4 = gr.Image(value=SAMPLE_CLOTHES["top"][3], 
+                                          label="Dress", 
+                                          height=150, width=100, 
+                                          interactive=False, show_label=True)
+                    top_btn4 = gr.Button("Select Dress", variant="secondary")
             
+            gr.HTML("<h4>Bottom Garments:</h4>")
             with gr.Row():
-                gr.HTML("<h4>Bottom Garments:</h4>")
-            with gr.Row():
-                bottom_btn1 = gr.Button("Pants", variant="secondary")
-                bottom_btn2 = gr.Button("Skirt", variant="secondary")
-                bottom_btn3 = gr.Button("Jeans", variant="secondary") 
-                bottom_btn4 = gr.Button("Shorts", variant="secondary")
+                with gr.Column():
+                    bottom_preview1 = gr.Image(value=SAMPLE_CLOTHES["bottom"][0], 
+                                             label="Pants", 
+                                             height=150, width=100, 
+                                             interactive=False, show_label=True)
+                    bottom_btn1 = gr.Button("Select Pants", variant="secondary")
+                
+                with gr.Column():
+                    bottom_preview2 = gr.Image(value=SAMPLE_CLOTHES["bottom"][1], 
+                                             label="Jeans", 
+                                             height=150, width=100, 
+                                             interactive=False, show_label=True)
+                    bottom_btn2 = gr.Button("Select Jeans", variant="secondary")
+                
+                with gr.Column():
+                    bottom_preview3 = gr.Image(value=SAMPLE_CLOTHES["bottom"][2], 
+                                             label="Skirt", 
+                                             height=150, width=100, 
+                                             interactive=False, show_label=True)
+                    bottom_btn3 = gr.Button("Select Skirt", variant="secondary")
+                
+                with gr.Column():
+                    bottom_preview4 = gr.Image(value=SAMPLE_CLOTHES["bottom"][3], 
+                                             label="Shorts", 
+                                             height=150, width=100, 
+                                             interactive=False, show_label=True)
+                    bottom_btn4 = gr.Button("Select Shorts", variant="secondary")
 
             with gr.Row():
                 garment_top = gr.Image(sources='upload', type="numpy", label="top garment")
@@ -174,6 +261,24 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
         with gr.Column():
             gallery = gr.Image()
 
+    # Connect model buttons to load model images
+    model_btn1.click(
+        lambda: load_model_image(MODEL_MAP["AI Model female_model"]),
+        outputs=[init_image]
+    )
+    model_btn2.click(
+        lambda: load_model_image(MODEL_MAP["AI Model female_model1"]),
+        outputs=[init_image]
+    )
+    model_btn3.click(
+        lambda: load_model_image(MODEL_MAP["AI Model male_model"]),
+        outputs=[init_image]
+    )
+    model_btn4.click(
+        lambda: load_model_image(MODEL_MAP["AI Model male_model2"]),
+        outputs=[init_image]
+    )
+
     # Connect buttons to load sample images
     top_btn1.click(
         lambda: load_sample_image(SAMPLE_CLOTHES["top"][0]),