Update app.py

app.py

@@ -2,12 +2,12 @@ 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"
 
 def generate_cloth(color_prompt):
@@ -32,92 +32,40 @@ def generate_design(design_prompt):
     else:
         raise Exception(f"Error generating design: {response.status_code}")
 
-segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
-
-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]
-    
-    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)
-
-def post_process_mask(mask):
-    mask_np = np.array(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
-    
-    processed_mask = Image.fromarray(mask_np).filter(ImageFilter.GaussianBlur(3))
-    return processed_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)
-
-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):
+def overlay_design_on_tshirt(cloth_image, design_image):
+    # Ensure images are in RGBA mode
     cloth_image = cloth_image.convert("RGBA")
     design_image = design_image.convert("RGBA")
 
-    mask = get_tshirt_mask(cloth_image)
-
-    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's dimensions
     scale_factor = 0.7
-    design_width = int(tshirt_width * scale_factor)
-    design_height = int(tshirt_height * scale_factor)
+    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)
 
-    x_offset = bbox[0] + (tshirt_width - design_width) // 2
-    y_offset = bbox[1] + (tshirt_height - design_height) // 2
+    # Calculate position to center the design on the T-shirt
+    x_offset = (width - design_width) // 2
+    y_offset = (height - design_height) // 2
 
+    # Create a transparent layer for the design
     transparent_layer = Image.new("RGBA", cloth_image.size, (0, 0, 0, 0))
     transparent_layer.paste(design_image, (x_offset, y_offset), design_image)
 
+    # Blend the T-shirt with the design
     final_image = Image.alpha_composite(cloth_image, transparent_layer)
     return final_image
 
-def debug_intermediate_outputs(cloth_image, mask):
-    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:
-        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)
+        cloth_image = generate_cloth(color_prompt)
+        design_image = generate_design(design_prompt)
+        final_image = overlay_design_on_tshirt(cloth_image, design_image)
         return final_image
     except Exception as e:
         raise Exception(f"Error in design process: {str(e)}")
 
+# Gradio UI
 with gr.Blocks() as interface:
     gr.Markdown("# **AI Cloth Designer**")
     gr.Markdown("Generate custom T-shirts by specifying a color and adding a design that perfectly fits the T-shirt.")
@@ -135,4 +83,4 @@ with gr.Blocks() as interface:
         outputs=output_image,
     )
 
-interface.launch(debug=True)
+interface.launch(debug=True)