File size: 4,011 Bytes
88cc045 d66f063 88cc045 d66f063 bc97f92 d66f063 ffe5741 d66f063 ffe5741 d66f063 716f6d9 d66f063 88cc045 d66f063 716f6d9 d66f063 ffe5741 d66f063 6a91eb0 bc97f92 eb3ffe1 d66f063 ffe5741 d66f063 ffe5741 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 98 99 100 101 |
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, top_left_x, top_left_y, bottom_right_x, bottom_right_y):
displacement_map = generate_displacement_map(image_a)
# Extract bounding box coordinates
top_left = (int(top_left_x), int(top_left_y))
bottom_right = (int(bottom_right_x), int(bottom_right_y))
# 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.convert("RGBA") # Ensure design image has an alpha channel
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) # Use image_b as the mask for pasting
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, top_left_x, top_left_y, bottom_right_x, bottom_right_y):
result = fit_and_warp_design(image_a, image_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y)
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")
top_left_x = gr.Slider(minimum=0, maximum=1000, label="Top-left X Coordinate", value=50)
top_left_y = gr.Slider(minimum=0, maximum=1000, label="Top-left Y Coordinate", value=100)
bottom_right_x = gr.Slider(minimum=0, maximum=1000, label="Bottom-right X Coordinate", value=300)
bottom_right_y = gr.Slider(minimum=0, maximum=1000, label="Bottom-right Y Coordinate", value=400)
# Define the Gradio interface
iface = gr.Interface(
fn=process_images,
inputs=[image_input_a, image_input_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y],
outputs="image",
title="Clothing Design Fitting with Adjustable Bounding Box",
description="Upload a clothing image and a design image. Adjust the design's position and size using sliders to fit the design onto the clothing.",
)
# Launch the interface
iface.launch()
|