Update app.py

app.py

@@ -7,6 +7,7 @@ import base64
 import requests
 import json
 import time
+from PIL import Image
 
 
 def tryon(person_img, garment_img, seed, randomize_seed):
@@ -33,20 +34,18 @@ def tryon(person_img, garment_img, seed, randomize_seed):
     }
     try:
         response = requests.post(url, headers=headers, data=json.dumps(data), timeout=50)
-        # print("post response code", response.status_code)
         if response.status_code == 200:
             result = response.json()['result']
             status = result['status']
             if status == "success":
                 uuid = result['result']
-                # print(uuid)
     except Exception as err:
         print(f"Post Exception Error: {err}")
         raise gr.Error("Too many users, please try again later")
     post_end_time = time.time()
     print(f"post time used: {post_end_time-post_start_time}")
 
-    get_start_time =time.time()
+    get_start_time = time.time()
     time.sleep(9)
     Max_Retry = 12
     result_img = None
@@ -56,7 +55,6 @@ def tryon(person_img, garment_img, seed, randomize_seed):
         try:
             url = "http://" + os.environ['tryon_url'] + "Query?taskId=" + uuid
             response = requests.get(url, headers=headers, timeout=20)
-            # print("get response code", response.status_code)
             if response.status_code == 200:
                 result = response.json()['result']
                 status = result['status']
@@ -72,9 +70,8 @@ def tryon(person_img, garment_img, seed, randomize_seed):
                     info = "Error"
                     break
             else:
-                # print(response.text)
-                err_log = "URL error, pleace contact the admin"
-                info = "URL error, pleace contact the admin"
+                err_log = "URL error, please contact the admin"
+                info = "URL error, please contact the admin"
                 break
         except requests.exceptions.ReadTimeout:
             err_log = "Http Timeout"
@@ -94,6 +91,7 @@ def tryon(person_img, garment_img, seed, randomize_seed):
 
     return result_img, seed, info
 
+
 def start_tryon(person_img, garment_img, seed, randomize_seed):
     start_time = time.time()
     if person_img is None or garment_img is None:
@@ -132,44 +130,96 @@ def start_tryon(person_img, garment_img, seed, randomize_seed):
                 result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)
                 info = "Success"
             else:
-                info = "Try again latter"
+                info = "Try again later"
         else:
             print(response.text)
-            info = "URL error, pleace contact the admin"
+            info = "URL error, please contact the admin"
     except requests.exceptions.ReadTimeout:
         print("timeout")
         info = "Too many users, please try again later"
         raise gr.Error("Too many users, please try again later")
     except Exception as err:
-        print(f"其他错误: {err}")
-        info = "Error, pleace contact the admin"
+        print(f"Other error: {err}")
+        info = "Error, please contact the admin"
     end_time = time.time()
     print(f"time used: {end_time-start_time}")
 
     return result_img, seed, info
 
+
+def capture_person_photo(webcam_image):
+    """Process captured person photo from webcam"""
+    if webcam_image is None:
+        gr.Warning("No image captured from camera")
+        return None
+    return webcam_image
+
+
+def capture_garment_photo(webcam_image):
+    """Process captured garment photo from webcam"""
+    if webcam_image is None:
+        gr.Warning("No image captured from camera")
+        return None
+    return webcam_image
+
+
+def clear_person_image():
+    """Clear the person image"""
+    return None
+
+
+def clear_garment_image():
+    """Clear the garment image"""
+    return None
+
+
+def switch_to_upload_person():
+    """Switch to upload mode for person image"""
+    return gr.update(visible=True), gr.update(visible=False)
+
+
+def switch_to_camera_person():
+    """Switch to camera mode for person image"""
+    return gr.update(visible=False), gr.update(visible=True)
+
+
+def switch_to_upload_garment():
+    """Switch to upload mode for garment image"""
+    return gr.update(visible=True), gr.update(visible=False)
+
+
+def switch_to_camera_garment():
+    """Switch to camera mode for garment image"""
+    return gr.update(visible=False), gr.update(visible=True)
+
+
 MAX_SEED = 999999
 
 example_path = os.path.join(os.path.dirname(__file__), 'assets')
 
