clean

app.py

@@ -1,8 +1,8 @@
 import argparse
-import spaces
+# import spaces
 from visualcloze import VisualClozeModel
 import gradio as gr
-import demo_tasks
+import examples
 from functools import partial
 from data.prefix_instruction import get_layout_instruction
 from huggingface_hub import snapshot_download
@@ -15,11 +15,8 @@ default_grid_w = 3
 default_upsampling_noise = 0.4
 default_steps = 30
 
-GUIDANCE = """
-
 
-## 📧 Contact: 
-Need help or have questions? Contact us at: lizhongyu [AT] mail.nankai.edu.cn.
+GUIDANCE = """
 
 ## 📋 Quick Start Guide:
 1. Adjust **Number of In-context Examples**, 0 disables in-context learning.
@@ -41,6 +38,7 @@ When generating three images in a 3x4 grid, i.e., Image to Depth + Normal + Hed,
 the runtime is approximately **110s**. 
 **Deploying locally with an 80G A100 can reduce the runtime by more than half.**
 
+### Note: For better quality, you can deploy the demo locally using the [model](https://huggingface.co/VisualCloze/VisualCloze/blob/main/visualcloze-512-lora.pth), which supports a higher resolution than this online demo, by following the instructions in the [GitHub repository](https://github.com/lzyhha/VisualCloze/tree/main?tab=readme-ov-file#2-web-demo-gradio).
 """
 
 CITATION = r"""
