parokshsaxena commited on
Commit
1dddd5f
β€’
1 Parent(s): 9065906

using shein sizes

Browse files
Files changed (2) hide show
  1. app.py +7 -5
  2. src/background_processor.py +200 -2
app.py CHANGED
@@ -122,10 +122,12 @@ pipe = TryonPipeline.from_pretrained(
122
  )
123
  pipe.unet_encoder = UNet_Encoder
124
 
125
- #WIDTH = int(4160/4) # 768
126
- #HEIGHT = int(6240/4) # 1024
127
- WIDTH = int(768)
128
- HEIGHT = int(1024)
 
 
129
  POSE_WIDTH = int(WIDTH/2) # int(WIDTH/2)
130
  POSE_HEIGHT = int(HEIGHT/2) #int(HEIGHT/2)
131
 
@@ -259,7 +261,7 @@ def start_tryon(dict,garm_img,garment_des, background_img, is_checked,is_checked
259
  # apply background to final image
260
  if background_img:
261
  logging.info("Adding background")
262
- final_image = BackgroundProcessor.add_background(final_image, background_img)
263
  return final_image, mask_gray
264
  # return images[0], mask_gray
265
 
 
122
  )
123
  pipe.unet_encoder = UNet_Encoder
124
 
125
+ # Standard size of shein images
126
+ WIDTH = int(4160/5)
127
+ HEIGHT = int(6240/5)
128
+ # Standard size on which model is trained
129
+ #WIDTH = int(768)
130
+ #HEIGHT = int(1024)
131
  POSE_WIDTH = int(WIDTH/2) # int(WIDTH/2)
132
  POSE_HEIGHT = int(HEIGHT/2) #int(HEIGHT/2)
133
 
 
261
  # apply background to final image
262
  if background_img:
263
  logging.info("Adding background")
264
+ final_image = BackgroundProcessor.add_background_v3(final_image, background_img)
265
  return final_image, mask_gray
266
  # return images[0], mask_gray
267
 
src/background_processor.py CHANGED
@@ -1,4 +1,5 @@
1
- from PIL import Image
 
2
  import numpy as np
3
  from preprocess.humanparsing.run_parsing import Parsing
4
 
@@ -35,4 +36,201 @@ class BackgroundProcessor:
35
  result_img = Image.fromarray(human_with_background.astype('uint8'))
36
 
37
  # Return or save the result
38
- return result_img
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image, ImageEnhance
2
+ import cv2
3
  import numpy as np
4
  from preprocess.humanparsing.run_parsing import Parsing
5
 
 
36
  result_img = Image.fromarray(human_with_background.astype('uint8'))
37
 
38
  # Return or save the result
