Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from PIL import Image, ImageFilter
+from PIL import Image, ImageDraw, ImageFilter
 import requests
 from io import BytesIO
 import torch
@@ -11,121 +11,152 @@ import cv2
 # AI model repo for design generation
 repo = "artificialguybr/TshirtDesignRedmond-V2"
 
-# ---------- Step 1: Generate Plain T-shirt ----------
 def generate_cloth(color_prompt):
-    prompt = f"A plain {color_prompt.strip().lower()} colored T-shirt hanging on a plain wall."
+    prompt = f"A plain {color_prompt} colored T-shirt hanging on a plain wall."
     api_url = f"https://api-inference.huggingface.co/models/{repo}"
+    headers = {}
     payload = {"inputs": prompt, "parameters": {"num_inference_steps": 30}}
-    
-    try:
-        response = requests.post(api_url, json=payload, timeout=60)
-        response.raise_for_status()
-        image_data = response.content
-        return Image.open(BytesIO(image_data)).convert("RGB")
-    except Exception as e:
-        print(f"Error generating cloth: {e}")
-        raise Exception("Failed to generate the T-shirt. Check your API key and internet connection.")
+    response = requests.post(api_url, headers=headers, json=payload)
+    if response.status_code == 200:
+        return Image.open(BytesIO(response.content)).convert("RGB")
+    else:
+        raise Exception(f"Error generating cloth: {response.status_code}")
 
-# ---------- Step 2: Generate Design ----------
 def generate_design(design_prompt):
-    prompt = f"A bold {design_prompt.strip().lower()} design with vibrant colors, highly detailed."
+    prompt = f"A bold {design_prompt} design with vibrant colors, highly detailed."
     api_url = f"https://api-inference.huggingface.co/models/{repo}"
+    headers = {}
     payload = {"inputs": prompt, "parameters": {"num_inference_steps": 30}}
-    
-    try:
-        response = requests.post(api_url, json=payload, timeout=60)
-        response.raise_for_status()
-        image_data = response.content
-        return Image.open(BytesIO(image_data)).convert("RGBA")
-    except Exception as e:
-        print(f"Error generating design: {e}")
-        raise Exception("Failed to generate the design. Check your API key and internet connection.")
+    response = requests.post(api_url, headers=headers, json=payload)
+    if response.status_code == 200:
+        return Image.open(BytesIO(response.content)).convert("RGBA")
+    else:
+        raise Exception(f"Error generating design: {response.status_code}")
 
-# ---------- Step 3: Load Segmentation Model ----------
-try:
-    segmentation_model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
-except Exception as e:
-    print(f"Error loading segmentation model: {e}")
-    raise Exception("Failed to load the segmentation model. Check your PyTorch installation.")
+# 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):
-    transform = T.Compose([
-        T.Resize(520),
+    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])
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
-    input_tensor = transform(image).unsqueeze(0)
-    
+    input_tensor = preprocess(image).unsqueeze(0)
     with torch.no_grad():
-        output = segmentation_model(input_tensor)['out'][0]
+        output = segmentation_model(input_tensor)["out"][0]
     
+    # Extract T-shirt mask (class 15 in COCO dataset)
     mask = output.argmax(0).byte().cpu().numpy()
-    mask_img = Image.fromarray((mask == 0).astype('uint8') * 255).convert("L")
-    
-    return mask_img.filter(ImageFilter.GaussianBlur(10)).resize(image.size)
+    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)
 
-# ---------- Step 4: Fit Design to T-shirt Area ----------
-def fit_design_to_tshirt(design_image, mask):
+# Post-process mask to improve quality
+def post_process_mask(mask):
+    # Convert mask to NumPy array
     mask_np = np.array(mask)
-    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     
-    if contours:
-        # Largest contour represents the T-shirt
-        c = max(contours, key=cv2.contourArea)
-        x, y, w, h = cv2.boundingRect(c)
-
-        # Log measurements of the T-shirt area
-        print(f"T-shirt Area Detected at (x: {x}, y: {y}), Width: {w}px, Height: {h}px")
-
-        # Resize design to fit within the T-shirt area
-        design_resized = design_image.resize((w, h), Image.Resampling.LANCZOS)
-        
-        # Overlay design on a transparent canvas of the same size
-        canvas = Image.new("RGBA", mask.size, (0, 0, 0, 0))
-        canvas.paste(design_resized, (x, y), design_resized)
-        
-        return canvas
-    else:
-        print("No contour found on the mask.")
-        return design_image
-        
-# ---------- Step 5: Overlay Design ----------
+    # 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
+
+# Overlay design on the T-shirt
 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)
-    fitted_design = fit_design_to_tshirt(design_image, mask)
-    return Image.alpha_composite(cloth_image.convert("RGBA"), fitted_design)
+    
+    # Extract bounding box for precise placement
+    bbox = get_bounding_box(mask)
+    tshirt_width = bbox[2] - bbox[0]
+    tshirt_height = bbox[3] - bbox[1]
+
+    # Resize the design to fit the T-shirt
+    design_width = int(tshirt_width * 0.6)
+    design_height = int(tshirt_height * 0.6)
+    resized_design = design_image.resize((design_width, design_height))
+
+    # Position the design in the center of the T-shirt
+    design_position = (
+        bbox[0] + (tshirt_width - design_width) // 2,
+        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(resized_design, design_position, resized_design)
+
+    # Mask the design to the T-shirt area
+    masked_design = Image.composite(transparent_layer, Image.new("RGBA", cloth_image.size), mask)
+
+    # Combine the cloth image with the masked design
+    final_image = Image.alpha_composite(cloth_image, masked_design)
+    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")
 
-# ---------- Step 6: Final Design Pipeline ----------
 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)
+        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 cloth_image, design_image, final_image
+        return final_image
     except Exception as e:
-        print(f"Error: {e}")
-        return None, None, None
+        raise Exception(f"Error in design process: {str(e)}")
 
-# ---------- Step 7: Gradio Interface ----------
+# Gradio UI
 with gr.Blocks() as interface:
-    gr.Markdown("# **AI T-Shirt Designer (Debug Mode)**")
-    gr.Markdown("Generate a T-shirt with a design that perfectly fits inside the T-shirt boundaries.")
-
+    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():
             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():
-            plain_tshirt_output = gr.Image(label="Generated Plain T-Shirt")
-            design_output = gr.Image(label="Generated Design")
-            final_output = gr.Image(label="Final Fitted T-Shirt Design")
+            output_image = gr.Image(label="Final T-Shirt Design")
 
     generate_button.click(
         design_tshirt,
         inputs=[color_prompt, design_prompt],
-        outputs=[plain_tshirt_output, design_output, final_output],
+        outputs=output_image,
     )
 
-interface.launch(debug=True, share=True)
+interface.launch(debug=True)