-garm_list = os.listdir(os.path.join(example_path,"cloth"))
-garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]
+# Check if assets directory exists, if not create dummy lists
+if os.path.exists(example_path):
+    garm_list = os.listdir(os.path.join(example_path, "cloth")) if os.path.exists(os.path.join(example_path, "cloth")) else []
+    garm_list_path = [os.path.join(example_path, "cloth", garm) for garm in garm_list]
+    
+    human_list = os.listdir(os.path.join(example_path, "human")) if os.path.exists(os.path.join(example_path, "human")) else []
+    human_list_path = [os.path.join(example_path, "human", human) for human in human_list]
+else:
+    garm_list_path = []
+    human_list_path = []
 
-human_list = os.listdir(os.path.join(example_path,"human"))
-human_list_path = [os.path.join(example_path,"human",human) for human in human_list]
-
-css="""
+css = """
 #col-left {
     margin: 0 auto;
-    max-width: 430px;
+    max-width: 450px;
 }
 #col-mid {
     margin: 0 auto;
-    max-width: 430px;
+    max-width: 450px;
 }
 #col-right {
     margin: 0 auto;
-    max-width: 430px;
+    max-width: 450px;
 }
 #col-showcase {
     margin: 0 auto;
@@ -178,61 +228,149 @@ css="""
 #button {
     color: blue;
 }
+.camera-section {
+    border: 2px dashed #ccc;
+    border-radius: 10px;
+    padding: 10px;
+    margin: 5px 0;
+}
+.mode-buttons {
+    display: flex;
+    gap: 10px;
+    margin-bottom: 10px;
+}
 """
 
 def load_description(fp):
-    with open(fp, 'r', encoding='utf-8') as f:
-        content = f.read()
-    return content
+    """Load description from file if it exists, otherwise return default"""
+    try:
+        with open(fp, 'r', encoding='utf-8') as f:
+            content = f.read()
+        return content
+    except FileNotFoundError:
+        return """
+        # 🎯 Virtual Try-On Application
+        
+        **Experience the future of fashion with AI-powered virtual try-on technology!**
+        
+        Upload or capture photos of yourself and garments to see how they look together.
+        """
 
 def change_imgs(image1, image2):
     return image1, image2
 
 with gr.Blocks(css=css) as Tryon:
     gr.HTML(load_description("assets/title.md"))
+    
     with gr.Row():
-        with gr.Column(elem_id = "col-left"):
+        with gr.Column(elem_id="col-left"):
             gr.HTML("""
             <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                 <div>
-                Step 1.  Upload a person image ⬇️
+                📸 Step 1. Get Person Image ⬇️
                 </div>
             </div>
             """)
-        with gr.Column(elem_id = "col-mid"):
+        with gr.Column(elem_id="col-mid"):
             gr.HTML("""
             <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                 <div>
-                Step 2. Upload a garment image ⬇️
+                👔 Step 2. Get Garment Image ⬇️
                 </div>
             </div>
             """)
-        with gr.Column(elem_id = "col-right"):
+        with gr.Column(elem_id="col-right"):
             gr.HTML("""
             <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                 <div>
-                Step 3. Press “Run” to get try-on results
+                ✨ Step 3. Generate Try-On Result
                 </div>
             </div>
             """)
+    
     with gr.Row():
