|
import gradio as gr |
|
from PIL import Image, ImageOps |
|
import numpy as np |
|
import cv2 |
|
|
|
|
|
tshirt_img = None |
|
design_img = None |
|
vertices = [] |
|
polygon_drawn = False |
|
design_pos = (250, 250) |
|
|
|
def load_tshirt(image): |
|
global tshirt_img |
|
tshirt_img = image.resize((600, 600), Image.Resampling.LANCZOS) |
|
return tshirt_img |
|
|
|
def load_design(image): |
|
global design_img, design_pos |
|
design_img = image.resize((100, 100), Image.Resampling.LANCZOS) |
|
design_pos = (250, 250) |
|
return tshirt_img, design_img, design_pos |
|
|
|
def add_polygon_vertex(x, y): |
|
global vertices, polygon_drawn |
|
if tshirt_img and not polygon_drawn: |
|
vertices.append((x, y)) |
|
if len(vertices) > 2 and np.linalg.norm(np.array(vertices[-1]) - np.array(vertices[0])) < 10: |
|
vertices[-1] = vertices[0] |
|
polygon_drawn = True |
|
return tshirt_img, design_img, design_pos |
|
|
|
def fit_design(): |
|
global design_img, tshirt_img, vertices, polygon_drawn |
|
if design_img and tshirt_img and polygon_drawn: |
|
src_pts = np.array([[0, 0], [design_img.width, 0], [design_img.width, design_img.height], [0, design_img.height]], dtype="float32") |
|
dst_pts = np.array(vertices, dtype="float32") |
|
|
|
M = cv2.getPerspectiveTransform(src_pts, dst_pts) |
|
design_np = np.array(design_img) |
|
transformed_design = cv2.warpPerspective(design_np, M, (tshirt_img.width, tshirt_img.height)) |
|
transformed_design_img = Image.fromarray(transformed_design) |
|
|
|
design_img = transformed_design_img |
|
return tshirt_img, design_img, design_pos |
|
return tshirt_img, design_img, design_pos |
|
|
|
def rotate_design(angle): |
|
global design_img |
|
if design_img: |
|
angle = int(angle) |
|
rotated_design = design_img.rotate(angle, expand=True) |
|
design_img = rotated_design |
|
return tshirt_img, design_img, design_pos |
|
return tshirt_img, design_img, design_pos |
|
|
|
def display_image(x=0, y=0, angle=0, image_type=None, image=None): |
|
if image_type == "tshirt": |
|
return load_tshirt(image) |
|
elif image_type == "design": |
|
return load_design(image) |
|
elif image_type == "polygon": |
|
return add_polygon_vertex(x, y) |
|
elif image_type == "fit": |
|
return fit_design() |
|
elif image_type == "rotate": |
|
return rotate_design(angle) |
|
return tshirt_img, design_img, design_pos |
|
|
|
with gr.Blocks() as demo: |
|
tshirt_image = gr.Image(label="T-Shirt Image") |
|
design_image = gr.Image(label="Design Image") |
|
rotate_slider = gr.Slider(minimum=0, maximum=360, step=1, label="Rotate Design") |
|
x_coord = gr.Number(label="X Coordinate") |
|
y_coord = gr.Number(label="Y Coordinate") |
|
|
|
tshirt_output = gr.Image(label="Output T-Shirt with Design") |
|
design_output = gr.Image(label="Output Design on T-Shirt") |
|
|
|
|
|
tshirt_image.upload(lambda image: display_image(image_type="tshirt", image=image), inputs=tshirt_image, outputs=tshirt_output) |
|
design_image.upload(lambda image: display_image(image_type="design", image=image), inputs=design_image, outputs=[tshirt_output, design_output]) |
|
|
|
|
|
gr.Button("Add Polygon Vertex").click(lambda x, y: display_image(x=x, y=y, image_type="polygon"), inputs=[x_coord, y_coord], outputs=[tshirt_output, design_output]) |
|
gr.Button("Fit Design").click(lambda: display_image(image_type="fit"), None, [tshirt_output, design_output]) |
|
rotate_slider.change(lambda angle: display_image(angle=angle, image_type="rotate"), inputs=rotate_slider, outputs=[tshirt_output, design_output]) |
|
|
|
demo.launch(share=True) |
|
|