Update app.py

app.py

@@ -1,62 +1,107 @@
+import os
 import requests
+import json
+import time
+import cv2
 import base64
-
-# --- Configuration ---
-API_KEY = "sk_7c19e4f2ad434a3ebc70fdd85ade5309"
-PERSON_IMAGE_PATH = "path_to_person_image.jpg"
-GARMENT_IMAGE_PATH = "path_to_garment_image.jpg"
-
-# --- Helper Function: Read Image and Convert to Base64 ---
-def encode_image(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
-
-# --- Step 1: Remove Background from Person Image ---
-def remove_background(image_path):
-    with open(image_path, "rb") as image_file:
-        response = requests.post(
-            "https://api.developer.pixelcut.ai/v1/remove-background",
-            headers={
-                "Accept": "application/json",
-                "X-API-KEY": API_KEY,
-            },
-            files={"image": image_file},
-            data={"format": "png"},
-        )
-
-    if response.status_code == 200:
-        return response.json().get("url")
-    else:
-        print("Error:", response.json())
-        return None
-
-# --- Step 2: Try-On Function ---
-def tryonn(person_image_path, garment_image_path):
-    # Get background-removed URLs
-    person_url = remove_background(person_image_path)
-    garment_url = remove_background(garment_image_path)
-
-    if not person_url or not garment_url:
-        print("Failed to process images!")
-        return
-
-    # Make the try-on API call
-    tryon_response = requests.post(
-        "https://api.example.com/v1/tryon",
-        headers={"Accept": "application/json", "X-API-KEY": API_KEY},
-        json={"person_image_url": person_url, "garment_image_url": garment_url},
-    )
-
-    if tryon_response.status_code == 200:
-        tryon_result_url = tryon_response.json().get("result_url")
-        print("Try-On Success! Result:", tryon_result_url)
-    else:
-        print("Try-On Failed:", tryon_response.json())
-
-# --- Execute the try-on function ---
-tryonn(PERSON_IMAGE_PATH, GARMENT_IMAGE_PATH)
-
-
+import random
+import numpy as np
+import gradio as gr
+
+MAX_SEED = 999999
+
+# ✅ Hardcoded Pixelcut API key
+pixelcut_api_key = "sk_7c19e4f2ad434a3ebc70fdd85ade5309"
+
+
+# Convert images to PNG format
+def convert_to_png(image):
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+
+# Encode images to base64 PNG format
+def encode_image(image):
+    image = convert_to_png(image)
+    _, buffer = cv2.imencode('.png', image)
+    return base64.b64encode(buffer).decode('utf-8')
+
+
+# Main try-on function using Pixelcut API
+def tryon(person_img, garment_img, seed, randomize_seed):
+    start_time = time.time()
+
+    # 🛑 Check if images are provided
+    if person_img is None or garment_img is None:
+        return None, None, "Empty image"
+
+    # 🎲 Handle seed randomization
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    # 🔥 Convert images to base64 PNG
+    encoded_person_img = encode_image(person_img)
+    encoded_garment_img = encode_image(garment_img)
+
+    # 🌟 Setup Pixelcut API endpoint and headers
+    url = "https://api.developer.pixelcut.ai/v1/remove-background"
+    headers = {
+        'Content-Type': 'application/json',
+        'Accept': 'application/json',
+        'X-API-KEY': pixelcut_api_key
+    }
+
+    # 🔧 Send person image for background removal
+    person_data = {"image": encoded_person_img, "format": "png"}
+    garment_data = {"image": encoded_garment_img, "format": "png"}
+
+    try:
+        # 🚀 Process Person Image
+        person_response = requests.post(url, headers=headers, json=person_data, timeout=60)
+        if person_response.status_code == 200:
+            person_result_url = person_response.json().get('url', '')
+        else:
+            return None, None, f"Person image error: {person_response.status_code}"
+
+        # 🚀 Process Garment Image
+        garment_response = requests.post(url, headers=headers, json=garment_data, timeout=60)
+        if garment_response.status_code == 200:
+            garment_result_url = garment_response.json().get('url', '')
+        else:
+            return None, None, f"Garment image error: {garment_response.status_code}"
+
+        # 🔥 Call the try-on API with both processed images
+        tryon_url = "https://api.example.com/v1/tryon"
+        tryon_payload = {
+            "person_image_url": person_result_url,
+            "garment_image_url": garment_result_url
+        }
+        tryon_response = requests.post(tryon_url, headers=headers, json=tryon_payload, timeout=60)
+
+        # ✅ Handle the try-on result
+        if tryon_response.status_code == 200:
+            result_url = tryon_response.json().get('result_url', '')
+            result_img_data = requests.get(result_url).content
+            result_np = np.frombuffer(result_img_data, np.uint8)
+            result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
+            info = "Success"
+        else:
+            info = f"Try-On error: {tryon_response.status_code}"
+
+    # 🛑 Handle timeout or API errors
+    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"❌ Other error: {err}")
+        info = "Error, please contact the admin"
+        result_img = None
+
+    end_time = time.time()
+    print(f"⏳ Time used: {end_time - start_time:.2f} seconds")
+
+    return result_img, seed, info
 
 example_path = os.path.join(os.path.dirname(__file__), 'assets')