add color harmonization

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+import torch
+from utils import *
+
+torch.hub.download_url_to_file(
+    'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/wikipedia.org_website.png',
+    'wikipedia.org_website.png')
+torch.hub.download_url_to_file(
+    'https://github.com/aalto-ui/aim/raw/aim2/backend/data/tests/input_values/aalto.fi_website.png',
+    'aalto.fi_website.png')
+
+
+def inference(img, template, angel):
+    color_image = cv2.imread(img.name, cv2.IMREAD_COLOR)
+    height, width, _ = color_image.shape
+
+    # Resize if it is bigeer than 960 * 800
+    if width > height:
+        if width > 960:  # 3/4 * 1280
+            coef_div = width / 960.0
+            color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)),
+                                     interpolation=cv2.INTER_CUBIC)
+    else:
+        if height > 800:  # 800
+            coef_div = height / 800.0
+            color_image = cv2.resize(color_image, dsize=(int(width / coef_div), int(height / coef_div)),
+                                     interpolation=cv2.INTER_CUBIC)
+
+    HSV_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
+    selected_harmomic_scheme = HarmonicScheme(str(template), int(angel))
+    new_HSV_image = best_harmomic_scheme.hue_shifted(HSV_image, num_superpixels=-1)
+
+    # Compute shifted histogram
+    histo_1 = count_hue_histogram(HSV_image)
+    histo_2 = count_hue_histogram(new_HSV_image)
+
+    # Create Hue Plots
+    fig1 = plothis(histo_1, best_harmomic_scheme, "Source Hue")
+    fig_1_cv = get_img_from_fig(fig1)
+    fig2 = plothis(histo_2, best_harmomic_scheme, "Target Hue")
+    fig_2_cv = get_img_from_fig(fig2)
+
+    # Stack Hue Plots
+    vis = np.concatenate((fig_1_cv, fig_2_cv), axis=0)
+    # Convert HSV to BGR
+    result_image = cv2.cvtColor(new_HSV_image, cv2.COLOR_HSV2BGR)
+
+    # Final output
+    canvas = np.full((800, 960, 3), (255, 255, 255), dtype=np.uint8)
+    # compute center offset
+    x_center = (960 - width) // 2
+    y_center = (800 - height) // 2
+    # copy img image into center of result image
+    canvas[y_center:y_center + height, x_center:x_center + width] = result_image
+
+    # Combine
+    output = np.concatenate((vis, canvas), axis=1)
+    cv2.imwrite('output.png', output)
+
+    return ['output.png']
+
+
+title = 'Color Harmonization'
+description = 'Compute Color Harmonization with Different Templates'
+article = "<p style='text-align: center'></p>"
+examples = [['wikipedia.org_website.png'], ['aalto.fi_website.png']]
+css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
+
+gr.Interface(
+    inference,
+    [gr.inputs.Image(type='file', label='Input'),
+     gr.inputs.Dropdown(["X", "Y", "T", "I", "mirror_L", "L", "V", "i"],
+                        default="X",
+                        label="Template"),
+     gr.inputs.Slider(0, 359, label="Angle")],
+    [gr.outputs.Image(type='file', label='Color Harmonization of Output Image')],
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    css=css,
+).launch(debug=True, enable_queue=True)

packages.txt

@@ -0,0 +1,6 @@
+libgl1
+libglib2.0-0
+libgl1-mesa-glx
+ffmpeg 
+libsm6 
+libxext6

requirements.txt

@@ -0,0 +1,6 @@
+opencv-python-headless
+torch
+gradio
+opencv-python==4.1.2
+numpy==1.21.6
+matplotlib==3.2.2

utils.py

