farm-recolor / recolorReinhardAlgo.py
vettorazi's picture
copied local files. docker initial setup
52cbb9c
raw
history blame
3.46 kB
import cv2
import numpy as np
from PIL import Image
def reinhards_color_transfer(source, target):
# Convert the images from the RGB to the Lab color space
source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2Lab).astype(np.float64)
target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2Lab).astype(np.float64)
# Compute mean and standard deviation for each channel in both images
l_mean_src, l_std_src = np.mean(source_lab[:, :, 0]), np.std(source_lab[:, :, 0])
a_mean_src, a_std_src = np.mean(source_lab[:, :, 1]), np.std(source_lab[:, :, 1])
b_mean_src, b_std_src = np.mean(source_lab[:, :, 2]), np.std(source_lab[:, :, 2])
l_mean_tar, l_std_tar = np.mean(target_lab[:, :, 0]), np.std(target_lab[:, :, 0])
a_mean_tar, a_std_tar = np.mean(target_lab[:, :, 1]), np.std(target_lab[:, :, 1])
b_mean_tar, b_std_tar = np.mean(target_lab[:, :, 2]), np.std(target_lab[:, :, 2])
# Subtract the means from the source image
source_lab[:, :, 0] -= l_mean_src
source_lab[:, :, 1] -= a_mean_src
source_lab[:, :, 2] -= b_mean_src
# Scale by the standard deviations
source_lab[:, :, 0] = (l_std_tar / l_std_src) * source_lab[:, :, 0]
source_lab[:, :, 1] = (a_std_tar / a_std_src) * source_lab[:, :, 1]
source_lab[:, :, 2] = (b_std_tar / b_std_src) * source_lab[:, :, 2]
# Add the target means
source_lab[:, :, 0] += l_mean_tar
source_lab[:, :, 1] += a_mean_tar
source_lab[:, :, 2] += b_mean_tar
# Clip pixel values to ensure they fall within the valid Lab range
source_lab[:, :, 0] = np.clip(source_lab[:, :, 0], 0, 255)
source_lab[:, :, 1] = np.clip(source_lab[:, :, 1], 0, 255)
source_lab[:, :, 2] = np.clip(source_lab[:, :, 2], 0, 255)
# Convert back to RGB
transferred_rgb = cv2.cvtColor(source_lab.astype(np.uint8), cv2.COLOR_Lab2BGR)
return transferred_rgb
def rgb_to_hex(rgb):
return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2])
def hex_to_rgb(hex):
hex = hex.lstrip('#')
return tuple(int(hex[i:i+2], 16) for i in (0, 2, 4))
def create_color_palette(colors, palette_width=800, palette_height=200):
"""
Receives a list of colors in hex format and creates a palette image
"""
pixels = []
n_colors = len(colors)
for i in range(n_colors):
color = hex_to_rgb(colors[i])
for j in range(palette_width//n_colors * palette_height):
pixels.append(color)
img = Image.new('RGB', (palette_height, palette_width))
img.putdata(pixels)
# img.show()
return img
# if __name__ == "__main__":
# source = cv2.imread("estampa.jpg")
# colors = ['#6b3d68', '#6d2055', '#695977', '#6b7988', '#6f9b9b']
# # Generate palette image
# palette_img = create_color_palette(colors)
# # Convert the palette image to BGR format
# palette_bgr = cv2.cvtColor(np.array(palette_img), cv2.COLOR_RGB2BGR)
# # Save the palette image
# # cv2.imwrite("palette.jpg", palette_bgr)
# target = palette_bgr#cv2.imread("palette.jpg")
# transferred = reinhards_color_transfer(source, target)
# cv2.imwrite("transferred_reinhard.jpg", transferred)
def recolor(source, colors):
palette_img = create_color_palette(colors)
palette_bgr = cv2.cvtColor(np.array(palette_img), cv2.COLOR_RGB2BGR)
recolored = reinhards_color_transfer(source, palette_bgr)
recoloredFile = cv2.imwrite("result.jpg", recolored)
return recolored