Update app.py

app.py

@@ -7,7 +7,6 @@ import torch
 from torchvision import transforms
 from torchvision.models.segmentation import deeplabv3_resnet101
 
-# Load Pretrained DeepLabV3 Model
 model = deeplabv3_resnet101(pretrained=True)
 model.eval()
 
@@ -25,40 +24,33 @@ def segment_clothing(image):
         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]))
     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')
+    displacement_map[mask != 15] = 0  
     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 / 50.0  # Adjust the factor for intensity
+    x_displacement = x + displacement_map / 50.0 
     y_displacement = y + displacement_map / 50.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 overlay_text(image, text, font_size, color, mask):
-    """Overlay text onto the detected clothing region."""
     pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
     draw = ImageDraw.Draw(pil_image)
 
-    # Find the bounding box of the mask (clothing area)
-    y_indices, x_indices = np.where(mask == 15)  # Class 15 corresponds to 'person' in DeepLabV3
+    y_indices, x_indices = np.where(mask == 15)  
     if len(x_indices) == 0 or len(y_indices) == 0:
         return None, "No clothing region detected."
 
@@ -68,7 +60,6 @@ def overlay_text(image, text, font_size, color, mask):
     clothing_width = x_max - x_min
     clothing_height = y_max - y_min
 
-    # Load font and adjust size dynamically
     font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
     if not os.path.exists(font_path):
         return None, "Font file not found. Please provide a valid font path."
@@ -82,36 +73,34 @@ def overlay_text(image, text, font_size, color, mask):
         font = ImageFont.truetype(font_path, font_size)
         text_width, text_height = font.getbbox(text)[2:]
 
-    # Calculate position to center the text
     text_x = x_min + (clothing_width - text_width) // 2
     text_y = y_min + (clothing_height - text_height) // 2
 
-    # Draw the text on a transparent overlay
     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 + (255,))
+
+    try:
+        rgba_color = tuple(color) + (255,)  
+        text_draw.text((text_x, text_y), text, font=font, fill=rgba_color)
+    except Exception as e:
+        return None, f"Error applying color: {str(e)}"
 
     return text_overlay, None
 
 def process_image(image, text, font_size, color):
     try:
-        # Segment the 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 the text
         text_overlay, error = overlay_text(image, text, font_size, color, mask)
         if error:
             return error
 
-        # Warp text using displacement map
         warped_text = warp_text(image, text_overlay, displacement_map)
 
-        # Blend the warped text back onto the original image
         pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).convert("RGBA")
         final_image = Image.alpha_composite(pil_image, warped_text)
 
@@ -120,14 +109,13 @@ def process_image(image, text, font_size, color):
         print(f"Error processing image: {str(e)}")
         return f"Error: {str(e)}"
 
-# Gradio Interface
 gr.Interface(
     fn=process_image,
     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")  # Default color: black
+        gr.ColorPicker(label="Text Color", value="#000000")
     ],
     outputs=gr.Image(type="pil", label="Final Image with Warped Text"),
     title="Warped Text Overlay on Clothing",