Add virtual try-on code

app.py

@@ -0,0 +1,56 @@
+import gradio as gr
+from PIL import Image
+import os
+import torch
+import mediapipe as mp
+from transformers import SamModel, SamProcessor
+from diffusers.utils import load_image
+from torchvision import transforms
+
+# Load model once
+model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-50")
+processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-50")
+
+def get_shoulder_coordinates(image: Image.Image):
+    mp_pose = mp.solutions.pose
+    pose = mp_pose.Pose()
+    
+    image_rgb = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    results = pose.process(image_rgb)
+    if results.pose_landmarks:
+        height, width, _ = image_rgb.shape
+        landmarks = results.pose_landmarks.landmark
+        left = (int(landmarks[11].x * width), int(landmarks[11].y * height))
+        right = (int(landmarks[12].x * width), int(landmarks[12].y * height))
+        return left, right
+    return None
+
+def try_on(person_img, tshirt_img):
+    coordinates = get_shoulder_coordinates(person_img)
+    if coordinates is None:
+        return "No pose detected", None
+    
+    left_shoulder, right_shoulder = coordinates
+    input_points = [[[left_shoulder[0], left_shoulder[1]], [right_shoulder[0], right_shoulder[1]]]]
+
+    inputs = processor(person_img, input_points=input_points, return_tensors="pt")
+    outputs = model(**inputs)
+
+    masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(),
+                                                          inputs["original_sizes"].cpu(),
+                                                          inputs["reshaped_input_sizes"].cpu())
+
+    mask_tensor = masks[0][0][2].to(dtype=torch.uint8)
+    mask = transforms.ToPILImage()(mask_tensor * 255)
+
+    tshirt_img = tshirt_img.resize(person_img.size, Image.LANCZOS)
+    result = Image.composite(tshirt_img, person_img, mask)
+    return result
+
+demo = gr.Interface(fn=try_on, 
+                    inputs=["image", "image"], 
+                    outputs="image",
+                    title="Virtual Try-On using SlimSAM",
+                    description="Upload a person image and a t-shirt image.")
+
+demo.launch()

requirements.txt

@@ -0,0 +1,10 @@
+Flask==2.3.3
+gunicorn==21.2.0
+Pillow==10.1.0
+opencv-python==4.9.0.80
+torch==2.2.0
+torchvision==0.17.0
+mediapipe==0.10.21
+transformers==4.40.1
+diffusers==0.27.2
+safetensors==0.4.2