app.py

import os
import cv2
import numpy as np
import json
import random
from PIL import Image, ImageDraw, ImageFont
import asyncio

import requests
import base64
import gradio as gr

# Set the machine number and model path
machine_number = 0
model = os.path.join(os.path.dirname(__file__), "models", "eva", "Eva_0.png")

# Define a mapping of model names to file paths
MODEL_MAP = {
    "AI Model Rouyan_0": os.path.join("models", "rouyan_new", "rouyan_new\\Rouyan_0.png"),
    "AI Model Rouyan_1": os.path.join("models", "rouyan_new", "rouyan_new\\Rouyan_1.png"),
    "AI Model Rouyan_2": os.path.join("models", "rouyan_new", "rouyan_new\\Rouyan_2.png"),
    "AI Model Eva_0": os.path.join("models", "eva", "Eva_0.png"),
    "AI Model Eva_1": os.path.join("models", "eva", "Eva_1.png"),
    "AI Model Simon_0": os.path.join("models", "simon_online", "Simon_0.png"),
    "AI Model Simon_1": os.path.join("models", "simon_online", "Simon_1.png"),
    "AI Model Xuanxuan_0": os.path.join("models", "xiaoxuan_online", "Xuanxuan_0.png"),
    "AI Model Xuanxuan_1": os.path.join("models", "xiaoxuan_online", "Xuanxuan_1.png"),
    "AI Model Xuanxuan_2": os.path.join("models", "xiaoxuan_online", "Xuanxuan_2.png"),
    "AI Model Yaqi_0": os.path.join("models", "yaqi", "Yaqi_0.png"),
    "AI Model Yaqi_1": os.path.join("models", "yaqi", "Yaqi_1.png"),
    "AI Model Yaqi_2": os.path.join("models", "yaqi", "Yaqi_2.png"),
    "AI Model Yaqi_3": os.path.join("models", "yaqi", "Yaqi_3.png"),
    "AI Model Yifeng_0": os.path.join("models", "yifeng_online", "Yifeng_0.png"),
    "AI Model Yifeng_1": os.path.join("models", "yifeng_online", "Yifeng_1.png"),
    "AI Model Yifeng_2": os.path.join("models", "yifeng_online", "Yifeng_2.png"),
    "AI Model Yifeng_3": os.path.join("models", "yifeng_online", "Yifeng_3.png"),
}

# Function to add watermark text to image
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

# Function to process try-on results
def get_tryon_result(model_name, garment1, garment2, seed=1234):
    if isinstance(model_name, np.ndarray):
       model_name = model_name[0] 
    
    model_name = "AI Model " + model_name.split("\\")[-1].split(".")[0]  # Handle Windows path
    print(type(model_name))


    # Directly load the model image from the disk, no need for Gradio file upload
    model_image = cv2.imread(MODEL_MAP.get(model_name))  # Load model image from disk
    if model_image is None:
        raise ValueError(f"Model image {model_name} could not be loaded.")

    # Encode garments as base64
    encoded_garment1 = cv2.imencode('.jpg', garment1)[1].tobytes()
    encoded_garment1 = base64.b64encode(encoded_garment1).decode('utf-8')

    if garment2 is not None:
        encoded_garment2 = cv2.imencode('.jpg', garment2)[1].tobytes()
        encoded_garment2 = base64.b64encode(encoded_garment2).decode('utf-8')
    else:
        encoded_garment2 = ''

    # Get the IP address from environment variable or default to localhost
    url = os.environ.get('OA_IP_ADDRESS', 'http://localhost:5000')
    headers = {'Content-Type': 'application/json'}
    seed = random.randint(0, 1222222222)

    # Prepare data for POST request
    data = {
        "garment1": encoded_garment1,
        "garment2": encoded_garment2,
        "model_name": model_name,
        "seed": seed
    }

    # Send POST request to server
    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!')

    final_img = add_waterprint(result_img)

    return final_img


with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 400px !important}") as demo:
    
    # Header Section
    gr.HTML(
        """
        <div style="text-align: center; padding: 20px;">
            <h1 style="font-size: 2.5rem; color: #2c3e50;">Outfit Anyone</h1>
            <h2 style="color: #34495e;">Ultra-high quality virtual try-on for any clothing and any person</h2>
        </div>
        """
    )

    # UI Layout for Image Inputs and Text Description
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Upload Your Model Image")
            init_image = gr.Image(sources='upload', type="numpy", label="Select a Model Image", value=None)
            example = gr.Examples(inputs=init_image,
                                  examples_per_page=4,
                                  examples=[os.path.join(os.path.dirname(__file__), MODEL_MAP.get('AI Model Rouyan_0'))])
        
        with gr.Column():
            gr.Markdown(
                """
                <h3 style="color: #2c3e50;">Instructions</h3>
                <p style="font-size: 1.1rem; color: #7f8c8d;">Please upload your model image and garment images (top and bottom). 
                The models are pre-loaded and cannot be modified. 
                For a dress or coat, you only need to upload the image for the 'Top Garment' section and leave the 'Bottom Garment' section empty.</p>
                """
            )
            with gr.Row():
                garment_top = gr.Image(sources='upload', type="numpy", label="Top Garment")
                example_top = gr.Examples(inputs=garment_top,
                                          examples_per_page=5,
                                          examples=[os.path.join(os.path.dirname(__file__), "garments", "top222.JPG")])
                
                garment_down = gr.Image(sources='upload', type="numpy", label="Bottom Garment")
                example_down = gr.Examples(inputs=garment_down,
                                           examples_per_page=5,
                                           examples=[os.path.join(os.path.dirname(__file__), "garments", "bottom1.png")])

            run_button = gr.Button(value="Run Try-On")
        
        with gr.Column():
            gallery = gr.Image(label="Try-On Result")

            run_button.click(fn=get_tryon_result, 
                             inputs=[init_image, garment_top, garment_down], 
                             outputs=[gallery], 
                             concurrency_limit=2)

    # Example Section
    gr.Markdown("## Example Try-On Results")
    with gr.Row():
        reference_image1 = gr.Image(label="Model Example", scale=1, value="examples\\examples_basemodel.png")
        reference_image2 = gr.Image(label="Garment Example", scale=1, value="examples\\examples_garment1.jpg")
        reference_image3 = gr.Image(label="Result Example", scale=1, value="examples\\examples_result1.png")

    gr.Examples(
        examples=[["examples\\examples_basemodel.png", "examples\\examples_garment1.png", "examples\\examples_result1.png"]],
        inputs=[reference_image1, reference_image2, reference_image3],
        label="Check out our example outfits!",
    )

if __name__ == "__main__":
    ip = requests.get('http://ifconfig.me/ip', timeout=1).text.strip()
    print("IP address", ip)
    demo.queue(max_size=10)
    demo.launch()