Create app.py

app.py

@@ -0,0 +1,159 @@
+import os
+import numpy as np
+import base64
+import json
+import random
+import cv2
+from PIL import Image
+import requests
+import gradio as gr
+
+# In-memory user storage for demonstration
+users_db = {}
+
+# Ensure the model path is correct
+current_dir = os.path.dirname(os.path.abspath(__file__))
+default_model = os.path.join(current_dir, "models/eva/Eva_0.png")
+
+# Map of AI models with their corresponding file paths
+MODEL_MAP = {
+    "AI Model Rouyan_0": os.path.join(current_dir, 'models/rouyan_new/Rouyan_0.png'),
+    "AI Model Eva_0": os.path.join(current_dir, 'models/eva/Eva_0.png'),
+    "AI Model Eva_1": os.path.join(current_dir, 'models/eva/Eva_1.png'),
+    "AI Model Simon_0": os.path.join(current_dir, 'models/simon_online/Simon_0.png'),
+    "AI Model Simon_1": os.path.join(current_dir, 'models/simon_online/Simon_1.png'),
+    "AI Model Xuanxuan_0": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_0.png'),
+    "AI Model Xuanxuan_1": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_1.png'),
+    "AI Model Xuanxuan_2": os.path.join(current_dir, 'models/xiaoxuan_online/Xuanxuan_2.png'),
+    "AI Model Yaqi_0": os.path.join(current_dir, 'models/yaqi/Yaqi_0.png'),
+    "AI Model Yaqi_1": os.path.join(current_dir, 'models/yaqi/Yaqi_1.png'),
+    "AI Model Yaqi_2": os.path.join(current_dir, 'models/yaqi/Yaqi_2.png'),
+    "AI Model Yaqi_3": os.path.join(current_dir, 'models/yaqi/Yaqi_3.png'),
+    "AI Model Yifeng_0": os.path.join(current_dir, 'models/yifeng_online/Yifeng_0.png'),
+    "AI Model Yifeng_1": os.path.join(current_dir, 'models/yifeng_online/Yifeng_1.png'),
+    "AI Model Yifeng_2": os.path.join(current_dir, 'models/yifeng_online/Yifeng_2.png'),
+    "AI Model Yifeng_3": os.path.join(current_dir, 'models/yifeng_online/Yifeng_3.png'),
+}
+
+# Registration function
+def register(username, password, confirm_password):
+    if username in users_db:
+        return "Username already exists", gr.update(visible=True), gr.update(visible=False)
+    if password != confirm_password:
+        return "Passwords do not match", gr.update(visible=True), gr.update(visible=False)
+    users_db[username] = password
+    return "Registration successful! Please login.", gr.update(visible=True), gr.update(visible=False)
+
+# Login function
+def login(username, password):
+    if users_db.get(username) == password:
+        return gr.update(visible=False), gr.update(visible=True)
+    else:
+        return gr.update(visible=True), gr.update(visible=False)
+
+# Watermark function
+def add_watermark(image):
+    height, width, _ = image.shape
+    cv2.putText(image, 'Powered by OutfitAnyone', (int(0.3 * width), height - 20),
+                cv2.FONT_HERSHEY_PLAIN, 2, (128, 128, 128), 2, cv2.LINE_AA)
+    return image
+
+# Try-on result processing
+def get_tryon_result(model_name, top_garment, bottom_garment=None):
+    model_key = "AI Model " + model_name.split("/")[-1].split(".")[0]
+    print(f"Selected model: {model_key}")
+
+    encoded_top = base64.b64encode(cv2.imencode('.jpg', top_garment)[1].tobytes()).decode('utf-8')
+    
+    # Check if bottom_garment is not None and is not an empty array
+    encoded_bottom = ""
+    if bottom_garment is not None and bottom_garment.size > 0:
+        encoded_bottom = base64.b64encode(cv2.imencode('.jpg', bottom_garment)[1].tobytes()).decode('utf-8')
+
+    server_url = os.environ.get('OA_IP_ADDRESS', 'http://localhost:5000')
+    headers = {'Content-Type': 'application/json'}
+    payload = {
+        "garment1": encoded_top,
+        "garment2": encoded_bottom,
+        "model_name": model_key,
+        "seed": random.randint(0, 99999999)
+    }
+
+    try:
+        response = requests.post(server_url, headers=headers, data=json.dumps(payload))
+
+        if response.status_code == 200:
+            result = response.json()
+            result_img = cv2.imdecode(np.frombuffer(base64.b64decode(result['images'][0]), np.uint8), cv2.IMREAD_UNCHANGED)
+            final_img = add_watermark(result_img)
+            return final_img
+        else:
+            print(f"Error: Server responded with status code {response.status_code}")
+            return None
+    except requests.exceptions.ConnectionError as e:
+        print(f"ConnectionError: Could not connect to server at {server_url}. Please ensure the server is running.")
+        print(f"Details: {e}")
+        return None
+    except requests.exceptions.RequestException as e:
+        print(f"RequestException: An error occurred during the request.")
+        print(f"Details: {e}")
+        return None
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.HTML("<script>console.log('Custom JavaScript');</script>")
+    login_block = gr.Column(visible=True)
+    register_block = gr.Column(visible=False)
+    main_app = gr.Column(visible=False)
+    
+    # Registration Page
+    with register_block:
+        gr.HTML("<h2>Register for Outfit Anyone</h2>")
+        username = gr.Textbox(label="Username")
+        password = gr.Textbox(label="Password", type="password", elem_id="reg_pass")
+        confirm_password = gr.Textbox(label="Confirm Password", type="password", elem_id="reg_conf_pass")
+        register_button = gr.Button("Register")
+        back_to_login_button = gr.Button("Back to Login")
+        register_message = gr.Textbox(label="", visible=False)
+
+    # Login Page
+    with login_block:
+        gr.HTML("<h2>Login to Outfit Anyone</h2>")
+        username_login = gr.Textbox(label="Username")
+        password_login = gr.Textbox(label="Password", type="password", elem_id="login_pass")
+        login_button = gr.Button("Login")
+        go_to_register_button = gr.Button("Register")
+        login_message = gr.Textbox(label="", visible=False)
+    
+    # Main App Interface
+    with main_app:
+        gr.HTML("""
+            <div style="text-align: center;">
+                <h1>Outfit Anyone: Virtual Try-On</h1>
+                <h4>v1.0</h4>
+                <p>Upload your garments and choose a model to see the virtual try-on.</p>
+            </div>
+        """)
+        with gr.Row():
+            with gr.Column():
+                model_selector = gr.Image(sources='upload', type="filepath", label="Model", value=default_model)
+                example_models = gr.Examples(inputs=model_selector, examples=[MODEL_MAP['AI Model Rouyan_0'], MODEL_MAP['AI Model Eva_0']], examples_per_page=4)
+            with gr.Column():
+                gr.HTML("<h3>Select Garments for Virtual Try-On</h3>")
+                top_garment_input = gr.Image(sources='upload', type="numpy", label="Top Garment")
+                bottom_garment_input = gr.Image(sources='upload', type="numpy", label="Bottom Garment (Optional)")
+                generate_button = gr.Button(value="Generate Outfit")
+            with gr.Column():
+                result_display = gr.Image()
+
+                generate_button.click(fn=get_tryon_result, inputs=[model_selector, top_garment_input, bottom_garment_input], outputs=[result_display])
+
+    # Button Click Actions
+    register_button.click(fn=register, inputs=[username, password, confirm_password], outputs=[register_message, register_block, login_block])
+    back_to_login_button.click(lambda: (gr.update(visible=False), gr.update(visible=True)), None, [register_block, login_block])
+    go_to_register_button.click(lambda: (gr.update(visible=True), gr.update(visible=False)), None, [register_block, login_block])
+    login_button.click(fn=login, inputs=[username_login, password_login], outputs=[login_block, main_app])
+    
+if __name__ == "__main__":
+    demo.queue(max_size=10)
+    demo.launch(server_name="0.0.0.0", server_port=8080)