Update app.py

app.py

@@ -29,51 +29,35 @@ def segment_clothing(image):
     mask = cv2.resize(output_predictions, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
     return mask
 
-def generate_displacement_map(image, mask):
-    """Generate a displacement map from the clothing region."""
-    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  # Apply mask (class 15 corresponds to 'person')
-    return displacement_map
-
-def warp_text(image, text_overlay, displacement_map):
-    """Warp the text overlay based on the displacement map."""
-    text_overlay_array = np.array(text_overlay)
-    displacement_map = cv2.GaussianBlur(displacement_map, (15, 15), 0)
-
-    # Create an x, y distortion map
-    h, w = displacement_map.shape
-    x, y = np.meshgrid(np.arange(w), np.arange(h))
-    x_displacement = x + displacement_map / 100.0
-    y_displacement = y + displacement_map / 100.0
-
-    # Warp text overlay using remap
-    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 apply_displacement_map(text_img, clothing_img, strength=20):
+    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 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):
-    """Overlay text onto the detected clothing region."""
     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
-
-    # Load font and create transparent overlay for text
     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)
@@ -82,28 +66,17 @@ def overlay_text(image, text, font_path, font_size, color, mask, manual_coords=N
 
 def process_image(image, text, font_size, color, font_path, placement_x, placement_y):
     try:
-        # Segment clothing using DeepLabV3
         mask = segment_clothing(image)
         if mask.sum() == 0:
             return "No clothing detected. Try another image."
-
-        # Generate displacement map
-        displacement_map = generate_displacement_map(image, mask)
-
-        # Overlay text with manual coordinates (if provided)
         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
-
-        # Warp text using displacement map
-        warped_text = warp_text(image, text_overlay, displacement_map)
-
-        # Blend warped text with the original image
-        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
-        final_image = Image.alpha_composite(pil_image, warped_text).convert("RGB")
-
-        return final_image
+        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)}"