Update app.py

app.py

@@ -2,10 +2,11 @@ import gradio as gr
 from PIL import Image, ImageDraw, ImageFilter
 import requests
 from io import BytesIO
-import numpy as np
 import torch
 import torchvision.transforms as T
 from torchvision import models
+import numpy as np
+import cv2
 
 # AI model repo for design generation
 repo = "artificialguybr/TshirtDesignRedmond-V2"
@@ -32,21 +33,83 @@ def generate_design(design_prompt):
     else:
         raise Exception(f"Error generating design: {response.status_code}")
 
-def overlay_design_on_tshirt(cloth_image, design_image):
+# Load pretrained DeepLabV3 model for T-shirt segmentation
+segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
+
+# Apply segmentation to extract T-shirt mask
+def get_tshirt_mask(image):
+    image = image.convert("RGB")  # Ensure 3 channels
+    preprocess = T.Compose([
+        T.Resize((520, 520)),  # Resize to avoid distortion
+        T.ToTensor(),
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    input_tensor = preprocess(image).unsqueeze(0)
+    with torch.no_grad():
+        output = segmentation_model(input_tensor)["out"][0]
+    
+    # Extract T-shirt mask (class 15 in COCO dataset)
+    mask = output.argmax(0).byte().cpu().numpy()
+    raw_mask = Image.fromarray((mask == 15).astype("uint8") * 255)  # Binary mask
+    processed_mask = post_process_mask(raw_mask)  # Apply post-processing
+    return processed_mask.resize(image.size)
+
+# Post-process mask to improve quality
+def post_process_mask(mask):
+    # Convert mask to NumPy array
+    mask_np = np.array(mask)
+    
+    # Morphological operations to refine mask
+    kernel = np.ones((5, 5), np.uint8)
+    mask_np = cv2.dilate(mask_np, kernel, iterations=2)  # Expand mask
+    mask_np = cv2.erode(mask_np, kernel, iterations=1)   # Remove noise
+    
+    # Convert back to PIL image and smooth
+    processed_mask = Image.fromarray(mask_np).filter(ImageFilter.GaussianBlur(3))
+    return processed_mask
+
+# Get bounding box from mask
+def get_bounding_box(mask):
+    mask_np = np.array(mask)
+    coords = np.column_stack(np.where(mask_np > 0))
+    if coords.size == 0:
+        raise Exception("No T-shirt detected in the image.")
+    x_min, y_min = coords.min(axis=0)
+    x_max, y_max = coords.max(axis=0)
+    return (x_min, y_min, x_max, y_max)
+
+# Visualize mask and bounding box on the image for debugging
+def visualize_mask(image, mask):
+    overlay = image.copy().convert("RGBA")
+    draw = ImageDraw.Draw(overlay)
+    bbox = get_bounding_box(mask)
+    draw.rectangle(bbox, outline="red", width=3)  # Draw bounding box
+    blended = Image.blend(image.convert("RGBA"), overlay, alpha=0.5)  # Overlay mask
+    blended.save("debug_visualization.png")  # Save debug image
+    return blended
+
+def overlay_design(cloth_image, design_image):
     # Ensure images are in RGBA mode
     cloth_image = cloth_image.convert("RGBA")
     design_image = design_image.convert("RGBA")
 
-    # Resize the design to fit within 70% of the T-shirt's dimensions
+    # Generate T-shirt mask
+    mask = get_tshirt_mask(cloth_image)
+
+    # Extract bounding box and T-shirt coordinates
+    bbox = get_bounding_box(mask)
+    tshirt_width = bbox[2] - bbox[0]
+    tshirt_height = bbox[3] - bbox[1]
+
+    # Resize the design to fit within 70% of the T-shirt bounding box
     scale_factor = 0.7
-    width, height = cloth_image.size
-    design_width = int(width * scale_factor)
-    design_height = int(height * scale_factor)
-    design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
+    design_width = int(tshirt_width * scale_factor)
+    design_height = int(tshirt_height * scale_factor)
+    design_image = design_image.resize((design_width, design_height), Image.ANTIALIAS)
 
     # Calculate position to center the design on the T-shirt
-    x_offset = (width - design_width) // 2
-    y_offset = (height - design_height) // 2
+    x_offset = bbox[0] + (tshirt_width - design_width) // 2
+    y_offset = bbox[1] + (tshirt_height - design_height) // 2
 
     # Create a transparent layer for the design
     transparent_layer = Image.new("RGBA", cloth_image.size, (0, 0, 0, 0))
@@ -56,11 +119,19 @@ def overlay_design_on_tshirt(cloth_image, design_image):
     final_image = Image.alpha_composite(cloth_image, transparent_layer)
     return final_image
 
+def debug_intermediate_outputs(cloth_image, mask):
+    # Save debug images
+    cloth_image.save("debug_cloth_image.png")
+    mask.save("debug_tshirt_mask.png")
+
 def design_tshirt(color_prompt, design_prompt):
+    cloth_image = generate_cloth(color_prompt)
+    design_image = generate_design(design_prompt)
     try:
-        cloth_image = generate_cloth(color_prompt)
-        design_image = generate_design(design_prompt)
-        final_image = overlay_design_on_tshirt(cloth_image, design_image)
+        mask = get_tshirt_mask(cloth_image)
+        debug_intermediate_outputs(cloth_image, mask)  # Debugging
+        visualize_mask(cloth_image, mask)  # Save visualization
+        final_image = overlay_design(cloth_image, design_image)
         return final_image
     except Exception as e:
         raise Exception(f"Error in design process: {str(e)}")
@@ -83,4 +154,4 @@ with gr.Blocks() as interface:
         outputs=output_image,
     )
 
-interface.launch(debug=True)
+interface.launch(debug=True)