Spaces:
Running
Running
| 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) | |