Update app.py

app.py

@@ -8,7 +8,6 @@ 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):
@@ -33,10 +32,8 @@ def generate_design(design_prompt):
     else:
         raise Exception(f"Error generating design: {response.status_code}")
 
-# 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([
@@ -48,27 +45,21 @@ def get_tshirt_mask(image):
     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))
@@ -78,7 +69,6 @@ def get_bounding_box(mask):
     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)
@@ -89,38 +79,30 @@ def visualize_mask(image, mask):
     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")
 
-    # 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
     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)
+    design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
 
-    # Calculate position to center the design on the T-shirt
     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))
     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):
-    # Save debug images
     cloth_image.save("debug_cloth_image.png")
     mask.save("debug_tshirt_mask.png")
 
@@ -136,17 +118,16 @@ def design_tshirt(color_prompt, design_prompt):
     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.")
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
             color_prompt = gr.Textbox(label="Cloth Color", placeholder="E.g., Red, Blue")
             design_prompt = gr.Textbox(label="Design Details", placeholder="E.g., Abstract art, Nature patterns")
             generate_button = gr.Button("Generate T-Shirt")
-        with gr.Column(scale=4):
-            output_image = gr.Image(label="Final T-Shirt Design", type="pil", height=500)
+        with gr.Column():
+            output_image = gr.Image(label="Final T-Shirt Design")
 
     generate_button.click(
         design_tshirt,