parokshsaxena commited on
Commit
658ee20
β€’
1 Parent(s): 6093a65

doing color transfer for image harmonization

Browse files
Files changed (1) hide show
  1. src/background_processor.py +26 -1
src/background_processor.py CHANGED
@@ -194,6 +194,9 @@ class BackgroundProcessor:
194
  #height = foreground_img_pil.height
195
  #background_image_pil = Image.open(background_img_path)
196
  #background_image_pil = background_image_pil.resize((width, height))
 
 
 
197
 
198
  foreground_binary = ImageFormatConvertor.pil_image_to_binary_data(foreground_img_pil)
199
  background_binary = ImageFormatConvertor.pil_image_to_binary_data(background_image_pil)
@@ -229,4 +232,26 @@ class BackgroundProcessor:
229
  logging.error(f"failed to use remove bg. Status: {remove_bg_request.status_code}. Resp: {remove_bg_request.content}")
230
  return None
231
 
232
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
194
  #height = foreground_img_pil.height
195
  #background_image_pil = Image.open(background_img_path)
196
  #background_image_pil = background_image_pil.resize((width, height))
197
+
198
+ # Do color transfer of background to foreground to adjust lighting condition
199
+ foreground_img_pil = cls.color_transfer(foreground_img_pil, background_image_pil)
200
 
201
  foreground_binary = ImageFormatConvertor.pil_image_to_binary_data(foreground_img_pil)
202
  background_binary = ImageFormatConvertor.pil_image_to_binary_data(background_image_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
+ @classmethod
236
+ def color_transfer(cls, source_pil: Image, target_pil: Image) -> Image:
237
+ source = ImageFormatConvertor.pil_to_cv2(source_pil)
238
+ target = ImageFormatConvertor.pil_to_cv2(target_pil)
239
+ source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB)
240
+ target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB)
241
+
242
+ # Compute the mean and standard deviation of the source and target images
243
+ source_mean, source_std = cv2.meanStdDev(source)
244
+ target_mean, target_std = cv2.meanStdDev(target)
245
+
246
+ #Reshape the mean and std to (1, 1, 3) so they can be broadcast correctly
247
+ source_mean = source_mean.reshape((1, 1, 3))
248
+ source_std = source_std.reshape((1, 1, 3))
249
+ target_mean = target_mean.reshape((1, 1, 3))
250
+ target_std = target_std.reshape((1, 1, 3))
251
+ # Subtract the mean from the source image
252
+ result = (source - source_mean) * (target_std / source_std) + target_mean
253
+ result = np.clip(result, 0, 255).astype(np.uint8)
254
+
255
+ res = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
256
+ res_pil = ImageFormatConvertor.cv2_to_pil(res)
257
+ return res_pil