Skip to content

Commit

Permalink
Create transfer.py
Browse files Browse the repository at this point in the history
  • Loading branch information
Adam Oswald authored and huggingface-web committed Nov 7, 2022
1 parent e7af5d9 commit 7a1112b
Showing 1 changed file with 31 additions and 0 deletions.
31 changes: 31 additions & 0 deletions transfer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
import numpy as np
import cv2 as cv


def transfer_color(source_image: np.ndarray, target_image: np.ndarray) -> np.ndarray:
"""Color transfer between images
Args:
source_image (np.ndarray): Color source image
target_image (np.ndarray): Target image
Returns:
np.ndarray: The result of the color transfer
Reference:
doi: 10.1109/38.946629
"""
# RGB -> L*a*b*
src_img = cv.cvtColor(source_image, cv.COLOR_RGB2Lab)
dst_img = cv.cvtColor(target_image, cv.COLOR_RGB2Lab)

# Calculate mean and std
src_means, src_stds = src_img.mean(axis=(0, 1)), src_img.std(axis=(0, 1))
dst_means, dst_stds = dst_img.mean(axis=(0, 1)), dst_img.std(axis=(0, 1))

# Transfer
dst_img = dst_img - dst_means.reshape((1, 1, 3))
dst_img *= (dst_stds / src_stds).reshape((1, 1, 3))
dst_img += src_means.reshape((1, 1, 3))

# L*a*b* -> RGB
dst_img = np.clip(dst_img, 0, 255).astype(np.uint8)
dst_img = cv.cvtColor(dst_img, cv.COLOR_LAB2RGB)
return dst_img

0 comments on commit 7a1112b

Please sign in to comment.