Update app.py

app.py

@@ -1,126 +1,72 @@
 import cv2
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
-import torch
-import torchvision.transforms as T
-from torchvision.models import resnet50
-from scipy.ndimage import gaussian_filter
 import gradio as gr
 
-class TextClothBlender:
-    def __init__(self, font_path: str):
-        self.font_path = font_path
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model = resnet50(pretrained=True).to(self.device)
-        self.model.eval()
-
-    def preprocess_cloth_image(self, cloth_image_path: str):
-        # Load the cloth image
-        cloth_image = cv2.imread(cloth_image_path)
-        gray_image = cv2.cvtColor(cloth_image, cv2.COLOR_BGR2GRAY)
-
-        # Detect texture using edge detection
-        edges = cv2.Canny(gray_image, 50, 150)
-
-        # Extract features using ResNet
-        preprocess = T.Compose([
-            T.ToPILImage(),
-            T.Resize((224, 224)),
-            T.ToTensor(),
-            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
-        input_tensor = preprocess(cloth_image).unsqueeze(0).to(self.device)
-        with torch.no_grad():
-            features = self.model(input_tensor)
-
-        return cloth_image, edges, features
-
-    def render_text(self, text: str, font_size: int, text_color: tuple):
-        # Create a blank image for text
-        font = ImageFont.truetype(self.font_path, font_size)
-        text_size = font.getsize(text)
-        text_image = Image.new('RGBA', text_size, (255, 255, 255, 0))
-        draw = ImageDraw.Draw(text_image)
-        draw.text((0, 0), text, font=font, fill=text_color)
-        
-        return np.array(text_image)
-
-    def apply_perspective_transform(self, text_image: np.ndarray, cloth_image_shape: tuple):
-        # Define points for perspective transformation
-        h, w, _ = cloth_image_shape
-        src_points = np.float32([[0, 0], [text_image.shape[1], 0], [0, text_image.shape[0]], [text_image.shape[1], text_image.shape[0]]])
-        dst_points = np.float32([[50, 50], [w - 50, 30], [50, h - 100], [w - 50, h - 120]])
-
-        matrix = cv2.getPerspectiveTransform(src_points, dst_points)
-        warped_text = cv2.warpPerspective(text_image, matrix, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
-
-        return warped_text
-
-    def blend_text_with_cloth(self, cloth_image: np.ndarray, text_image: np.ndarray, edges: np.ndarray):
-        # Convert cloth and text images to the same size
-        h, w, _ = cloth_image.shape
-        text_resized = cv2.resize(text_image, (w, h), interpolation=cv2.INTER_AREA)
-
-        # Convert text to grayscale for masking
-        text_gray = cv2.cvtColor(text_resized, cv2.COLOR_RGBA2GRAY)
-        _, text_mask = cv2.threshold(text_gray, 1, 255, cv2.THRESH_BINARY)
-
-        # Apply displacement mapping using the edges
-        displace_map = gaussian_filter(edges, sigma=5)
-        displaced_text = cv2.addWeighted(text_resized, 0.5, displace_map[..., None], 0.5, 0)
-
-        # Blend text and cloth using overlay mode
-        blended = cv2.addWeighted(cloth_image, 0.7, displaced_text[..., :3], 0.3, 0)
-
-        return blended
-
-    def refine_output(self, blended_image: np.ndarray):
-        # Apply Gaussian blur for smooth edges
-        refined = cv2.GaussianBlur(blended_image, (5, 5), 0)
-        return refined
-
-    def process(self, cloth_image_path: str, text: str, font_size: int, text_color: tuple):
-        # Step 1: Preprocess the cloth image
-        cloth_image, edges, _ = self.preprocess_cloth_image(cloth_image_path)
-
-        # Step 2: Render the text
-        text_image = self.render_text(text, font_size, text_color)
-
-        # Step 3: Apply perspective transformation
-        warped_text = self.apply_perspective_transform(text_image, cloth_image.shape)
-
-        # Step 4: Blend text with the cloth
-        blended_image = self.blend_text_with_cloth(cloth_image, warped_text, edges)
-
-        # Step 5: Refine the output
-        final_image = self.refine_output(blended_image)
-
-        return final_image
-
-# Define the Gradio interface
-def blend_text_on_cloth(cloth_image, text, font_size, text_color):
-    font_path = "path/to/font.ttf"  # Ensure the font file exists in the deployed environment
-    blender = TextClothBlender(font_path)
-    text_color = tuple(map(int, text_color.strip('()').split(',')))  # Convert string to tuple
-    cloth_image_path = "temp_cloth_image.jpg"
-
-    # Save the uploaded cloth image temporarily
-    cv2.imwrite(cloth_image_path, cv2.cvtColor(np.array(cloth_image), cv2.COLOR_RGB2BGR))
-    
-    # Process the image
-    result = blender.process(cloth_image_path, text, int(font_size), text_color)
-    return result[:, :, ::-1]  # Convert BGR to RGB for display
-
-iface = gr.Interface(
-    fn=blend_text_on_cloth,
+
+# Function to detect fabric texture and generate a displacement map
+def generate_displacement_map(image):
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    laplacian = cv2.Laplacian(gray, cv2.CV_32F)  # Detect edges and texture
+    displacement_map = cv2.normalize(laplacian, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)
+    return displacement_map
+
+
+# Function to blend text with the cloth image
+def blend_text_with_cloth(cloth_image, text, font_size=32, font_color=(255, 0, 0), x=50, y=50):
+    # Convert cloth image to BGR for OpenCV processing
+    cloth_image = np.array(cloth_image)
+    cloth_bgr = cv2.cvtColor(cloth_image, cv2.COLOR_RGBA2BGRA)
+
+    # Generate displacement map
+    displacement_map = generate_displacement_map(cloth_bgr)
+
+    # Render text
+    img_pil = Image.fromarray(cloth_image)
+    draw = ImageDraw.Draw(img_pil)
+
+    # Load a font (use a default font if custom font isn't available)
+    try:
+        font = ImageFont.truetype("arial.ttf", font_size)
+    except:
+        font = ImageFont.load_default()
+
+    draw.text((x, y), text, font=font, fill=font_color)
+    rendered_text = np.array(img_pil)
+
+    # Convert rendered text image to BGR
+    text_bgr = cv2.cvtColor(rendered_text, cv2.COLOR_RGBA2BGRA)
+
+    # Blend text with cloth image using displacement map
+    blended = cv2.addWeighted(cloth_bgr, 0.8, text_bgr, 0.5, 0)
+    final_image = cv2.cvtColor(blended, cv2.COLOR_BGRA2RGBA)
+
+    return Image.fromarray(final_image)
+
+
+# Gradio function
+def process_image(image, text, font_size, font_color, x, y):
+    font_color = tuple(map(int, font_color.strip("()").split(",")))  # Convert string to tuple
+    result = blend_text_with_cloth(image, text, font_size, font_color, x, y)
+    return result
+
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=process_image,
     inputs=[
         gr.Image(type="pil", label="Upload Cloth Image"),
-        gr.Textbox(label="Text to Blend"),
-        gr.Slider(10, 100, step=1, label="Font Size"),
-        gr.Textbox(label="Text Color (R,G,B,A)", placeholder="e.g., 255,0,0,255")
+        gr.Textbox(label="Text to Blend", default="Sample Text"),
+        gr.Slider(10, 100, step=2, label="Font Size", default=32),
+        gr.Textbox(label="Font Color (RGB)", default="(255, 0, 0)"),
+        gr.Slider(0, 1000, step=10, label="X Coordinate", default=50),
+        gr.Slider(0, 1000, step=10, label="Y Coordinate", default=50),
     ],
-    outputs=gr.Image(type="numpy", label="Blended Output")
+    outputs=gr.Image(type="pil", label="Blended Output"),
+    title="Text-Cloth Blending",
+    description="Upload a cloth image and add text that blends naturally into the fabric. Adjust parameters like font size, color, and position.",
 )
 
+# Launch the app
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()