39
+ return result_img
40
+
41
+ @classmethod
42
+ def temp_v2(cls, human_img_path, background_img_path, mask_img_path):
43
+ # Load the images
44
+ foreground_img = cv2.imread(human_img_path).resize((768,1024)) # The segmented person image
45
+ background_img = cv2.imread(background_img_path) # The new background image
46
+ mask_img = cv2.imread(mask_img_path, cv2.IMREAD_GRAYSCALE) # The mask image from the human parser model
47
+
48
+ # Ensure the foreground image and the mask are the same size
49
+ if foreground_img.shape[:2] != mask_img.shape[:2]:
50
+ raise ValueError("Foreground image and mask must be the same size")
51
+
52
+ # Resize background image to match the size of the foreground image
53
+ background_img = cv2.resize(background_img, (foreground_img.shape[1], foreground_img.shape[0]))
54
+
55
+ # Create an inverted mask
56
+ mask_inv = cv2.bitwise_not(mask_img)
57
+
58
+ # Convert mask to 3 channels
59
+ mask_3ch = cv2.cvtColor(mask_img, cv2.COLOR_GRAY2BGR)
60
+ mask_inv_3ch = cv2.cvtColor(mask_inv, cv2.COLOR_GRAY2BGR)
61
+
62
+ # Extract the person from the foreground image using the mask
63
+ person = cv2.bitwise_and(foreground_img, mask_3ch)
64
+
65
+ # Extract the background where the person is not present
66
+ background = cv2.bitwise_and(background_img, mask_inv_3ch)
67
+
68
+ # Combine the person and the new background
69
+ combined_img = cv2.add(person, background)
70
+
71
+ # Refine edges using Gaussian Blur (feathering technique)
72
+ blurred_combined_img = cv2.GaussianBlur(combined_img, (5, 5), 0)
73
+
74
+ # Post-processing: Adjust brightness, contrast, etc. (optional)
75
+ alpha = 1.2 # Contrast control (1.0-3.0)
76
+ beta = 20 # Brightness control (0-100)
77
+
78
+ post_processed_img = cv2.convertScaleAbs(blurred_combined_img, alpha=alpha, beta=beta)
79
+
80
+ # Save the final image
81
+ # cv2.imwrite('path_to_save_final_image.png', post_processed_img)
82
+
83
+ # Display the images (optional)
84
+ cv2.imshow('Foreground', foreground_img)
85
+ cv2.imshow('Background', background_img)
86
+ cv2.imshow('Mask', mask_img)
87
+ cv2.imshow('Combined', combined_img)
88
+ cv2.imshow('Post Processed', post_processed_img)
89
+ cv2.waitKey(0)
90
+ cv2.destroyAllWindows()
91
+ return post_processed_img
92
+
93
+
94
+ @classmethod
95
+ def add_background_v3(cls, foreground_pil: Image, background_pil: Image):
96
+ foreground_pil= foreground_pil.convert("RGB")
97
+ width = foreground_pil.width
98
+ height = foreground_pil.height
99
+
100
+ # Create mask image
101
+ parsed_img, _ = parsing_model(foreground_pil)
102
+ mask_pil = parsed_img.convert("L")
103
+ # Apply a threshold to convert to binary image
104
+ # mask_pil = mask_pil.point(lambda p: 1 if p > 127 else 0, mode='1')
105
+ mask_pil = mask_pil.resize((width, height))
106
+
107
+ # Resize background image
108
+ background_pil = background_pil.convert("RGB")
109
+ background_pil = background_pil.resize((width, height))
110
+
111
+ # Load the images using PIL
112
+ #foreground_pil = Image.open(human_img_path).convert("RGB") # The segmented person image
113
+ #background_pil = Image.open(background_img_path).convert("RGB") # The new background image
114
+ #mask_pil = Image.open(mask_img_path).convert('L') # The mask image from the human parser model
115
+
116
+ # Resize the background to match the size of the foreground
117
+ #background_pil = background_pil.resize(foreground_pil.size)
118
+
119
+ # Resize mask
120
+ #mask_pil = mask_pil.resize(foreground_pil.size)
121
+
122
+ # Convert PIL images to OpenCV format
123
+ foreground_cv2 = cls.pil_to_cv2(foreground_pil)
124
+ background_cv2 = cls.pil_to_cv2(background_pil)
125
+ #mask_cv2 = pil_to_cv2(mask_pil)
126
+ mask_cv2 = np.array(mask_pil) # Directly convert to NumPy array without color conversion
127
+
128
+ # Ensure the mask is a single channel image
129
+ if len(mask_cv2.shape) == 3:
130
+ mask_cv2 = cv2.cvtColor(mask_cv2, cv2.COLOR_BGR2GRAY)
131
+
132
+ # Threshold the mask to convert it to pure black and white
133
+ _, mask_cv2 = cv2.threshold(mask_cv2, 0, 255, cv2.THRESH_BINARY)
134
+
135
+ # Ensure the mask is a single channel image
136
+ #if len(mask_cv2.shape) == 3:
137
+ # mask_cv2 = cv2.cvtColor(mask_cv2, cv2.COLOR_BGR2GRAY)
138
+
139
+ # Create an inverted mask
140
+ mask_inv_cv2 = cv2.bitwise_not(mask_cv2)
141
+
142
+ # Convert mask to 3 channels
143
+ mask_3ch_cv2 = cv2.cvtColor(mask_cv2, cv2.COLOR_GRAY2BGR)
144
+ mask_inv_3ch_cv2 = cv2.cvtColor(mask_inv_cv2, cv2.COLOR_GRAY2BGR)
145
+
146
+ # Extract the person from the foreground image using the mask
147
+ person_cv2 = cv2.bitwise_and(foreground_cv2, mask_3ch_cv2)
148
+
149
+ # Extract the background where the person is not present
150
+ background_extracted_cv2 = cv2.bitwise_and(background_cv2, mask_inv_3ch_cv2)
151
+
152
+ # Combine the person and the new background
153
+ combined_cv2 = cv2.add(person_cv2, background_extracted_cv2)
154
+
155
+ # Refine edges using Gaussian Blur (feathering technique)
156
+ blurred_combined_cv2 = cv2.GaussianBlur(combined_cv2, (5, 5), 0)
157
+
158
+ # Convert the result back to PIL format
159
+ combined_pil = cls.cv2_to_pil(blurred_combined_cv2)
160
+
161
+
162
+ """
163
+ # Post-processing: Adjust brightness, contrast, etc. (optional)
164
+ enhancer = ImageEnhance.Contrast(combined_pil)
165
+ post_processed_pil = enhancer.enhance(1.2) # Adjust contrast
166
+ enhancer = ImageEnhance.Brightness(post_processed_pil)
167
+ post_processed_pil = enhancer.enhance(1.2) # Adjust brightness
168
+ """
169
+
170
+
171
+ # Save the final image
172
+ # post_processed_pil.save('path_to_save_final_image_1.png')
173
+
174
+ # Display the images (optional)
175
+ #foreground_pil.show(title="Foreground")
176
+ #background_pil.show(title="Background")
177
+ #mask_pil.show(title="Mask")
178
+ #combined_pil.show(title="Combined")
179
+ # post_processed_pil.show(title="Post Processed")
180
+
181
+ return combined_pil
182
+
183
+ @classmethod
184
+ def replace_background(cls, foreground_img_path: str, background_img_path: str):
185
+ # Load the input image (with alpha channel) and the background image
186
+ #input_image = cv2.imread(foreground_img_path, cv2.IMREAD_UNCHANGED)
187
+ input_image = cv2.imread(foreground_img_path)
188
+ background_image = cv2.imread(background_img_path)
189
+
190
+ # Ensure the input image has an alpha channel
191
+ if input_image.shape[2] != 4:
192
+ raise ValueError("Input image must have an alpha channel")
193
+
194
+ # Extract the alpha channel
195
+ alpha_channel = input_image[:, :, 3]
196
+
197
+ # Resize the background image to match the input image dimensions
198
+ background_image = cv2.resize(background_image, (input_image.shape[1], input_image.shape[0]))
199
+
200
+ # Convert alpha channel to 3 channels
201
+ alpha_channel_3ch = cv2.cvtColor(alpha_channel, cv2.COLOR_GRAY2BGR)
202
+ alpha_channel_3ch = alpha_channel_3ch / 255.0 # Normalize to 0-1
203
+
204
+ # Extract the BGR channels of the input image
205
+ input_bgr = input_image[:, :, :3]
206
+
207
+ # Blend the images using the alpha channel
208
+ foreground = cv2.multiply(alpha_channel_3ch, input_bgr.astype(float))
209
+ background = cv2.multiply(1.0 - alpha_channel_3ch, background_image.astype(float))
210
+ combined_image = cv2.add(foreground, background).astype(np.uint8)
211
+
212
+ # Save and display the result
213
+ cv2.imwrite('path_to_save_combined_image.png', combined_image)
214
+ cv2.imshow('Combined Image', combined_image)
215
+ cv2.waitKey(0)
216
+ cv2.destroyAllWindows()
217
+
218
+
219
+
220
+ # Function to convert PIL Image to OpenCV format
221
+ @classmethod
222
+ def pil_to_cv2(cls, pil_image):
223
+ open_cv_image = np.array(pil_image)
224
+ # Convert RGB to BGR if it's a 3-channel image
225
+ if len(open_cv_image.shape) == 3:
226
+ open_cv_image = open_cv_image[:, :, ::-1].copy()
227
+ return open_cv_image
228
+
229
+ # Function to convert OpenCV format to PIL Image
230
+ @classmethod
231
+ def cv2_to_pil(cls, cv2_image):
232
+ # Convert BGR to RGB if it's a 3-channel image
233
+ if len(cv2_image.shape) == 3:
234
+ cv2_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)
235
+ pil_image = Image.fromarray(cv2_image)
236
+ return pil_image