Gopalag commited on
Commit
670fde4
·
verified ·
1 Parent(s): dfb11c1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +70 -81
app.py CHANGED
@@ -1,4 +1,4 @@
1
- import gradio as gr
2
  import numpy as np
3
  import random
4
  import spaces
@@ -6,86 +6,96 @@ import torch
6
  from diffusers import DiffusionPipeline
7
  from PIL import Image
8
  import io
 
9
 
10
- dtype = torch.bfloat16
11
- device = "cuda" if torch.cuda.is_available() else "cpu"
12
-
13
- pipe = DiffusionPipeline.from_pretrained(
14
- "black-forest-labs/FLUX.1-schnell",
15
- torch_dtype=dtype
16
- ).to(device)
17
-
18
- MAX_SEED = np.iinfo(np.int32).max
19
- MAX_IMAGE_SIZE = 2048
20
-
21
- import numpy as np
22
- from collections import Counter
23
-
24
- def get_prominent_colors(image, num_colors=5):
25
  """
26
- Get the most prominent colors from an image, focusing on edges
27
  """
28
  # Convert to numpy array
29
  img_array = np.array(image)
30
 
31
- # Create a simple edge mask using gradient magnitude
32
- gradient_x = np.gradient(img_array.mean(axis=2))[1]
33
- gradient_y = np.gradient(img_array.mean(axis=2))[0]
34
- gradient_magnitude = np.sqrt(gradient_x**2 + gradient_y**2)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
35
 
36
- # Threshold to get edge pixels
37
- edge_threshold = np.percentile(gradient_magnitude, 90) # Adjust percentile as needed
38
- edge_mask = gradient_magnitude > edge_threshold
 
 
 
 
 
39
 
40
- # Get colors from edge pixels
41
- edge_colors = img_array[edge_mask]
 
 
42
 
43
- # Convert colors to tuples for counting
44
- colors = [tuple(color) for color in edge_colors]
 
 
 
45
 
46
- # Count occurrences of each color
47
- color_counts = Counter(colors)
 
48
 
49
- # Get most common colors
50
- prominent_colors = color_counts.most_common(num_colors)
 
51
 
52
- return prominent_colors
53
 
54
  def create_tshirt_preview(design_image, tshirt_color="white"):
55
  """
56
- Overlay the design onto the existing t-shirt template and color match
57
  """
58
  # Load the template t-shirt image
59
  tshirt = Image.open('image.jpeg')
60
  tshirt_width, tshirt_height = tshirt.size
61
 
62
- # Convert design to PIL Image if it's not already
63
- if not isinstance(design_image, Image.Image):
64
- design_image = Image.fromarray(design_image)
65
-
66
- # Get prominent colors from the design
67
- prominent_colors = get_prominent_colors(design_image)
68
- if prominent_colors:
69
- # Use the most prominent color for the t-shirt
70
- main_color = prominent_colors[0][0] # RGB tuple of most common color
71
- else:
72
- # Fallback to white if no colors found
73
- main_color = (255, 255, 255)
74
 
75
- # Convert design to PIL Image if it's not already
76
- if not isinstance(design_image, Image.Image):
77
- design_image = Image.fromarray(design_image)
78
 
79
- # Resize design to fit nicely on shirt (40% of shirt width)
80
- design_width = int(tshirt_width * 0.35) # Adjust this percentage as needed
81
  design_height = int(design_width * design_image.size[1] / design_image.size[0])
82
  design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
83
 
84
- # Calculate position to center design on shirt
85
  x = (tshirt_width - design_width) // 2
86
- y = int(tshirt_height * 0.2) # Adjust this value based on your template
87
 
88
- # If design has transparency (RGBA), create mask
89
  if design_image.mode == 'RGBA':
90
  mask = design_image.split()[3]
91
  else:
@@ -94,34 +104,10 @@ def create_tshirt_preview(design_image, tshirt_color="white"):
94
  # Paste design onto shirt
95
  tshirt.paste(design_image, (x, y), mask)
96
 
97
- return tshirt
98
-
99
- def enhance_prompt_for_tshirt(prompt, style=None):
100
- """Add specific terms to ensure good t-shirt designs."""
101
- style_terms = {
102
- "minimal": ["simple geometric shapes", "clean lines", "minimalist illustration"],
103
- "vintage": ["distressed effect", "retro typography", "vintage illustration"],
104
- "artistic": ["hand-drawn style", "watercolor effect", "artistic illustration"],
105
- "geometric": ["abstract shapes", "geometric patterns", "modern design"],
106
- "typography": ["bold typography", "creative lettering", "text-based design"],
107
- "realistic": ["realistic", "cinematic", "photograph"]
108
- }
109
-
110
- base_terms = [
111
- "create t-shirt design",
112
- "with centered composition",
113
- "high quality",
114
- "professional design",
115
- "clear background"
116
- ]
117
-
118
- enhanced_prompt = f"{prompt}, {', '.join(base_terms)}"
119
 
