Add color correction on loopback for img2img (add new dependency)

frontend/frontend.py

@@ -256,12 +256,18 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
 
                 img2img_image_editor_mode.change(
                     uifn.change_image_editor_mode,
-                    [img2img_image_editor_mode, img2img_image_editor, img2img_resize, img2img_width, img2img_height],
+                    [img2img_image_editor_mode,
+                     img2img_image_editor,
+                     img2img_image_mask,
+                     img2img_resize,
+                     img2img_width,
+                     img2img_height
+                     ],
                     [img2img_image_editor, img2img_image_mask, img2img_btn_editor, img2img_btn_mask,
                      img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength]
                 )
 
-                # img2img_image_editor.edit(
+                # img2img_image_editor_mode.change(
                 #     uifn.update_image_mask,
                 #     [img2img_image_editor, img2img_resize, img2img_width, img2img_height],
                 #     img2img_image_mask

frontend/ui_functions.py

@@ -6,14 +6,17 @@
 import re
 
 
-def change_image_editor_mode(choice, cropped_image, resize_mode, width, height):
+def change_image_editor_mode(choice, cropped_image, masked_image, resize_mode, width, height):
     if choice == "Mask":
-        return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)]
-    return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)]
+        update_image_result = update_image_mask(cropped_image, resize_mode, width, height)
+        return [gr.update(visible=False), update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)]
+
+    update_image_result = update_image_mask(masked_image["image"], resize_mode, width, height)
+    return [update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)]
 
 def update_image_mask(cropped_image, resize_mode, width, height):
     resized_cropped_image = resize_image(resize_mode, cropped_image, width, height) if cropped_image else None
-    return gr.update(value=resized_cropped_image)
+    return gr.update(value=resized_cropped_image, visible=True)
 
 def toggle_options_gfpgan(selection):
     if 0 in selection:

webui.py

@@ -1,5 +1,7 @@
 import argparse, os, sys, glob, re
 
+import cv2
+
 from frontend.frontend import draw_gradio_ui
 from frontend.job_manager import JobManager, JobInfo
 from frontend.ui_functions import resize_image
@@ -1388,7 +1390,17 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name)
         history = []
         initial_seed = None
 
+        do_color_correction = False
+        try:
+            from skimage import exposure
+            do_color_correction = True
+        except:
+            print("Install scikit-image to perform color correction on loopback")
+
         for i in range(n_iter):
+            if do_color_correction and i == 0:
+                correction_target = cv2.cvtColor(np.asarray(init_img.copy()), cv2.COLOR_RGB2LAB)
+
             output_images, seed, info, stats = process_images(
                 outpath=outpath,
                 func_init=init,
@@ -1430,6 +1442,17 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name)
                 initial_seed = seed
 
             init_img = output_images[0]
+
+            if do_color_correction and correction_target is not None:
+                init_img = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
+                    cv2.cvtColor(
+                        np.asarray(init_img),
+                        cv2.COLOR_RGB2LAB
+                    ),
+                    correction_target,
+                    channel_axis=2
+                ), cv2.COLOR_LAB2RGB).astype("uint8"))
+
             if not random_seed_loopback:
                 seed = seed + 1
             else:

webui_playground.py

@@ -65,13 +65,13 @@ def txt2img(*args, **kwargs):
 def img2img(*args, **kwargs):
     return mock_processing(
         prompt=args[0],
-        seed=args[12],
-        width=args[14],
-        height=args[13],
-        steps=args[5],
-        cfg_scale=args[10],
-        sampler=args[6],
-        batch_count=args[9]
+        seed=args[11],
+        width=args[13],
+        height=args[12],
+        steps=args[4],
+        cfg_scale=args[9],
+        sampler=args[5],
+        batch_count=args[8]
     )