Update app.py

app.py

@@ -18,14 +18,60 @@ pipe = DiffusionPipeline.from_pretrained(
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+import numpy as np
+from collections import Counter
+
+def get_prominent_colors(image, num_colors=5):
+    """
+    Get the most prominent colors from an image, focusing on edges
+    """
+    # Convert to numpy array
+    img_array = np.array(image)
+    
+    # Create a simple edge mask using gradient magnitude
+    gradient_x = np.gradient(img_array.mean(axis=2))[1]
+    gradient_y = np.gradient(img_array.mean(axis=2))[0]
+    gradient_magnitude = np.sqrt(gradient_x**2 + gradient_y**2)
+    
+    # Threshold to get edge pixels
+    edge_threshold = np.percentile(gradient_magnitude, 90)  # Adjust percentile as needed
+    edge_mask = gradient_magnitude > edge_threshold
+    
+    # Get colors from edge pixels
+    edge_colors = img_array[edge_mask]
+    
+    # Convert colors to tuples for counting
+    colors = [tuple(color) for color in edge_colors]
+    
+    # Count occurrences of each color
+    color_counts = Counter(colors)
+    
+    # Get most common colors
+    prominent_colors = color_counts.most_common(num_colors)
+    
+    return prominent_colors
+
 def create_tshirt_preview(design_image, tshirt_color="white"):
     """
-    Overlay the design onto the existing t-shirt template
+    Overlay the design onto the existing t-shirt template and color match
     """
     # Load the template t-shirt image
     tshirt = Image.open('image.jpeg')
     tshirt_width, tshirt_height = tshirt.size
     
+    # Convert design to PIL Image if it's not already
+    if not isinstance(design_image, Image.Image):
+        design_image = Image.fromarray(design_image)
+    
+    # Get prominent colors from the design
+    prominent_colors = get_prominent_colors(design_image)
+    if prominent_colors:
+        # Use the most prominent color for the t-shirt
+        main_color = prominent_colors[0][0]  # RGB tuple of most common color
+    else:
+        # Fallback to white if no colors found
+        main_color = (255, 255, 255)
+    
     # Convert design to PIL Image if it's not already
     if not isinstance(design_image, Image.Image):
         design_image = Image.fromarray(design_image)
@@ -37,7 +83,7 @@ def create_tshirt_preview(design_image, tshirt_color="white"):
     
     # Calculate position to center design on shirt
     x = (tshirt_width - design_width) // 2
-    y = int(tshirt_height * 0.20)  # Adjust this value based on your template
+    y = int(tshirt_height * 0.2)  # Adjust this value based on your template
     
     # If design has transparency (RGBA), create mask
     if design_image.mode == 'RGBA':
@@ -50,41 +96,6 @@ def create_tshirt_preview(design_image, tshirt_color="white"):
     
     return tshirt
 
-# def create_tshirt_preview(design_image, tshirt_color="white"):
-#     """
-#     Overlay the design onto a t-shirt template
-#     """
-#     # Create a base t-shirt shape
-#     tshirt_width = 800
-#     tshirt_height = 1000
-    
-#     # Create base t-shirt image
-#     tshirt = Image.new('RGB', (tshirt_width, tshirt_height), tshirt_color)
-    
-#     # Convert design to PIL Image if it's not already
-#     if not isinstance(design_image, Image.Image):
-#         design_image = Image.fromarray(design_image)
-    
-#     # Resize design to fit nicely on shirt (30% of shirt width)
-#     design_width = int(tshirt_width * 0.3)
-#     design_height = int(design_width * design_image.size[1] / design_image.size[0])
-#     design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
-    
-#     # Calculate position to center design on shirt (top third of shirt)
-#     x = (tshirt_width - design_width) // 2
-#     y = int(tshirt_height * 0.25)  # Position in top third
-    
-#     # If design has transparency (RGBA), create mask
-#     if design_image.mode == 'RGBA':
-#         mask = design_image.split()[3]
-#     else:
-#         mask = None
-    
-#     # Paste design onto shirt
-#     tshirt.paste(design_image, (x, y), mask)
-    
-#     return tshirt
-
 def enhance_prompt_for_tshirt(prompt, style=None):
     """Add specific terms to ensure good t-shirt designs."""
     style_terms = {
@@ -97,9 +108,9 @@ def enhance_prompt_for_tshirt(prompt, style=None):
     }
     
     base_terms = [
-        "create a t-shirt design",
+        "create t-shirt design",
         "with centered composition",
-        "4k high quality",
+        "high quality",
         "professional design",
         "clear background"
     ]
@@ -201,14 +212,14 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(
             """
-            # 👕 Deradh's T-Shirt Design Generator
+            # 👕Deradh's T-Shirt Design Generator
             """,
             elem_classes=["main-title"]
         )
         
         gr.Markdown(
             """
-            Create unique t-shirt designs using AI. 
+            Create unique t-shirt designs using Deradh's AI. 
             Describe your design idea and select a style to generate professional-quality artwork 
             perfect for custom t-shirts.
             """,