120
- if style and style in style_terms:
121
- style_specific_terms = style_terms[style]
122
- enhanced_prompt = f"{enhanced_prompt}, {', '.join(style_specific_terms)}"
123
-
124
- return enhanced_prompt
125
 
126
  @spaces.GPU()
127
  def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=False,
@@ -143,6 +129,9 @@ def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=Fals
143
  guidance_scale=0.0
144
  ).images[0]
145
 
 
 
 
146
  # Create t-shirt preview
147
  tshirt_preview = create_tshirt_preview(design_image, tshirt_color)
148
 
 
1
+ mport gradio as gr
2
  import numpy as np
3
  import random
4
  import spaces
 
6
  from diffusers import DiffusionPipeline
7
  from PIL import Image
8
  import io
9
+ from PIL import ImageEnhance
10
 
11
+ def get_edge_color(image):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12
  """
13
+ Get a random color from the edge of the image
14
  """
15
  # Convert to numpy array
16
  img_array = np.array(image)
17
 
18
+ # Get pixels from all edges
19
+ top_edge = img_array[0, :, :]
20
+ bottom_edge = img_array[-1, :, :]
21
+ left_edge = img_array[:, 0, :]
22
+ right_edge = img_array[:, -1, :]
23
+
24
+ # Combine all edge pixels
25
+ edge_pixels = np.concatenate([top_edge, bottom_edge, left_edge, right_edge])
26
+
27
+ # Pick a random edge pixel
28
+ random_edge_color = tuple(edge_pixels[random.randint(0, len(edge_pixels)-1)])
29
+
30
+ return random_edge_color
31
+
32
+ def color_match_tshirt(tshirt_image, target_color, threshold=30):
33
+ """
34
+ Change white/near-white areas of the t-shirt to the target color
35
+ """
36
+ # Convert to numpy array
37
+ img_array = np.array(tshirt_image)
38
+
39
+ # Create a mask for near-white pixels
40
+ white_mask = np.all(np.abs(img_array - [255, 255, 255]) < threshold, axis=2)
41
+
42
+ # Apply the new color to masked areas
43
+ img_array[white_mask] = target_color
44
 
45
+ return Image.fromarray(img_array)
46
+
47
+ def add_logo_watermark(image, logo_path='logo.png', size_percentage=0.2):
48
+ """
49
+ Add a logo watermark to the bottom right corner
50
+ """
51
+ # Open and resize logo
52
+ logo = Image.open(logo_path)
53
 
54
+ # Calculate new logo size (20% of main image width)
55
+ new_width = int(image.size[0] * size_percentage)
56
+ new_height = int(new_width * logo.size[1] / logo.size[0])
57
+ logo = logo.resize((new_width, new_height), Image.Resampling.LANCZOS)
58
 
59
+ # If logo has transparency, use it as mask
60
+ if logo.mode == 'RGBA':
61
+ mask = logo.split()[3]
62
+ else:
63
+ mask = None
64
 
65
+ # Calculate position (bottom right corner with padding)
66
+ position = (image.size[0] - logo.size[0] - 10,
67
+ image.size[1] - logo.size[1] - 10)
68
 
69
+ # Create a copy of the image
70
+ result = image.copy()
71
+ result.paste(logo, position, mask)
72
 
73
+ return result
74
 
75
  def create_tshirt_preview(design_image, tshirt_color="white"):
76
  """
77
+ Overlay the design onto the t-shirt template with color matching
78
  """
79
  # Load the template t-shirt image
80
  tshirt = Image.open('image.jpeg')
81
  tshirt_width, tshirt_height = tshirt.size
82
 
83
+ # Get a random edge color from the design
84
+ edge_color = get_edge_color(design_image)
 
 
 
 
 
 
 
 
 
 
85
 
86
+ # Color match the t-shirt
87
+ tshirt = color_match_tshirt(tshirt, edge_color)
 
88
 
89
+ # Resize design (35% of shirt width)
90
+ design_width = int(tshirt_width * 0.35)
91
  design_height = int(design_width * design_image.size[1] / design_image.size[0])
92
  design_image = design_image.resize((design_width, design_height), Image.Resampling.LANCZOS)
93
 
94
+ # Calculate position to center design
95
  x = (tshirt_width - design_width) // 2
96
+ y = int(tshirt_height * 0.2)
97
 
98
+ # Create mask if design has transparency
99
  if design_image.mode == 'RGBA':
100
  mask = design_image.split()[3]
101
  else:
 
104
  # Paste design onto shirt
105
  tshirt.paste(design_image, (x, y), mask)
106
 
107
+ # Add logo watermark
108
+ tshirt = add_logo_watermark(tshirt)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
109
 
110
+ return tshirt
 
 
 
 
111
 
112
  @spaces.GPU()
113
  def infer(prompt, style=None, tshirt_color="white", seed=42, randomize_seed=False,
 
129
  guidance_scale=0.0
130
  ).images[0]
131
 
132
+ # Add logo to design
133
+ design_image = add_logo_watermark(design_image)
134
+
135
  # Create t-shirt preview
136
  tshirt_preview = create_tshirt_preview(design_image, tshirt_color)
137