parokshsaxena commited on
Commit
0c774a3
β€’
1 Parent(s): fbd5743

removing naize harmonization, adding functions for creating masks

Browse files
Files changed (2) hide show
  1. app.py +3 -2
  2. src/background_processor.py +57 -0
app.py CHANGED
@@ -185,9 +185,10 @@ def start_tryon(human_img_dict,garm_img,garment_des, background_img, is_checked,
185
  else:
186
  human_img = human_img_orig.resize((WIDTH, HEIGHT))
187
 
 
188
  # Do color transfer from background image for better image harmonization
189
- if background_img:
190
- human_img = BackgroundProcessor.intensity_transfer(human_img, background_img)
191
 
192
 
193
  if is_checked:
 
185
  else:
186
  human_img = human_img_orig.resize((WIDTH, HEIGHT))
187
 
188
+ # Commenting out naize harmonization for now. We will have to integrate with Deep Learning based Harmonization methods
189
  # Do color transfer from background image for better image harmonization
190
+ #if background_img:
191
+ # human_img = BackgroundProcessor.intensity_transfer(human_img, background_img)
192
 
193
 
194
  if is_checked:
src/background_processor.py CHANGED
@@ -232,6 +232,63 @@ class BackgroundProcessor:
232
  logging.error(f"failed to use remove bg. Status: {remove_bg_request.status_code}. Resp: {remove_bg_request.content}")
233
  return None
234
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
235
 
236
  @classmethod
237
  def color_transfer(cls, source_pil: Image, target_pil: Image) -> Image:
 
232
  logging.error(f"failed to use remove bg. Status: {remove_bg_request.status_code}. Resp: {remove_bg_request.content}")
233
  return None
234
 
235
+ @classmethod
236
+ def create_mask(cls, foreground_path: str, mask_path: str):
237
+ """
238
+ Given foreground image path with background removed, create a maska and save it in mask_path
239
+ """
240
+ # Load the foreground image with alpha channel
241
+ foreground = Image.open(foreground_path)
242
+
243
+ # Convert to RGBA if not already
244
+ foreground = foreground.convert("RGBA")
245
+
246
+ # Create the mask from the alpha channel
247
+ alpha_channel = np.array(foreground.split()[-1])
248
+
249
+ # Create a binary mask where alpha > 0 is white (255) and alpha == 0 is black (0)
250
+ mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
251
+
252
+ # Save the mask to a file
253
+ Image.fromarray(mask).save(mask_path)
254
+
255
+ @classmethod
256
+ def get_minimal_bounding_box(cls, foreground_pil: Image):
257
+ """
258
+ Result x1,y1,x2,y2 ie cordinate of bottom left and top right
259
+ """
260
+ # convert to cv2
261
+ foreground = ImageFormatConvertor.pil_to_cv2(foreground_pil)
262
+ # Ensure the image has an alpha channel (transparency)
263
+ if foreground.shape[2] == 4:
264
+ # Extract the alpha channel
265
+ alpha_channel = foreground[:, :, 3]
266
+ # Create a binary image from the alpha channel
267
+ _, binary_image = cv2.threshold(alpha_channel, 1, 255, cv2.THRESH_BINARY)
268
+ else:
269
+ # If there is no alpha channel, convert the image to grayscale
270
+ gray_image = cv2.cvtColor(foreground, cv2.COLOR_BGR2GRAY)
271
+ # Apply binary thresholding
272
+ _, binary_image = cv2.threshold(gray_image, 1, 255, cv2.THRESH_BINARY)
273
+
274
+ # Find all non-zero points (non-background)
275
+ non_zero_points = cv2.findNonZero(binary_image)
276
+
277
+ # Get the minimal bounding rectangle
278
+ if non_zero_points is not None:
279
+ x, y, w, h = cv2.boundingRect(non_zero_points)
280
+ """
281
+ # Optionally, draw the bounding box on the image for visualization
282
+ output_image = foreground.copy()
283
+ cv2.rectangle(output_image, (x, y), (x+w, y+h), (0, 255, 0, 255), 2)
284
+ # Save or display the output image
285
+ output_image_pil = ImageFormatConvertor.cv2_to_pil(output_image)
286
+ output_image_pil.save('output_with_bounding_box.png')
287
+ """
288
+
289
+ return (x, y, x + w, y + h)
290
+ else:
291
+ return 0,0,w,h
292
 
293
  @classmethod
294
  def color_transfer(cls, source_pil: Image, target_pil: Image) -> Image: