Update app.py

app.py

@@ -20,11 +20,40 @@ MODEL_MAP = {
     "AI Model male_model2": 'models/male_model2.png',
 }
 
+# Sample clothing items from your clothes folder
+SAMPLE_CLOTHES = {
+    "top": [
+        "clothes/shirt1.jpg",
+        "clothes/shirt2.jpg", 
+        "clothes/dress1.jpg",
+        "clothes/jacket1.jpg"
+    ],
+    "bottom": [
+        "clothes/pants1.jpg",
+        "clothes/skirt1.jpg",
+        "clothes/jeans1.jpg", 
+        "clothes/shorts1.jpg"
+    ]
+}
+
 def add_waterprint(img):
     h, w, _ = img.shape
     img = cv2.putText(img, 'Powered by OutfitAnyone', (int(0.3*w), h-20), cv2.FONT_HERSHEY_PLAIN, 2, (128, 128, 128), 2, cv2.LINE_AA)
     return img
 
+def load_sample_image(image_path):
+    """Load a sample clothing image and return it as numpy array"""
+    try:
+        if os.path.exists(image_path):
+            img = cv2.imread(image_path)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            return img
+        else:
+            # Return a placeholder if image doesn't exist
+            return np.zeros((400, 400, 3), dtype=np.uint8)
+    except Exception as e:
+        print(f"Error loading image {image_path}: {e}")
+        return np.zeros((400, 400, 3), dtype=np.uint8)
 
 def get_tryon_result(model_name, garment1, garment2, seed=1234):
     # model_name = "AI Model " + model_name.split("\\")[-1].split(".")[0] # windows
@@ -40,7 +69,17 @@ def get_tryon_result(model_name, garment1, garment2, seed=1234):
     else:
         encoded_garment2 = ''
 
-    url = os.environ['OA_IP_ADDRESS']
+    # Fix for the missing environment variable
+    try:
+        url = os.environ['OA_IP_ADDRESS']
+    except KeyError:
+        print("Error: OA_IP_ADDRESS environment variable not set!")
+        # Return a dummy image or handle the error appropriately
+        dummy_img = np.zeros((512, 512, 3), dtype=np.uint8)
+        dummy_img = cv2.putText(dummy_img, 'Error: OA_IP_ADDRESS not set', (50, 256), 
+                               cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+        return add_waterprint(dummy_img)
+    
     headers = {'Content-Type': 'application/json'}
     seed = random.randint(0, 1222222222)
     data = {
@@ -49,15 +88,26 @@ def get_tryon_result(model_name, garment1, garment2, seed=1234):
         "model_name": model_name,
         "seed": seed
     }
-    response = requests.post(url, headers=headers, data=json.dumps(data))
-    print("response code", response.status_code)
-    if response.status_code == 200:
-        result = response.json()
-        result = base64.b64decode(result['images'][0])
-        result_np = np.frombuffer(result, np.uint8)
-        result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
-    else:
-        print('server error!')
+    
+    try:
+        response = requests.post(url, headers=headers, data=json.dumps(data))
+        print("response code", response.status_code)
+        if response.status_code == 200:
+            result = response.json()
+            result = base64.b64decode(result['images'][0])
+            result_np = np.frombuffer(result, np.uint8)
+            result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
+        else:
+            print('server error!')
+            # Return error image
+            result_img = np.zeros((512, 512, 3), dtype=np.uint8)
+            result_img = cv2.putText(result_img, f'Server Error: {response.status_code}', 
+                                   (50, 256), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
+    except Exception as e:
+        print(f"Request error: {e}")
+        result_img = np.zeros((512, 512, 3), dtype=np.uint8)
+        result_img = cv2.putText(result_img, 'Connection Error', (50, 256), 
+                               cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
 
     final_img = add_waterprint(result_img)
     return final_img
@@ -95,6 +145,26 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
                 </div>
                 </div>
                 """)
+            
+            # Sample clothing buttons section
+            gr.HTML("<h3>Quick Select Sample Clothes:</h3>")
+            
+            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.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.Row():
                 garment_top = gr.Image(sources='upload', type="numpy", label="top garment")
                 garment_down = gr.Image(sources='upload', type="numpy", label="lower garment")
@@ -104,14 +174,49 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
         with gr.Column():
             gallery = gr.Image()
 
-            run_button.click(fn=get_tryon_result, 
-                             inputs=[
-                                    init_image,
-                                    garment_top,
-                                    garment_down,
-                                    ], 
-                             outputs=[gallery],
-                             concurrency_limit=2)
+    # Connect buttons to load sample images
+    top_btn1.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["top"][0]),
+        outputs=[garment_top]
+    )
+    top_btn2.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["top"][1]), 
+        outputs=[garment_top]
+    )
+    top_btn3.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["top"][2]),
+        outputs=[garment_top]
+    )
+    top_btn4.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["top"][3]),
+        outputs=[garment_top]
+    )
+    
+    bottom_btn1.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][0]),
+        outputs=[garment_down]
+    )
+    bottom_btn2.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][1]),
+        outputs=[garment_down] 
+    )
+    bottom_btn3.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][2]),
+        outputs=[garment_down]
+    )
+    bottom_btn4.click(
+        lambda: load_sample_image(SAMPLE_CLOTHES["bottom"][3]),
+        outputs=[garment_down]
+    )
+
+    run_button.click(fn=get_tryon_result, 
+                     inputs=[
+                            init_image,
+                            garment_top,
+                            garment_down,
+                            ], 
+                     outputs=[gallery],
+                     concurrency_limit=2)
 
 if __name__ == "__main__":
     ip = requests.get('http://ifconfig.me/ip', timeout=1).text.strip()