Update app.py

app.py

@@ -1,107 +1,126 @@
 import cv2
 import numpy as np
-import gradio as gr
 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
 
-
-def generate_text_image(text, style, width, height, font_scale):
-    """Generate an image of the input text with the selected style and scale."""
-    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)
-
-    # Dynamically calculate font size based on user-defined font scale
-    font_size = int(min(width, height) * font_scale)
-    try:
-        font = ImageFont.truetype(font_path[style], size=font_size)
-    except:
-        font = ImageFont.load_default()
-
-    text_bbox = draw.textbbox((0, 0), text, font=font)
-    text_width = text_bbox[2] - text_bbox[0]
-    text_height = text_bbox[3] - text_bbox[1]
-
-    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
-    )
-
-    return text_warped
-
-
-def overlay_text_on_clothing(clothing_image, text_input, style, font_scale, strength=20, alpha=0.7, photo=None, blend_alpha=0.5):
-    """Blend generated text and optional photo onto the clothing image."""
-    clothing_img = cv2.imread(clothing_image)  # Read image from file path
-
-    # Generate text image dynamically
-    text_img_pil = generate_text_image(text_input, style, clothing_img.shape[1], clothing_img.shape[0], font_scale)
-    text_img = cv2.cvtColor(np.array(text_img_pil), cv2.COLOR_RGBA2BGRA)
-
-    alpha_channel = text_img[:, :, 3] / 255.0
-    text_img = text_img[:, :, :3]
-
-    # Apply displacement map
-    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]
-
-    # Blend optional photo if provided
-    if photo:
-        photo_img = Image.open(photo).resize((clothing_img.shape[1], clothing_img.shape[0]))
-        photo_img = cv2.cvtColor(np.array(photo_img), cv2.COLOR_RGBA2BGR)
-        clothing_img = cv2.addWeighted(clothing_img, 1 - blend_alpha, photo_img, blend_alpha, 0)
-
-    # Convert the final blended image back to a PIL Image and return
-    blended_img = Image.fromarray(cv2.cvtColor(clothing_img, cv2.COLOR_BGR2RGB))
-    return blended_img
-
-
-interface = gr.Interface(
-    fn=overlay_text_on_clothing,
+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,
     inputs=[
-        gr.Image(type="filepath", label="Upload Clothing Image", interactive=True),
-        gr.Textbox(label="Enter Text for Design"),
-        gr.Radio(["Bold", "Italic", "Graffiti", "Calligraphy"], label="Select Style", value="Bold"),
-        gr.Slider(0.05, 0.5, step=0.05, value=0.1, label="Text Size (Scale)"),
-        gr.Slider(10, 50, step=5, value=20, label="Displacement Strength"),
-        gr.Slider(0.1, 1.0, step=0.1, value=0.7, label="Alpha Blending"),
-        gr.Image(type="filepath", label="Optional Photo for Blending", interactive=True),
-        gr.Slider(0.0, 1.0, step=0.1, value=0.5, label="Photo Blend Alpha")
+        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")
     ],
-    outputs=gr.Image(type="pil", label="Final Design"),
-    title="AI-Powered Clothing Text and Photo Overlay",
-    description="Upload a clothing image, enter a text design, and select a style. Adjust text size, blending, and optional photo overlay for enhanced designs.",
-    allow_flagging="never"
+    outputs=gr.Image(type="numpy", label="Blended Output")
 )
 
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()