app.py

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import torch
import torchvision.transforms as transforms
from skimage.filters import sobel
from skimage.restoration import denoise_tv_chambolle
from scipy.interpolate import Rbf


# Function to estimate a normal map from cloth texture
def estimate_normal_map(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    sobel_x = sobel(gray)
    sobel_y = sobel(gray)
    
    normal_map = np.stack([sobel_x, sobel_y, np.ones_like(sobel_x)], axis=-1)
    normal_map /= np.linalg.norm(normal_map, axis=-1, keepdims=True)

    return (normal_map * 255).astype(np.uint8)


def apply_tps_warping(design, normal_map):
    # Resize normal map to match the design size
    normal_map = cv2.resize(normal_map, (design.shape[1], design.shape[0]))

    h, w = design.shape[:2]
    x, y = np.meshgrid(np.arange(w), np.arange(h))

    # Generate warp offsets from normal map
    control_x = x + (normal_map[:, :, 0] - 128) * 0.5
    control_y = y + (normal_map[:, :, 1] - 128) * 0.5

    # Apply Radial Basis Function (RBF) interpolation
    rbf_x = Rbf(x.flatten(), y.flatten(), control_x.flatten(), function='thin_plate')
    rbf_y = Rbf(x.flatten(), y.flatten(), control_y.flatten(), function='thin_plate')

    warped_x = rbf_x(x, y).astype(np.float32)
    warped_y = rbf_y(x, y).astype(np.float32)

    # Warp the design
    warped_design = cv2.remap(design, warped_x, warped_y, interpolation=cv2.INTER_LINEAR)

    return warped_design


# Function to blend design onto the cloth using Poisson Editing
def blend_design_cloth(cloth, design, x=50, y=50):
    cloth_bgr = np.array(cloth)
    design_bgr = np.array(design)

    normal_map = estimate_normal_map(cloth_bgr)

    # Resize design to fit the center of the cloth
    design_resized = cv2.resize(design_bgr, (cloth_bgr.shape[1] // 2, cloth_bgr.shape[0] // 5))
    
    # Convert to grayscale and create a mask
    design_gray = cv2.cvtColor(design_resized, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(design_gray, 1, 255, cv2.THRESH_BINARY)

    # Warp design using normal map
    warped_design = apply_tps_warping(design_resized, normal_map)

    # Blend using Poisson seamless cloning
    center = (x + design_resized.shape[1] // 2, y + design_resized.shape[0] // 2)
    blended = cv2.seamlessClone(warped_design, cloth_bgr, mask, center, cv2.MIXED_CLONE)

    return Image.fromarray(blended)


# Gradio function
def process_image(cloth_image, design_image, x=50, y=50):
    # Blend design onto cloth
    result = blend_design_cloth(cloth_image, design_image, x, y)
    return result


# Gradio Interface
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Cloth Image"),
        gr.Image(type="pil", label="Upload Design"),
        gr.Slider(0, 1000, step=10, label="X Coordinate", value=50),
        gr.Slider(0, 1000, step=10, label="Y Coordinate", value=50),
    ],
    outputs=gr.Image(type="pil", label="Blended Output"),
    title="Advanced Cloth Design Blending",
    description="Upload a cloth image and a design to blend them naturally using advanced warping & Poisson blending.",
)

# Launch the app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)