i

README.md

@@ -121,10 +121,10 @@ This may also result in several interesting applications.
 ![p](github_misc/12.png)
 
 
-# Luminance Removal
+# Color Reshuffle
 
-Removing the luminance in images. Will result in a useful ControlNet for relighting.
+Randomly shuffle the color of each pixel but preserve the image structure. 
 
-We recommend to multiply this result image with random color to synthesize different albedo if they are used in your training.
+Useful in training re-coloring ControlNets.
 
-![p](github_misc/13b.png)
+![p](github_misc/14.png)

annotator/shuffle/__init__.py

@@ -1,13 +1,28 @@
 import cv2
 import numpy as np
-from annotator.util import make_noise_disk
+from annotator.util import make_noise_disk, min_max_norm
 
 
-class ShuffleDetector:
+class ContentShuffleDetector:
     def __call__(self, img):
         H, W, C = img.shape
         F = 256
         x = make_noise_disk(H, W, 1, F) * float(W - 1)
         y = make_noise_disk(H, W, 1, F) * float(H - 1)
         flow = np.stack([x, y], axis=2).astype(np.float32)
         return cv2.remap(img, flow, None, cv2.INTER_LINEAR)
+
+
+class ColorShuffleDetector:
+    def __call__(self, img):
+        H, W, C = img.shape
+        F = 256
+        A = make_noise_disk(H, W, 3, F)
+        B = make_noise_disk(H, W, 3, F)
+        C = (A + B) / 2.0
+        A = (C + (A - C) * 2.0).clip(0, 1)
+        B = (C + (B - C) * 2.0).clip(0, 1)
+        L = img.astype(np.float32) / 255.0
+        Y = A * L + B * (1 - L)
+        Y = min_max_norm(Y) * 255.0
+        return Y.clip(0, 255).astype(np.uint8)

annotator/util.py

@@ -57,9 +57,15 @@ def nms(x, t, s):
 
 
 def make_noise_disk(H, W, C, F):
-    noise = np.random.uniform(low=0, high=1, size=(2 * H // F, 2 * W // F, C))
-    noise = cv2.resize(noise, (W * 2, H * 2), interpolation=cv2.INTER_CUBIC)
-    noise = noise[H // 2: H // 2 + H, W // 2: W // 2 + W]
+    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F: F + H, F: F + W]
     noise -= np.min(noise)
     noise /= np.max(noise)
     return noise
+
+
+def min_max_norm(x):
+    x -= np.min(x)
+    x /= np.maximum(np.max(x), 1e-5)
+    return x

gradio_annotator.py

@@ -172,29 +172,29 @@ def oneformer_ade20k(img, res):
     return [result]
 
 
-model_shuffler = None
+model_content_shuffler = None
 
 
-def shuffler(img, res):
+def content_shuffler(img, res):
     img = resize_image(HWC3(img), res)
-    global model_shuffler
-    if model_shuffler is None:
-        from annotator.shuffle import ShuffleDetector
-        model_shuffler = ShuffleDetector()
-    result = model_shuffler(img)
+    global model_content_shuffler
+    if model_content_shuffler is None:
+        from annotator.shuffle import ContentShuffleDetector
+        model_content_shuffler = ContentShuffleDetector()
+    result = model_content_shuffler(img)
     return [result]
 
 
-model_deluma = None
+model_color_shuffler = None
 
 
-def deluma(img, res):
+def color_shuffler(img, res):
     img = resize_image(HWC3(img), res)
-    global model_deluma
-    if model_deluma is None:
-        from annotator.deluma import DelumaDetector
-        model_deluma = DelumaDetector()
-    result = model_deluma(img)
+    global model_color_shuffler
+    if model_color_shuffler is None:
+        from annotator.shuffle import ColorShuffleDetector
+        model_color_shuffler = ColorShuffleDetector()
+    result = model_color_shuffler(img)
     return [result]
 
 
@@ -359,18 +359,18 @@ def deluma(img, res):
             run_button = gr.Button(label="Run")
         with gr.Column():
             gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
-    run_button.click(fn=shuffler, inputs=[input_image, resolution], outputs=[gallery])
+    run_button.click(fn=content_shuffler, inputs=[input_image, resolution], outputs=[gallery])
 
     with gr.Row():
-        gr.Markdown("## Luminance Removal")
+        gr.Markdown("## Color Shuffle")
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(source='upload', type="numpy")
             resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
             run_button = gr.Button(label="Run")
         with gr.Column():
             gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
-    run_button.click(fn=deluma, inputs=[input_image, resolution], outputs=[gallery])
+    run_button.click(fn=color_shuffler, inputs=[input_image, resolution], outputs=[gallery])
 
 
 block.launch(server_name='0.0.0.0')