Update app.py

app.py

@@ -1,48 +1,110 @@
-import torch
-import gradio as gr
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-from PIL import Image
 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
+from scipy.ndimage import map_coordinates
 
-# Load ControlNet pre-trained model for depth-based warping
-controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11f1p_sd15_depth", torch_dtype=torch.float16)
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    controlnet=controlnet,
-    torch_dtype=torch.float16
-).to("cuda")
-
-# Function to generate depth map of the cloth
-def generate_depth_map(image_path):
-    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
-    depth_map = cv2.Laplacian(image, cv2.CV_64F)
-    depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
-    return Image.fromarray(depth_map)
-
-# Function to blend design onto fabric using ControlNet
-def blend_design_on_cloth(fabric_image, design_image, prompt="T-shirt with embedded design"):
-    # Generate depth map for fabric
-    depth_map = generate_depth_map(fabric_image)
-
-    # Generate realistic blended output
-    result = pipe(prompt=prompt, image=fabric_image, control_image=depth_map, num_inference_steps=30).images[0]
-    return result
 
-# Gradio interface
-def process_image(fabric_image, design_image):
-    result = blend_design_on_cloth(fabric_image, design_image)
+# Function to estimate a normal map from the 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 = normal_map / np.linalg.norm(normal_map, axis=-1, keepdims=True)
+    return (normal_map * 255).astype(np.uint8)
+
+
+# Function to apply Thin Plate Spline (TPS) warping
+def apply_tps_warping(text_image, normal_map):
+    h, w = text_image.shape[:2]
+    x, y = np.meshgrid(np.arange(w), np.arange(h))
+    
+    # Generate control points from the normal map
+    control_x = x + (normal_map[:, :, 0] - 128) * 0.5
+    control_y = y + (normal_map[:, :, 1] - 128) * 0.5
+
+    # Interpolate using Radial Basis Function (RBF)
+    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)
+
+    # Apply warping
+    warped_text = cv2.remap(text_image, warped_x, warped_y, interpolation=cv2.INTER_LINEAR)
+    
+    return warped_text
+
+
+# Function to blend text onto cloth using Poisson editing
+def blend_text_cloth(cloth, text, x=50, y=50):
+    cloth_bgr = np.array(cloth)
+    text_bgr = np.array(text)
+
+    normal_map = estimate_normal_map(cloth_bgr)
+
+    # Resize text to fit on cloth
+    text_resized = cv2.resize(text_bgr, (cloth_bgr.shape[1] // 2, cloth_bgr.shape[0] // 5))
+    
+    # Convert to grayscale and create a mask
+    text_gray = cv2.cvtColor(text_resized, cv2.COLOR_BGR2GRAY)
+    _, mask = cv2.threshold(text_gray, 1, 255, cv2.THRESH_BINARY)
+
+    # Warp text using normal map
+    warped_text = apply_tps_warping(text_resized, normal_map)
+
+    # Blend the text using Poisson editing
+    center = (x + text_resized.shape[1] // 2, y + text_resized.shape[0] // 2)
+    blended = cv2.seamlessClone(warped_text, cloth_bgr, mask, center, cv2.MIXED_CLONE)
+
+    return Image.fromarray(blended)
+
+
+# Gradio function
+def process_image(cloth_image, text, font_size=32, font_color=(255, 0, 0), x=50, y=50):
+    # Convert font color input
+    font_color = tuple(map(int, font_color.strip("()").split(",")))
+
+    # Create a blank image with text
+    text_img = Image.new('RGB', (400, 200), (0, 0, 0, 0))
+    draw = ImageDraw.Draw(text_img)
+
+    try:
+        font = ImageFont.truetype("arial.ttf", font_size)
+    except:
+        font = ImageFont.load_default()
+
+    draw.text((50, 50), text, font=font, fill=font_color)
+    text_img = np.array(text_img)
+
+    # Blend text onto cloth
+    result = blend_text_cloth(cloth_image, text_img, 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/Text Image")],
+    inputs=[
+        gr.Image(type="pil", label="Upload Cloth Image"),
+        gr.Textbox(label="Text to Blend", value="Sample Text"),
+        gr.Slider(10, 100, step=2, label="Font Size", value=32),
+        gr.Textbox(label="Font Color (RGB)", value="(255, 0, 0)"),
+        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="AI-Powered Text & Design Blending on Cloth",
-    description="Upload a T-shirt or fabric image and a design. AI will blend it naturally into the fabric!",
+    title="Advanced Text-Cloth Blending",
+    description="Upload a cloth image and blend text naturally using advanced warping & blending techniques.",
 )
 
-# Launch app for Hugging Face Spaces
+# Launch the app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch(server_name="0.0.0.0", server_port=7860)