app.py

import os
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import torch
from torchvision import transforms
from torchvision.models.segmentation import deeplabv3_resnet101

model = deeplabv3_resnet101(pretrained=True)
model.eval()

def segment_clothing(image):
    """Segment clothing region using DeepLabV3."""
    preprocess = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((512, 512)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
    input_tensor = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        output = model(input_tensor)['out'][0]
    output_predictions = output.argmax(0).byte().cpu().numpy()

    mask = cv2.resize(output_predictions, (image.shape[1], image.shape[0]))
    return mask

def generate_displacement_map(image, mask):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (15, 15), 0)
    displacement_map = cv2.normalize(blurred, None, 0, 255, cv2.NORM_MINMAX)
    displacement_map[mask != 15] = 0  
    return displacement_map

def warp_text(image, text_overlay, displacement_map):
    text_overlay_array = np.array(text_overlay)
    displacement_map = cv2.GaussianBlur(displacement_map, (15, 15), 0)

    h, w = displacement_map.shape
    x, y = np.meshgrid(np.arange(w), np.arange(h))
    x_displacement = x + displacement_map / 50.0 
    y_displacement = y + displacement_map / 50.0

    warped = cv2.remap(text_overlay_array, x_displacement.astype(np.float32), y_displacement.astype(np.float32), interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
    return Image.fromarray(warped)

def overlay_text(image, text, font_size, color, mask):
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
    draw = ImageDraw.Draw(pil_image)

    y_indices, x_indices = np.where(mask == 15)  
    if len(x_indices) == 0 or len(y_indices) == 0:
        return None, "No clothing region detected."

    x_min, x_max = x_indices.min(), x_indices.max()
    y_min, y_max = y_indices.min(), y_indices.max()

    clothing_width = x_max - x_min
    clothing_height = y_max - y_min

    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    if not os.path.exists(font_path):
        return None, "Font file not found. Please provide a valid font path."

    font = ImageFont.truetype(font_path, font_size)
    text_width, text_height = font.getbbox(text)[2:]
    while text_width > clothing_width or text_height > clothing_height:
        font_size -= 1
        if font_size <= 5:
            return None, "Text too large to fit on the clothing. Try smaller text or font size."
        font = ImageFont.truetype(font_path, font_size)
        text_width, text_height = font.getbbox(text)[2:]

    text_x = x_min + (clothing_width - text_width) // 2
    text_y = y_min + (clothing_height - text_height) // 2

    text_overlay = Image.new("RGBA", pil_image.size, (255, 255, 255, 0))
    text_draw = ImageDraw.Draw(text_overlay)

    try:
        rgba_color = tuple(color) + (255,)  
        text_draw.text((text_x, text_y), text, font=font, fill=rgba_color)
    except Exception as e:
        return None, f"Error applying color: {str(e)}"

    return text_overlay, None

def process_image(image, text, font_size, color):
    try:
        mask = segment_clothing(image)
        if mask.sum() == 0:
            return "No clothing detected. Try another image."

        displacement_map = generate_displacement_map(image, mask)

        text_overlay, error = overlay_text(image, text, font_size, color, mask)
        if error:
            return error

        warped_text = warp_text(image, text_overlay, displacement_map)

        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
        final_image = Image.alpha_composite(pil_image, warped_text)

        return final_image
    except Exception as e:
        print(f"Error processing image: {str(e)}")
        return f"Error: {str(e)}"

gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Clothing Image"),
        gr.Textbox(label="Enter Text"),
        gr.Slider(10, 150, step=5, label="Font Size"),
        gr.ColorPicker(label="Text Color", value="#000000")
    ],
    outputs=gr.Image(type="pil", label="Final Image with Warped Text"),
    title="Warped Text Overlay on Clothing",
    description="Upload a clothing image and add warped text that conforms to folds and curves."
).launch()