Spaces:
Runtime error
Runtime error
File size: 2,123 Bytes
9585106 |
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 |
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()
|