-        with gr.Column(elem_id = "col-left"):
-            imgs = gr.Image(label="Person image", sources='upload', type="numpy")
-            # category = gr.Dropdown(label="Garment category", choices=['upper_body', 'lower_body', 'dresses'],  value="upper_body")
-            example = gr.Examples(
-                inputs=imgs,
-                examples_per_page=12,
-                examples=human_list_path
-            )
-        with gr.Column(elem_id = "col-mid"):
-            garm_img = gr.Image(label="Garment image", sources='upload', type="numpy")
-            example = gr.Examples(
-                inputs=garm_img,
-                examples_per_page=12,
-                examples=garm_list_path
-            )
-        with gr.Column(elem_id = "col-right"):
-            image_out = gr.Image(label="Result", show_share_button=False)
+        # Person Image Column
+        with gr.Column(elem_id="col-left"):
+            # Mode selection buttons for person
+            with gr.Row(elem_classes="mode-buttons"):
+                upload_person_btn = gr.Button("📁 Upload Photo", variant="secondary", size="sm")
+                camera_person_btn = gr.Button("📷 Take Photo", variant="secondary", size="sm")
+                clear_person_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
+            
+            # Upload interface for person (visible by default)
+            with gr.Group(visible=True) as person_upload_group:
+                person_img_upload = gr.Image(
+                    label="Person Image (Upload)", 
+                    sources=['upload'], 
+                    type="numpy",
+                    height=400
+                )
+            
+            # Camera interface for person (hidden by default)
+            with gr.Group(visible=False, elem_classes="camera-section") as person_camera_group:
+                person_img_camera = gr.Image(
+                    label="Person Image (Camera)", 
+                    sources=['webcam'], 
+                    type="numpy",
+                    height=400
+                )
+                capture_person_btn = gr.Button("📸 Capture Person Photo", variant="primary")
+            
+            # Combined person image state
+            person_img_final = gr.State()
+            
+            # Examples for person images
+            if human_list_path:
+                example_person = gr.Examples(
+                    inputs=person_img_upload,
+                    examples_per_page=8,
+                    examples=human_list_path,
+                    label="Example Person Images"
+                )
+        
+        # Garment Image Column
+        with gr.Column(elem_id="col-mid"):
+            # Mode selection buttons for garment
+            with gr.Row(elem_classes="mode-buttons"):
+                upload_garment_btn = gr.Button("📁 Upload Photo", variant="secondary", size="sm")
+                camera_garment_btn = gr.Button("📷 Take Photo", variant="secondary", size="sm")
+                clear_garment_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
+            
+            # Upload interface for garment (visible by default)
+            with gr.Group(visible=True) as garment_upload_group:
+                garment_img_upload = gr.Image(
+                    label="Garment Image (Upload)", 
+                    sources=['upload'], 
+                    type="numpy",
+                    height=400
+                )
+            
+            # Camera interface for garment (hidden by default)
+            with gr.Group(visible=False, elem_classes="camera-section") as garment_camera_group:
+                garment_img_camera = gr.Image(
+                    label="Garment Image (Camera)", 
+                    sources=['webcam'], 
+                    type="numpy",
+                    height=400
+                )
+                capture_garment_btn = gr.Button("📸 Capture Garment Photo", variant="primary")
+            
+            # Combined garment image state
+            garment_img_final = gr.State()
+            
+            # Examples for garment images
+            if garm_list_path:
+                example_garment = gr.Examples(
+                    inputs=garment_img_upload,
+                    examples_per_page=8,
+                    examples=garm_list_path,
+                    label="Example Garment Images"
+                )
+        
+        # Results Column
+        with gr.Column(elem_id="col-right"):
+            image_out = gr.Image(label="🎯 Try-On Result", show_share_button=False, height=400)
+            
             with gr.Row():
                 seed = gr.Slider(
                     label="Seed",
@@ -242,34 +380,135 @@ with gr.Blocks(css=css) as Tryon:
                     value=0,
                 )
                 randomize_seed = gr.Checkbox(label="Random seed", value=True)
+            
             with gr.Row():
                 seed_used = gr.Number(label="Seed used")
