File size: 5,022 Bytes
968edaf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import os
import cv2
import random
import gradio as gr
from gradio_client import Client
machine_number = 0
model = os.path.join(os.path.dirname(__file__), "models/simon_online/Simon_0.png")
url = os.environ['OA_IP_ADDRESS']
print("API:", url)
client = Client(url)
MODEL_MAP = {
"AI Model Simon_0": 'models/simon_online/Simon_0.png',
"AI Model Xuanxuan_0": 'models/xiaoxuan_online/Xuanxuan_0.png',
"AI Model Yifeng_0": 'models/yifeng_online/Yifeng_0.png'
}
def add_waterprint(img):
h, w, _ = img.shape
img = cv2.putText(img, 'AI VTON', (int(0.3 * w), h - 20), cv2.FONT_HERSHEY_PLAIN, 2,
(128, 128, 128), 2, cv2.LINE_AA)
return img
def get_tryon_result(model_name, garment1, garment2, seed=1234):
# _model = "AI Model " + model_name.split("\\")[-1].split(".")[0] # windows
_model = "AI Model " + model_name.split("/")[-1].split(".")[0] # linux
print("Use Model:", _model)
seed = random.randint(0, 1222222222)
result = client.predict(
model_name,
garment1,
garment2,
api_name="/get_tryon_result",
fn_index=seed
)
final_img = remove_watermark2(result)
return final_img
def remove_watermark2(path):
img = cv2.imread(path)
img_ = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w, _ = img.shape
y_start = max(h - 45, 0)
y_end = h
x_start = max(int(0.3 * w), 0)
x_end = w
img_[y_start:y_end, x_start:x_end, :] = [255, 255, 255]
return img_
with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 400px !important} ") as demo:
# gr.Markdown("# Outfit Anyone v0.9")
with gr.Row():
with gr.Column():
init_image = gr.Image(sources='clipboard', type="filepath", label="model", value=model)
example = gr.Examples(inputs=init_image,
examples_per_page=4,
examples=[
os.path.join(os.path.dirname(__file__), MODEL_MAP.get('AI Model Simon_0')),
os.path.join(os.path.dirname(__file__),
MODEL_MAP.get('AI Model Xuanxuan_0')),
os.path.join(os.path.dirname(__file__), MODEL_MAP.get('AI Model Yifeng_0')),
])
with gr.Column():
gr.HTML(
"""
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
<div>
</div>
</div>
""")
with gr.Row():
garment_top = gr.Image(sources='upload', type="filepath", 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"),
os.path.join(os.path.dirname(__file__), "garments/top5.png"),
os.path.join(os.path.dirname(__file__), "garments/top333.png"),
os.path.join(os.path.dirname(__file__), "garments/dress1.png"),
os.path.join(os.path.dirname(__file__), "garments/dress2.png"),
])
garment_down = gr.Image(sources='upload', type="filepath", label="lower garment")
example_down = gr.Examples(inputs=garment_down,
examples_per_page=5,
examples=[os.path.join(os.path.dirname(__file__), "garments/bottom1.png"),
os.path.join(os.path.dirname(__file__), "garments/bottom2.PNG"),
os.path.join(os.path.dirname(__file__), "garments/bottom3.JPG"),
os.path.join(os.path.dirname(__file__), "garments/bottom4.PNG"),
os.path.join(os.path.dirname(__file__), "garments/bottom5.png"),
])
run_button = gr.Button(value="Run")
with gr.Column():
gallery = gr.Image()
run_button.click(fn=get_tryon_result,
inputs=[
init_image,
garment_top,
garment_down,
],
outputs=[gallery],
concurrency_limit=4)
if __name__ == "__main__":
demo.queue(max_size=10)
demo.launch(share=False, server_name='0.0.0.0',server_port=7860)
|