Spaces:
Running
on
Zero
Running
on
Zero
parokshsaxena
commited on
Commit
β’
658ee20
1
Parent(s):
6093a65
doing color transfer for image harmonization
Browse files- 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
|