|
import numpy as np |
|
import cv2 |
|
import ot |
|
from PIL import Image |
|
|
|
def transfer_channel(source_channel, target_channel): |
|
source_hist, _ = np.histogram(source_channel, bins=256, range=(0, 256)) |
|
target_hist, _ = np.histogram(target_channel, bins=256, range=(0, 256)) |
|
|
|
source_hist = source_hist.astype(np.float64) / source_hist.sum() |
|
target_hist = target_hist.astype(np.float64) / target_hist.sum() |
|
|
|
r = np.arange(256).reshape((-1, 1)) |
|
c = np.arange(256).reshape((1, -1)) |
|
M = (r - c) ** 2 |
|
M = M / M.max() |
|
|
|
P = ot.emd(source_hist, target_hist, M) |
|
|
|
transferred_channel = np.zeros_like(source_channel) |
|
for i in range(256): |
|
transferred_channel[source_channel == i] = P[i].argmax() |
|
|
|
return transferred_channel |
|
|
|
def optimal_transport_color_transfer(source, target): |
|
source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2Lab).astype(np.float64) |
|
target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2Lab).astype(np.float64) |
|
|
|
|
|
for ch in range(3): |
|
source_lab[:, :, ch] = transfer_channel(source_lab[:, :, ch], target_lab[:, :, ch]) |
|
|
|
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) |
|
|
|
return img |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def recolor(source, colors): |
|
pallete_img = create_color_palette(colors) |
|
palette_bgr = cv2.cvtColor(np.array(pallete_img), cv2.COLOR_RGB2BGR) |
|
recolored = optimal_transport_color_transfer(source, palette_bgr) |
|
smooth = True |
|
if smooth: |
|
|
|
diameter = 10 |
|
sigma_color = 25 |
|
sigma_space = 15 |
|
smoothed = cv2.bilateralFilter(recolored, diameter, sigma_color, sigma_space) |
|
recoloredFile = cv2.imwrite("result.jpg", smoothed, [cv2.IMWRITE_JPEG_QUALITY, 100]) |
|
return recoloredFile |
|
else: |
|
recoloredFile = cv2.imwrite("result.jpg", recolored) |
|
return recoloredFile |
|
|
|
|
|
|
|
|