@@ -0,0 +1,258 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Color Harmonization utility functions.
+Some Codes are imported and adopted from https://github.com/tartarskunk/ColorHarmonization
+"""
+
+# Import Libraries
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+import io
+
+# Constants
+HueTemplates = {
+    "i": [(0.00, 0.05)],
+    "V": [(0.00, 0.26)],
+    "L": [(0.00, 0.05), (0.25, 0.22)],
+    "mirror_L": [(0.00, 0.05), (-0.25, 0.22)],
+    "I": [(0.00, 0.05), (0.50, 0.05)],
+    "T": [(0.25, 0.50)],
+    "Y": [(0.00, 0.26), (0.50, 0.05)],
+    "X": [(0.00, 0.26), (0.50, 0.26)],
+}
+template_types = list(HueTemplates.keys())
+M = len(template_types)
+A = 360
+
+
+def deg_distance(a, b):
+    d1 = np.abs(a - b)
+    d2 = np.abs(360 - d1)
+    d = np.minimum(d1, d2)
+    return d
+
+
+def normalized_gaussian(X, mu, S):
+    X = np.asarray(X).astype(np.float64)
+    S = np.asarray(S).astype(np.float64)
+    D = np.deg2rad(X - mu)
+    S = np.deg2rad(S)
+    D2 = np.multiply(D, D)
+    S2 = np.multiply(S, S)
+    return np.exp(-D2 / (2 * S2))
+
+
+class HueSector:
+
+    def __init__(self, center, width):
+        # In Degree [0,2 pi)
+        self.center = center
+        self.width = width
+        self.border = [(self.center - self.width / 2), (self.center + self.width / 2)]
+
+    def is_in_sector(self, H):
+        # True/False matrix if hue resides in the sector
+        return deg_distance(H, self.center) < self.width / 2
+
+    def distance_to_border(self, H):
+        H_1 = deg_distance(H, self.border[0])
+        H_2 = deg_distance(H, self.border[1])
+        H_dist2bdr = np.minimum(H_1, H_2)
+        return H_dist2bdr
+
+    def closest_border(self, H):
+        H_1 = deg_distance(H, self.border[0])
+        H_2 = deg_distance(H, self.border[1])
+        H_cls_bdr = np.argmin((H_1, H_2), axis=0)
+        H_cls_bdr = 2 * (H_cls_bdr - 0.5)
+        return H_cls_bdr
+
+    def distance_to_center(self, H):
+        H_dist2ctr = deg_distance(H, self.center)
+        return H_dist2ctr
+
+
+class HarmonicScheme:
+
+    def __init__(self, m, alpha):
+        self.m = m
+        self.alpha = alpha
+        self.reset_sectors()
+
+    def reset_sectors(self):
+        self.sectors = []
+        for t in HueTemplates[self.m]:
+            center = t[0] * 360 + self.alpha
+            width = t[1] * 360
+            sector = HueSector(center, width)
+            self.sectors.append(sector)
+
+    def harmony_score(self, X):
+        # Opencv store H as [0, 180) --> [0, 360)
+        H = X[:, :, 0].astype(np.int32) * 2
+        # Opencv store S as [0, 255] --> [0, 1]
+        S = X[:, :, 1].astype(np.float32) / 255.0
+
+        H_dis = self.hue_distance(H)
+        H_dis = np.deg2rad(H_dis)
+        return np.sum(np.multiply(H_dis, S))
+
+    def hue_distance(self, H):
+        H_dis = []
+        for i in range(len(self.sectors)):
+            sector = self.sectors[i]
+            H_dis.append(sector.distance_to_border(H))
+            H_dis[i][sector.is_in_sector(H)] = 0
+        H_dis = np.asarray(H_dis)
+        H_dis = H_dis.min(axis=0)
+        return H_dis
+
+    def hue_shifted(self, X, num_superpixels=-1):
+        Y = X.copy()
+        H = X[:, :, 0].astype(np.int32) * 2
+        S = X[:, :, 1].astype(np.float32) / 255.0
+
+        H_d2b = [sector.distance_to_border(H) for sector in self.sectors]
+        H_d2b = np.asarray(H_d2b)
+
+        H_cls = np.argmin(H_d2b, axis=0)
+        if num_superpixels != -1:
+            SEEDS = cv2.ximgproc.createSuperpixelSEEDS(X.shape[1], X.shape[0], X.shape[2], num_superpixels, 10)
+            SEEDS.iterate(X, 4)
+
+            V = np.zeros(H.shape).reshape(-1)
+            N = V.shape[0]
+
+            H_ctr = np.zeros((H.shape))
+            grid_num = SEEDS.getNumberOfSuperpixels()
+            labels = SEEDS.getLabels()
+            for i in range(grid_num):
+
+                P = [[], []]
+                s = np.average(H_cls[labels == i])
+                if s > 0.5:
+                    s = 1
+                else:
+                    s = 0
+                H_cls[labels == i] = s
+
+        H_ctr = np.zeros((H.shape))
+        H_wid = np.zeros((H.shape))
+        H_d2c = np.zeros((H.shape))
+        H_dir = np.zeros((H.shape))
+
+        for i in range(len(self.sectors)):
+            sector = self.sectors[i]
+            mask = (H_cls == i)
+            H_ctr[mask] = sector.center
+            H_wid[mask] = sector.width
+            H_dir += sector.closest_border(H) * mask
+            H_dist2ctr = sector.distance_to_center(H)
+            H_d2c += H_dist2ctr * mask
+
+        H_sgm = H_wid / 2
+        H_gau = normalized_gaussian(H_d2c, 0, H_sgm)
+        H_tmp = np.multiply(H_wid / 2, 1 - H_gau)
+        H_shf = np.multiply(H_dir, H_tmp)
+        H_new = (H_ctr + H_shf).astype(np.int32)
+
+        for i in range(len(self.sectors)):
+            sector = self.sectors[i]
+            mask = sector.is_in_sector(H)
+            np.copyto(H_new, H, where=sector.is_in_sector(H))
+
+        H_new = np.remainder(H_new, 360)
+        H_new = (H_new / 2).astype(np.uint8)
+        Y[:, :, 0] = H_new
+        return Y
+
+
+def count_hue_histogram(X):
+    N = 360
+    H = X[:, :, 0].astype(np.int32) * 2
+    S = X[:, :, 1].astype(np.float64) / 255.0
+    H_flat = H.flatten()
+    S_flat = S.flatten()
+
+    histo = np.zeros(N)
+    for i in range(len(H_flat)):
+        histo[H_flat[i]] += S_flat[i]
+    return histo
+
+
+def plothis(hue_histo, harmonic_scheme, caption: str):
+    N = 360
+
+    # Compute pie slices
+    theta = np.linspace(0.0, 2 * np.pi, N, endpoint=False)
+    width = np.pi / 180
+
+    # Compute colors, RGB values for the hue
+    hue_colors = np.zeros((N, 4))
+    for i in range(hue_colors.shape[0]):
+        color_HSV = np.zeros((1, 1, 3), dtype=np.uint8)
+        color_HSV[0, 0, :] = [int(i / 2), 255, 255]
+        color_BGR = cv2.cvtColor(color_HSV, cv2.COLOR_HSV2BGR)
+        B = int(color_BGR[0, 0, 0]) / 255.0
+        G = int(color_BGR[0, 0, 1]) / 255.0
+        R = int(color_BGR[0, 0, 2]) / 255.0
+        hue_colors[i] = (R, G, B, 1.0)
+
+    # Compute colors, for the shadow
+    shadow_colors = np.zeros((N, 4))
+    for i in range(shadow_colors.shape[0]):
+        shadow_colors[i] = (0.0, 0.0, 0.0, 1.0)
+
+    # Create hue, guidline and shadow arrays
+    hue_histo = hue_histo.astype(float)
+
+    hue_histo_msx = float(np.max(hue_histo))
+    if hue_histo_msx != 0.0:
+        hue_histo /= np.max(hue_histo)
+    guide_histo = np.array([0.05] * N)
+    shadow_histo = np.array([0.0] * N)
+
+    # Compute angels of shadow, template types
+    for sector in harmonic_scheme.sectors:
+        sector_center = sector.center
+        sector_width = sector.width
+        end = int((sector_center + sector_width / 2) % 360)
+        start = int((sector_center - sector_width / 2) % 360)
+
+        if start < end:
+            shadow_histo[start: end] = 1.0
+        else:
+            shadow_histo[start: 360] = 1.0
+            shadow_histo[0: end] = 1.0
+
+    # Plot, 1280 * 800
+    fig = plt.figure(figsize=(3.2, 4))
+    ax = fig.add_subplot(111, projection='polar')
+    # add hue histogram
+    ax.bar(theta, hue_histo, width=width, bottom=0.0, color=hue_colors, alpha=1.0)
+    # add guidline
+    ax.bar(theta, guide_histo, width=width, bottom=1.0, color=hue_colors, alpha=1.0)
+    # add shadow angels for the template types
+    ax.bar(theta, shadow_histo, width=width, bottom=0.0, color=shadow_colors, alpha=0.1)
+    ax.set_title(caption, pad=15)
+
+    plt.close()
+
+    return fig
+
+
+# https://stackoverflow.com/questions/7821518/matplotlib-save-plot-to-numpy-array
+def get_img_from_fig(fig, dpi=100):
+    """
+    a function which returns an image as numpy array from figure
+    """
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", dpi=dpi)
+    buf.seek(0)
+    img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
+    buf.close()
+    img = cv2.imdecode(img_arr, 1)
+    return img