-                result_info = gr.Text(label="Response")
-            # try_button = gr.Button(value="Run", elem_id="button")
-            test_button = gr.Button(value="Run", elem_id="button")
-
-
-    # try_button.click(fn=start_tryon, inputs=[imgs, garm_img, seed, randomize_seed], outputs=[image_out, seed_used, result_info], api_name='tryon',concurrency_limit=10)
-    test_button.click(fn=tryon, inputs=[imgs, garm_img, seed, randomize_seed], outputs=[image_out, seed_used, result_info], api_name=False, concurrency_limit=45)
-
-    with gr.Column(elem_id = "col-showcase"):
+                result_info = gr.Text(label="Status")
+            
+            # Main try-on button
+            tryon_button = gr.Button("🚀 Generate Try-On", variant="primary", size="lg")
+    
+    # Event handlers for person image mode switching
+    upload_person_btn.click(
+        switch_to_upload_person,
+        outputs=[person_upload_group, person_camera_group]
+    )
+    
+    camera_person_btn.click(
+        switch_to_camera_person,
+        outputs=[person_upload_group, person_camera_group]
+    )
+    
+    clear_person_btn.click(
+        clear_person_image,
+        outputs=[person_img_upload]
+    )
+    
+    # Event handlers for garment image mode switching
+    upload_garment_btn.click(
+        switch_to_upload_garment,
+        outputs=[garment_upload_group, garment_camera_group]
+    )
+    
+    camera_garment_btn.click(
+        switch_to_camera_garment,
+        outputs=[garment_upload_group, garment_camera_group]
+    )
+    
+    clear_garment_btn.click(
+        clear_garment_image,
+        outputs=[garment_img_upload]
+    )
+    
+    # Capture button handlers
+    capture_person_btn.click(
+        capture_person_photo,
+        inputs=[person_img_camera],
+        outputs=[person_img_final]
+    )
+    
+    capture_garment_btn.click(
+        capture_garment_photo,
+        inputs=[garment_img_camera],
+        outputs=[garment_img_final]
+    )
+    
+    # Function to get the current person image
+    def get_current_person_image(upload_img, camera_img, final_img):
+        if final_img is not None:
+            return final_img
+        elif upload_img is not None:
+            return upload_img
+        elif camera_img is not None:
+            return camera_img
+        return None
+    
+    # Function to get the current garment image
+    def get_current_garment_image(upload_img, camera_img, final_img):
+        if final_img is not None:
+            return final_img
+        elif upload_img is not None:
+            return upload_img
+        elif camera_img is not None:
+            return camera_img
+        return None
+    
+    # Modified try-on function to handle multiple image sources
+    def run_tryon(person_upload, person_camera, person_final, 
+                  garment_upload, garment_camera, garment_final, 
+                  seed, randomize_seed):
+        
+        # Get the current person image
+        person_img = get_current_person_image(person_upload, person_camera, person_final)
+        
+        # Get the current garment image
+        garment_img = get_current_garment_image(garment_upload, garment_camera, garment_final)
+        
+        # Run the try-on process
+        return tryon(person_img, garment_img, seed, randomize_seed)
+    
+    # Main try-on button click handler
+    tryon_button.click(
+        fn=run_tryon,
+        inputs=[
+            person_img_upload, person_img_camera, person_img_final,
+            garment_img_upload, garment_img_camera, garment_img_final,
+            seed, randomize_seed
+        ],
+        outputs=[image_out, seed_used, result_info],
+        api_name=False,
+        concurrency_limit=45
+    )
+    
+    # Showcase section
+    with gr.Column(elem_id="col-showcase"):
         gr.HTML("""
-        <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
-            <div> </div>
-            <br>
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px; margin-top: 30px;">
             <div>
-            Virtual try-on examples in pairs of person and garment images
+            🌟 Virtual Try-On Examples
+            <br><br>
+            <em>See how different garments look on various models</em>
             </div>
         </div>
         """)
-        show_case = gr.Examples(
-            examples=[
-                ["assets/examples/model2.png", "assets/examples/garment2.png", "assets/examples/result2.png"],
-                ["assets/examples/model3.png", "assets/examples/garment3.png", "assets/examples/result3.png"],
-                ["assets/examples/model1.png", "assets/examples/garment1.png", "assets/examples/result1.png"],
-            ],
-            inputs=[imgs, garm_img, image_out],
-            label=None
-        )
-
-Tryon.queue(api_open=False).launch(show_api=False)
+        
+        # Only show examples if they exist
+        if os.path.exists("assets/examples"):
+            show_case = gr.Examples(
+                examples=[
+                    ["assets/examples/model2.png", "assets/examples/garment2.png", "assets/examples/result2.png"],
+                    ["assets/examples/model3.png", "assets/examples/garment3.png", "assets/examples/result3.png"],
+                    ["assets/examples/model1.png", "assets/examples/garment1.png", "assets/examples/result1.png"],
+                ],
+                inputs=[person_img_upload, garment_img_upload, image_out],
+                label="Example Try-On Results"
+            )
+
+# Launch the app
+if __name__ == "__main__":
+    Tryon.queue(api_open=False).launch(
+        show_api=False,
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )