import os import random import textwrap import cv2 import gradio as gr import numpy as np from PIL import Image from cv2.ximgproc import guidedFilter from imgutils.data import load_image from imgutils.restore import restore_with_nafnet, restore_with_scunet def dynamic_clean_adverse( input_image: Image.Image, diameter_min: int = 4, diameter_max: int = 6, sigma_color_min: float = 6.0, sigma_color_max: float = 10.0, sigma_space_min: float = 6.0, sigma_space_max: float = 10.0, radius_min: int = 3, radius_max: int = 6, eps_min: float = 16.0, eps_max: float = 24.0, b_iters: int = 64, g_iters: int = 8, ): img = np.array(input_image).astype(np.float32) y = img.copy() for _ in range(b_iters): diameter = random.randint(diameter_min, diameter_max) sigma_color = random.random() * (sigma_color_max - sigma_color_min) + sigma_color_min sigma_space = random.random() * (sigma_space_max - sigma_space_min) + sigma_space_min y = cv2.bilateralFilter(y, diameter, sigma_color, sigma_space) for _ in range(g_iters): radius = random.randint(radius_min, radius_max) eps = random.random() * (eps_max - eps_min) + eps_min y = guidedFilter(img, y, radius, eps) output_image = Image.fromarray(y.clip(0, 255).astype(np.uint8)) return output_image def clean( image: Image.Image, diameter_min: int = 4, diameter_max: int = 6, sigma_color_min: float = 6.0, sigma_color_max: float = 10.0, sigma_space_min: float = 6.0, sigma_space_max: float = 10.0, radius_min: int = 3, radius_max: int = 6, eps_min: float = 16.0, eps_max: float = 24.0, b_iters: int = 64, g_iters: int = 8, use_scunet_clean: bool = False, use_nafnet_clean: bool = False ) -> Image.Image: image = load_image(image) image = dynamic_clean_adverse( image, diameter_min, diameter_max, sigma_color_min, sigma_color_max, sigma_space_min, sigma_space_max, radius_min, radius_max, eps_min, eps_max, b_iters, g_iters ) if use_scunet_clean: image = restore_with_scunet(image) if use_nafnet_clean: image = restore_with_nafnet(image) return image if __name__ == '__main__': with gr.Blocks() as demo: with gr.Row(): gr_markdown = gr.Markdown(textwrap.dedent(""" Cleaner for [MIST](https://github.com/mist-project/mist-v2)(**M**IST **I**s **S**tupid **T**rash) noises. Inspired by https://github.com/lllyasviel/AdverseCleaner * **Update 2023.12.18**, allow random dynamic adversarial clean and iterate steps. """).strip()) with gr.Row(): with gr.Column(): gr_input_image = gr.Image(label='Input Image', type="pil") gr_submit = gr.Button(value='MIST = MIST is Stupid Trash', variant='primary') with gr.Accordion("Advanced Config", open=False): with gr.Row(): gr_diameter_min = gr.Slider( minimum=1, maximum=30, step=1, value=4, label="Diameter Min (default = 4)", interactive=True, ) gr_diameter_max = gr.Slider( minimum=1, maximum=30, step=1, value=6, label="Diameter Max (default = 6)", interactive=True, ) with gr.Row(): gr_sigma_color_min = gr.Slider( minimum=1, maximum=30, step=1, value=6, label="SigmaColor Min (default = 6)", interactive=True, ) gr_sigma_color_max = gr.Slider( minimum=1, maximum=30, step=1, value=10, label="SigmaColor Max (default = 10)", interactive=True, ) with gr.Row(): gr_sigma_space_min = gr.Slider( minimum=1, maximum=30, step=1, value=6, label="SigmaSpace Min (default = 6)", interactive=True, ) gr_sigma_space_max = gr.Slider( minimum=1, maximum=30, step=1, value=10, label="SigmaSpace Max (default = 10)", interactive=True, ) with gr.Row(): gr_radius_min = gr.Slider( minimum=1, maximum=30, step=1, value=3, label="Radius Min (default = 3)", interactive=True, ) gr_radius_max = gr.Slider( minimum=1, maximum=30, step=1, value=6, label="Radius Max (default = 6)", interactive=True, ) with gr.Row(): gr_eps_min = gr.Slider( minimum=1, maximum=30, step=1, value=16, label="Accuracy Min (default = 16)", interactive=True, ) gr_eps_max = gr.Slider( minimum=1, maximum=30, step=1, value=24, label="Accuracy Max (default = 24)", interactive=True, ) with gr.Row(): gr_b_iters = gr.Slider( minimum=1, maximum=256, step=1, value=64, label="Bilateral Filter Iters (default = 64)", interactive=True, ) gr_g_iters = gr.Slider( minimum=1, maximum=32, step=1, value=8, label="Guided Filter Iters (default = 8)", interactive=True, ) with gr.Accordion("Extra Restoration", open=False): with gr.Row(): gr_scunet = gr.Checkbox(label='Use SCUNET', value=False) gr_nafnet = gr.Checkbox(label='Use NAFNET', value=False) with gr.Column(): gr_output_image = gr.Image(label='Output Image', type="pil") gr_submit.click( fn=clean, inputs=[ gr_input_image, gr_diameter_min, gr_diameter_max, gr_sigma_color_min, gr_sigma_color_max, gr_sigma_space_min, gr_sigma_space_max, gr_radius_min, gr_radius_max, gr_eps_min, gr_eps_max, gr_b_iters, gr_g_iters, gr_scunet, gr_nafnet, ], outputs=[gr_output_image], ) demo.queue(os.cpu_count()).launch()