improvement in the Recolor OTA algorithim

recolorOTAlgo.py

@@ -1,52 +1,44 @@
 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))
+def color_transfer(source, target):
+    source = cv2.cvtColor(source, cv2.COLOR_BGR2Lab).astype("float32")
+    target = cv2.cvtColor(target, cv2.COLOR_BGR2Lab).astype("float32")
 
-    source_hist = source_hist.astype(np.float64) / source_hist.sum()
-    target_hist = target_hist.astype(np.float64) / target_hist.sum()
+    lMeanSrc, aMeanSrc, bMeanSrc = (source[..., i].mean() for i in range(3))
+    lStdSrc, aStdSrc, bStdSrc = (source[..., i].std() for i in range(3))
 
-    r = np.arange(256).reshape((-1, 1))
-    c = np.arange(256).reshape((1, -1))
-    M = (r - c) ** 2
-    M = M / M.max()
+    lMeanTar, aMeanTar, bMeanTar = (target[..., i].mean() for i in range(3))
+    lStdTar, aStdTar, bStdTar = (target[..., i].std() for i in range(3))
 
-    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
+    (l, a, b) = cv2.split(source)
+    l -= lMeanSrc
+    a -= aMeanSrc
+    b -= bMeanSrc
 
-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)
+    l = (lStdTar / lStdSrc) * l
+    a = (aStdTar / aStdSrc) * a
+    b = (bStdTar / bStdSrc) * b
 
-    # Transfer all channels L, a, and b
-    for ch in range(3):
-        source_lab[:, :, ch] = transfer_channel(source_lab[:, :, ch], target_lab[:, :, ch])
+    l += lMeanTar
+    a += aMeanTar
+    b += bMeanTar
 
-    transferred_rgb = cv2.cvtColor(source_lab.astype(np.uint8), cv2.COLOR_Lab2BGR)
-    return transferred_rgb
-    
+    l = np.clip(l, 0, 255)
+    a = np.clip(a, 0, 255)
+    b = np.clip(b, 0, 255)
 
-def rgb_to_hex(rgb):
-    return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2])
+    transfer = cv2.merge([l, a, b])
+    transfer = cv2.cvtColor(transfer.astype("uint8"), cv2.COLOR_Lab2BGR)
 
-def hex_to_rgb(hex):
-    hex = hex.lstrip('#')
-    return tuple(int(hex[i:i+2], 16) for i in (0, 2, 4))
+    return transfer
 
+def hex_to_rgb(hex_val):
+    hex_val = hex_val.lstrip('#')
+    return tuple(int(hex_val[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):
@@ -54,44 +46,53 @@ def create_color_palette(colors, palette_width=800, palette_height=200):
         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()
+    img.putdata(pixels)
     return img
 
-
-# if __name__ == "__main__":
-#     source = cv2.imread("estampa-test.png")
-#     target = cv2.imread("color.png")
-#     transferred = optimal_transport_color_transfer(source, target)
-
-#     smooth = False#test with true to have a different result.
-#     if smooth:
-#         # Apply bilateral filtering
-#         diameter = 30  # diameter of each pixel neighborhood, adjust based on your image size
-#         sigma_color = 25  # larger value means colors farther to each other will mix together
-#         sigma_space = 25  # larger values means farther pixels will influence each other if their colors are close enough
-#         smoothed = cv2.bilateralFilter(transferred, diameter, sigma_color, sigma_space)
-#         cv2.imwrite("result_OTA.jpg", smoothed)
-#     else:
-#         cv2.imwrite("result_OTA.jpg", transferred)
-
-
-def recolor(source, colors):
+def recolor_statistical(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#test with true for different results.
-    if smooth:
-        # Apply bilateral filtering
-        diameter = 10  # diameter of each pixel neighborhood, adjust based on your image size
-        sigma_color = 25  # larger value means colors farther to each other will mix together
-        sigma_space = 15  # larger values means farther pixels will influence each other if their colors are close enough
-        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
+    recolored = color_transfer(source, palette_bgr)
+    
+    # Adjust bilateral filtering parameters
+    diameter = 15  
+    sigma_color = 30  
+    sigma_space = 20  
+    smoothed = cv2.bilateralFilter(recolored, diameter, sigma_color, sigma_space)
+    
+    # Save image at maximum quality
+    filename = "result.jpg"
+    cv2.imwrite(filename, smoothed, [cv2.IMWRITE_JPEG_QUALITY, 100])
+    return filename
 
+# Usage
+# source_img = cv2.imread("path_to_your_source_image.jpg")
+# hexcolors = ["#db5a1e", "#555115", "#9a690e", "#1f3a19", "#da8007", 
+#              "#9a0633", "#b70406", "#d01b4b", "#e20b0f", "#f7515d"]
+# result_path = recolor_statistical(source_img, hexcolors)
+
+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#test with true for different results.
+    # if smooth:
+    #     # Apply bilateral filtering
+    #     diameter = 10  # diameter of each pixel neighborhood, adjust based on your image size
+    #     sigma_color = 25  # larger value means colors farther to each other will mix together
+    #     sigma_space = 15  # larger values means farther pixels will influence each other if their colors are close enough
+    #     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
+    source = source.astype(np.uint8)
+    source = cv2.cvtColor(source, cv2.COLOR_RGB2BGR)
+    
+    source_img = source
+    hexcolors = ["#db5a1e", "#555115", "#9a690e", "#1f3a19", "#da8007", 
+                "#9a0633", "#b70406", "#d01b4b", "#e20b0f", "#f7515d"]
+    result_path = recolor_statistical(source_img, colors)