@@ -53,13 +51,17 @@ If our work is useful for your research, please consider citing:
 @article{li2025visualcloze,
   title={VisualCloze: A Universal Image Generation Framework via Visual In-Context Learning},
   author={Li, Zhong-Yu and Du, ruoyi and Yan, Juncheng and Zhuo, Le and Li, Zhen and Gao, Peng and Ma, Zhanyu and Cheng, Ming-Ming},
-  booktitle={arXiv preprint arxiv:},
+  journal={arXiv preprint arxiv:},
   year={2025}
 }
 ```
 📋 **License**
 <br>
-This project is licensed under xxx.
+This project is licensed under apache-2.0.
+
+📧 **Contact**
+<br>
+Need help or have questions? Contact us at: lizhongyu [AT] mail.nankai.edu.cn.
 """
 
 NOTE = r"""
@@ -150,7 +152,7 @@ def create_demo(model):
                     steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=default_steps, step=1)
                     cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=50.0, value=30, step=1)
                     upsampling_steps = gr.Slider(label="Upsampling steps (SDEdit)", minimum=1, maximum=100.0, value=10, step=1)
-                    upsampling_noise = gr.Slider(label="Upsampling noise (SDEdit)", minimum=0, maximum=1.0, value=default_upsampling_noise, step=0.01)
+                    upsampling_noise = gr.Slider(label="Upsampling noise (SDEdit)", minimum=0, maximum=1.0, value=default_upsampling_noise, step=0.05)
 
                 gr.Markdown(CITATION)
 
@@ -170,70 +172,70 @@ def create_demo(model):
                 gr.Markdown("# Task Examples")
                 text_dense_prediction_tasks = gr.Textbox(label="Task", visible=False)
                 dense_prediction_tasks = gr.Dataset(
-                    samples=demo_tasks.dense_prediction_text, 
+                    samples=examples.dense_prediction_text, 
                     label='Dense Prediction', 
                     samples_per_page=1000, 
                     components=[text_dense_prediction_tasks])
                 
                 text_conditional_generation_tasks = gr.Textbox(label="Task", visible=False)
                 conditional_generation_tasks = gr.Dataset(
-                    samples=demo_tasks.conditional_generation_text, 
+                    samples=examples.conditional_generation_text, 
                     label='Conditional Generation',
                     samples_per_page=1000, 
                     components=[text_conditional_generation_tasks])
                 
                 text_image_restoration_tasks = gr.Textbox(label="Task", visible=False)
                 image_restoration_tasks = gr.Dataset(
-                    samples=demo_tasks.image_restoration_text, 
+                    samples=examples.image_restoration_text, 
                     label='Image Restoration',
                     samples_per_page=1000, 
                     components=[text_image_restoration_tasks])
                 
                 text_style_transfer_tasks = gr.Textbox(label="Task", visible=False)
                 style_transfer_tasks = gr.Dataset(
-                    samples=demo_tasks.style_transfer_text, 
+                    samples=examples.style_transfer_text, 
                     label='Style Transfer', 
                     samples_per_page=1000, 
                     components=[text_style_transfer_tasks])
                 
                 text_style_condition_fusion_tasks = gr.Textbox(label="Task", visible=False)
                 style_condition_fusion_tasks = gr.Dataset(
-                    samples=demo_tasks.style_condition_fusion_text, 
+                    samples=examples.style_condition_fusion_text, 
                     label='Style Condition Fusion', 
                     samples_per_page=1000, 
                     components=[text_style_condition_fusion_tasks]) 
                 
                 text_tryon_tasks = gr.Textbox(label="Task", visible=False)
                 tryon_tasks = gr.Dataset(
-                    samples=demo_tasks.tryon_text, 
+                    samples=examples.tryon_text, 
                     label='Virtual Try-On', 
                     samples_per_page=1000, 
                     components=[text_tryon_tasks])
 
                 text_relighting_tasks = gr.Textbox(label="Task", visible=False)
                 relighting_tasks = gr.Dataset(
-                    samples=demo_tasks.relighting_text, 
+                    samples=examples.relighting_text, 
                     label='Relighting', 
                     samples_per_page=1000, 
                     components=[text_relighting_tasks])
                 
                 text_photodoodle_tasks = gr.Textbox(label="Task", visible=False)
                 photodoodle_tasks = gr.Dataset(
-                    samples=demo_tasks.photodoodle_text, 
+                    samples=examples.photodoodle_text, 
                     label='Photodoodle', 
                     samples_per_page=1000, 
                     components=[text_photodoodle_tasks])   
 
                 text_editing_tasks = gr.Textbox(label="Task", visible=False)
                 editing_tasks = gr.Dataset(
-                    samples=demo_tasks.editing_text, 
+                    samples=examples.editing_text, 
                     label='Editing', 
                     samples_per_page=1000, 
                     components=[text_editing_tasks])
 
                 text_unseen_tasks = gr.Textbox(label="Task", visible=False)
                 unseen_tasks = gr.Dataset(
-                    samples=demo_tasks.unseen_tasks_text, 
+                    samples=examples.unseen_tasks_text, 
                     label='Unseen Tasks (May produce unstable effects)', 
                     samples_per_page=1000, 
                     components=[text_unseen_tasks]) 
@@ -241,42 +243,42 @@ def create_demo(model):
                 gr.Markdown("# Subject-driven Tasks Examples")
                 text_subject_driven_tasks = gr.Textbox(label="Task", visible=False)
                 subject_driven_tasks = gr.Dataset(
-                    samples=demo_tasks.subject_driven_text, 
+                    samples=examples.subject_driven_text, 
                     label='Subject-driven Generation', 
                     samples_per_page=1000, 
                     components=[text_subject_driven_tasks])
                 
                 text_condition_subject_fusion_tasks = gr.Textbox(label="Task", visible=False)
                 condition_subject_fusion_tasks = gr.Dataset(
-                    samples=demo_tasks.condition_subject_fusion_text, 
+                    samples=examples.condition_subject_fusion_text, 
                     label='Condition+Subject Fusion', 
                     samples_per_page=1000, 
                     components=[text_condition_subject_fusion_tasks])
                 
                 text_style_transfer_with_subject_tasks = gr.Textbox(label="Task", visible=False)
                 style_transfer_with_subject_tasks = gr.Dataset(
-                    samples=demo_tasks.style_transfer_with_subject_text, 
+                    samples=examples.style_transfer_with_subject_text, 
                     label='Style Transfer with Subject', 
                     samples_per_page=1000, 
                     components=[text_style_transfer_with_subject_tasks])
 
                 text_condition_subject_style_fusion_tasks = gr.Textbox(label="Task", visible=False)
                 condition_subject_style_fusion_tasks = gr.Dataset(
-                    samples=demo_tasks.condition_subject_style_fusion_text, 
+                    samples=examples.condition_subject_style_fusion_text, 
                     label='Condition+Subject+Style Fusion', 
                     samples_per_page=1000, 
                     components=[text_condition_subject_style_fusion_tasks])
 
                 text_editing_with_subject_tasks = gr.Textbox(label="Task", visible=False)
                 editing_with_subject_tasks = gr.Dataset(
-                    samples=demo_tasks.editing_with_subject_text, 
+                    samples=examples.editing_with_subject_text, 
                     label='Editing with Subject', 
                     samples_per_page=1000, 
                     components=[text_editing_with_subject_tasks])
 
                 text_image_restoration_with_subject_tasks = gr.Textbox(label="Task", visible=False)
                 image_restoration_with_subject_tasks = gr.Dataset(
-                    samples=demo_tasks.image_restoration_with_subject_text, 
+                    samples=examples.image_restoration_with_subject_text, 
                     label='Image Restoration with Subject', 
                     samples_per_page=1000, 
                     components=[text_image_restoration_with_subject_tasks])
@@ -318,7 +320,6 @@ def create_demo(model):
         def generate_image(*inputs):
             images = []
             if grid_h.value + 1 != model.grid_h or grid_w.value != model.grid_w:
-                print(grid_h.value, grid_w.value, model.grid_h, model.grid_w, type(grid_h.value), type(model.grid_h))
                 raise gr.Error('Please wait for the loading to complete.')
             for i in range(model.grid_h):
                 images.append([])
@@ -337,7 +338,7 @@ def create_demo(model):
                     upsampling_steps=upsampling_steps, upsampling_noise=upsampling_noise
                 )
             except Exception as e:
-                raise gr.Error('Process error. Possible that the task examples have not finished loading yet. Error: ' + str(e))
+                raise gr.Error('Process error. Possible that the task examples have not finished loading yet. Error: ' + e)
 
             output = gr.update(
                 elem_id='output_gallery', 
@@ -375,104 +376,104 @@ def create_demo(model):
             update_grid(cur_hrid_h, cur_hrid_w)
             output = gr.update(
                 elem_id='output_gallery', 
-                value=output, 
-                columns=min(len(output), 2),
-                rows=int(len(output) / 2 + 0.5))
+                value=[o for o, m in zip(output, mask) if m == 1], 
+                columns=min(sum(mask), 2),
+                rows=int(sum(mask) / 2 + 0.5))
             return [output] + current_example + state
         
         dense_prediction_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_dense_prediction_tasks), 
+            partial(process_tasks, func=examples.process_dense_prediction_tasks), 
             inputs=[dense_prediction_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full", 
             show_progress_on=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps] + [generate_btn])
 
         conditional_generation_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_conditional_generation_tasks), 
+            partial(process_tasks, func=examples.process_conditional_generation_tasks), 
             inputs=[conditional_generation_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         image_restoration_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_image_restoration_tasks), 
+            partial(process_tasks, func=examples.process_image_restoration_tasks), 
             inputs=[image_restoration_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         style_transfer_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_style_transfer_tasks), 
+            partial(process_tasks, func=examples.process_style_transfer_tasks), 
             inputs=[style_transfer_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         style_condition_fusion_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_style_condition_fusion_tasks), 
+            partial(process_tasks, func=examples.process_style_condition_fusion_tasks), 
             inputs=[style_condition_fusion_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
             
         relighting_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_relighting_tasks), 
+            partial(process_tasks, func=examples.process_relighting_tasks), 
             inputs=[relighting_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         tryon_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_tryon_tasks), 
+            partial(process_tasks, func=examples.process_tryon_tasks), 
             inputs=[tryon_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
         
         photodoodle_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_photodoodle_tasks), 
+            partial(process_tasks, func=examples.process_photodoodle_tasks), 
             inputs=[photodoodle_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         editing_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_editing_tasks), 
+            partial(process_tasks, func=examples.process_editing_tasks), 
             inputs=[editing_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
         
         unseen_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_unseen_tasks), 
+            partial(process_tasks, func=examples.process_unseen_tasks), 
             inputs=[unseen_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         subject_driven_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_subject_driven_tasks), 
+            partial(process_tasks, func=examples.process_subject_driven_tasks), 
             inputs=[subject_driven_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         style_transfer_with_subject_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_style_transfer_with_subject_tasks), 
+            partial(process_tasks, func=examples.process_style_transfer_with_subject_tasks), 
             inputs=[style_transfer_with_subject_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         condition_subject_fusion_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_condition_subject_fusion_tasks), 
+            partial(process_tasks, func=examples.process_condition_subject_fusion_tasks), 
             inputs=[condition_subject_fusion_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         condition_subject_style_fusion_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_condition_subject_style_fusion_tasks), 
+            partial(process_tasks, func=examples.process_condition_subject_style_fusion_tasks), 
             inputs=[condition_subject_style_fusion_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         editing_with_subject_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_editing_with_subject_tasks), 
+            partial(process_tasks, func=examples.process_editing_with_subject_tasks), 
             inputs=[editing_with_subject_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
 
         image_restoration_with_subject_tasks.click(
-            partial(process_tasks, func=demo_tasks.process_image_restoration_with_subject_tasks), 
+            partial(process_tasks, func=examples.process_image_restoration_with_subject_tasks), 
             inputs=[image_restoration_with_subject_tasks], 
             outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
             show_progress="full")
@@ -496,7 +497,7 @@ def create_demo(model):
 
 
 # @spaces.GPU(duration=120)
-@spaces.GPU(duration=10)
+# @spaces.GPU(duration=10)
 def generate(
     images, 
     prompts, 
@@ -514,7 +515,7 @@ def generate(
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default="models/visualcloze-384-lora.pth")
+    parser.add_argument("--model_path", type=str, default="checkpoints/visualcloze-384-lora.pth")
     parser.add_argument("--precision", type=str, choices=["fp32", "bf16", "fp16"], default="bf16")
     parser.add_argument("--resolution", type=int, default=384)
     return parser.parse_args()
@@ -523,7 +524,7 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
     
-    # snapshot_download(repo_id="VisualCloze/VisualCloze", repo_type="model", local_dir="models")
+    # snapshot_download(repo_id="VisualCloze/VisualCloze", repo_type="model", local_dir="checkpoints")
 
     # Initialize model
     model = VisualClozeModel(resolution=args.resolution, model_path=args.model_path, precision=args.precision)
@@ -532,5 +533,5 @@ if __name__ == "__main__":
     demo = create_demo(model)
     
     # Start Gradio server
-    demo.launch()
-    # demo.launch(share=False, server_port=10050, server_name="0.0.0.0")
+    # demo.launch()
+    demo.launch(share=False, server_port=10050, server_name="0.0.0.0")

data/degradation_toolkit/add_degradation_various.py

@@ -0,0 +1,401 @@
+import os
+import numpy as np
+import random
+import cv2
+import math
+from scipy import special
+from skimage import restoration
+
+import torch
+from torch.nn import functional as F
+from torchvision.utils import make_grid
+
+
+def uint2single(img):
+    return np.float32(img/255.)
+
+
+def single2uint(img):
+    return np.uint8((img.clip(0, 1)*255.).round())
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    """Convert torch Tensors into image numpy arrays.
+    After clamping to [min, max], values will be normalized to [0, 1].
+    Args:
+        tensor (Tensor or list[Tensor]): Accept shapes:
+            1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
+            2) 3D Tensor of shape (3/1 x H x W);
+            3) 2D Tensor of shape (H x W).
+            Tensor channel should be in RGB order.
+        rgb2bgr (bool): Whether to change rgb to bgr.
+        out_type (numpy type): output types. If ``np.uint8``, transform outputs
+            to uint8 type with range [0, 255]; otherwise, float type with
+            range [0, 1]. Default: ``np.uint8``.
+        min_max (tuple[int]): min and max values for clamp.
+    Returns:
+        (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
+        shape (H x W). The channel order is BGR.
+    """
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
+
+    if torch.is_tensor(tensor):
+        tensor = [tensor]
+    result = []
+    for _tensor in tensor:
+        _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+        _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
+
+        n_dim = _tensor.dim()
+        if n_dim == 4:
+            img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if rgb2bgr:
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 3:
+            img_np = _tensor.numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if img_np.shape[2] == 1:  # gray image
+                img_np = np.squeeze(img_np, axis=2)
+            else:
+                if rgb2bgr:
+                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 2:
+            img_np = _tensor.numpy()
+        else:
+            raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
+        if out_type == np.uint8:
+            # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
+            img_np = (img_np * 255.0).round()
+        img_np = img_np.astype(out_type)
+        result.append(img_np)
+    if len(result) == 1:
+        result = result[0]
+    return result
+
+
+def get_noise(img, value=10):
+
+    noise = np.random.uniform(0, 256, img.shape[0:2])
+    
+    v = value * 0.01
+    noise[np.where(noise < (256 - v))] = 0
+
+    k = np.array([[0, 0.1, 0],
+                  [0.1, 8, 0.1],
+                  [0, 0.1, 0]])
+
+    noise = cv2.filter2D(noise, -1, k)
+
+    '''cv2.imshow('img',noise)
+    cv2.waitKey()
+    cv2.destroyWindow('img')'''
+    return noise
+
+
+def rain_blur(noise, length=10, angle=0, w=1):
+
+    trans = cv2.getRotationMatrix2D((length / 2, length / 2), angle - 45, 1 - length / 100.0)
+    dig = np.diag(np.ones(length))  
+    k = cv2.warpAffine(dig, trans, (length, length))  
+    k = cv2.GaussianBlur(k, (w, w), 0)  
+
+    blurred = cv2.filter2D(noise, -1, k)  
+
+    cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
+    blurred = np.array(blurred, dtype=np.uint8)
+
+    rain = np.expand_dims(blurred, 2)
+    blurred = np.repeat(rain, 3, 2)
+
+    return blurred
+
+
+def add_rain(img,value):
+    if np.max(img) > 1:
+        pass
+    else:
+        img = img*255
+
+
+    w, h, c = img.shape
+    h = h - (h % 4)
+    w = w - (w % 4)
+    img = img[0:w, 0:h, :]
+
+
+    w = np.random.choice([3, 5, 7, 9, 11], p=[0.2, 0.2, 0.2, 0.2, 0.2])
+    length = np.random.randint(30, 41) 
+    angle = np.random.randint(-45, 45) 
+
+    noise = get_noise(img, value=value)
+    rain = rain_blur(noise, length=length, angle=angle, w=w)
+
+    img = img.astype('float32') + rain
+    np.clip(img, 0, 255, out=img)
+    img = img/255.0
+    return img
+
+
+def add_rain_range(img, value_min, value_max):
+    value = np.random.randint(value_min, value_max)
+    if np.max(img) > 1:
+        pass
+    else:
+        img = img*255
+
+
+    w, h, c = img.shape
+    h = h - (h % 4)
+    w = w - (w % 4)
+    img = img[0:w, 0:h, :]
+
+
+    w = np.random.choice([3, 5, 7, 9, 11], p=[0.2, 0.2, 0.2, 0.2, 0.2])
+    length = np.random.randint(30, 41) 
+    angle = np.random.randint(-45, 45) 
+
+    noise = get_noise(img, value=value)
+    rain = rain_blur(noise, length=length, angle=angle, w=w)
+
+    img = img.astype('float32') + rain
+    np.clip(img, 0, 255, out=img)
+    img = img/255.0
+    return img
+
+
+def add_Poisson_noise(img, level=2):
+    # input range[0, 1]
+    vals = 10**(level)  
+    img = np.random.poisson(img * vals).astype(np.float32) / vals
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_Gaussian_noise(img, level=20):
+    # input range[0, 1]
+    noise_level = level / 255.0
+    noise_map = np.random.normal(loc=0.0, scale=1.0, size=img.shape)*noise_level
+    img += noise_map
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_Gaussian_noise_range(img, min_level=10, max_level=50):
+    # input range[0, 1]
+    level = random.uniform(min_level, max_level)
+    noise_level = level / 255.0
+    noise_map = np.random.normal(loc=0.0, scale=1.0, size=img.shape)*noise_level
+    img += noise_map
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_sp_noise(img, snr=0.95, salt_pro=0.5):
+    # input range[0, 1]
+    output = np.copy(img)
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            rdn = random.random()
+            if rdn < snr:
+                output[i][j] = img[i][j]
+            else:
+                rdn = random.random()
+                if rdn < salt_pro:
+                    output[i][j] = 1
+                else:
+                    output[i][j] = 0
+    
+    return output
+
+
+def add_JPEG_noise(img, level):
+
+    quality_factor = level
+    img = single2uint(img)
+    _, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
+    img = cv2.imdecode(encimg, 1)
+    img = uint2single(img)
+    
+    return img
+
+
+def add_JPEG_noise_range(img, level_min, level_max):
+
+    quality_factor = random.randint(level_min, level_max)
+    img = single2uint(img)
+    _, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
+    img = cv2.imdecode(encimg, 1)
+    img = uint2single(img)
+    
+    return img
+
+
+def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):
+    """2D sinc filter, ref: https://dsp.stackexchange.com/questions/58301/2-d-circularly-symmetric-low-pass-filter
+
+    Args:
+        cutoff (float): cutoff frequency in radians (pi is max)
+        kernel_size (int): horizontal and vertical size, must be odd.
+        pad_to (int): pad kernel size to desired size, must be odd or zero.
+    """
+    assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
+    kernel = np.fromfunction(
+        lambda x, y: cutoff * special.j1(cutoff * np.sqrt(
+            (x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2)) / ((2 * np.pi * np.sqrt(
+            (x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2)) + 1e-9), [kernel_size, kernel_size])
+    kernel[(kernel_size - 1) // 2, (kernel_size - 1) // 2] = cutoff ** 2 / (4 * np.pi)
+    kernel = kernel / np.sum(kernel)
+    if pad_to > kernel_size:
+        pad_size = (pad_to - kernel_size) // 2
+        kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
+    return kernel
+
+
+def filter2D(img, kernel):
+    """PyTorch version of cv2.filter2D
+    Args:
+        img (Tensor): (b, c, h, w)
+        kernel (Tensor): (b, k, k)
+    """
+    k = kernel.size(-1)
+    b, c, h, w = img.size()
+    if k % 2 == 1:
+        img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
+    else:
+        raise ValueError('Wrong kernel size')
+
+    ph, pw = img.size()[-2:]
+
+    if kernel.size(0) == 1:
+        # apply the same kernel to all batch images
+        img = img.view(b * c, 1, ph, pw)
+        kernel = kernel.view(1, 1, k, k)
+        return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
+    else:
+        img = img.view(1, b * c, ph, pw)
+        kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
+        return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
+
+
+def sinc(img, kernel_size,omega_c):
+
+    sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
+    sinc_kernel = torch.FloatTensor(sinc_kernel)
+
+    img = filter2D(img,sinc_kernel)
+
+    return img
+
+
+def add_ringing(img):
+    # input: [0, 1]
+    img = img2tensor([img])[0].unsqueeze(0)
+    ks = 15
+    omega_c = round(1.2, 2)
+    img = sinc(img, ks, omega_c)
+    img = torch.clamp((img * 255.0).round(), 0, 255) / 255.
+    img = tensor2img(img, min_max=(0, 1))
+    img = img/255.0
+    return img
+
+
+def low_light(img, lum_scale):
+    img = img*lum_scale
+    return img
+
+
+def low_light_range(img):
+    lum_scale = random.uniform(0.1, 0.5)
+    img = img*lum_scale
+    return img
+
+
+def iso_GaussianBlur(img, window, sigma):
+    img = cv2.GaussianBlur(img.copy(), (window, window), sigma)
+    return img
+
+
+def iso_GaussianBlur_range(img, window, min_sigma=2, max_sigma=4):
+    sigma = random.uniform(min_sigma, max_sigma)
+    img = cv2.GaussianBlur(img.copy(), (window, window), sigma)
+    return img
+
+
+def add_resize(img):
+    ori_H, ori_W = img.shape[0], img.shape[1]
+    rnum = np.random.rand()
+    if rnum > 0.8:  # up
+        sf1 = random.uniform(1, 2)
+    elif rnum < 0.7:  # down
+        sf1 = random.uniform(0.2, 1)
+    else:
+        sf1 = 1.0
+    img = cv2.resize(img, (int(sf1*img.shape[1]), int(sf1*img.shape[0])), interpolation=random.choice([1, 2, 3]))
+    img = cv2.resize(img, (int(ori_W), int(ori_H)), interpolation=random.choice([1, 2, 3]))
+    
+    img = np.clip(img, 0.0, 1.0)
+
+    return img
+
+
+def r_l(img):
+    img = img2tensor([img],bgr2rgb=False)[0].unsqueeze(0)
+    psf = np.ones((1, 1, 5, 5))
+    psf = psf / psf.sum()
+    img = img.numpy()
+    img = np.pad(img, ((0, 0), (0, 0), (7, 7), (7, 7)), 'linear_ramp')
+    img = restoration.richardson_lucy(img, psf, 1)
+    img = img[:, :, 7:-7, 7:-7]
+    img = torch.from_numpy(img)
+    img = img.squeeze(0).numpy().transpose(1, 2, 0)
+    return img
+
+
+def inpainting(img,l_num,l_thick):
+
+    ori_h, ori_w = img.shape[0], img.shape[1]
+    mask = np.zeros((ori_h, ori_w, 3), np.uint8)
+    col = random.choice(['white', 'black'])
+    while (l_num):
+        x1, y1 = random.randint(0, ori_w), random.randint(0, ori_h)
+        x2, y2 = random.randint(0, ori_w), random.randint(0, ori_h)
+        pts = np.array([[x1, y1], [x2, y2]], np.int32)
+        pts = pts.reshape((-1, 1, 2))
+        mask = cv2.polylines(mask, [pts], 0, (1, 1, 1), l_thick)
+        l_num -= 1
+
+    if col == 'white':
+        img = np.clip(img + mask, 0, 1)
+    else:
+        img = np.clip(img - mask, 0, 1)
+
+    return img

data/degradation_toolkit/degradation_utils.py

@@ -0,0 +1,232 @@
+import numpy as np
+import cv2
+import random
+from PIL import Image
+
+from data.degradation_toolkit.add_degradation_various import *
+from data.degradation_toolkit.image_operators import *
+from data.degradation_toolkit.x_distortion import *
+
+
+degradation_list1 = [
+    'blur',
+    'noise',
+    'compression',
+    'brighten',
+    'darken',
+    'spatter',
+    'contrast_strengthen',
+    'contrast_weaken',
+    'saturate_strengthen',
+    'saturate_weaken',
+    'oversharpen',
+    'pixelate',
+    'quantization',
+]
+
+
+degradation_list2 = [
+    'Rain', 
+    'Ringing', 
+    'r_l', 
+    'Inpainting', 
+    'mosaic', 
+    'SRx2', 
+    'SRx4',
+    'GaussianNoise',
+    'GaussianBlur',
+    'JPEG',
+    'Resize',
+    'SPNoise',
+    'LowLight',
+    'PoissonNoise',
+    'gray',
+    'ColorDistortion',
+]
+
+
+degradation_list3 = [
+    'Laplacian', 
+    'Canny',
+    'Sobel',
+    'Defocus',
+    'Mosaic',
+    'Barrel',
+    'Pincushion',
+    'Spatter',
+    'Elastic',
+    'Frost',
+    'Contrast',
+]
+
+
+degradation_list4 = [
+    'flip',
+    'rotate90',
+    'rotate180',
+    'rotate270',
+    'identity',
+]
+
+
+all_degradation_types = degradation_list1 + degradation_list2 + degradation_list3 + degradation_list4
+
+
+def single2uint(img):
+    return np.uint8((img.clip(0, 1) * 255.0).round())
+
+
+def uint2single(img):
+    return np.float32(img / 255.0)
+
+
+def add_x_distortion_single_images(img_gt1, deg_type):
+    # np.uint8, BGR
+    x_distortion_dict = distortions_dict
+    severity = random.choice([1, 2, 3, 4, 5])
+    if deg_type == 'compression' or deg_type == "quantization":
+        severity = min(3, severity)
+    deg_type = random.choice(x_distortion_dict[deg_type])
+
+    img_gt1 = cv2.cvtColor(img_gt1, cv2.COLOR_BGR2RGB)
+    img_lq1 = globals()[deg_type](img_gt1, severity)
+
+    img_gt1 = cv2.cvtColor(img_gt1, cv2.COLOR_RGB2BGR)
+    img_lq1 = cv2.cvtColor(img_lq1, cv2.COLOR_RGB2BGR)
+
+    return img_lq1, img_gt1, deg_type
+
+
+def add_degradation_single_images(img_gt1, deg_type):
+    if deg_type == 'Rain':
+        value = random.uniform(40, 200)
+        img_lq1 = add_rain(img_gt1, value=value)
+    elif deg_type == 'Ringing':
+        img_lq1 = add_ringing(img_gt1)
+    elif deg_type == 'r_l':
+        img_lq1 = r_l(img_gt1)
+    elif deg_type == 'Inpainting':
+        l_num = random.randint(20, 50)
+        l_thick = random.randint(10, 20)
+        img_lq1 = inpainting(img_gt1, l_num=l_num, l_thick=l_thick)
+    elif deg_type == 'mosaic':
+        img_lq1 = mosaic_CFA_Bayer(img_gt1)
+    elif deg_type == 'SRx2':
+        H, W, _ = img_gt1.shape
+        img_lq1 = cv2.resize(img_gt1, (W//2, H//2), interpolation=cv2.INTER_CUBIC)
+        img_lq1 = cv2.resize(img_lq1, (W, H), interpolation=cv2.INTER_CUBIC)
+    elif deg_type == 'SRx4':
+        H, W, _ = img_gt1.shape
+        img_lq1 = cv2.resize(img_gt1, (W//4, H//4), interpolation=cv2.INTER_CUBIC)
+        img_lq1 = cv2.resize(img_lq1, (W, H), interpolation=cv2.INTER_CUBIC)
+
+    elif deg_type == 'GaussianNoise':
+        level = random.uniform(10, 50)
+        img_lq1 = add_Gaussian_noise(img_gt1, level=level)
+    elif deg_type == 'GaussianBlur':
+        sigma = random.uniform(2, 4)
+        img_lq1 = iso_GaussianBlur(img_gt1, window=15, sigma=sigma)
+    elif deg_type == 'JPEG':
+        level = random.randint(10, 40)
+        img_lq1 = add_JPEG_noise(img_gt1, level=level)
+    elif deg_type == 'Resize':
+        img_lq1 = add_resize(img_gt1)
+    elif deg_type == 'SPNoise':
+        img_lq1 = add_sp_noise(img_gt1)
+    elif deg_type == 'LowLight':
+        lum_scale = random.uniform(0.3, 0.4)
+        img_lq1 = low_light(img_gt1, lum_scale=lum_scale)
+    elif deg_type == 'PoissonNoise':
+        img_lq1 = add_Poisson_noise(img_gt1, level=2)
+    elif deg_type == 'gray':
+        img_lq1 = cv2.cvtColor(img_gt1, cv2.COLOR_BGR2GRAY)
+        img_lq1 = np.expand_dims(img_lq1, axis=2)
+        img_lq1 = np.concatenate((img_lq1, img_lq1, img_lq1), axis=2)
+    elif deg_type == 'None':
+        img_lq1 = img_gt1
+    elif deg_type == 'ColorDistortion':
+        if random.random() < 0.5:
+            channels = list(range(3))
+            random.shuffle(channels) 
+            img_lq1 = img_gt1[..., channels]
+        else:
+            channel = random.randint(0, 2)
+            img_lq1 = img_gt1.copy()
+            if random.random() < 0.5:
+                img_lq1[..., channel] = 0
+            else:
+                img_lq1[..., channel] = 1
+    else:
+        print('Error!', '-', deg_type, '-')
+        exit()
+    img_lq1 = np.clip(img_lq1 * 255, 0, 255).round().astype(np.uint8)
+    img_lq1 = img_lq1.astype(np.float32) / 255.0
+    img_gt1 = np.clip(img_gt1 * 255, 0, 255).round().astype(np.uint8)
+    img_gt1 = img_gt1.astype(np.float32) / 255.0
+
+    return img_lq1, img_gt1
+
+
+def calculate_operators_single_images(img_gt1, deg_type):
+    img_gt1 = img_gt1.copy()
+    
+    if deg_type == 'Laplacian':
+        img_lq1 = Laplacian_edge_detector(img_gt1)
+    elif deg_type == 'Canny':
+        img_lq1 = Canny_edge_detector(img_gt1)
+    elif deg_type == 'Sobel':
+        img_lq1 = Sobel_edge_detector(img_gt1)
+    elif deg_type == 'Defocus':
+        img_lq1 = defocus_blur(img_gt1, level=(3, 0.2))
+    elif deg_type == 'Mosaic':
+        img_lq1 = mosaic_CFA_Bayer(img_gt1)
+    elif deg_type == 'Barrel':
+        img_lq1 = simulate_barrel_distortion(img_gt1, k1=0.1, k2=0.05)
+    elif deg_type == 'Pincushion':
+        img_lq1 = simulate_pincushion_distortion(img_gt1, k1=-0.1, k2=-0.05)
+    elif deg_type == 'Spatter':
+        img_lq1 = uint2single(spatter((img_gt1), severity=1))
+    elif deg_type == 'Elastic':
+        img_lq1 = elastic_transform((img_gt1), severity=4)
+    elif deg_type == 'Frost':
+        img_lq1 = uint2single(frost(img_gt1, severity=4))
+    elif deg_type == 'Contrast':
+        img_lq1 = adjust_contrast(img_gt1, clip_limit=4.0, tile_grid_size=(4, 4))
+        
+    if np.mean(img_lq1).astype(np.float16) == 0:
+        print(deg_type, 'prompt&query zero images.')
+        img_lq1 = img_gt1.copy()
+        
+    return img_lq1, img_gt1
+
+
+def add_degradation(image, deg_type):
+    if deg_type in degradation_list1:
+        list_idx = 1
+        img_lq1, _, _ = add_x_distortion_single_images(np.copy(image), deg_type)
+        img_lq1 = uint2single(img_lq1)
+    elif deg_type in degradation_list2:
+        list_idx = 2
+        img_lq1, _ = add_degradation_single_images(np.copy(uint2single(image)), deg_type)
+    elif deg_type in degradation_list3:
+        list_idx = 3
+        if deg_type in ['Laplacian', 'Canny', 'Sobel', 'Frost']:
+            img_lq1, _ = calculate_operators_single_images(np.copy(image), deg_type)
+        else:
+            img_lq1, _ = calculate_operators_single_images(np.copy(uint2single(image)), deg_type)
+        if img_lq1.max() > 1:
+            img_lq1 = uint2single(img_lq1)
+    elif deg_type in degradation_list4:
+        list_idx = 4
+        img_lq1 = np.copy(uint2single(image))
+        if deg_type == 'flip':
+            img_lq1 = np.flip(img_lq1, axis=1)
+        elif deg_type == 'rotate90':
+            img_lq1 = np.rot90(img_lq1, k=1)
+        elif deg_type == 'rotate180':
+            img_lq1 = np.rot90(img_lq1, k=2)
+        elif deg_type == 'rotate270':
+            img_lq1 = np.rot90(img_lq1, k=3)
+        elif deg_type == 'identity':
+            pass
+    return Image.fromarray(single2uint(img_lq1)), list_idx

data/degradation_toolkit/image_operators.py

@@ -0,0 +1,420 @@
+import os
+import cv2
+import numpy as np
+import argparse 
+from skimage.filters import gaussian
+from scipy.ndimage.interpolation import map_coordinates
+from tqdm import tqdm
+from PIL import Image
+
+
+def single2uint(img):
+    return np.uint8((img.clip(0, 1)*255.).round())
+
+
+def uint2single(img):
+    return np.float32(img/255.)
+
+
+def Laplacian_edge_detector(img):
+    # input: [0, 1]
+    # return: [0, 1] (H, W, 3)
+    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
+    img = cv2.GaussianBlur(img, (3, 3), 0)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img = cv2.Laplacian(img, cv2.CV_16S) # (H, W)
+    img = cv2.convertScaleAbs(img)
+    img = img.astype(np.float32) / 255.
+    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
+    return img
+
+
+def Laplacian_edge_detector_uint8(img):
+    # input: [0, 255]
+    # return: [0, 255] (H, W, 3)
+    img = cv2.GaussianBlur(img, (3, 3), 0)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img = cv2.Laplacian(img, cv2.CV_16S) # (H, W)
+    img = cv2.convertScaleAbs(img)
+    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
+    return img
+
+
+def Canny_edge_detector(img):
+    # input: [0, 1]
+    # return: [0, 1] (H, W, 3)
+    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
+    img = cv2.GaussianBlur(img, (3, 3), 0)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img = cv2.Canny(img, 50, 200) # (H, W)
+    img = cv2.convertScaleAbs(img)
+    img = img.astype(np.float32) / 255.
+    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
+    return img
+
+
+def Canny_edge_detector_uint8(img):
+    # input: [0, 255]
+    # return: [0, 255] (H, W, 3)
+    img = cv2.GaussianBlur(img, (3, 3), 0)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img = cv2.Canny(img, 50, 200) # (H, W)
+    img = cv2.convertScaleAbs(img)
+    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
+    return img
+
+
+def Sobel_edge_detector(img):
+    # input: [0, 1]
+    # return: [0, 1] (H, W, 3)
+    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
+    img = cv2.GaussianBlur(img, (3, 3), 0)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img = cv2.Sobel(img, cv2.CV_16S, 1, 1) # (H, W)
+    img = cv2.convertScaleAbs(img)
+    img = img.astype(np.float32) / 255.
+    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
+    return img
+
+
+def erosion(img, kernel_size=5):
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    img = cv2.erode(img, kernel, iterations=1)
+    return img
+
+
+def dilatation(img, kernel_size=5):
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    img = cv2.dilate(img, kernel, iterations=1)
+    return img
+
+
+def opening(img):
+    return dilatation(erosion(img))
+
+
+def closing(img):
+    return erosion(dilatation(img))
+
+
+def morphological_gradient(img):
+    return dilatation(img) - erosion(img)
+
+
+def top_hat(img):
+    return img - opening(img)
+
+
+def black_hat(img):
+    return closing(img) - img
+
+
+def adjust_contrast(image, clip_limit=2.0, tile_grid_size=(8, 8)):
+    
+    image = single2uint(image)
+    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+
+    l, a, b = cv2.split(lab)
+
+    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
+    l_eq = clahe.apply(l)
+
+    lab_eq = cv2.merge((l_eq, a, b))
+    result = cv2.cvtColor(lab_eq, cv2.COLOR_LAB2BGR)
+
+    result = uint2single(result)
+    return result
+
+
+def embossing(img):
+    kernel = np.array([[0, -1, -1],
+                       [1, 0, -1],
+                       [1, 1, 0]])
+    return cv2.filter2D(img, -1, kernel)
+
+
+def hough_transform_line_detection(img):
+    img = single2uint(img)
+    dst = cv2.Canny(img, 50, 200, apertureSize=3)
+    cdst = cv2.cvtColor(dst, cv2.COLOR_GRAY2BGR)
+    lines = cv2.HoughLinesP(dst, 1, np.pi / 180, 230, None, 0, 0)
+    if lines is not None:
+        for i in range(0, len(lines)):
+            rho = lines[i][0][0]
+            theta = lines[i][0][1]
+            a = np.cos(theta)
+            b = np.sin(theta)
+
+            x0 = a * rho
+            y0 = b * rho
+            pt1 = (int(x0 + 1000*(-b)), int(y0 + 1000*(a)))
+
+            pt2 = (int(x0 - 1000*(-b)), int(y0 - 1000*(a)))
+            cv2.line(img, pt1, pt2, (0, 0, 255), 3, cv2.LINE_AA)
+    
+    return uint2single(img)
+
+
+def hough_circle_detection(img):
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, 100, param1=100, param2=30, minRadius=50, maxRadius=200)
+    circles = np.uint16(np.around(circles))
+    for i in circles[0, :]:
+        cv2.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
+    return img
+
+
+def disk(radius, alias_blur=0.1, dtype=np.float32):
+    if radius <= 8:
+        L = np.arange(-8, 8 + 1)
+        ksize = (3, 3)
+    else:
+        L = np.arange(-radius, radius + 1)
+        ksize = (5, 5)
+    X, Y = np.meshgrid(L, L)
+    aliased_disk = np.array((X ** 2 + Y ** 2) <= radius ** 2, dtype=dtype)
+    aliased_disk /= np.sum(aliased_disk)
+
+    # supersample disk to antialias
+    return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)
+
+
+def defocus_blur(image, level=(1, 0.1)):
+    c = level
+    kernel = disk(radius=c[0], alias_blur=c[1])
+
+    channels = []
+    for d in range(3):
+        channels.append(cv2.filter2D(image[:, :, d], -1, kernel))
+    channels = np.array(channels).transpose((1, 2, 0))  # 3x64x64 -> 64x64x3
+
+    return np.clip(channels, 0, 1)
+
+
+def masks_CFA_Bayer(shape):
+    pattern = "RGGB"
+    channels = dict((channel, np.zeros(shape)) for channel in "RGB")
+    for channel, (y, x) in zip(pattern, [(0, 0), (0, 1), (1, 0), (1, 1)]):
+        channels[channel][y::2, x::2] = 1
+    return tuple(channels[c].astype(bool) for c in "RGB")
+
+
+def cfa4_to_rgb(CFA4):
+    RGB = np.zeros((CFA4.shape[0]*2, CFA4.shape[1]*2, 3), dtype=np.uint8)
+    RGB[0::2, 0::2, 0] = CFA4[:, :, 0]  # R
+    RGB[0::2, 1::2, 1] = CFA4[:, :, 1]  # G on R row
+    RGB[1::2, 0::2, 1] = CFA4[:, :, 2]  # G on B row
+    RGB[1::2, 1::2, 2] = CFA4[:, :, 3]  # B
+
+    return RGB
+
+
+def mosaic_CFA_Bayer(RGB):
+    RGB = single2uint(RGB)
+    R_m, G_m, B_m = masks_CFA_Bayer(RGB.shape[0:2])
+    mask = np.concatenate(
+        (R_m[..., np.newaxis], G_m[..., np.newaxis], B_m[..., np.newaxis]), axis=-1
+    )
+    mosaic = np.multiply(mask, RGB)  # mask*RGB
+    CFA = mosaic.sum(2).astype(np.uint8)
+
+    CFA4 = np.zeros((RGB.shape[0] // 2, RGB.shape[1] // 2, 4), dtype=np.uint8)
+    CFA4[:, :, 0] = CFA[0::2, 0::2]
+    CFA4[:, :, 1] = CFA[0::2, 1::2]
+    CFA4[:, :, 2] = CFA[1::2, 0::2]
+    CFA4[:, :, 3] = CFA[1::2, 1::2]
+
+    rgb = cfa4_to_rgb(CFA4)
+    rgb = uint2single(rgb)
+    return rgb
+
+
+def simulate_barrel_distortion(image, k1=0.02, k2=0.01):
+    height, width = image.shape[:2]
+    mapx, mapy = np.meshgrid(np.arange(width), np.arange(height))
+    mapx = 2 * mapx / (width - 1) - 1
+    mapy = 2 * mapy / (height - 1) - 1
+    r = np.sqrt(mapx**2 + mapy**2)
+    mapx = mapx * (1 + k1 * r**2 + k2 * r**4)
+    mapy = mapy * (1 + k1 * r**2 + k2 * r**4)
+    mapx = (mapx + 1) * (width - 1) / 2
+    mapy = (mapy + 1) * (height - 1) / 2
+    distorted_image = cv2.remap(image, mapx.astype(np.float32), mapy.astype(np.float32), cv2.INTER_LINEAR)
+    return distorted_image
+
+
+def simulate_pincushion_distortion(image, k1=-0.02, k2=-0.01):
+    height, width = image.shape[:2]
+    mapx, mapy = np.meshgrid(np.arange(width), np.arange(height))
+    mapx = 2 * mapx / (width - 1) - 1
+    mapy = 2 * mapy / (height - 1) - 1
+    r = np.sqrt(mapx**2 + mapy**2)
+    mapx = mapx * (1 + k1 * r**2 + k2 * r**4)
+    mapy = mapy * (1 + k1 * r**2 + k2 * r**4)
+    mapx = (mapx + 1) * (width - 1) / 2
+    mapy = (mapy + 1) * (height - 1) / 2
+    distorted_image = cv2.remap(image, mapx.astype(np.float32), mapy.astype(np.float32), cv2.INTER_LINEAR)
+    return distorted_image
+
+
+def rgb2gray(rgb):
+    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
+
+
+def spatter(x, severity=1):
+    c = [(0.65, 0.3, 4, 0.69, 0.6, 0),
+         (0.65, 0.3, 3, 0.68, 0.6, 0),
+         (0.65, 0.3, 2, 0.68, 0.5, 0),
+         (0.65, 0.3, 1, 0.65, 1.5, 1),
+         (0.67, 0.4, 1, 0.65, 1.5, 1)][severity - 1]
+    x_PIL = x
+    x = np.array(x, dtype=np.float32) / 255.
+
+    liquid_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])
+
+    liquid_layer = gaussian(liquid_layer, sigma=c[2])
+    liquid_layer[liquid_layer < c[3]] = 0
+    if c[5] == 0:
+        liquid_layer = (liquid_layer * 255).astype(np.uint8)
+        dist = 255 - cv2.Canny(liquid_layer, 50, 150)
+        dist = cv2.distanceTransform(dist, cv2.DIST_L2, 5)
+        _, dist = cv2.threshold(dist, 20, 20, cv2.THRESH_TRUNC)
+        dist = cv2.blur(dist, (3, 3)).astype(np.uint8)
+        dist = cv2.equalizeHist(dist)
+        ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
+        dist = cv2.filter2D(dist, cv2.CV_8U, ker)
+        dist = cv2.blur(dist, (3, 3)).astype(np.float32)
+
+        m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)
+        m /= np.max(m, axis=(0, 1))
+        m *= c[4]
+        # water is pale turqouise
+        color = np.concatenate((175 / 255. * np.ones_like(m[..., :1]),
+                                238 / 255. * np.ones_like(m[..., :1]),
+                                238 / 255. * np.ones_like(m[..., :1])), axis=2)
+
+        color = cv2.cvtColor(color, cv2.COLOR_BGR2BGRA)
+
+        if len(x.shape) < 3 or x.shape[2] < 3:
+            add_spatter_color = cv2.cvtColor(np.clip(m * color, 0, 1),
+                                             cv2.COLOR_BGRA2BGR)
+            add_spatter_gray = rgb2gray(add_spatter_color)
+
+            return (np.clip(x + add_spatter_gray, 0, 1) * 255).astype(np.uint8)
+
+        else:
+
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2BGRA)
+
+            return (cv2.cvtColor(np.clip(x + m * color, 0, 1),
+                                cv2.COLOR_BGRA2BGR) * 255).astype(np.uint8)
+    else:
+        m = np.where(liquid_layer > c[3], 1, 0)
+        m = gaussian(m.astype(np.float32), sigma=c[4])
+        m[m < 0.8] = 0
+
+        x_rgb = np.array(x_PIL)
+
+        # mud brown
+        color = np.concatenate((63 / 255. * np.ones_like(x_rgb[..., :1]),
+                                42 / 255. * np.ones_like(x_rgb[..., :1]),
+                                20 / 255. * np.ones_like(x_rgb[..., :1])),
+                               axis=2)
+        color *= m[..., np.newaxis]
+        if len(x.shape) < 3 or x.shape[2] < 3:
+            x *= (1 - m)
+            return (np.clip(x + rgb2gray(color), 0, 1) * 255).astype(np.uint8)
+
+        else:
+            x *= (1 - m[..., np.newaxis])
+            return (np.clip(x + color, 0, 1) * 255).astype(np.uint8)
+
+
+# mod of https://gist.github.com/erniejunior/601cdf56d2b424757de5
+def elastic_transform(image, severity=3):
+    image = np.array(image, dtype=np.float32) / 255.
+    shape = image.shape
+    shape_size = shape[:2]
+
+    sigma = np.array(shape_size) * 0.01
+    alpha = [250 * 0.05, 250 * 0.065, 250 * 0.085, 250 * 0.1, 250 * 0.12][
+        severity - 1]
+    max_dx = shape[0] * 0.005
+    max_dy = shape[0] * 0.005
+
+    dx = (gaussian(np.random.uniform(-max_dx, max_dx, size=shape[:2]),
+                   sigma, mode='reflect', truncate=3) * alpha).astype(
+        np.float32)
+    dy = (gaussian(np.random.uniform(-max_dy, max_dy, size=shape[:2]),
+                   sigma, mode='reflect', truncate=3) * alpha).astype(
+        np.float32)
+
+    if len(image.shape) < 3 or image.shape[2] < 3:
+        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1))
+    else:
+        dx, dy = dx[..., np.newaxis], dy[..., np.newaxis]
+        x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]),
+                              np.arange(shape[2]))
+        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx,
+                                                          (-1, 1)), np.reshape(
+            z, (-1, 1))
+    return np.clip(
+        map_coordinates(image, indices, order=1, mode='reflect').reshape(
+            shape), 0, 1) * 255
+
+
+def frost(x, severity=2):
+    c = [(1, 0.4),
+         (0.8, 0.6),
+         (0.7, 0.7),
+         (0.65, 0.7),
+         (0.6, 0.75)][severity - 1]
+
+    idx = np.random.randint(5)
+    filename = [os.path.join("degradation_toolkit/frost", 'frost1.png'),
+                os.path.join("degradation_toolkit/frost", 'frost2.png'),
+                os.path.join("degradation_toolkit/frost", 'frost3.png'),
+                os.path.join("degradation_toolkit/frost", 'frost4.jpg'),
+                os.path.join("degradation_toolkit/frost", 'frost5.jpg'),
+                os.path.join("degradation_toolkit/frost", 'frost6.jpg')][idx]
+    frost = Image.open(filename)
+    frost = frost.convert("RGB")
+    frost = np.array(frost)
+    # frost = cv2.imread(filename)
+    frost = uint2single(frost)
+    frost_shape = frost.shape
+    x_shape = np.array(x).shape
+
+    # resize the frost image so it fits to the image dimensions
+    scaling_factor = 1
+    if frost_shape[0] >= x_shape[0] and frost_shape[1] >= x_shape[1]:
+        scaling_factor = 1
+    elif frost_shape[0] < x_shape[0] and frost_shape[1] >= x_shape[1]:
+        scaling_factor = x_shape[0] / frost_shape[0]
+    elif frost_shape[0] >= x_shape[0] and frost_shape[1] < x_shape[1]:
+        scaling_factor = x_shape[1] / frost_shape[1]
+    elif frost_shape[0] < x_shape[0] and frost_shape[1] < x_shape[
+        1]:  # If both dims are too small, pick the bigger scaling factor
+        scaling_factor_0 = x_shape[0] / frost_shape[0]
+        scaling_factor_1 = x_shape[1] / frost_shape[1]
+        scaling_factor = np.maximum(scaling_factor_0, scaling_factor_1)
+
+    scaling_factor *= 1.1
+    new_shape = (int(np.ceil(frost_shape[1] * scaling_factor)),
+                 int(np.ceil(frost_shape[0] * scaling_factor)))
+    frost_rescaled = cv2.resize(frost, dsize=new_shape,
+                                interpolation=cv2.INTER_CUBIC)
+
+    # randomly crop
+    x_start, y_start = np.random.randint(0, frost_rescaled.shape[0] - x_shape[
+        0]), np.random.randint(0, frost_rescaled.shape[1] - x_shape[1])
+
+    if len(x_shape) < 3 or x_shape[2] < 3:
+        frost_rescaled = frost_rescaled[x_start:x_start + x_shape[0],
+                         y_start:y_start + x_shape[1]]
+        frost_rescaled = rgb2gray(frost_rescaled)
+    else:
+        frost_rescaled = frost_rescaled[x_start:x_start + x_shape[0],
+                         y_start:y_start + x_shape[1]][..., [2, 1, 0]]
+    return c[0] * np.array(x) + c[1] * frost_rescaled

data/degradation_toolkit/x_distortion/__init__.py

@@ -0,0 +1,120 @@
+from .blur import *
+from .brightness import *
+from .quantization import *
+from .compression import *
+from .contrast import *
+from .noise import *
+from .oversharpen import *
+from .pixelate import *
+from .saturate import *
+
+
+def add_distortion(img, severity=1, distortion_name=None):
+    """This function returns a distorted version of the given image.
+
+    @param img (np.ndarray, unit8): Input image, H x W x 3, RGB, [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @distortion_name: 
+    @return: Degraded image (np.ndarray, unit8), H x W x 3, RGB, [0, 255]
+    """
+
+    if not isinstance(img, np.ndarray):
+        raise AttributeError('Expecting type(img) to be numpy.ndarray')
+    if not (img.dtype.type is np.uint8):
+        raise AttributeError('Expecting img.dtype.type to be numpy.uint8')
+
+    if not (img.ndim in [2, 3]):
+        raise AttributeError('Expecting img.shape to be either (h x w) or (h x w x c)')
+    if img.ndim == 2:
+        img = np.stack((img,) * 3, axis=-1)
+
+    h, w, c = img.shape
+    if (h < 32 or w < 32):
+        raise AttributeError('The (w, h) must be at least 32 pixels')
+    if not (c in [1, 3]):
+        raise AttributeError('Expecting img to have either 1 or 3 chennels')
+    if c == 1:
+        img = np.stack((np.squeeze(img),) * 3, axis=-1)
+
+    if severity not in [1, 2, 3, 4, 5]:
+        raise AttributeError('The severity must be an integer in [1, 5]')
+
+    if distortion_name:
+        img_lq = globals()[distortion_name](img, severity)
+    else:
+        raise ValueError("The distortion_name must be passed")
+
+    return np.uint8(img_lq)
+
+
+distortions_dict = {
+    "blur": [
+        "blur_gaussian", 
+        "blur_motion", 
+        "blur_glass", 
+        "blur_lens", 
+        "blur_zoom", 
+        "blur_jitter", 
+    ],
+    "noise": [
+        "noise_gaussian_RGB", 
+        "noise_gaussian_YCrCb", 
+        "noise_speckle", 
+        "noise_spatially_correlated", 
+        "noise_poisson", 
+        "noise_impulse", 
+    ], 
+    "compression": [
+        "compression_jpeg", 
+        "compression_jpeg_2000", 
+    ], 
+    "brighten": [
+        "brightness_brighten_shfit_HSV", 
+        "brightness_brighten_shfit_RGB", 
+        "brightness_brighten_gamma_HSV", 
+        "brightness_brighten_gamma_RGB", 
+    ], 
+    "darken": [
+        "brightness_darken_shfit_HSV", 
+        "brightness_darken_shfit_RGB", 
+        "brightness_darken_gamma_HSV", 
+        "brightness_darken_gamma_RGB", 
+    ], 
+    "contrast_strengthen": [
+        "contrast_strengthen_scale",
+        "contrast_strengthen_stretch", 
+    ],
+    "contrast_weaken": [
+        "contrast_weaken_scale",
+        "contrast_weaken_stretch",
+    ],
+    "saturate_strengthen": [
+        "saturate_strengthen_HSV", 
+        "saturate_strengthen_YCrCb", 
+    ], 
+    "saturate_weaken": [
+        "saturate_weaken_HSV", 
+        "saturate_weaken_YCrCb", 
+    ], 
+    "oversharpen": [
+        "oversharpen", 
+    ], 
+    "pixelate": [
+        "pixelate", 
+    ], 
+    "quantization": [
+        "quantization_otsu", 
+        "quantization_median", 
+        "quantization_hist", 
+    ], 
+    "spatter": [
+        "spatter", 
+    ], 
+}
+
+
+def get_distortion_names(subset=None):
+    if subset in distortions_dict:
+        print(distortions_dict[subset])
+    else:
+        print(distortions_dict)

data/degradation_toolkit/x_distortion/blur.py

@@ -0,0 +1,155 @@
+import cv2
+import numpy as np
+
+from skimage.filters import gaussian
+from .helper import (
+    _motion_blur,
+    shuffle_pixels_njit, 
+    clipped_zoom, 
+    gen_disk, 
+    gen_lensmask, 
+)
+
+
+def blur_gaussian(img, severity=1):
+    """
+    Gaussian Blur. 
+    severity=[1, 2, 3, 4, 5] corresponding to sigma=[1, 2, 3, 4, 5].
+    severity mainly refer to KADID-10K and Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [1, 2, 3, 4, 5][severity - 1]
+    img = np.array(img) / 255.
+    img = gaussian(img, sigma=c, channel_axis=-1)
+    img = np.clip(img, 0, 1) * 255
+    return img.round().astype(np.uint8)
+
+
+def blur_gaussian_lensmask(img, severity=1):
+    """
+    Gaussian Blur with Lens Mask. 
+    severity=[1, 2, 3, 4, 5] corresponding to 
+    [gamma, sigma]=[[2.0, 2], [2.4, 4], [3.0, 6], [3.8, 8], [5.0, 10]].
+    severity mainly refer to PieAPP.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [(2.0, 2), (2.4, 4), (3.0, 6), (3.8, 8), (5.0, 10)][severity - 1]
+    img_orig = np.array(img) / 255.
+    h, w = img.shape[:2]
+    mask = gen_lensmask(h, w, gamma=c[0])[:, :, None]
+    img = gaussian(img_orig, sigma=c[1], channel_axis=-1)
+    img = mask * img_orig + (1 - mask) * img
+    img = np.clip(img, 0, 1) * 255
+    return img.round().astype(np.uint8)
+
+
+def blur_motion(img, severity=1):
+    """
+    Motion Blur. 
+    severity = [1, 2, 3, 4, 5] corresponding to radius=[5, 10, 15, 15, 20] and
+    sigma=[1, 2, 3, 4, 5].
+    severity mainly refer to Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [0, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [(5, 3), (10, 5), (15, 7), (15, 9), (20, 12)][severity - 1]
+    angle = np.random.uniform(-90, 90)
+    img = np.array(img)
+    img = _motion_blur(img, radius=c[0], sigma=c[1], angle=angle)
+    img = np.clip(img, 0, 255)
+    return img.round().astype(np.uint8)
+
+
+def blur_glass(img, severity=1):
+    """
+    Glass Blur. 
+    severity = [1, 2, 3, 4, 5] corresponding to 
+    [sigma, shift, iteration]=[(0.7, 1, 1), (0.9, 2, 1), (1.2, 2, 2), (1.4, 3, 2), (1.6, 4, 2)].
+    severity mainly refer to Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [0, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [(0.7, 1, 1), (0.9, 2, 1), (1.2, 2, 2), (1.4, 3, 2), (1.6, 4, 2)][severity - 1]
+    img = np.array(img) / 255.
+    img = gaussian(img, sigma=c[0], channel_axis=-1)
+    img = shuffle_pixels_njit(img, shift=c[1], iteration=c[2])
+    img = np.clip(gaussian(img, sigma=c[0], channel_axis=-1), 0, 1) * 255
+    return img.round().astype(np.uint8)
+
+
+def blur_lens(img, severity=1):
+    """
+    Lens Blur. 
+    severity = [1, 2, 3, 4, 5] corresponding to radius=[2, 3, 4, 6, 8].
+    severity mainly refer to KADID-10K.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [0, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [2, 3, 4, 6, 8][severity - 1]
+    img = np.array(img) / 255.
+    kernel = gen_disk(radius=c)
+    img_lq = []
+    for i in range(3):
+        img_lq.append(cv2.filter2D(img[:, :, i], -1, kernel))
+    img_lq = np.array(img_lq).transpose((1, 2, 0))
+    img_lq = np.clip(img_lq, 0, 1) * 255
+    return img_lq.round().astype(np.uint8)
+
+
+def blur_zoom(img, severity=1):
+    """
+    Zoom Blur. 
+    severity = [1, 2, 3, 4, 5] corresponding to radius=
+        [np.arange(1, 1.03, 0.02),
+         np.arange(1, 1.06, 0.02),
+         np.arange(1, 1.10, 0.02),
+         np.arange(1, 1.15, 0.02),
+         np.arange(1, 1.21, 0.02)].
+    severity mainly refer to Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [0, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [np.arange(1, 1.03, 0.02),
+         np.arange(1, 1.06, 0.02),
+         np.arange(1, 1.10, 0.02),
+         np.arange(1, 1.15, 0.02),
+         np.arange(1, 1.21, 0.02)][severity - 1]
+    img = (np.array(img) / 255.).astype(np.float32)
+    h, w = img.shape[:2]
+    img_lq = np.zeros_like(img)
+    for zoom_factor in c:
+        zoom_layer = clipped_zoom(img, zoom_factor)
+        img_lq += zoom_layer[:h, :w, :]
+    img_lq = (img + img_lq) / (len(c) + 1)
+    img_lq = np.clip(img_lq, 0, 1) * 255
+    return img_lq.round().astype(np.uint8)
+
+
+def blur_jitter(img, severity=1):
+    """
+    Jitter Blur.
+    severity = [1, 2, 3, 4, 5] corresponding to shift=[1, 2, 3, 4, 5]. 
+    severity mainly refer to KADID-10K.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [0, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [1, 2, 3, 4, 5][severity - 1]
+    img = np.array(img)
+    img_lq = shuffle_pixels_njit(img, shift=c, iteration=1)
+    return np.uint8(img_lq)

data/degradation_toolkit/x_distortion/brightness.py

@@ -0,0 +1,150 @@
+import numpy as np
+import cv2
+from .helper import gen_lensmask
+
+
+def brightness_brighten_shfit_HSV(img, severity=1):
+    """
+    The RGB image is mapping to HSV, and then enhance the brightness by V channel
+    severity=[1,2,3,4,5] is corresponding to c=[0.1, 0.2, 0.3, 0.4, 0.5]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    c = [0.1, 0.2, 0.3, 0.4, 0.5][severity-1]
+    img = np.float32(np.array(img) / 255.)
+    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    img_hsv[:, :, 2] += c
+    img_lq = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_brighten_shfit_RGB(img, severity=1):
+    """
+    The RGB image is directly enhanced by RGB mean shift
+    severity=[1,2,3,4,5] is corresponding to c=[0.1, 0.15, 0.2, 0.27, 0.35]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    c = [0.1, 0.15, 0.2, 0.27, 0.35][severity-1]
+    img = np.float32(np.array(img) / 255.)
+    img_lq = img + c
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_brighten_gamma_RGB(img, severity=1):
+    """
+    The RGB image is enhanced by V channel with a gamma function
+    severity=[1,2,3,4,5] is corresponding to gamma=[0.8, 0.7, 0.6, 0.45, 0.3]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    gamma = [0.8, 0.7, 0.6, 0.45, 0.3][severity-1]
+    img = np.array(img / 255.)
+    img_lq = img ** gamma
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_brighten_gamma_HSV(img, severity=1):
+    """
+    The RGB image is enhanced by V channel with a gamma function
+    severity=[1,2,3,4,5] is corresponding to gamma=[0.7, 0.55, 0.4, 0.25, 0.1]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    gamma = [0.7, 0.58, 0.47, 0.36, 0.25][severity-1]
+    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    img_hsv = np.array(img_hsv / 255.)
+    img_hsv[:, :, 2] = img_hsv[:, :, 2] ** gamma
+    img_lq = np.uint8(np.clip(img_hsv, 0, 1) * 255.)
+    img_lq = cv2.cvtColor(img_lq, cv2.COLOR_HSV2RGB)
+    return img_lq
+
+
+def brightness_darken_shfit_HSV(img, severity=1):
+    """
+    The RGB image is mapping to HSV, and then darken the brightness by V channel
+    severity=[1,2,3,4,5] is corresponding to c=[0.1, 0.2, 0.3, 0.4, 0.5]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    c = [0.1, 0.2, 0.3, 0.4, 0.5][severity-1]
+    img = np.float32(np.array(img) / 255.)
+    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    img_hsv[:, :, 2] -= c
+    img_lq = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_darken_shfit_RGB(img, severity=1):
+    """
+    The RGB image's brightness is directly reduced by RGB mean shift
+    severity=[1,2,3,4,5] is corresponding to c=[0.1, 0.15, 0.2, 0.27, 0.35]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    c = [0.1, 0.15, 0.2, 0.27, 0.35][severity-1]
+    img = np.float32(np.array(img)/255.)
+    img_lq = img - c
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_darken_gamma_RGB(img, severity=1):
+    """
+    The RGB image is darkened by V channel with a gamma function
+    severity=[1,2,3,4,5] is corresponding to gamma=[1.4, 1.7, 2.1, 2.6, 3.2]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    gamma = [1.4, 1.7, 2.1, 2.6, 3.2][severity-1]
+    img = np.array(img / 255.)
+    img_lq = img ** gamma
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def brightness_darken_gamma_HSV(img, severity=1):
+    """
+    The RGB image is enhanced by V channel with a gamma function
+    severity=[1,2,3,4,5] is corresponding to gamma=[1.5, 1.8, 2.2, 2.7, 3.5]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    gamma = [1.5, 1.8, 2.2, 2.7, 3.5][severity-1]
+    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
+    img_hsv = np.array(img_hsv / 255.)
+    img_hsv[:, :, 2] = img_hsv[:, :, 2] ** gamma
+    img_lq = np.uint8(np.clip(img_hsv, 0, 1) * 255.)
+    img_lq = cv2.cvtColor(img_lq, cv2.COLOR_HSV2RGB)
+    return img_lq
+
+
+def brightness_vignette(img, severity=1):
+    """
+    The RGB image is suffered from the vignette effect.
+    severity=[1,2,3,4,5] is corresponding to gamma=[0.5, 0.875, 1.25, 1.625, 2]
+
+    @param img: Input image, H x W x RGB, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x RGB, value range [0, 255]
+    """
+    gamma = [0.5, 0.875, 1.25, 1.625, 2][severity - 1]
+    img = np.array(img)
+    h, w = img.shape[:2]
+    mask = gen_lensmask(h, w, gamma=gamma)[:, :, None]
+    img_lq = mask * img
+    return np.uint8(np.clip(img_lq, 0, 255))

data/degradation_toolkit/x_distortion/compression.py

@@ -0,0 +1,78 @@
+import numpy as np
+from PIL import Image
+from io import BytesIO
+
+
+def compression_jpeg(img, severity=1):
+    """
+    JPEG compression on a NumPy array.
+    severity=[1,2,3,4,5] corresponding to quality=[25,18,15,10,7].
+    from https://github.com/bethgelab/imagecorruptions/blob/master/imagecorruptions/corruptions.py
+
+    @param img: Input image as NumPy array, H x W x C, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image as NumPy array, H x W x C, value range [0, 255]
+    """
+    assert img.dtype == np.uint8, "Image array should have dtype of np.uint8"
+    assert severity in [1, 2, 3, 4, 5], 'Severity must be an integer between 1 and 5.'
+
+    quality = [25, 18, 12, 8, 5][severity - 1]
+    output = BytesIO()
+    gray_scale = False
+    if img.shape[2] == 1:  # Check if the image is grayscale
+        gray_scale = True
+    # Convert NumPy array to PIL Image
+    img = Image.fromarray(img)
+    if gray_scale:
+        img = img.convert('L')
+    else:
+        img = img.convert('RGB')
+    # Save image to a bytes buffer using JPEG compression
+    img.save(output, 'JPEG', quality=quality)
+    output.seek(0)
+    # Load the compressed image from the bytes buffer
+    img_lq = Image.open(output)
+    # Convert PIL Image back to NumPy array
+    if gray_scale:
+        img_lq = np.array(img_lq.convert('L'))
+        img_lq = img_lq.reshape((img_lq.shape[0], img_lq.shape[1], 1))  # Maintaining the original shape (H, W, 1)
+    else:
+        img_lq = np.array(img_lq.convert('RGB'))
+    return img_lq
+
+
+def compression_jpeg_2000(img, severity=1):
+    """
+    JPEG2000 compression on a NumPy array.
+    severity=[1,2,3,4,5] corresponding to quality=[29,27.5,26,24.5,23], quality_mode='dB'.
+
+    @param x: Input image as NumPy array, H x W x C, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image as NumPy array, H x W x C, value range [0, 255]
+    """
+    assert img.dtype == np.uint8, "Image array should have dtype of np.uint8"
+    assert severity in [1, 2, 3, 4, 5], 'Severity must be an integer between 1 and 5.'
+
+    quality = [29, 27.5, 26, 24.5, 23][severity - 1]
+    output = BytesIO()
+    gray_scale = False
+    if img.shape[2] == 1:  # Check if the image is grayscale
+        gray_scale = True
+    # Convert NumPy array to PIL Image
+    img = Image.fromarray(img)
+    if gray_scale:
+        img = img.convert('L')
+    else:
+        img = img.convert('RGB')
+    # Save image to a bytes buffer using JPEG compression
+    img.save(output, 'JPEG2000', quality_mode='dB', quality_layers=[quality])
+    output.seek(0)
+    # Load the compressed image from the bytes buffer
+    img_lq = Image.open(output)
+    # Convert PIL Image back to NumPy array
+    if gray_scale:
+        img_lq = np.array(img_lq.convert('L'))
+        img_lq = img_lq.reshape((img_lq.shape[0], img_lq.shape[1], 1))  # Maintaining the original shape (H, W, 1)
+    else:
+        img_lq = np.array(img_lq.convert('RGB'))
+    return img_lq

data/degradation_toolkit/x_distortion/contrast.py

@@ -0,0 +1,74 @@
+import cv2
+import numpy as np
+from PIL import Image
+from PIL import ImageEnhance
+
+
+def contrast_weaken_scale(img, severity=1):
+    """
+    Contrast Weaken by scaling. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[0.75, 0.6, 0.45, 0.3, 0.2].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.75, 0.6, 0.45, 0.3, 0.2][severity - 1]
+    img = Image.fromarray(img)
+    enhancer = ImageEnhance.Contrast(img)
+    img = enhancer.enhance(c)
+    img = np.uint8(np.clip(np.array(img), 0, 255))
+    return img
+
+
+def contrast_weaken_stretch(img, severity=1):
+    """
+    Contrast Weaken by stretching. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[1.0, 0.9, 0.8, 0.6, 0.4].
+    severity mainly refer to PieAPP.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [1.0, 0.9, 0.8, 0.6, 0.4][severity - 1]
+    img = np.array(img) / 255.
+    img_mean = np.mean(img, axis=(0,1), keepdims=True)
+    img = 1. / (1 + (img_mean / (img + 1e-12)) ** c)
+    img = np.uint8(np.clip(img, 0, 1) * 255)
+    return img
+
+
+def contrast_strengthen_scale(img, severity=1):
+    """
+    Contrast Strengthen by scaling. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[1.4, 1.7, 2.1, 2.6, 4.0].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [1.4, 1.7, 2.1, 2.6, 4.0][severity - 1]
+    img = Image.fromarray(img)
+    enhancer = ImageEnhance.Contrast(img)
+    img = enhancer.enhance(c)
+    img = np.uint8(np.clip(np.array(img), 0, 255))
+    return img
+
+
+def contrast_strengthen_stretch(img, severity=1):
+    """
+    Contrast Strengthen by stretching. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[2.0, 4.0, 6.0, 8.0, 10.0].
+    severity mainly refer to PieAPP.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [2.0, 4.0, 6.0, 8.0, 10.0][severity - 1]
+    img = np.array(img) / 255.
+    img_mean = np.mean(img, axis=(0,1), keepdims=True)
+    img = 1. / (1 + (img_mean / (img + 1e-12)) ** c)
+    img = np.uint8(np.clip(img, 0, 1) * 255)
+    return img

data/degradation_toolkit/x_distortion/helper.py

@@ -0,0 +1,171 @@
+import cv2
+from scipy.ndimage import zoom as scizoom
+from numba import njit, prange
+import numpy as np
+import math
+
+
+def gen_lensmask(h, w, gamma):
+    """
+    Generate lens mask with shape (h, w). 
+    For point (i, j), 
+    distance = [(i - h // 2)^2 + (j - w // 2)^2] ^ (1/2) / [h // 2)^2 + (w // 2)^2] ^ (1/2)
+    mask = scale * (1 - distance) ^ gamma
+    
+    @param h: height
+    @param w: width
+    @param gamma: exponential factor
+    @return: Mask, H x W
+    """
+    dist1 = np.array([list(range(w))] * h) - w // 2
+    dist2 = np.array([list(range(h))] * w) - h // 2
+    dist2 = np.transpose(dist2, (1, 0))
+    dist = np.sqrt((dist1 ** 2 + dist2 ** 2)) / np.sqrt((w ** 2 + h ** 2) / 4)
+    mask = (1 - dist) ** gamma
+    return mask
+
+
+def gen_disk(radius, dtype=np.float32):
+    if radius <= 8:
+        L = np.arange(-8, 8 + 1)
+    else:
+        L = np.arange(-radius, radius + 1)
+    X, Y = np.meshgrid(L, L)
+    disk = np.array((X ** 2 + Y ** 2) <= radius ** 2, dtype=dtype)
+    disk /= np.sum(disk)
+    return disk
+
+
+# modification of https://github.com/FLHerne/mapgen/blob/master/diamondsquare.py
+def plasma_fractal(mapsize=256, wibbledecay=3):
+    """
+    Generate a heightmap using diamond-square algorithm.
+    Return square 2d array, side length 'mapsize', of floats in range 0-255.
+    'mapsize' must be a power of two.
+    """
+    assert (mapsize & (mapsize - 1) == 0)
+    maparray = np.empty((mapsize, mapsize), dtype=np.float_)
+    maparray[0, 0] = 0
+    stepsize = mapsize
+    wibble = 100
+
+    def wibbledmean(array):
+        return array / 4 + wibble * np.random.uniform(-wibble, wibble,
+                                                      array.shape)
+
+    def fillsquares():
+        """For each square of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        cornerref = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        squareaccum = cornerref + np.roll(cornerref, shift=-1, axis=0)
+        squareaccum += np.roll(squareaccum, shift=-1, axis=1)
+        maparray[stepsize // 2:mapsize:stepsize,
+        stepsize // 2:mapsize:stepsize] = wibbledmean(squareaccum)
+
+    def filldiamonds():
+        """For each diamond of points stepsize apart,
+           calculate middle value as mean of points + wibble"""
+        mapsize = maparray.shape[0]
+        drgrid = maparray[stepsize // 2:mapsize:stepsize,
+                 stepsize // 2:mapsize:stepsize]
+        ulgrid = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        ldrsum = drgrid + np.roll(drgrid, 1, axis=0)
+        lulsum = ulgrid + np.roll(ulgrid, -1, axis=1)
+        ltsum = ldrsum + lulsum
+        maparray[0:mapsize:stepsize,
+        stepsize // 2:mapsize:stepsize] = wibbledmean(ltsum)
+        tdrsum = drgrid + np.roll(drgrid, 1, axis=1)
+        tulsum = ulgrid + np.roll(ulgrid, -1, axis=0)
+        ttsum = tdrsum + tulsum
+        maparray[stepsize // 2:mapsize:stepsize,
+        0:mapsize:stepsize] = wibbledmean(ttsum)
+
+    while stepsize >= 2:
+        fillsquares()
+        filldiamonds()
+        stepsize //= 2
+        wibble /= wibbledecay
+
+    maparray -= maparray.min()
+    return maparray / maparray.max()
+
+
+def clipped_zoom(img, zoom_factor):
+    # clipping along the width dimension:
+    ch0 = int(np.ceil(img.shape[0] / float(zoom_factor)))
+    top0 = (img.shape[0] - ch0) // 2
+
+    # clipping along the height dimension:
+    ch1 = int(np.ceil(img.shape[1] / float(zoom_factor)))
+    top1 = (img.shape[1] - ch1) // 2
+
+    img = scizoom(img[top0:top0 + ch0, top1:top1 + ch1],
+                  (zoom_factor, zoom_factor, 1), order=1)
+
+    return img
+
+
+def getOptimalKernelWidth1D(radius, sigma):
+    return radius * 2 + 1
+
+
+def gauss_function(x, mean, sigma):
+    return (np.exp(- (x - mean)**2 / (2 * (sigma**2)))) / (np.sqrt(2 * np.pi) * sigma)
+
+
+def getMotionBlurKernel(width, sigma):
+    k = gauss_function(np.arange(width), 0, sigma)
+    Z = np.sum(k)
+    return k/Z
+
+
+def shift(image, dx, dy):
+    if(dx < 0):
+        shifted = np.roll(image, shift=image.shape[1]+dx, axis=1)
+        shifted[:,dx:] = shifted[:,dx-1:dx]
+    elif(dx > 0):
+        shifted = np.roll(image, shift=dx, axis=1)
+        shifted[:,:dx] = shifted[:,dx:dx+1]
+    else:
+        shifted = image
+
+    if(dy < 0):
+        shifted = np.roll(shifted, shift=image.shape[0]+dy, axis=0)
+        shifted[dy:,:] = shifted[dy-1:dy,:]
+    elif(dy > 0):
+        shifted = np.roll(shifted, shift=dy, axis=0)
+        shifted[:dy,:] = shifted[dy:dy+1,:]
+    return shifted
+
+
+def _motion_blur(x, radius, sigma, angle):
+    width = getOptimalKernelWidth1D(radius, sigma)
+    kernel = getMotionBlurKernel(width, sigma)
+    point = (width * np.sin(np.deg2rad(angle)), width * np.cos(np.deg2rad(angle)))
+    hypot = math.hypot(point[0], point[1])
+
+    blurred = np.zeros_like(x, dtype=np.float32)
+    for i in range(width):
+        dy = -math.ceil(((i*point[0]) / hypot) - 0.5)
+        dx = -math.ceil(((i*point[1]) / hypot) - 0.5)
+        if (np.abs(dy) >= x.shape[0] or np.abs(dx) >= x.shape[1]):
+            # simulated motion exceeded image borders
+            break
+        shifted = shift(x, dx, dy)
+        blurred = blurred + kernel[i] * shifted
+    return blurred
+
+
+# Numba nopython compilation to shuffle_pixles
+@njit()
+def shuffle_pixels_njit(img, shift, iteration):
+    height, width = img.shape[:2]
+    # locally shuffle pixels
+    for _ in range(iteration):
+        for h in range(height - shift, shift, -1):
+            for w in range(width - shift, shift, -1):
+                dx, dy = np.random.randint(-shift, shift, size=(2,))
+                h_prime, w_prime = h + dy, w + dx
+                # swap
+                img[h, w], img[h_prime, w_prime] = img[h_prime, w_prime], img[h, w]
+    return img

data/degradation_toolkit/x_distortion/noise.py

@@ -0,0 +1,117 @@
+import cv2
+import numpy as np
+import skimage as sk
+
+
+def noise_gaussian_RGB(img, severity=1):
+    """
+    Additive Gaussian noise in RGB channels. 
+    severity=[1, 2, 3, 4, 5] is corresponding to sigma=[0.05, 0.1, 0.15, 0.2, 0.25]. 
+    severity mainly refer to KADID-10K and Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    sigma = [0.05, 0.1, 0.15, 0.2, 0.25][severity-1]
+    img = np.array(img) / 255.
+    noise = np.random.normal(0, sigma, img.shape)
+    img_lq = img + noise
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def noise_gaussian_YCrCb(img, severity=1):
+    """
+    Additive Gaussian noise with higher noise in color channels. 
+    severity=[1, 2, 3, 4, 5] is corresponding to
+    sigma_l=[0.05, 0.06, 0.07, 0.08, 0.09], 
+    sigma_r=[1, 1.45, 1.9, 2.35, 2.8], 
+    sigma_b=[1, 1.45, 1.9, 2.35, 2.8]. 
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    sigma_l = [0.05, 0.06, 0.07, 0.08, 0.09][severity-1]
+    sigma_r = sigma_l * [1, 1.45, 1.9, 2.35, 2.8][severity - 1]
+    sigma_b = sigma_l * [1, 1.45, 1.9, 2.35, 2.8][severity - 1]
+    h, w = img.shape[:2]
+    img = np.float32(np.array(img) / 255.)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2YCR_CB)
+    noise_l = np.expand_dims(np.random.normal(0, sigma_l, (h, w)), 2)
+    noise_r = np.expand_dims(np.random.normal(0, sigma_r, (h, w)), 2)
+    noise_b = np.expand_dims(np.random.normal(0, sigma_b, (h, w)), 2)
+    noise = np.concatenate((noise_l, noise_r, noise_b), axis=2)
+    img_lq = np.float32(img + noise)
+    img_lq = cv2.cvtColor(img_lq, cv2.COLOR_YCR_CB2RGB)
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def noise_speckle(img, severity=1):
+    """
+    Multiplicative Gaussian noise. 
+    severity=[1, 2, 3, 4, 5] is corresponding to sigma=[0.14, 0.21, 0.28, 0.35, 0.42].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.14, 0.21, 0.28, 0.35, 0.42][severity - 1]
+    img = np.array(img) / 255.
+    noise = img * np.random.normal(size=img.shape, scale=c)
+    img_lq = img + noise
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def noise_spatially_correlated(img, severity=1):
+    """
+    Spatially correlated noise. 
+    severity=[1, 2, 3, 4, 5] is corresponding to sigma=[0.08, 0.11, 0.14, 0.18, 0.22].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    sigma = [0.08, 0.11, 0.14, 0.18, 0.22][severity - 1]
+    img = np.array(img) / 255.
+    noise = np.random.normal(0, sigma, img.shape)
+    img_lq = img + noise
+    img_lq = cv2.blur(img_lq, [3, 3])
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def noise_poisson(img, severity=1):
+    """
+    Poisson noise. 
+    PieAPP keeps this distortion free of additional parameters.
+    The default:
+    c =  vals = len(np.unique(image))
+    vals = 2 ** np.ceil(np.log2(vals))
+    But Imagecorruptions introduces a extra parameter c
+    ranging [60, 25, 12, 5, 3] for sigma = sqrt(I / c).
+    severity=[1, 2, 3, 4, 5] is corresponding to c=[80, 60, 40, 25, 15].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [80, 60, 40, 25, 15][severity - 1]
+    img = np.array(img) / 255.
+    img_lq = np.random.poisson(img * c) / float(c)
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)
+
+
+def noise_impulse(img, severity=1):
+    """
+    Impulse noise is also known as salt&pepper noise.
+    PieAPP introduce the range [1e-4, 0.045].
+    severity=[1, 2, 3, 4, 5] is corresponding to amount=[0.01, 0.03, 0.05, 0.07, 0.10].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.01, 0.03, 0.05, 0.07, 0.10][severity - 1]
+    img = np.array(img) / 255.
+    img_lq = sk.util.random_noise(img, mode='s&p', amount=c)
+    return np.uint8(np.clip(img_lq, 0, 1) * 255.)

data/degradation_toolkit/x_distortion/oversharpen.py

@@ -0,0 +1,31 @@
+import cv2
+import numpy as np
+
+
+def oversharpen(img, severity=1):
+    """
+    OverSharpening filter on a NumPy array.
+    severity = [1, 5] corresponding to amount = [2, 4, 6, 8, 10]
+    
+    @param x: Input image as NumPy array, H x W x C, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image as NumPy array, H x W x C, value range [0, 255]
+    """
+    assert img.dtype == np.uint8, "Image array should have dtype of np.uint8"
+    assert severity in [1, 2, 3, 4, 5], 'Severity must be an integer between 1 and 5.'
+    
+    amount = [2, 2.8, 4, 6, 8][severity - 1]
+
+    # Setting the kernel size and sigmaX value for Gaussian blur
+    # In OpenCV's Size(kernel_width, kernel_height), both kernel_width and kernel_height
+    # should be odd numbers; for example, we can use (2*radius+1, 2*radius+1)
+    blur_radius = 2  # The radius is the blur radius used to set the size of the Gaussian kernel
+    sigmaX = 0
+
+    # Create a blurred/smoothed version of the image
+    blurred = cv2.GaussianBlur(img, (2*blur_radius+1, 2*blur_radius+1), sigmaX)
+
+    # Compute the sharpened image with an enhancement factor of 'amount'
+    sharpened = cv2.addWeighted(img, 1 + amount, blurred, -amount, 0)
+
+    return sharpened

data/degradation_toolkit/x_distortion/pixelate.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from PIL import Image
+
+
+def pixelate(img, severity=1):
+    """
+    Pixelate. 
+    severity=[1, 2, 3, 4, 5] corresponding to sigma=[0.5, 0.4, 0.3, 0.25, 0.2].
+    severity mainly refer to Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.5, 0.4, 0.3, 0.25, 0.2][severity - 1]
+    h, w = np.array(img).shape[:2]
+    img = Image.fromarray(img)
+    img = img.resize((int(w * c), int(h * c)), Image.BOX)
+    img = img.resize((w, h), Image.NEAREST)
+    return np.array(img).astype(np.uint8)

data/degradation_toolkit/x_distortion/quantization.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+from PIL import Image
+from skimage.filters import threshold_multiotsu
+
+
+
+def quantization_otsu(img, severity=1):
+    """
+    Color Quantization using OTSU method.
+    severity=[1, 2, 3, 4, 5] corresponding to num_classes=[15, 11, 8, 5, 3].
+    severity mainly refer to KADID-10K and Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [15, 11, 8, 5, 3][severity - 1]
+    img = np.array(img).astype(np.float32)
+    for i in range(img.shape[2]):
+        img_gray = img[:, :, i]
+        thresholds = threshold_multiotsu(img_gray, classes=c, nbins=30) # modify skimage
+        v_max = img_gray.max()
+        v_min = img_gray.min()
+        img[:, :, i] = np.digitize(img[:, :, i], bins=thresholds) * (v_max - v_min) / c + v_min
+    img = np.clip(img, 0, 255)
+    return img
+
+
+def quantization_median(img, severity=1):
+    """
+    Color Quantization using Histogram Median.
+    severity=[1, 2, 3, 4, 5] corresponding to num_classes=[20, 15, 10, 6, 3].
+    severity mainly refer to KADID-10K and Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [20, 15, 10, 6, 3][severity - 1]
+    for i in range(img.shape[2]):
+        img_gray = Image.fromarray(img[:, :, i])
+        img_gray = img_gray.quantize(colors=c, method=Image.Quantize.MEDIANCUT).convert("L")
+        img[:, :, i] = np.array(img_gray)
+    img = np.clip(img, 0, 255)
+    return img
+
+
+def quantization_hist(img, severity=1):
+    """
+    Color Quantization using Histogram Equalization.
+    severity=[1, 2, 3, 4, 5] corresponding to num_classes=[24, 16, 8, 6, 4].
+    severity mainly refer to KADID-10K and Imagecorruptions.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [24, 16, 8, 6, 4][severity - 1]
+    hist, _ = np.histogram(img.flatten(), bins=c, range=[0, 255])
+    cdf = hist.cumsum()
+    cdf_m = np.ma.masked_equal(cdf, 0)
+    cdf_m = (cdf_m - cdf_m.min()) * 255 / (cdf_m.max() - cdf_m.min())
+    cdf = np.ma.filled(cdf_m, 0).astype('uint8')
+    img = np.uint8(np.round(img / 255 * (c - 1)))
+    img = cdf[img]
+    img = np.clip(img, 0, 255)
+    return img

data/degradation_toolkit/x_distortion/saturate.py

@@ -0,0 +1,75 @@
+import cv2
+import numpy as np
+
+
+def saturate_weaken_HSV(img, severity=1):
+    """
+    Saturate Weaken by scaling S channel in HSV. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[0.7, 0.55, 0.4, 0.2, 0.0].
+    severity mainly refer to KADID-10K.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.7, 0.55, 0.4, 0.2, 0.0][severity - 1]
+    hsv = np.array(cv2.cvtColor(img, cv2.COLOR_RGB2HSV), dtype=np.float32)
+    hsv[:, :, 1] = c * hsv[:, :, 1]
+    hsv = np.uint8(np.clip(hsv, 0, 255))
+    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+    return img
+
+
+def saturate_weaken_YCrCb(img, severity=1):
+    """
+    Saturate Weaken by scaling S channel in YCrCb. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[0.6, 0.4, 0.2, 0.1, 0.0].
+    severity mainly refer to PieAPP.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [0.6, 0.4, 0.2, 0.1, 0.0][severity - 1]
+    ycrcb = np.array(cv2.cvtColor(img, cv2.COLOR_RGB2YCR_CB), dtype=np.float32)
+    ycrcb[:, :, 1] = 128 + (ycrcb[:, :, 1] - 128) * c
+    ycrcb[:, :, 2] = 128 + (ycrcb[:, :, 2] - 128) * c
+    ycrcb = np.uint8(np.clip(ycrcb, 0, 255))
+    img = cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB)
+    return img
+
+
+def saturate_strengthen_HSV(img, severity=1):
+    """
+    Saturate Strengthen by scaling S channel in HSV. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[3.0, 6.0, 12.0, 20.0, 64.0].
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [3.0, 6.0, 12.0, 20.0, 64.0][severity - 1]
+    hsv = np.array(cv2.cvtColor(img, cv2.COLOR_RGB2HSV), dtype=np.float32)
+    hsv[:, :, 1] = c * hsv[:, :, 1]
+    hsv = np.uint8(np.clip(hsv, 0, 255))
+    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+    return img
+
+
+def saturate_strengthen_YCrCb(img, severity=1):
+    """
+    Saturate Strengthen by scaling S channel in YCrCb. 
+    severity=[1, 2, 3, 4, 5] corresponding to scale=[2.0, 3.0, 5.0, 8.0, 16.0].
+    severity mainly refer to PieAPP.
+
+    @param img: Input image, H x W x 3, value range [0, 255]
+    @param severity: Severity of distortion, [1, 5]
+    @return: Degraded image, H x W x 3, value range [0, 255]
+    """
+    c = [2.0, 3.0, 5.0, 8.0, 16.0][severity - 1]
+    ycrcb = np.array(cv2.cvtColor(img, cv2.COLOR_RGB2YCR_CB), dtype=np.float32)
+    ycrcb[:, :, 1] = 128 + (ycrcb[:, :, 1] - 128) * c
+    ycrcb[:, :, 2] = 128 + (ycrcb[:, :, 2] - 128) * c
+    ycrcb = np.uint8(np.clip(ycrcb, 0, 255))
+    img = cv2.cvtColor(ycrcb, cv2.COLOR_YCR_CB2RGB)
+    return img

data/degradation_toolkit/x_distortion/spatter.py

@@ -0,0 +1,74 @@
+import cv2
+import numpy as np
+from skimage.filters import gaussian
+
+def rgb2gray(rgb):
+    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])
+
+def spatter(x, severity=1):
+    c = [(0.65, 0.3, 4, 0.69, 0.6, 0),
+         (0.65, 0.3, 3, 0.68, 0.6, 0),
+         (0.65, 0.3, 2, 0.68, 0.5, 0),
+         (0.65, 0.3, 1, 0.65, 1.5, 1),
+         (0.67, 0.4, 1, 0.65, 1.5, 1)][severity - 1]
+    x_PIL = x
+    x = np.array(x, dtype=np.float32) / 255.
+
+    liquid_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])
+
+    liquid_layer = gaussian(liquid_layer, sigma=c[2])
+    liquid_layer[liquid_layer < c[3]] = 0
+    if c[5] == 0:
+        liquid_layer = (liquid_layer * 255).astype(np.uint8)
+        dist = 255 - cv2.Canny(liquid_layer, 50, 150)
+        dist = cv2.distanceTransform(dist, cv2.DIST_L2, 5)
+        _, dist = cv2.threshold(dist, 20, 20, cv2.THRESH_TRUNC)
+        dist = cv2.blur(dist, (3, 3)).astype(np.uint8)
+        dist = cv2.equalizeHist(dist)
+        ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
+        dist = cv2.filter2D(dist, cv2.CV_8U, ker)
+        dist = cv2.blur(dist, (3, 3)).astype(np.float32)
+
+        m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)
+        m /= np.max(m, axis=(0, 1))
+        m *= c[4]
+        # water is pale turqouise
+        color = np.concatenate((175 / 255. * np.ones_like(m[..., :1]),
+                                238 / 255. * np.ones_like(m[..., :1]),
+                                238 / 255. * np.ones_like(m[..., :1])), axis=2)
+
+        color = cv2.cvtColor(color, cv2.COLOR_BGR2BGRA)
+
+        if len(x.shape) < 3 or x.shape[2] < 3:
+            add_spatter_color = cv2.cvtColor(np.clip(m * color, 0, 1),
+                                             cv2.COLOR_BGRA2BGR)
+            add_spatter_gray = rgb2gray(add_spatter_color)
+
+            return np.clip(x + add_spatter_gray, 0, 1) * 255
+
+        else:
+
+            x = cv2.cvtColor(x, cv2.COLOR_BGR2BGRA)
+
+            return cv2.cvtColor(np.clip(x + m * color, 0, 1),
+                                cv2.COLOR_BGRA2BGR) * 255
+    else:
+        m = np.where(liquid_layer > c[3], 1, 0)
+        m = gaussian(m.astype(np.float32), sigma=c[4])
+        m[m < 0.8] = 0
+
+        x_rgb = np.array(x_PIL.convert('RGB'))
+
+        # mud brown
+        color = np.concatenate((63 / 255. * np.ones_like(x_rgb[..., :1]),
+                                42 / 255. * np.ones_like(x_rgb[..., :1]),
+                                20 / 255. * np.ones_like(x_rgb[..., :1])),
+                               axis=2)
+        color *= m[..., np.newaxis]
+        if len(x.shape) < 3 or x.shape[2] < 3:
+            x *= (1 - m)
+            return np.clip(x + rgb2gray(color), 0, 1) * 255
+
+        else:
+            x *= (1 - m[..., np.newaxis])
+            return np.clip(x + color, 0, 1) * 255

examples/__init__.py

@@ -0,0 +1,13 @@
+from .gradio_tasks import dense_prediction_text, conditional_generation_text, process_dense_prediction_tasks, process_conditional_generation_tasks
+from .gradio_tasks_restoration import image_restoration_text, process_image_restoration_tasks
+from .gradio_tasks_style import style_transfer_text, style_condition_fusion_text, process_style_transfer_tasks, process_style_condition_fusion_tasks
+from .gradio_tasks_tryon import tryon_text, process_tryon_tasks
+from .gradio_tasks_editing import editing_text, process_editing_tasks
+from .gradio_tasks_photodoodle import photodoodle_text, process_photodoodle_tasks
+from .gradio_tasks_editing_subject import editing_with_subject_text, process_editing_with_subject_tasks
+from .gradio_tasks_relighting import relighting_text, process_relighting_tasks
+from .gradio_tasks_unseen import unseen_tasks_text, process_unseen_tasks
+from .gradio_tasks_subject import subject_driven_text, condition_subject_fusion_text, condition_subject_style_fusion_text, style_transfer_with_subject_text, \
+    image_restoration_with_subject_text, \
+    process_subject_driven_tasks, process_image_restoration_with_subject_tasks, process_style_transfer_with_subject_tasks, process_condition_subject_style_fusion_tasks, \
+    process_condition_subject_fusion_tasks