File size: 3,403 Bytes
88cc045 d66f063 88cc045 d66f063 bc97f92 d66f063 88cc045 d66f063 6a91eb0 bc97f92 6a91eb0 d66f063 6a91eb0 d66f063 6a91eb0 |
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 |
import torch
import torchvision.transforms as transforms
from PIL import Image
import cv2
import numpy as np
import gradio as gr
# Load MiDaS model
midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large", trust_repo=True)
midas.eval()
# Preprocessing function
def preprocess_image(image):
transform = transforms.Compose([
transforms.Resize(384),
transforms.CenterCrop(384),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
return transform(image).unsqueeze(0)
# Function to generate the displacement map
def generate_displacement_map(image_a):
input_batch = preprocess_image(image_a)
with torch.no_grad():
depth_map = midas(input_batch)
depth_map = depth_map.squeeze().cpu().numpy()
depth_map = cv2.resize(depth_map, (image_a.width, image_a.height))
depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
displacement_map = depth_map * 30
return displacement_map
# Function to warp and fit Image-B onto Image-A
def fit_and_warp_design(image_a, image_b, design_bbox):
displacement_map = generate_displacement_map(image_a)
# Extract bounding box coordinates
top_left = (int(design_bbox[0]), int(design_bbox[1]))
bottom_right = (int(design_bbox[2]), int(design_bbox[3]))
# Resize the design to fit within the specified bounding box
design_width = bottom_right[0] - top_left[0]
design_height = bottom_right[1] - top_left[1]
image_b = image_b.resize((design_width, design_height))
# Create a blank canvas with the same size as Image-A
canvas = Image.new('RGBA', (displacement_map.shape[1], displacement_map.shape[0]), (0, 0, 0, 0))
canvas.paste(image_b, top_left, image_b)
canvas_np = np.array(canvas)
h, w = displacement_map.shape
y_indices, x_indices = np.indices((h, w), dtype=np.float32)
x_displacement = (x_indices + displacement_map).astype(np.float32)
y_displacement = (y_indices + displacement_map).astype(np.float32)
x_displacement = np.clip(x_displacement, 0, w - 1)
y_displacement = np.clip(y_displacement, 0, h - 1)
warped_canvas = cv2.remap(canvas_np, x_displacement, y_displacement, cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
image_a_rgba = image_a.convert("RGBA")
image_a_np = np.array(image_a_rgba)
non_transparent_pixels = warped_canvas[..., 3] > 0
image_a_np[non_transparent_pixels] = warped_canvas[non_transparent_pixels]
final_image = Image.fromarray(image_a_np)
return final_image
# Gradio interface function
def process_images(image_a, image_b, design_bbox):
result = fit_and_warp_design(image_a, image_b, design_bbox)
return result
# Gradio UI components
image_input_a = gr.Image(label="Upload Clothing Image", type="pil")
image_input_b = gr.Image(label="Upload Design Image", type="pil")
design_bbox_input = gr.Image(label="Adjust Design Position and Size", type="pil")
# Define the Gradio interface
iface = gr.Interface(
fn=process_images,
inputs=[image_input_a, image_input_b, design_bbox_input],
outputs="image",
title="Clothing Design Fitting with Drag-and-Drop",
description="Upload a clothing image and a design image. Drag and resize the design onto the clothing using the cursor.",
)
# Launch the interface
iface.launch()
|