Update app.py

app.py

@@ -1,12 +1,23 @@
 import gradio as gr
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageOps
 import requests
 from io import BytesIO
-import os
+import torch
+from torchvision import transforms
+from torchvision.models.segmentation import deeplabv3_resnet101
 
 # AI model repo for design generation
 repo = "artificialguybr/TshirtDesignRedmond-V2"
 
+# Load a pretrained segmentation model for masking
+segmentation_model = deeplabv3_resnet101(pretrained=True).eval()
+
+# Transform for input image
+transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor(),
+])
+
 # Generate plain cloth image with specified color
 def generate_cloth(color_prompt):
     prompt = f"A plain {color_prompt} colored T-shirt hanging on a plain wall."
@@ -41,49 +52,44 @@ def generate_design(design_prompt):
     else:
         raise Exception(f"Error generating design: {response.status_code}")
 
-# Overlay design on cloth with adjustable width and height
-def overlay_design(cloth_image, design_image, x, y, width, height):
-    # Ensure images are in RGBA mode
-    cloth_image = cloth_image.convert("RGBA")
-    design_image = design_image.convert("RGBA")
-
-    # Resize design based on user inputs
-    resized_design = design_image.resize((width, height))
-
-    # Overlay the design at specified coordinates
-    result = cloth_image.copy()
-    result.paste(resized_design, (x, y), resized_design)
+# Apply segmentation mask to fit design on T-shirt
+def apply_mask(cloth_image, design_image):
+    input_tensor = transform(cloth_image).unsqueeze(0)
+    with torch.no_grad():
+        output = segmentation_model(input_tensor)['out'][0]
+    mask = output.argmax(0).byte().numpy()
+    mask_image = Image.fromarray((mask == 15).astype('uint8') * 255)  # Class 15 for T-shirt
+    mask_resized = mask_image.resize(cloth_image.size)
 
-    return result
+    # Fit the design inside the T-shirt mask
+    design_resized = design_image.resize(cloth_image.size)
+    design_masked = Image.composite(design_resized, cloth_image, mask_resized)
+    return design_masked
 
 # Full workflow: Generate cloth, design, and combine them
-def design_tshirt(color_prompt, design_prompt, x, y, width, height):
+def design_tshirt(color_prompt, design_prompt):
     cloth_image = generate_cloth(color_prompt)
     design_image = generate_design(design_prompt)
-    final_image = overlay_design(cloth_image, design_image, x, y, width, height)
+    final_image = apply_mask(cloth_image, design_image)
     return final_image
 
 # Gradio interface
 with gr.Blocks() as interface:
     gr.Markdown("# **AI Cloth Designer**")
-    gr.Markdown("Generate custom T-shirts by specifying a color and adding a draggable design with adjustable size.")
+    gr.Markdown("Generate custom T-shirts by specifying a color and adding a perfectly fitted design.")
 
     with gr.Row():
         with gr.Column(scale=1):
             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")
-            x_coord = gr.Slider(label="X Coordinate", minimum=0, maximum=400, step=10, value=100)
-            y_coord = gr.Slider(label="Y Coordinate", minimum=0, maximum=600, step=10, value=100)
-            width_slider = gr.Slider(label="Design Width", minimum=100, maximum=500, step=10, value=200)
-            height_slider = gr.Slider(label="Design Height", minimum=100, maximum=500, step=10, value=300)
             generate_button = gr.Button("Generate T-Shirt")
         with gr.Column(scale=1):
             output_image = gr.Image(label="Final T-Shirt Design")
 
     generate_button.click(
         design_tshirt,
-        inputs=[color_prompt, design_prompt, x_coord, y_coord, width_slider, height_slider],
+        inputs=[color_prompt, design_prompt],
         outputs=output_image,
     )
 
-interface.launch(debug=True)
+interface.launch(debug=True)