parokshsaxena commited on
Commit
1993dd0
β€’
1 Parent(s): 5ca0873

match color function

Browse files
Files changed (1) hide show
  1. src/background_processor.py +31 -1
src/background_processor.py CHANGED
@@ -232,8 +232,10 @@ 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
  @classmethod
236
  def color_transfer(cls, source_pil: Image, target_pil: Image) -> Image:
 
237
  source = ImageFormatConvertor.pil_to_cv2(source_pil)
238
  # Resize background image
239
  width, height = source_pil.width, source_pil.height
@@ -263,6 +265,7 @@ class BackgroundProcessor:
263
 
264
  @classmethod
265
  def intensity_transfer(cls, source_pil: Image, target_pil: Image) -> Image:
 
266
  """
267
  Transfers the intensity distribution from the target image to the source image.
268
 
@@ -306,4 +309,31 @@ class BackgroundProcessor:
306
  # return cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
307
  res = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
308
  res_pil = ImageFormatConvertor.cv2_to_pil(res)
309
- return res_pil
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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:
238
+ # NOT IN USE as output color was not good
239
  source = ImageFormatConvertor.pil_to_cv2(source_pil)
240
  # Resize background image
241
  width, height = source_pil.width, source_pil.height
 
265
 
266
  @classmethod
267
  def intensity_transfer(cls, source_pil: Image, target_pil: Image) -> Image:
268
+
269
  """
270
  Transfers the intensity distribution from the target image to the source image.
271
 
 
309
  # return cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
310
  res = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR)
311
  res_pil = ImageFormatConvertor.cv2_to_pil(res)
312
+ return res_pil
313
+
314
+ @classmethod
315
+ def match_color(cls, source_pil: Image, target_pil: Image):
316
+ source = ImageFormatConvertor.pil_to_cv2(source_pil)
317
+ # Resize background image
318
+ width, height = source_pil.width, source_pil.height
319
+ target_pil = target_pil.convert("RGB")
320
+ target_pil = target_pil.resize((width, height))
321
+
322
+ target = ImageFormatConvertor.pil_to_cv2(target_pil)
323
+
324
+ matched_foreground = cv2.cvtColor(source, cv2.COLOR_BGR2LAB)
325
+ matched_background = cv2.cvtColor(target, cv2.COLOR_BGR2LAB)
326
+
327
+ # Match the histograms
328
+ for i in range(3):
329
+ matched_foreground[:, :, i] = cv2.equalizeHist(matched_foreground[:, :, i])
330
+ matched_background[:, :, i] = cv2.equalizeHist(matched_background[:, :, i])
331
+
332
+ matched_foreground = cv2.cvtColor(matched_foreground, cv2.COLOR_LAB2BGR)
333
+ matched_background = cv2.cvtColor(matched_background, cv2.COLOR_LAB2BGR)
334
+
335
+ matched_foreground_pil = ImageFormatConvertor.cv2_to_pil(matched_foreground)
336
+ matched_background_pil = ImageFormatConvertor.cv2_to_pil(matched_background)
337
+
338
+ return matched_foreground_pil, matched_background_pil
339
+