Update app.py

app.py

@@ -1,44 +1,126 @@
+import os
+import cv2
+import numpy as np
+import json
+import random
+from PIL import Image, ImageDraw, ImageFont
+import requests
+import base64
 import gradio as gr
-from transformers import CLIPProcessor, CLIPModel
-from PIL import Image
-import torch
-
-# Load the CLIP model and processor from Hugging Face
-model_name = "openai/clip-vit-base-patch32"
-model = CLIPModel.from_pretrained(model_name)
-processor = CLIPProcessor.from_pretrained(model_name)
-
-def generate_outfit(base_image, garment_image):
-    # Open the images using PIL
-    base_image = Image.open(base_image).convert("RGB")
-    garment_image = Image.open(garment_image).convert("RGB")
+
+# Define the path to the default model
+default_model = os.path.join(os.path.dirname(__file__), "models/eva/Eva_0.png")
+
+# Map of AI models with their corresponding file paths
+MODEL_MAP = {
+    "AI Model Rouyan_0": 'models/rouyan_new/Rouyan_0.png',
+    "AI Model Rouyan_1": 'models/rouyan_new/Rouyan_1.png',
+    "AI Model Rouyan_2": 'models/rouyan_new/Rouyan_2.png',
+    "AI Model Eva_0": 'models/eva/Eva_0.png',
+    "AI Model Eva_1": 'models/eva/Eva_1.png',
+    "AI Model Simon_0": 'models/simon_online/Simon_0.png',
+    "AI Model Simon_1": 'models/simon_online/Simon_1.png',
+    "AI Model Xuanxuan_0": 'models/xiaoxuan_online/Xuanxuan_0.png',
+    "AI Model Xuanxuan_1": 'models/xiaoxuan_online/Xuanxuan_1.png',
+    "AI Model Xuanxuan_2": 'models/xiaoxuan_online/Xuanxuan_2.png',
+    "AI Model Yaqi_0": 'models/yaqi/Yaqi_0.png',
+    "AI Model Yaqi_1": 'models/yaqi/Yaqi_1.png',
+    "AI Model Yaqi_2": 'models/yaqi/Yaqi_2.png',
+    "AI Model Yaqi_3": 'models/yaqi/Yaqi_3.png',
+    "AI Model Yifeng_0": 'models/yifeng_online/Yifeng_0.png',
+    "AI Model Yifeng_1": 'models/yifeng_online/Yifeng_1.png',
+    "AI Model Yifeng_2": 'models/yifeng_online/Yifeng_2.png',
+    "AI Model Yifeng_3": 'models/yifeng_online/Yifeng_3.png',
+}
+
+def add_watermark(image):
+    """
+    Adds a watermark to the provided 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
+
+def get_tryon_result(model_name, top_garment, bottom_garment=None):
+    """
+    Processes the virtual try-on result by sending a request to the server.
+    """
+    # Format the model name for the server request
+    model_key = "AI Model " + model_name.split("/")[-1].split(".")[0]
+    print(f"Selected model: {model_key}")
+
+    # Encode the garments as base64
+    encoded_top = base64.b64encode(cv2.imencode('.jpg', top_garment)[1].tobytes()).decode('utf-8')
+    encoded_bottom = base64.b64encode(cv2.imencode('.jpg', bottom_garment)[1].tobytes()).decode('utf-8') if bottom_garment else ''
+
+    # Server request setup
+    server_url = os.environ.get('OA_IP_ADDRESS', 'http://localhost:5000')  # Default to localhost if environment variable is not set
+    headers = {'Content-Type': 'application/json'}
+    payload = {
+        "garment1": encoded_top,
+        "garment2": encoded_bottom,
+        "model_name": model_key,
+        "seed": random.randint(0, 99999999)
+    }
+
+    # Send the request
+    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
+
+# Set up the Gradio interface
+with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 400px !important}") as demo:
+    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>
+    """)
     
-    # Process the images with the CLIP model
-    inputs = processor(
-        text=["a person wearing a garment"],
-        images=[base_image, garment_image],
-        return_tensors="pt",
-        padding=True  # Ensure padding is applied for batched inputs
+    with gr.Row():
+        with gr.Column():
+            model_selector = gr.Image(sources='clipboard', 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)")
+            example_top_garments = gr.Examples(inputs=top_garment_input, 
+                                               examples=[os.path.join(os.path.dirname(__file__), "garments/top1.jpg")],
+                                               examples_per_page=5)
+            example_bottom_garments = gr.Examples(inputs=bottom_garment_input,
+                                                  examples=[os.path.join(os.path.dirname(__file__), "garments/bottom1.jpg")],
+                                                  examples_per_page=5)
+            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])
+
+    gr.Markdown("## Example Outputs")
+    with gr.Row():
+        ref_image = gr.Image(label="Model Example", value="examples/model_example.jpg")
+        garment_example = gr.Image(label="Garment Example", value="examples/garment_example.jpg")
+        result_example = gr.Image(label="Result Example", value="examples/result_example.jpg")
+    gr.Examples(
+        examples=[["examples/model_example.jpg", "examples/garment_example.jpg", "examples/result_example.jpg"]],
+        inputs=[ref_image, garment_example, result_example],
     )
-    
-    # Perform inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Get the similarity score between the base image and the garment image
-    similarity_score = outputs.logits_per_image.item()
-    
-    # Return the similarity score
-    return f"Similarity Score: {similarity_score:.4f}"
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=generate_outfit,
-    inputs=[gr.Image(type="filepath", label="Base Image"), gr.Image(type="filepath", label="Garment Image")],
-    outputs="text",
-    title="Outfit Generator",
-    description="Upload a base image and a garment image to generate outfit suggestions using CLIP."
-)
 
 if __name__ == "__main__":
+    demo.queue(max_size=10)
     demo.launch()