init

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+batch*
+phoe*
+*gt*
+*__pycache__*
+*ocr*
+*.pyc
+*.pyo
+*.pyd
+*.txt
+*.json
+.gradio
+output
+gradio_cache
+__pycache__

INSTALL.md

@@ -0,0 +1,29 @@
+# Installation
+
+Downloading VisualCloze repo from github:
+
+```bash
+git clone xxx
+```
+
+### 1. Create a conda environment and install PyTorch
+
+Note: You may want to adjust the CUDA version [according to your driver version](https://docs.nvidia.com/deploy/cuda-compatibility/#default-to-minor-version).
+
+```bash
+conda create -n visualcloze -y
+conda activate visualcloze
+conda install python=3.11 pytorch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda=12.1 -c pytorch -c nvidia -y
+```
+
+### 2. Install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### 3. Install flash-attn
+
+```bash
+pip install flash-attn --no-build-isolation
+```

app.py

@@ -0,0 +1,491 @@
+import argparse
+import gradio as gr
+import demo_tasks
+from functools import partial
+from data.prefix_instruction import get_layout_instruction
+from visualcloze import VisualClozeModel
+
+
+max_grid_h = 5
+max_grid_w = 5
+default_grid_h = 2
+default_grid_w = 3
+default_upsampling_noise = 0.4
+default_steps = 30
+
+GUIDANCE = """
+
+## 📋 Quick Start Guide:
+1. Adjust **Number of In-context Examples**, 0 disables in-context learning.
+2. Set **Task Columns**, the number of images involved in a task.
+3. Upload Images. For in-context examples, upload all images. For the current query, upload images exclude the target.
+4. Click **Generate** to create the images.
+5. Parameters can be fine-tuned under **Advanced Options**.
+
+<div style='font-size: 24px; font-weight: bold; color: #FF9999;'>Click the task button in the right bottom to acquire examples of various tasks.</div>
+
+"""
+
+CITATION = r"""
+If you find VisualCloze is helpful, please consider to star ⭐ the <a href='https://github.com/lzyhha/VisualCloze' target='_blank'>Github Repo</a>. Thanks! 
+---
+📝 **Citation**
+<br>
+If our work is useful for your research, please consider citing:
+```bibtex
+@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:},
+  year={2025}
+}
+```
+📋 **License**
+<br>
+This project is licensed under xxx.
+"""
+
+INTRODUCTION = """
+## 📋 Introduction:
+VisualCloze utilizes in-context examples as visual demonstrations to clarify the desired task. 
+
+Through in-context learning, VisualCloze can:
+1. support various in-domain tasks, 
+2. generalize to **unseen tasks** through in-context learning, 
+3. unify multiple tasks into one step and generate not only the target image but also the intermediate results, 
+4. support reverse generation, i.e., reverse-engineering a set of conditions from a target.
+
+"""
+
+def create_demo(model):
+    with gr.Blocks(title="VisualCloze Demo") as demo:
+        gr.Markdown("# VisualCloze: A Universal Image Generation Framework via Visual In-Context Learning")
+
+        gr.HTML("""
+        <div style="display:flex;column-gap:4px;">
+            <a href="xxx">
+                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
+            </a> 
+            <a href="xxx">
+                <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
+            </a>
+            <a href="xxx">
+                <img src='https://img.shields.io/badge/VisualCloze%20checkpoint-HF%20Model-green?logoColor=violet&label=%F0%9F%A4%97%20Checkpoint'>
+            </a>
+            <a href="xxx">
+                <img src='https://img.shields.io/badge/VisualCloze%20datasets-HF%20Dataset-6B88E3?logoColor=violet&label=%F0%9F%A4%97%20Graph200k%20Dataset'>
+            </a>
+        </div>
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                gr.Markdown(INTRODUCTION)
+            with gr.Column(scale=2):
+                gr.Markdown(GUIDANCE)
+                    
+        # gr.Markdown("<div style='font-size: 24px; font-weight: bold; color: #FF9999;'>" + 
+        #             "Note: Click the task button in the right bottom to acquire examples of tasks." + 
+        #             "</div>", )
+
+        # Pre-create all possible image components
+        all_image_inputs = []
+        rows = []
+        row_texts = []
+        with gr.Row():
+
+            # 左侧列：图像网格和提示输入
+            with gr.Column(scale=2):
+                # 图像网格部分
+                for i in range(max_grid_h):
+                    # Add row label before each row
+                    row_texts.append(gr.Markdown(
+                        "<div style='font-size: 24px; font-weight: bold;'>" + 
+                        ("query" if i == default_grid_h - 1 else f"In-context Example {i + 1}") + 
+                        "</div>", 
+                        elem_id=f"row_text_{i}", 
+                        visible=i < default_grid_h
+                    ))
+                    with gr.Row(visible=i < default_grid_h, elem_id=f"row_{i}") as row:
+                        rows.append(row)
+                        for j in range(max_grid_w):
+                            img_input = gr.Image(
+                                label=f"In-context Example {i + 1}/{j + 1}" if i != default_grid_h - 1 else f"Query {j + 1}",
+                                type="pil",
+                                visible= i < default_grid_h and j < default_grid_w,
+                                interactive=True,
+                                elem_id=f"img_{i}_{j}"
+                            )
+                            all_image_inputs.append(img_input)
+
+                # 提示输入部分
+                layout_prompt = gr.Textbox(
+                    label="Layout Description (Auto-filled, Read-only)",
+                    placeholder="Layout description will be automatically filled based on grid size...",
+                    value=get_layout_instruction(default_grid_w, default_grid_h),
+                    elem_id="layout_prompt",
+                    interactive=False
+                )
+                
+                task_prompt = gr.Textbox(
+                    label="Task Description (Can be modified by referring to examples to perform custom tasks, but may lead to unstable results)",
+                    placeholder="Describe what task should be performed...",
+                    value="",
+                    elem_id="task_prompt"
+                )
+                
+                content_prompt = gr.Textbox(
+                    label="Content Description (Image caption, Editing instructions, etc.)",
+                    placeholder="Describe the content requirements...",
+                    value="",
+                    elem_id="content_prompt"
+                )
+                
+                generate_btn = gr.Button("Generate", elem_id="generate_btn")
+
+                grid_h = gr.Slider(minimum=0, maximum=max_grid_h-1, value=default_grid_h-1, step=1, label="Number of In-context Examples", elem_id="grid_h")
+                grid_w = gr.Slider(minimum=1, maximum=max_grid_w, value=default_grid_w, step=1, label="Task Columns", elem_id="grid_w")
+                
+                with gr.Accordion("Advanced options", open=False):
+                    seed = gr.Number(label="Seed (0 for random)", value=0, precision=0)
+                    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)
+
+                gr.Markdown(CITATION)
+
+            # 右侧列：输出图像
+            with gr.Column(scale=2):
+                output_gallery = gr.Gallery(
+                    label="Generated Results",
+                    show_label=True,
+                    elem_id="output_gallery",
+                    columns=None,  # 设为None以允许自动调整
+                    rows=None,     # 设为None以允许自动调整
+                    height="auto",
+                    allow_preview=True,
+                    object_fit="contain"  # 确保图片完整显示
+                )
+            
+                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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    label='Unseen Tasks (May produce unstable effects)', 
+                    samples_per_page=1000, 
+                    components=[text_unseen_tasks]) 
+                
+                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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    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, 
+                    label='Image Restoration with Subject', 
+                    samples_per_page=1000, 
+                    components=[text_image_restoration_with_subject_tasks])
+                
+        def update_grid(h, w):
+            actual_h = h + 1
+            model.set_grid_size(actual_h, w)
+            
+            updates = []
+            
+            # Update image component visibility
+            for i in range(max_grid_h * max_grid_w):
+                curr_row = i // max_grid_w
+                curr_col = i % max_grid_w
+                updates.append(
+                    gr.update(
+                        label=f"In-context Example {curr_row + 1}/{curr_col + 1}" if curr_row != actual_h - 1 else f"Query {curr_col + 1}",
+                        elem_id=f"img_{curr_row}_{curr_col}", 
+                        visible=(curr_row < actual_h and curr_col < w)))
+            
+            # Update row visibility and labels
+            updates_row = []
+            updates_row_text = []
+            for i in range(max_grid_h):
+                updates_row.append(gr.update(f"row_{i}", visible=(i < actual_h)))
+                updates_row_text.append(
+                    gr.update(
+                        elem_id=f"row_text_{i}", 
+                        visible=i < actual_h, 
+                        value="<div style='font-size: 24px; font-weight: bold;'>" + 
+                        ("Query" if i == actual_h - 1 else f"In-context Example {i + 1}") + 
+                        "</div>", 
+                    )
+                )
+            
+            updates.extend(updates_row)
+            updates.extend(updates_row_text)
+            updates.append(gr.update(elem_id="layout_prompt", value=get_layout_instruction(w, actual_h)))
+            return updates
+
+        def generate_image(*inputs):
+            images = []
+            for i in range(model.grid_h):
+                images.append([])
+                for j in range(model.grid_w):
+                    images[i].append(inputs[i * max_grid_w + j])
+            seed, cfg, steps, upsampling_steps, upsampling_noise, layout_text, task_text, content_text = inputs[-8:]
+                
+            results = model.process_images(
+                images, 
+                [layout_text, task_text, content_text], 
+                seed=seed, cfg=cfg, steps=steps, 
+                upsampling_steps=upsampling_steps, upsampling_noise=upsampling_noise
+            )
+
+            output = gr.update(
+                elem_id='output_gallery', 
+                value=results, 
+                columns=min(len(results), 2),
+                rows=int(len(results) / 2 + 0.5))
+            
+            return output
+
+        def process_tasks(task, func):
+            outputs = func(task)
+            mask = outputs[0]
+            state = outputs[1:8]
+            if state[5] is None:
+                state[5] = default_upsampling_noise
+            if state[6] is None:
+                state[6] = default_steps
+            images = outputs[8:-len(mask)]
+            output = outputs[-len(mask):]
+            for i in range(len(mask)):
+                if mask[i] == 1:
+                    images.append(None)
+                else:
+                    images.append(output[-len(mask) + i])
+            
+            state[0] = state[0] - 1
+            cur_hrid_h = state[0]
+            cur_hrid_w = state[1]
+
+            current_example = [None] * 25
+            for i, image in enumerate(images):
+                pos = (i // cur_hrid_w) * 5 + (i % cur_hrid_w)
+                if image is not None:
+                    current_example[pos] = image
+            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))
+            return [output] + current_example + state
+        
+        dense_prediction_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        conditional_generation_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        image_restoration_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        style_transfer_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        style_condition_fusion_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+            
+        relighting_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        tryon_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+        
+        photodoodle_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        editing_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+        
+        unseen_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        subject_driven_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        style_transfer_with_subject_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        condition_subject_fusion_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        condition_subject_style_fusion_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        editing_with_subject_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+
+        image_restoration_with_subject_tasks.click(
+            partial(process_tasks, func=demo_tasks.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=False, queue=False)
+        # Initialize grid
+        model.set_grid_size(default_grid_h, default_grid_w)
+        
+        # Connect event processing function to all components that need updating
+        output_components = all_image_inputs + rows + row_texts + [layout_prompt]
+
+        grid_h.change(fn=update_grid, inputs=[grid_h, grid_w], outputs=output_components)
+        grid_w.change(fn=update_grid, inputs=[grid_h, grid_w], outputs=output_components)
+        
+        # Modify generate button click event
+        generate_btn.click(
+            fn=generate_image,
+            inputs=all_image_inputs + [seed, cfg, steps, upsampling_steps, upsampling_noise] + [layout_prompt, task_prompt, content_prompt],
+            outputs=output_gallery
+        )
+        
+    return demo
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default=None)
+    parser.add_argument("--precision", type=str, choices=["fp32", "bf16", "fp16"], default="bf16")
+    parser.add_argument("--resolution", type=int, default=384)
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    
+    # Initialize model
+    model = VisualClozeModel(resolution=args.resolution, model_path=args.model_path, precision=args.precision)
+    
+    # Create Gradio demo
+    demo = create_demo(model)
+    
+    # Start Gradio server
+    demo.launch(share=False, server_port=10050, server_name="0.0.0.0")

data/__init__.py

@@ -0,0 +1 @@
+from .prefix_instruction import *

data/prefix_instruction.py

@@ -0,0 +1,1086 @@
+import random
+
+condition_list = ["canny", "depth", "hed", "normal", "mlsd", "openpose", "sam2_mask", "mask", "foreground", "background", "uniformer"]
+style_list = ["InstantStyle", "ReduxStyle"]
+editing_list = ["DepthEdit", "FillEdit"]
+degradation_list = [
+    # blur
+    "blur",
+    "compression",
+    "SRx2",
+    "SRx4",
+    "pixelate",
+    "Defocus",
+    "GaussianBlur",
+    # sharpen
+    "oversharpen",
+    # nosie
+    "GaussianNoise",
+    "PoissonNoise",
+    "SPNoise",
+    # mosaic
+    "mosaic",
+    # contrast
+    "contrast_strengthen",
+    "contrast_weaken",
+    # quantization
+    "quantization",
+    "JPEG",
+    # light
+    "brighten",
+    "darken",
+    "LowLight",
+    # color
+    "saturate_strengthen",
+    "saturate_weaken",
+    "gray",
+    "ColorDistortion",
+    # infilling
+    "Inpainting",
+    # rotate
+    "rotate90",
+    "rotate180",
+    "rotate270",
+    # other
+    "Barrel",
+    "Pincushion",
+    "Elastic",
+    # spacial effect
+    "Rain",
+    "Frost",
+    ]
+
+
+def get_image_prompt(image_type):
+    image_prompts = {
+        "target": [
+            "a high-quality image",
+            "an aesthetically pleasing photograph",
+            "a high-resolution image",
+            "an image with vivid details",
+            "a visually striking and clear picture",
+            "a high-definition image",
+            "an image with artistic appeal",
+            "a sharp and beautifully composed photograph",
+            "a high-aesthetic image",
+            "an image with flawless clarity",
+            "a vibrant and professionally captured photo",
+            "a crystal-clear image",
+            "an image with artistic quality"
+            "a high-quality image with exceptional detail",
+            "a photo realistic image",
+        ],
+        "reference": [
+            "a reference image",
+            "an image featuring the primary object"
+            "a reference for the main object",
+            "a reference image highlighting the central object",
+            "an image containing the key object",
+            "a reference image with the main subject included",
+            "an image providing the main object",
+            "a reference image showcasing the dominant object",
+            "an image that includes the main object",
+            "a reference image capturing the primary subject",
+            "an image containing the main subject",
+        ],
+        # condition
+        "canny": [
+            "canny edge map with sharp black-and-white contours",
+            "black-and-white edge map highlighting crisp boundaries",
+            "canny result showing stark white edges on black",
+            "edge map with clean white lines on a dark background",
+            "canny output featuring precise white object outlines",
+            "black background with white edge-detected contours",
+            "canny edge map displaying clear white structural edges",
+            "white edge lines on black from canny detection",
+            "canny map with sharp white edges and dark voids",
+            "edge map revealing white outlines of object shapes",
+        ],
+        "depth": [
+            "depth map showing gray-scale object contours",
+            "gray-toned depth map with layered outlines",
+            "depth map featuring gradient-gray surfaces",
+            "gray-shaded depth map with distinct edges",
+            "depth map displaying soft gray gradients",
+            "gray-scale depth map with clear object boundaries",
+            "depth map highlighting gray-level depth variations",
+            "gray-textured depth map with smooth transitions",
+            "depth map revealing gray-toned spatial layers",
+            "gray-based depth map with detailed object contours",
+        ],
+        "hed": [
+            "hed edge map with smooth flowing contours",
+            "soft-edged map from hed detection",
+            "hed result showing refined continuous edges",
+            "edge map with natural well-connected outlines",
+            "hed output featuring smooth detailed boundaries",
+            "elegant edge map with seamless transitions",
+            "hed edge map displaying clean holistic contours",
+            "refined edge lines from hed detection",
+            "hed map with flowing natural object outlines",
+            "edge map revealing smooth interconnected shapes",
+        ],
+        "normal": [
+            "normal map showing surface orientation details",
+            "rgb-coded normal map for 3D lighting",
+            "normal map with encoded surface normals",
+            "detailed normal map for texture shading",
+            "normal map highlighting surface curvature",
+            "rgb normal map for bump mapping effects",
+            "normal map capturing fine geometric details",
+            "surface normal visualization in rgb colors",
+            "normal map for realistic lighting interaction",
+            "normal map displaying directional surface data",
+        ],
+        "mlsd": [
+            "mlsd detected straight line segments",
+            "line segments extracted using mlsd",
+            "mlsd output showing precise straight lines",
+            "straight edges detected by mlsd algorithm",
+            "mlsd result with clean line segment boundaries",
+            "line segment map generated by mlsd",
+            "mlsd-detected straight structural lines",
+            "straight line visualization from mlsd",
+            "mlsd-based line segment detection output",
+            "line segments highlighted by mlsd method",
+        ],
+        "openpose": [
+            "openpose skeleton with colorful connecting lines",
+            "body keypoints linked by bright colored lines",
+            "openpose output showing joints and vibrant skeleton",
+            "human pose with colored lines for bone structure",
+            "openpose-detected keypoints and colorful limbs",
+            "skeletal lines in vivid colors from openpose",
+            "body joints connected by multicolored straight lines",
+            "openpose visualization with colorful skeletal links",
+            "keypoints and bright lines forming body skeleton",
+            "human pose mapped with colored lines by openpose",
+        ],
+        "sam2_mask": [
+            "sam 2 generated colorful segmentation masks",
+            "color-coded masks from sam 2 segmentation",
+            "sam 2 output with vibrant object masks",
+            "segmentation masks in bright colors by sam 2",
+            "colorful object masks from sam 2 detection",
+            "sam 2 result showing multicolored regions",
+            "masks with distinct colors from sam 2",
+            "sam 2 segmentation with vivid mask overlays",
+            "colorful masks highlighting objects via sam 2",
+            "sam 2-generated masks with rich color coding",
+        ],
+        "uniformer": [
+            "color-coded objects in uniformer segmentation",
+            "uniformer map with colored object blocks",
+            "objects as distinct color patches by uniformer",
+            "color blocks representing objects in uniformer",
+            "uniformer output with colored object regions",
+            "objects highlighted as color zones in uniformer",
+            "uniformer segmentation showing color-divided objects",
+            "color patches for objects in uniformer result",
+            "uniformer map with objects as solid color areas",
+            "objects segmented as colored blocks by uniformer",
+            "uniformer map with objects as solid color areas",
+        ],
+        "mask": [
+            "Color-coded objects in open-world segmentation",
+            "Distinct colors marking different objects",
+            "Objects highlighted as unique color patches",
+            "Color blocks representing diverse objects",
+            "Segmented image with varied color zones",
+            "Objects visualized as solid color regions",
+            "Colorful map of open-world object segmentation",
+            "Objects divided by vibrant color boundaries",
+            "Color-coded segmentation of diverse items",
+            "Objects mapped as distinct colored areas",
+        ],
+        "foreground": [
+            "Foreground on solid color canvas",
+            "Image with foreground on plain backdrop",
+            "Foreground placed on monochrome background",
+            "Objects on solid color base",
+            "Foreground isolated on uniform color",
+            "Segmented subject on plain color field",
+            "Foreground displayed on solid color",
+            "Image with foreground on solid backdrop",
+            "Foreground on a clean color canvas",
+            "Objects on a solid color background",
+        ],
+        "background": [
+            "Background-only image with foreground masked",
+            "Photo showing background after masking foreground",
+            "Image with foreground removed leaving background",
+            "Background revealed by masking the foreground",
+            "Foreground masked to expose background",
+            "Picture with background visible after masking",
+            "Image displaying background without foreground",
+            "Foreground erased leaving only background",
+            "Background isolated by masking the foreground",
+            "Photo with foreground hidden showing background",
+        ],
+        # Style
+        "style_source": [
+            "Image in a distinct artistic style",
+            "Artistically styled picture with unique flair",
+            "Photo showcasing a specific art style",
+            "Image with a clear artistic aesthetic",
+            "Art-style influenced visual composition",
+            "Picture reflecting a particular art movement",
+            "Image with bold artistic characteristics",
+            "Artistically rendered visual content",
+            "Photo with a strong artistic theme",
+            "Image embodying a defined art style",
+        ],
+        "style_target": [
+            "High-quality image with striking artistic style",
+            "Crisp photo showcasing bold artistic flair",
+            "Visually stunning image with artistic influence",
+            "High-definition picture in a unique art style",
+            "Artistically styled image with exceptional clarity",
+            "High-quality visual with distinct artistic touch",
+            "Sharp photo reflecting a clear artistic theme",
+            "Artistically crafted image with high resolution",
+            "Vibrant picture blending quality and art style",
+            "High-aesthetic image with artistic precision",
+        ],
+        # Editing
+        "DepthEdit": [
+            "a high-quality image",
+            "an aesthetically pleasing photograph",
+            "a high-resolution image",
+            "an image with vivid details",
+            "a visually striking and clear picture",
+            "a high-definition image",
+            "an image with artistic appeal",
+            "a sharp and beautifully composed photograph",
+            "a high-aesthetic image",
+            "an image with flawless clarity",
+            "a vibrant and professionally captured photo",
+            "a crystal-clear image",
+            "an image with artistic quality",
+            "a high-quality image with exceptional detail",
+            "a photo realistic image",
+        ],
+        "FillEdit": [
+            "a high-quality image",
+            "an aesthetically pleasing photograph",
+            "a high-resolution image",
+            "an image with vivid details",
+            "a visually striking and clear picture",
+            "a high-definition image",
+            "an image with artistic appeal",
+            "a sharp and beautifully composed photograph",
+            "a high-aesthetic image",
+            "an image with flawless clarity",
+            "a vibrant and professionally captured photo",
+            "a crystal-clear image",
+            "an image with artistic quality",
+            "a high-quality image with exceptional detail",
+            "a photo realistic image",
+        ],
+        # degradation
+        # Blur
+        "blur": [
+            "a softly blurred image with smooth transitions",
+            "a photograph with a gentle motion blur effect",
+            "an image exhibiting subtle Gaussian blur",
+            "a picture with a light and even blurring",
+            "a softly defocused photograph with reduced sharpness",
+            "an image featuring a mild blur for artistic effect",
+            "a photograph with a gentle out-of-focus appearance",
+            "a softly smeared image with smooth edges",
+            "a picture with a light blur enhancing the mood",
+            "an image with a delicate blur creating a dreamy effect",
+        ],
+        "compression": [
+            "a highly compressed image with noticeable artifacts",
+            "a photograph showing compression-induced quality loss",
+            "an image with visible compression artifacts and reduced clarity",
+            "a picture exhibiting blocky artifacts from compression",
+            "a compressed photo with color banding and loss of detail",
+            "an image displaying noticeable compression noise",
+            "a photograph with degraded quality due to high compression",
+            "a picture showing pixelation from aggressive compression",
+            "an image with artifacts and reduced resolution from compression",
+            "a compressed image featuring loss of sharpness and detail",
+        ],
+        "SRx2": [
+            "an image downsampled by a factor of 2 with enhanced details",
+            "a photograph resized to half its original resolution",
+            "an downscaled image (2x) maintaining image quality",
+            "a picture downsized by 2x with preserved sharpness",
+            "an image scaled half its size with clear details",
+            "a low-resolution version of the original image (2x)",
+            "a half-resolution photograph with maintained clarity",
+            "an image decreased in size by 2x with minimal quality loss",
+            "a 2x downscaled picture retaining original details",
+            "an image resized to half its original dimensions with enhanced quality",
+        ],
+        "SRx4": [
+            "an image downsampled by a factor of 4 with enhanced details",
+            "a photograph resized to quarter its original resolution",
+            "an downscaled image (4x) maintaining image quality",
+            "a picture downsized by 4x with preserved sharpness",
+            "an image scaled four times its size with clear details",
+            "a low-resolution version of the original image (4x)",
+            "a quadruple-resolution photograph with maintained clarity",
+            "an image decreased in size by 4x with minimal quality loss",
+            "a 4x downscaled picture retaining original details",
+            "an image resized to quarter its original dimensions with enhanced quality",
+        ],
+        "pixelate": [
+            "a heavily pixelated image with large blocks",
+            "a picture showing strong pixelation effects",
+            "an image with noticeable pixel blocks obscuring details",
+            "a pixelated photograph with reduced image clarity",
+            "an image exhibiting coarse pixelation for a stylized look",
+            "a picture with large pixel squares creating a mosaic effect",
+            "a highly pixelated photo obscuring fine details",
+            "an image featuring prominent pixelation and blockiness",
+            "a pixelated image with distinct square blocks",
+            "a photograph with exaggerated pixelation for artistic effect",
+        ],
+        "Defocus": [
+            "a defocused image with soft and blurry regions",
+            "a photograph with intentional defocus creating a shallow depth of field",
+            "an image exhibiting a defocused background with a clear subject",
+            "a picture with selective defocus enhancing the main object",
+            "a defocused photo with smooth out-of-focus areas",
+            "an image showing a defocused effect for artistic blurring",
+            "a photograph with a softly defocused foreground",
+            "a picture with partial defocus creating a dreamy appearance",
+            "an image featuring defocus to highlight specific areas",
+            "a defocused photograph with gentle blurring around the subject",
+        ],
+        "GaussianBlur": [
+            "an image with Gaussian blurring creating a soft focus effect",
+            "a photograph with a Gaussian blur enhancing the subject",
+            "a picture with Gaussian blurring to highlight the main object",
+            "an image featuring Gaussian blur to soften the background",
+            "a Gaussian-blurred photograph with a soft focus",
+            "a Gaussian-blurred image with a gentle focus on the subject",
+            "a picture with Gaussian blurring to emphasize the main subject",
+            "an image with Gaussian blurring to create a dreamy effect",
+            "a Gaussian-blurred photograph with a soft focus on the main object",
+        ],
+        # Sharpen
+        "oversharpen": [
+            "an image with excessive sharpening creating halos around edges",
+            "a photograph overly sharpened with exaggerated edge contrast",
+            "an oversharpened picture showing unnatural edge highlights",
+            "a highly sharpened image with pronounced texture details",
+            "a picture exhibiting over-sharpening with visible artifacts",
+            "an image with extreme sharpening enhancing all details sharply",
+            "a photograph with oversharpened edges and increased contrast",
+            "an overly sharpened image causing unnatural texture emphasis",
+            "a picture with excessive sharpening effects on all elements",
+            "an image displaying over-sharpened features with enhanced edges",
+        ],
+        # Noise
+        "GaussianNoise": [
+            "an image with subtle Gaussian noise adding grain",
+            "a photograph exhibiting Gaussian noise for a textured look",
+            "a picture with light Gaussian noise enhancing realism",
+            "an image featuring Gaussian noise with smooth distribution",
+            "a photo with added Gaussian noise creating a grainy effect",
+            "an image showing gentle Gaussian noise for artistic texture",
+            "a photograph with mild Gaussian noise increasing depth",
+            "a picture with soft Gaussian noise enhancing the image",
+            "an image displaying Gaussian noise for a vintage feel",
+            "a photo with Gaussian noise subtly integrated into the image",
+        ],
+        "PoissonNoise": [
+            "an image with Poisson noise creating photon distribution effects",
+            "a photograph exhibiting Poisson noise for realistic grain",
+            "a picture with added Poisson noise enhancing texture",
+            "an image featuring Poisson noise with natural variance",
+            "a photo with Poisson noise simulating low-light conditions",
+            "an image showing Poisson noise for authentic grain patterns",
+            "a photograph with mild Poisson noise increasing image depth",
+            "a picture with Poisson noise adding subtle texture",
+            "an image displaying Poisson noise for a realistic appearance",
+            "a photo with Poisson noise integrated for enhanced realism",
+        ],
+        "SPNoise": [
+            "an image with salt and pepper noise introducing random pixels",
+            "a photograph exhibiting SP noise with black and white speckles",
+            "a picture with added salt and pepper noise creating scattered dots",
+            "an image featuring SP noise with random pixel disruptions",
+            "a photo with SP noise simulating transmission errors",
+            "an image showing salt and pepper noise for a gritty effect",
+            "a photograph with mild SP noise adding texture variation",
+            "a picture with SP noise introducing random black and white pixels",
+            "an image displaying salt and pepper noise for a distressed look",
+            "a photo with SP noise integrated for a speckled appearance",
+        ],
+        # Mosaic
+        "mosaic": [
+            "an image with a strong mosaic effect obscuring details",
+            "a photograph exhibiting mosaic patterns with large tiles",
+            "a picture with applied mosaic effect creating a tiled appearance",
+            "an image featuring mosaic blocks for privacy masking",
+            "a photo with mosaic segmentation highlighting regions",
+            "an image showing a mosaic overlay for abstract effect",
+            "a photograph with mosaic patterns simplifying the image",
+            "a picture with a mosaic filter creating geometric tiles",
+            "an image displaying a mosaic effect for stylistic purposes",
+            "a photo with mosaic segmentation emphasizing specific areas",
+        ],
+        # Contrast
+        "contrast_strengthen": [
+            "an image with enhanced contrast making colors pop",
+            "a photograph exhibiting strengthened contrast for vividness",
+            "a picture with increased contrast highlighting details",
+            "an image featuring heightened contrast for dramatic effect",
+            "a photo with boosted contrast enhancing visual depth",
+            "an image showing strengthened contrast with pronounced shadows and highlights",
+            "a photograph with amplified contrast for greater clarity",
+            "a picture with enhanced contrast making elements stand out",
+            "an image displaying increased contrast for a striking appearance",
+            "a photo with reinforced contrast improving overall image impact",
+        ],
+        "contrast_weaken": [
+            "an image with reduced contrast creating a softer look",
+            "a photograph exhibiting weakened contrast for a muted effect",
+            "a picture with decreased contrast making colors more subtle",
+            "an image featuring lowered contrast for a gentle appearance",
+            "a photo with diminished contrast softening the overall image",
+            "an image showing weakened contrast with less pronounced shadows and highlights",
+            "a photograph with reduced contrast for a flatter visual tone",
+            "a picture with softened contrast creating a delicate atmosphere",
+            "an image displaying decreased contrast for a subdued look",
+            "a photo with lowered contrast enhancing a calm and serene feel",
+        ],
+        # Quantization
+        "quantization": [
+            "an image with quantization artifacts reducing color depth",
+            "a photograph exhibiting quantization leading to banding effects",
+            "a picture with applied quantization simplifying color gradients",
+            "an image featuring quantized color levels creating discrete steps",
+            "a photo with quantization reducing the number of distinct colors",
+            "an image showing quantization leading to posterization effects",
+            "a photograph with quantized color palette for a stylized look",
+            "a picture with quantization introducing color banding and loss of detail",
+            "an image displaying quantization effects on smooth color transitions",
+            "a photo with quantization artifacts simplifying the overall color scheme",
+        ],
+        "JPEG": [
+            "a JPEG-compressed image with noticeable compression artifacts",
+            "a photograph saved in JPEG format showing quality loss",
+            "an image exhibiting JPEG artifacts like blockiness and blurring",
+            "a picture with JPEG compression leading to reduced clarity",
+            "an image featuring JPEG-induced artifacts affecting image quality",
+            "a photo with visible JPEG compression effects on details",
+            "an image showing JPEG artifacts such as color banding and pixelation",
+            "a photograph with degraded quality due to JPEG compression",
+            "a picture with JPEG compression artifacts impacting the overall appearance",
+            "an image displaying JPEG-induced quality loss with blurred edges",
+        ],
+        # Light
+        "brighten": [
+            "a brightly lit image with enhanced luminosity",
+            "a photograph exhibiting increased brightness for a vibrant look",
+            "a picture with boosted brightness making the scene more radiant",
+            "an image featuring heightened brightness illuminating all areas",
+            "a photo with amplified brightness creating a sunny appearance",
+            "an image showing increased brightness enhancing visibility",
+            "a photograph with enhanced brightness making colors more vivid",
+            "a picture with boosted luminosity brightening the overall image",
+            "an image displaying heightened brightness for a luminous effect",
+            "a photo with increased brightness adding warmth and clarity",
+        ],
+        "darken": [
+            "a darkened image with reduced luminosity creating a moody atmosphere",
+            "a photograph exhibiting decreased brightness for a subdued look",
+            "a picture with lowered brightness making the scene more somber",
+            "an image featuring diminished brightness enhancing shadows",
+            "a photo with reduced brightness creating a twilight appearance",
+            "an image showing decreased brightness adding depth and contrast",
+            "a photograph with darkened tones making colors more muted",
+            "a picture with lowered luminosity creating a dramatic effect",
+            "an image displaying reduced brightness for a darker aesthetic",
+            "a photo with decreased brightness enhancing the mysterious mood",
+        ],
+        "LowLight": [
+            "an image with low light conditions creating a dim and shadowy appearance",
+            "a photograph exhibiting low light to simulate night-time conditions",
+            "a picture with reduced illumination to create a night-time ambiance",
+            "an image featuring low light to emphasize the subject in darkness",
+            "a photo with low light conditions creating a mysterious mood",
+            "an image showing low light to enhance the dramatic lighting of the scene",
+            "a photograph with dim lighting to create a soft and dreamy effect",
+            "a picture with low light to emphasize the texture and details of the image",
+            "an image displaying low light conditions for a serene and peaceful feel",
+        ],
+        # Color
+        "saturate_strengthen": [
+            "an image with enhanced saturation making colors more vivid",
+            "a photograph exhibiting strengthened saturation for vibrant hues",
+            "a picture with boosted color saturation enhancing visual appeal",
+            "an image featuring heightened saturation creating rich color tones",
+            "a photo with amplified saturation making colors pop",
+            "an image showing increased saturation for a lively appearance",
+            "a photograph with saturated colors enhancing the overall image",
+            "a picture with strengthened color saturation adding vibrancy",
+            "an image displaying enhanced saturation for a dynamic look",
+            "a photo with boosted color intensity making the scene more colorful",
+        ],
+        "saturate_weaken": [
+            "an image with reduced saturation creating a muted color palette",
+            "a photograph exhibiting weakened saturation for subdued tones",
+            "a picture with lowered color saturation making colors more subtle",
+            "an image featuring diminished saturation creating a pastel look",
+            "a photo with decreased saturation softening the overall colors",
+            "an image showing reduced saturation for a faded appearance",
+            "a photograph with desaturated colors enhancing a minimalist aesthetic",
+            "a picture with weakened color saturation adding a calm feel",
+            "an image displaying lowered saturation for a gentle color scheme",
+            "a photo with diminished color intensity creating a subdued look",
+        ],
+        "gray": [
+            "a grayscale image with varying shades of gray",
+            "a black and white photograph emphasizing contrast and texture",
+            "a gray-toned picture highlighting light and shadow",
+            "an image converted to grayscale showcasing structural details",
+            "a monochromatic photo with rich gray gradients",
+            "a grayscale image emphasizing form and composition",
+            "a black and white picture with balanced gray tones",
+            "an image in gray scale enhancing depth and dimension",
+            "a monochrome photograph focusing on texture and contrast",
+            "a gray-toned image presenting a classic black and white aesthetic",
+        ],
+        "ColorDistortion": [
+            "an image with distorted and surreal colors",
+            "a picture featuring unnatural color tones",
+            "a visually striking image with altered hues",
+            "a photo showcasing disrupted color balance",
+            "an image with vibrant and unexpected colors",
+            "a picture displaying shifted color spectrums",
+            "an artwork-like image with perturbed colors",
+            "a photo with dreamlike and distorted hues",
+            "an image with unconventional color variations",
+            "a visually unique picture with color shifts",
+        ],
+        # Infilling
+        "Inpainting": [
+            "an inpainted image seamlessly filling missing areas",
+            "a photograph with inpainting repairing damaged regions",
+            "a picture featuring inpainting to restore obscured parts",
+            "an image using inpainting to complete incomplete areas",
+            "a photo with inpainting blending filled regions naturally",
+            "an image showing inpainting techniques removing unwanted objects",
+            "a photograph with inpainting reconstructing missing details",
+            "a picture utilizing inpainting to enhance image continuity",
+            "an image with inpainting seamlessly integrating filled sections",
+            "a photo using inpainting to mend and complete the visual content",
+        ],
+        # Rotate
+        "rotate90": [
+            "an image rotated 90 degrees clockwise for a new perspective",
+            "a photograph turned 90 degrees to the right altering the orientation",
+            "a picture rotated a quarter turn clockwise enhancing composition",
+            "an image featuring a 90-degree rotation adjusting the viewpoint",
+            "a photo with a 90-degree clockwise rotation changing the layout",
+            "an image showing a rotated view at 90 degrees for a fresh angle",
+            "a photograph rotated right by 90 degrees for dynamic framing",
+            "a picture with a 90-degree turn clockwise modifying the scene",
+            "an image displaying a 90-degree rotated orientation for visual interest",
+            "a photo rotated ninety degrees to enhance the composition",
+        ],
+        "rotate180": [
+            "an image rotated 180 degrees flipping it upside down",
+            "a photograph turned completely around with a 180-degree rotation",
+            "a picture rotated halfway, creating an inverted perspective",
+            "an image featuring a 180-degree turn altering the original orientation",
+            "a photo with an upside-down view due to 180-degree rotation",
+            "an image showing a flipped perspective with a 180-degree rotation",
+            "a photograph rotated twice around, changing the viewpoint",
+            "a picture with a half-turn rotation modifying the scene layout",
+            "an image displaying a 180-degree rotated orientation for a unique angle",
+            "a photo rotated one full half-circle to invert the composition",
+        ],
+        "rotate270": [
+            "an image rotated 270 degrees clockwise for a new angle",
+            "a photograph turned 270 degrees to the right altering the orientation",
+            "a picture rotated three quarters turn clockwise enhancing composition",
+            "an image featuring a 270-degree rotation adjusting the viewpoint",
+            "a photo with a 270-degree clockwise rotation changing the layout",
+            "an image showing a rotated view at 270 degrees for a fresh angle",
+            "a photograph rotated right by 270 degrees for dynamic framing",
+            "a picture with a 270-degree turn clockwise modifying the scene",
+            "an image displaying a 270-degree rotated orientation for visual interest",
+            "a photo rotated two and a half turns clockwise to enhance the composition",
+        ],
+        # Other
+        "Barrel": [
+            "an image with barrel distortion bending the edges outward",
+            "a photograph exhibiting barrel distortion creating a convex effect",
+            "a picture with barrel distortion warping the image edges",
+            "an image featuring barrel distortion causing peripheral stretching",
+            "a photo with barrel distortion curving the sides outward",
+            "an image showing barrel distortion for a fisheye lens effect",
+            "a photograph with warped edges due to barrel distortion",
+            "a picture with barrel distortion altering the straight lines",
+            "an image displaying barrel distortion creating a rounded appearance",
+            "a photo with barrel distortion enhancing the central focus",
+        ],
+        "Pincushion": [
+            "an image with pincushion distortion bending the edges inward",
+            "a photograph exhibiting pincushion distortion creating a concave effect",
+            "a picture with pincushion distortion warping the image edges inward",
+            "an image featuring pincushion distortion causing peripheral compression",
+            "a photo with pincushion distortion curving the sides inward",
+            "an image showing pincushion distortion for a telephoto lens effect",
+            "a photograph with warped edges due to pincushion distortion",
+            "a picture with pincushion distortion altering the straight lines inward",
+            "an image displaying pincushion distortion creating a pinched appearance",
+            "a photo with pincushion distortion enhancing the central focus inward",
+        ],
+        "Elastic": [
+            "an image with elastic deformation creating fluid distortions",
+            "a photograph exhibiting elastic transformations warping the structure",
+            "a picture with elastic effects bending and stretching elements",
+            "an image featuring elastic distortions for a dynamic appearance",
+            "a photo with elastic transformations altering the image geometry",
+            "an image showing elastic deformation for a fluid, wavy effect",
+            "a photograph with elastic warping adding motion-like distortions",
+            "a picture with elastic effects creating flexible and dynamic shapes",
+            "an image displaying elastic transformations enhancing creative distortion",
+            "a photo with elastic deformation modifying the original image structure",
+        ],
+        # Spatial Effect
+        "Rain": [
+            "an image with realistic rain effects adding dynamic streaks",
+            "a photograph exhibiting rain overlays creating a wet atmosphere",
+            "a picture with rain effects enhancing the scene with falling droplets",
+            "an image featuring rain streaks adding motion and mood",
+            "a photo with simulated rain creating a rainy day ambiance",
+            "an image showing rain effects with dynamic water droplets",
+            "a photograph with rain overlays adding a sense of movement",
+            "a picture with rain effects enhancing the visual texture",
+            "an image displaying rain streaks for a dramatic weather effect",
+            "a photo with realistic rain adding depth and atmosphere",
+        ],
+        "Frost": [
+            "an image with frost overlays creating icy textures",
+            "a photograph exhibiting frost effects adding a chilly ambiance",
+            "a picture with frost patterns enhancing the scene with icy details",
+            "an image featuring frost overlays creating a frozen appearance",
+            "a photo with simulated frost adding a wintry atmosphere",
+            "an image showing frost effects with delicate ice patterns",
+            "a photograph with frost overlays adding a sense of coldness",
+            "a picture with frost effects enhancing the visual texture with ice",
+            "an image displaying frost patterns for a frosty weather effect",
+            "a photo with realistic frost adding depth and a chilly mood",
+        ],
+    }
+    if image_type in style_list:
+        return [random.choice(image_prompts["style_source"]), random.choice(image_prompts["style_target"])]
+    elif image_type == 'clothing':
+        return [random.choice(image_prompts["clothing"]), random.choice(image_prompts["fullbody"])]
+    else:
+        return [random.choice(image_prompts[image_type])]
+
+
+def get_layout_instruction(cols, rows):
+    layout_instruction = [
+        f"A grid layout with {rows} rows and {cols} columns, displaying {cols*rows} images arranged side by side.",
+        f"{cols*rows} images are organized into a grid of {rows} rows and {cols} columns, evenly spaced.",
+        f"A {rows}x{cols} grid containing {cols*rows} images, aligned in a clean and structured layout.",
+        f"{cols*rows} images are placed in a grid format with {rows} horizontal rows and {cols} vertical columns.",
+        f"A visual grid composed of {rows} rows and {cols} columns, showcasing {cols*rows} images in a balanced arrangement.",
+        f"{cols*rows} images form a structured grid, with {rows} rows and {cols} columns, neatly aligned.",
+        f"A {rows}x{cols} grid layout featuring {cols*rows} images, arranged side by side in a precise pattern.",
+        f"{cols*rows} images are displayed in a grid of {rows} rows and {cols} columns, creating a uniform visual structure.",
+        f"A grid with {rows} rows and {cols} columns, containing {cols*rows} images arranged in a symmetrical layout.",
+        f"{cols*rows} images are organized into a {rows}x{cols} grid, forming a cohesive and orderly display.",
+    ]
+    return random.choice(layout_instruction)
+
+
+def get_task_instruction(condition_prompt, target_prompt):
+    task_instruction = [
+        f"Each row outlines a logical process, starting from {condition_prompt}, to achieve {target_prompt}.",
+        f"In each row, a method is described to use {condition_prompt} for generating {target_prompt}.",
+        f"Each row presents a task that leverages {condition_prompt} to produce {target_prompt}.",
+        f"Every row demonstrates how to transform {condition_prompt} into {target_prompt} through a logical approach.",
+        f"Each row details a strategy to derive {target_prompt} based on the provided {condition_prompt}.",
+        f"In each row, a technique is explained to convert {condition_prompt} into {target_prompt}.",
+        f"Each row illustrates a pathway from {condition_prompt} to {target_prompt} using a clear logical task.",
+        f"Every row provides a step-by-step guide to evolve {condition_prompt} into {target_prompt}.",
+        f"Each row describes a process that begins with {condition_prompt} and results in {target_prompt}.",
+        f"In each row, a logical task is demonstrated to achieve {target_prompt} based on {condition_prompt}.",
+    ]
+    return random.choice(task_instruction)
+
+
+def get_content_instruction():
+    content_instruction = [
+        "The content of the last image in the final row is: ",
+        "The last image of the last row depicts: ",
+        "In the final row, the last image shows: ",
+        "The last image in the bottom row illustrates: ",
+        "The content of the bottom-right image is: ",
+        "The final image in the last row portrays: ",
+        "The last image of the final row displays: ",
+        "In the last row, the final image captures: ",
+        "The bottom-right corner image presents: ",
+        "The content of the last image in the concluding row is: ",
+    ]
+    return random.choice(content_instruction)
+
+
+graph200k_task_dicts = [
+    {
+        "task_name": "conditional generation",
+        "sample_weight": 1,
+        "image_list": [
+            ["canny", "target"],
+            ["depth", "target"],
+            ["hed", "target"],
+            ["normal", "target"],
+            ["mlsd", "target"],
+            ["openpose", "target"],
+            ["sam2_mask", "target"],
+            ["uniformer", "target"],
+            ["mask", "target"],
+            ["foreground", "target"],
+            ["background", "target"],
+        ],
+    },
+    {
+        "task_name": "conditional generation with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "canny", "target"],
+            ["reference", "depth", "target"],
+            ["reference", "hed", "target"],
+            ["reference", "normal", "target"],
+            ["reference", "mlsd", "target"],
+            ["reference", "openpose", "target"],
+            ["reference", "sam2_mask", "target"],
+            ["reference", "uniformer", "target"],
+            ["reference", "mask", "target"],
+            ["reference", "background", "target"],
+        ],
+    },
+    {
+        "task_name": "conditional generation with style",
+        "sample_weight": 1,
+        "image_list": [
+            # instant style
+            ["canny", "InstantStyle"],
+            ["depth", "InstantStyle"],
+            ["hed", "InstantStyle"],
+            ["normal", "InstantStyle"],
+            ["mlsd", "InstantStyle"],
+            ["openpose", "InstantStyle"],
+            ["sam2_mask", "InstantStyle"],
+            ["uniformer", "InstantStyle"],
+            ["mask", "InstantStyle"],
+            # redux style
+            ["canny", "ReduxStyle"],
+            ["depth", "ReduxStyle"],
+            ["hed", "ReduxStyle"],
+            ["normal", "ReduxStyle"],
+            ["mlsd", "ReduxStyle"],
+            ["openpose", "ReduxStyle"],
+            ["sam2_mask", "ReduxStyle"],
+            ["uniformer", "ReduxStyle"],
+            ["mask", "ReduxStyle"],
+        ],
+    },
+    {
+        "task_name": "image generation with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "target"],
+        ],
+    },
+    {
+        "task_name": "subject extraction", 
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "reference"],
+        ],
+    },
+    {
+        "task_name": "style transfer",
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "InstantStyle"],
+            ["target", "ReduxStyle"],
+            ["reference", "InstantStyle"],
+        ],
+    },
+    {
+        "task_name": "style transfer with condition",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "canny", "InstantStyle"],
+            ["reference", "depth", "InstantStyle"],
+            ["reference", "hed", "InstantStyle"],
+            ["reference", "normal", "InstantStyle"],
+            ["reference", "mlsd", "InstantStyle"],
+            ["reference", "openpose", "InstantStyle"],
+            ["reference", "sam2_mask", "InstantStyle"],
+            ["reference", "uniformer", "InstantStyle"],
+            ["reference", "mask", "InstantStyle"],
+        ],
+    },
+    {
+        "task_name": "image editing",
+        "sample_weight": 1,
+        "image_list": [
+            ["DepthEdit", "target"],
+            ["FillEdit", "target"],
+        ],
+    },
+    {
+        "task_name": "image editing with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "DepthEdit", "target"],
+            ["reference", "FillEdit", "target"],
+        ],
+    },
+    {
+        "task_name": "dense prediction",
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "canny"],
+            ["target", "depth"],
+            ["target", "hed"],
+            ["target", "normal"],
+            ["target", "mlsd"],
+            ["target", "openpose"],
+            ["target", "sam2_mask"],
+            ["target", "uniformer"],
+        ],
+    },
+    {
+        "task_name": "restoration",
+        "sample_weight": 1,
+        "image_list": [
+            # blur related
+            ["blur", "target"],
+            ["compression", "target"], 
+            ["SRx2", "target"],
+            ["SRx4", "target"],
+            ["pixelate", "target"],
+            ["Defocus", "target"],
+            ["GaussianBlur", "target"],
+            
+            # sharpen related
+            ["oversharpen", "target"],
+            
+            # noise related
+            ["GaussianNoise", "target"],
+            ["PoissonNoise", "target"],
+            ["SPNoise", "target"],
+            
+            # mosaic
+            ["mosaic", "target"],
+            
+            # contrast related
+            ["contrast_strengthen", "target"],
+            ["contrast_weaken", "target"],
+            
+            # quantization related
+            ["quantization", "target"],
+            ["JPEG", "target"],
+            
+            # light related
+            ["brighten", "target"],
+            ["darken", "target"],
+            ["LowLight", "target"],
+            
+            # color related
+            ["saturate_strengthen", "target"],
+            ["saturate_weaken", "target"],
+            ["gray", "target"],
+            ["ColorDistortion", "target"],
+            
+            # infilling
+            ["Inpainting", "target"],
+            
+            # rotation related
+            ["rotate90", "target"],
+            ["rotate180", "target"],
+            ["rotate270", "target"],
+            
+            # distortion related
+            ["Barrel", "target"],
+            ["Pincushion", "target"],
+            ["Elastic", "target"],
+            
+            # special effects
+            ["Rain", "target"],
+            ["Frost", "target"]
+        ],
+    },
+    {
+        "task_name": "restoration with reference",
+        "sample_weight": 1,
+        "image_list": [
+            # blur related
+            ["reference", "blur", "target"],
+            ["reference", "compression", "target"],
+            ["reference", "SRx2", "target"],
+            ["reference", "SRx4", "target"],
+            ["reference", "pixelate", "target"],
+            ["reference", "Defocus", "target"],
+            ["reference", "GaussianBlur", "target"], # new
+            # sharpen related
+            ["reference", "oversharpen", "target"], 
+            # noise related
+            ["reference", "GaussianNoise", "target"],
+            ["reference", "PoissonNoise", "target"],
+            ["reference", "SPNoise", "target"],
+            # mosaic
+            ["reference", "mosaic", "target"],
+            # contrast related
+            ["reference", "contrast_strengthen", "target"],
+            ["reference", "contrast_weaken", "target"],
+            # quantization related
+            ["reference", "quantization", "target"],
+            ["reference", "JPEG", "target"],
+            # light related
+            ["reference", "brighten", "target"],
+            ["reference", "darken", "target"],
+            ["reference", "LowLight", "target"], # new
+            # color related
+            ["reference", "saturate_strengthen", "target"],
+            ["reference", "saturate_weaken", "target"],
+            ["reference", "gray", "target"],
+            ["reference", "ColorDistortion", "target"],
+            # infilling
+            ["reference", "Inpainting", "target"],
+            # rotation related
+            ["reference", "rotate90", "target"],
+            ["reference", "rotate180", "target"],
+            ["reference", "rotate270", "target"],
+            # distortion related
+            ["reference", "Barrel", "target"],
+            ["reference", "Pincushion", "target"],
+            ["reference", "Elastic", "target"],
+            # special effects
+            ["reference", "Rain", "target"],
+            ["reference", "Frost", "target"]
+        ],
+    }
+]
+
+
+test_task_dicts = [
+    {
+        "task_name": "conditional generation",
+        "sample_weight": 1,
+        "image_list": [
+            ["canny", "target"],
+            ["depth", "target"],
+            ["hed", "target"],
+            ["normal", "target"],
+            ["mlsd", "target"],
+            ["openpose", "target"],
+            ["sam2_mask", "target"],
+            ["uniformer", "target"],
+            ["mask", "target"],
+            ["foreground", "target"],
+            ["background", "target"],
+        ],
+    },
+    {
+        "task_name": "image generation with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "target"],
+        ],
+    },
+    {
+        "task_name": "conditional generation with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "depth", "target"],
+            ["reference", "openpose", "target"],
+        ],
+    },
+    {
+        "task_name": "subject extraction", 
+        "sample_weight": 0.2,
+        "image_list": [
+            ["target", "reference"],
+        ],
+    },
+    {
+        "task_name": "dense prediction",
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "depth"],
+            ["target", "openpose"],
+        ],
+    },
+    {
+        "task_name": "restoration",
+        "sample_weight": 1,
+        "image_list": [
+            # blur related
+            ["GaussianBlur", "target"],
+            
+            # infilling
+            ["Inpainting", "target"],
+            
+            # rotation related
+            ["rotate90", "target"],
+            
+            # distortion related
+            ["Elastic", "target"],
+        ],
+    },
+    {
+        "task_name": "restoration with reference",
+        "sample_weight": 1,
+        "image_list": [
+            # infilling
+            ["reference", "Inpainting", "target"],
+        ],
+    },
+    {
+        "task_name": "image editing with reference",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "DepthEdit", "target"],
+            ["reference", "FillEdit", "target"],
+        ],
+    },
+    {
+        "task_name": "style transfer",
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "InstantStyle"],
+            ["target", "ReduxStyle"],
+            ["reference", "InstantStyle"],
+        ],
+    },
+    {
+        "task_name": "style transfer with condition",
+        "sample_weight": 1,
+        "image_list": [
+            ["reference", "canny", "InstantStyle"],
+            ["reference", "depth", "InstantStyle"],
+            ["reference", "hed", "InstantStyle"],
+            ["reference", "normal", "InstantStyle"],
+            ["reference", "mlsd", "InstantStyle"],
+            ["reference", "openpose", "InstantStyle"],
+            ["reference", "sam2_mask", "InstantStyle"],
+            ["reference", "uniformer", "InstantStyle"],
+            ["reference", "mask", "InstantStyle"],
+        ],
+    },
+    {
+        "task_name": "subject extraction", 
+        "sample_weight": 1,
+        "image_list": [
+            ["target", "reference"],
+        ],
+    },
+]

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

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

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)

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)

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))

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

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

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

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.)

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

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)

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

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

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

degradation_utils.py

@@ -0,0 +1,232 @@
+import numpy as np
+import cv2
+import random
+from PIL import Image
+
+from degradation_toolkit.add_degradation_various import *
+from degradation_toolkit.image_operators import *
+from 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

demo_tasks/__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