OutfitAnyone / app.py
ali2780's picture
Update app.py
58f2376 verified
raw
history blame
7.37 kB
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()