Merge pull request SpenserCai#34 from SpenserCai/dev

Support sd-webui 1.6.0 add,Pre-Pre-Decolorization
ratribiana · Sep 7, 2023 · 49380c1 · 49380c1
2 parents 2fabd0e + 88fb3fa
commit 49380c1
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/scripts/deoldify_base.py b/scripts/deoldify_base.py
@@ -3,12 +3,14 @@
 Date: 2023-08-09 09:58:27
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2023-08-09 10:00:43
+LastEditTime: 2023-09-07 10:32:51
 Description: file content
 '''
 from modules import shared
 from deoldify import device as deoldify_device
 from deoldify.device_id import DeviceId
+import cv2
+from PIL.Image import Image
 
 device_id = shared.cmd_opts.device_id
 
@@ -23,4 +25,10 @@
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
 warnings.filterwarnings("ignore", category=UserWarning, message="The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.")
-warnings.filterwarnings("ignore", category=FutureWarning, message="Arguments other than a weight enum or `None`.*?")
+warnings.filterwarnings("ignore", category=FutureWarning, message="Arguments other than a weight enum or `None`.*?")
+
+def Decolorization(img:Image)->Image:
+    # 通过opencv吧图片褪色成黑白
+    img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2GRAY)
+    img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    return Image.fromarray(img)
diff --git a/scripts/postprocessing_deoldify.py b/scripts/postprocessing_deoldify.py
@@ -3,7 +3,7 @@
 Date: 2023-07-28 14:41:28
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2023-08-09 10:11:23
+LastEditTime: 2023-09-07 10:49:32
 Description: file content
 '''
 # DeOldify UI & Processing
@@ -25,21 +25,26 @@ def ui(self):
             # 一个名为artistic的复选框，初始值是False
             artistic = gr.Checkbox(label="Artistic")
             artistic.value = False
+            pre_decolorization = gr.Checkbox(label="Pre-Decolorization",info="For yellowed photos, this option can be used to fade to black and white before coloring")
+            pre_decolorization.value = False
 
         return {
             "is_enabled": is_enabled,
             "render_factor": render_factor,
             "artistic": artistic,
+            "pre_decolorization": pre_decolorization
         }
 
-    def process_image(self, image, render_factor, artistic):
+    def process_image(self, image, render_factor, artistic, pre_decolorization):
+        if pre_decolorization:
+            image = Decolorization(image)
         vis = get_image_colorizer(root_folder=Path(paths_internal.models_path),render_factor=render_factor, artistic=artistic)
         outImg = vis.get_transformed_image_from_image(image, render_factor=render_factor)
         return outImg
 
-    def process(self, pp: scripts_postprocessing.PostprocessedImage, is_enabled, render_factor, artistic):
+    def process(self, pp: scripts_postprocessing.PostprocessedImage, is_enabled, render_factor, artistic, pre_decolorization):
         if not is_enabled or is_enabled is False:
             return
 
-        pp.image = self.process_image(pp.image, render_factor, artistic)
-        pp.info["deoldify"] = f"render_factor={render_factor}, artistic={artistic}"
+        pp.image = self.process_image(pp.image, render_factor, artistic, pre_decolorization)
+        pp.info["deoldify"] = f"render_factor={render_factor}, artistic={artistic}, pre_decolorization={pre_decolorization}"
diff --git a/scripts/ui_deoldify.py b/scripts/ui_deoldify.py
@@ -3,7 +3,7 @@
 Date: 2023-08-06 20:15:12
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2023-08-21 17:29:20
+LastEditTime: 2023-09-07 10:21:59
 Description: file content
 '''
 from modules import script_callbacks, paths_internal
@@ -40,12 +40,12 @@ def deoldify_tab():
         with gr.Tab("Video"):
             with gr.Row():
                 with gr.Column():
-                    video_input = gr.inputs.Video(label="原视频")
+                    video_input = gr.Video(label="原视频")
                     # 一个名为render_factor的滑块，范围是1-50，初始值是35，步长是1
                     render_factor = gr.Slider(minimum=1, maximum=50, step=1, label="Render Factor")
                     render_factor.value = 35
                 with gr.Column():
-                    video_output = gr.outputs.Video(label="修复后的视频")
+                    video_output = gr.Video(label="修复后的视频",interactive=False)
             run_button = gr.Button(label="Run")
             run_button.click(inputs=[video_input,render_factor],outputs=[video_output],fn=process_image)