Update app.py

app.py

@@ -1,100 +1,86 @@
-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
-
-# Load Pretrained DeepLabV3 Model
-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()
-
-    # Scale back to original size
-    mask = cv2.resize(output_predictions, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
-    return mask
+from PIL import Image, ImageDraw, ImageFont
+import io
+
+def generate_text_image(text, style, width, height):
+    """Generate an image of the input text with the selected style."""
+    font_path = {
+        "Bold": "arialbd.ttf",
+        "Italic": "ariali.ttf",
+        "Graffiti": "graffiti.ttf",  # Provide custom font if needed
+        "Calligraphy": "calligraphy.ttf"  # Provide custom font if needed
+    }
+
+    img = Image.new("RGBA", (width, height), (255, 255, 255, 0))
+    draw = ImageDraw.Draw(img)
+
+    try:
+        font = ImageFont.truetype(font_path[style], size=80)
+    except:
+        font = ImageFont.load_default()
+
+    text_width, text_height = draw.textsize(text, font=font)
+    position = ((width - text_width) // 2, (height - text_height) // 2)
+
+    draw.text(position, text, font=font, fill=(0, 0, 0, 255))
+
+    return img
 
 def apply_displacement_map(text_img, clothing_img, strength=20):
+    """Apply displacement map to blend text onto clothing."""
     gray = cv2.cvtColor(clothing_img, cv2.COLOR_BGR2GRAY)
     grad_x = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=5)
     grad_y = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=5)
+
     grad_x = cv2.normalize(grad_x, None, 0, 1, cv2.NORM_MINMAX)
     grad_y = cv2.normalize(grad_y, None, 0, 1, cv2.NORM_MINMAX)
+
     displacement_map = np.zeros_like(clothing_img, dtype=np.float32)
     displacement_map[:, :, 0] = grad_x * strength
     displacement_map[:, :, 1] = grad_y * strength
-    text_warped = cv2.remap(text_img, displacement_map[:, :, 0].astype(np.float32), displacement_map[:, :, 1].astype(np.float32), interpolation=cv2.INTER_LINEAR)
+
+    text_warped = cv2.remap(text_img, 
+                            displacement_map[:, :, 0].astype(np.float32),
+                            displacement_map[:, :, 1].astype(np.float32), 
+                            interpolation=cv2.INTER_LINEAR)
+
     return text_warped
 
-def blend_images(foreground, background, alpha=0.7):
-    foreground = foreground.astype(float)
-    background = background.astype(float)
-    blended = cv2.addWeighted(foreground, alpha, background, 1 - alpha, 0)
-    return blended.astype(np.uint8)
-
-def overlay_text(image, text, font_path, font_size, color, mask, manual_coords=None):
-    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
-    y_indices, x_indices = np.where(mask == 15)
-    if len(x_indices) == 0 or len(y_indices) == 0:
-        return None, "No clothing region detected."
-    if manual_coords:
-        text_x, text_y = manual_coords
-    else:
-        x_min, x_max = x_indices.min(), x_indices.max()
-        y_min, y_max = y_indices.min(), y_indices.max()
-        text_x, text_y = x_min, y_min
-    font = ImageFont.truetype(font_path, font_size)
-    text_overlay = Image.new("RGBA", pil_image.size, (255, 255, 255, 0))
-    text_draw = ImageDraw.Draw(text_overlay)
-    text_draw.text((text_x, text_y), text, font=font, fill=color)
-    return text_overlay, None
-
-def process_image(image, text, font_size, color, font_path, placement_x, placement_y):
-    try:
-        mask = segment_clothing(image)
-        if mask.sum() == 0:
-            return "No clothing detected. Try another image."
-        manual_coords = (placement_x, placement_y)
-        text_overlay, error = overlay_text(image, text, font_path, font_size, color, mask, manual_coords)
-        if error:
-            return error
-        text_img = np.array(text_overlay.convert("RGB"))
-        text_warped = apply_displacement_map(text_img, image)
-        blended_image = blend_images(text_warped, image)
-        return Image.fromarray(blended_image)
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Gradio Interface
-font_options = ["/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", "/path/to/another/font.ttf"]
-
-gr.Interface(
-    fn=process_image,
+def overlay_text_on_clothing(clothing_image, text_input, style, strength=20, alpha=0.7):
+    """Blend generated text onto the clothing image."""
+    clothing_img = cv2.imdecode(np.frombuffer(clothing_image, np.uint8), cv2.IMREAD_COLOR)
+
+    # Generate text image dynamically
+    text_img_pil = generate_text_image(text_input, style, clothing_img.shape[1], clothing_img.shape[0])
+    text_img = cv2.cvtColor(np.array(text_img_pil), cv2.COLOR_RGBA2BGRA)
+
+    alpha_channel = text_img[:, :, 3] / 255.0
+    text_img = text_img[:, :, :3]
+
+    text_warped = apply_displacement_map(text_img, clothing_img, strength)
+
+    for c in range(3):
+        clothing_img[:, :, c] = (1 - alpha_channel) * clothing_img[:, :, c] + alpha_channel * text_warped[:, :, c]
+
+    _, buffer = cv2.imencode('.jpg', clothing_img)
+    return buffer.tobytes()
+
+interface = gr.Interface(
+    fn=overlay_text_on_clothing,
     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"),
-        gr.Dropdown(choices=font_options, label="Select Font", value=font_options[0]),
-        gr.Number(label="Text Placement X", value=0),
-        gr.Number(label="Text Placement Y", value=0),
+        gr.Image(type="binary", label="Upload Clothing Image", tool="editor"),
+        gr.Textbox(label="Enter Text for Design"),
+        gr.Radio(["Bold", "Italic", "Graffiti", "Calligraphy"], label="Select Style", value="Bold"),
+        gr.Slider(10, 50, step=5, default=20, label="Displacement Strength"),
+        gr.Slider(0.1, 1.0, step=0.1, default=0.7, label="Alpha Blending")
     ],
-    outputs=gr.Image(type="pil", label="Final Image with Warped Text"),
-    title="Warped Text Overlay on Clothing",
-    description="Upload a clothing image and customize text placement, font, and style."
-).launch()
+    outputs=gr.Image(type="pil", label="Final Design"),
+    title="AI-Powered Clothing Text Overlay",
+    description="Upload a clothing image, enter a text design, and select a style to blend them onto clothing with displacement mapping.",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    interface.launch()