full_files

app.py

@@ -1,213 +1,158 @@
 import spaces
-import json
-import subprocess
 import os
-import sys
-
-def run_command(command):
-    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
-    output, error = process.communicate()
-    if process.returncode != 0:
-        print(f"Error executing command: {command}")
-        print(error.decode('utf-8'))
-        exit(1)
-    return output.decode('utf-8')
-
-# Download CUDA installer
-download_command = "wget https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
-result = run_command(download_command)
-if result is None:
-    print("Failed to download CUDA installer.")
-    exit(1)
-
-# Run CUDA installer in silent mode
-install_command = "sh cuda_12.2.0_535.54.03_linux.run --silent --toolkit --samples --override"
-result = run_command(install_command)
-if result is None:
-    print("Failed to run CUDA installer.")
-    exit(1)
-
-print("CUDA installation process completed.")
-
-def install_packages():
-    
-    # Clone the repository with submodules
-    run_command("git clone --recurse-submodules https://github.com/abetlen/llama-cpp-python.git")
-    
-    # Change to the cloned directory
-    os.chdir("llama-cpp-python")
-    
-    # Checkout the specific commit in the llama.cpp submodule
-    os.chdir("vendor/llama.cpp")
-    run_command("git checkout 50e0535")
-    os.chdir("../..")
-    
-    # Upgrade pip
-    run_command("pip install --upgrade pip")
-
-
-
-    # Install all optional dependencies with CUDA support
-    run_command('CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS -DCUDA_PATH=/usr/local/cuda-12.2 -DCUDAToolkit_ROOT=/usr/local/cuda-12.2 -DCUDAToolkit_INCLUDE_DIR=/usr/local/cuda-12/include -DCUDAToolkit_LIBRARY_DIR=/usr/local/cuda-12.2/lib64" FORCE_CMAKE=1 pip install -e .')
-    
-    run_command("make clean && GGML_OPENBLAS=1 make -j")
-        
-    # Reinstall the package with CUDA support
-    run_command('CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS -DCUDA_PATH=/usr/local/cuda-12.2 -DCUDAToolkit_ROOT=/usr/local/cuda-12.2 -DCUDAToolkit_INCLUDE_DIR=/usr/local/cuda-12/include -DCUDAToolkit_LIBRARY_DIR=/usr/local/cuda-12.2/lib64" FORCE_CMAKE=1 pip install -e .')
-
-    # Install llama-cpp-agent
-    run_command("pip install llama-cpp-agent")
-    
-    run_command("export PYTHONPATH=$PYTHONPATH:$(pwd)")
-    
-    print("Installation complete!")
-
-try:
-    install_packages()
-    
-    # Add a delay to allow for package registration
-    import time
-    time.sleep(5)
-    
-    # Force Python to reload the site packages
-    import site
-    import importlib
-    importlib.reload(site)
-    
-    # Now try to import the libraries
-    from llama_cpp import Llama
-    from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
-    from llama_cpp_agent.providers import LlamaCppPythonProvider
-    from llama_cpp_agent.chat_history import BasicChatHistory
-    from llama_cpp_agent.chat_history.messages import Roles
-    
-    print("Libraries imported successfully!")
-except Exception as e:
-    print(f"Installation failed or libraries couldn't be imported: {str(e)}")
-    sys.exit(1)
-    
+import requests
+import yaml
+import torch
 import gradio as gr
+from PIL import Image
+import sys
+sys.path.append(os.path.abspath('./'))
+from inference.utils import *
+from core.utils import load_or_fail
+from train import WurstCoreB
+from gdf import VPScaler, CosineTNoiseCond, DDPMSampler, P2LossWeight, AdaptiveLossWeight
+from train import WurstCore_t2i as WurstCoreC
+import torch.nn.functional as F
+from core.utils import load_or_fail
+import numpy as np
+import random
+import math
+from einops import rearrange
 from huggingface_hub import hf_hub_download
 
-hf_hub_download(
-    repo_id="MaziyarPanahi/Mistral-Nemo-Instruct-2407-GGUF",
-    filename="Mistral-Nemo-Instruct-2407.Q5_K_M.gguf",
-    local_dir="./models"
-)
 
-# Initialize LLM outside the respond function
-llm = Llama(
-    model_path="models/Mistral-Nemo-Instruct-2407.Q5_K_M.gguf",
-    flash_attn=True,
-    n_gpu_layers=81,
-    n_batch=1024,
-    n_ctx=32768,
+def download_file(url, folder_path, filename):
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+    file_path = os.path.join(folder_path, filename)
+
+    if os.path.isfile(file_path):
+        print(f"File already exists: {file_path}")
+    else:
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(file_path, 'wb') as file:
+                for chunk in response.iter_content(chunk_size=1024):
+                    file.write(chunk)
+            print(f"File successfully downloaded and saved: {file_path}")
+        else:
+            print(f"Error downloading the file. Status code: {response.status_code}")
+
+def download_models():
+    models = {
+        "STABLEWURST_A": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_a.safetensors?download=true", "models/StableWurst", "stage_a.safetensors"),
+        "STABLEWURST_PREVIEWER": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/previewer.safetensors?download=true", "models/StableWurst", "previewer.safetensors"),
+        "STABLEWURST_EFFNET": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/effnet_encoder.safetensors?download=true", "models/StableWurst", "effnet_encoder.safetensors"),
+        "STABLEWURST_B_LITE": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_b_lite_bf16.safetensors?download=true", "models/StableWurst", "stage_b_lite_bf16.safetensors"),
+        "STABLEWURST_C": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_c_bf16.safetensors?download=true", "models/StableWurst", "stage_c_bf16.safetensors"),
+        "ULTRAPIXEL_T2I": ("https://huggingface.co/roubaofeipi/UltraPixel/resolve/main/ultrapixel_t2i.safetensors?download=true", "models/UltraPixel", "ultrapixel_t2i.safetensors"),
+        "ULTRAPIXEL_LORA_CAT": ("https://huggingface.co/roubaofeipi/UltraPixel/resolve/main/lora_cat.safetensors?download=true", "models/UltraPixel", "lora_cat.safetensors"),
+    }
+
+    for model, (url, folder, filename) in models.items():
+        download_file(url, folder, filename)
+
+download_models()
+
+# Global variables
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16
+
+# Load configs and setup models
+with open("configs/training/t2i.yaml", "r", encoding="utf-8") as file:
+    config_c = yaml.safe_load(file)
+
+with open("configs/inference/stage_b_1b.yaml", "r", encoding="utf-8") as file:
+    config_b = yaml.safe_load(file)
+
+core = WurstCoreC(config_dict=config_c, device=device, training=False)
+core_b = WurstCoreB(config_dict=config_b, device=device, training=False)
+
+extras = core.setup_extras_pre()
+models = core.setup_models(extras)
+models.generator.eval().requires_grad_(False)
+
+extras_b = core_b.setup_extras_pre()
+models_b = core_b.setup_models(extras_b, skip_clip=True)
+models_b = WurstCoreB.Models(
+   **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
 )
-
-provider = LlamaCppPythonProvider(llm)
-
-@spaces.GPU(duration=120)
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    top_k,
-    repeat_penalty,
-):
-    chat_template = MessagesFormatterType.MISTRAL
-
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=f"{system_message}",
-        predefined_messages_formatter_type=chat_template,
-        debug_output=True
-    )
+models_b.generator.bfloat16().eval().requires_grad_(False)
+
+# Load pretrained model
+pretrained_path = "models/ultrapixel_t2i.safetensors"
+sdd = torch.load(pretrained_path, map_location='cpu')
+collect_sd = {k[7:]: v for k, v in sdd.items()}
+models.train_norm.load_state_dict(collect_sd)
+models.generator.eval()
+models.train_norm.eval()
+
+# Set up sampling configurations
+extras.sampling_configs.update({
+    'cfg': 4,
+    'shift': 1,
+    'timesteps': 20,
+    't_start': 1.0,
+    'sampler': DDPMSampler(extras.gdf)
+})
+
+extras_b.sampling_configs.update({
+    'cfg': 1.1,
+    'shift': 1,
+    'timesteps': 10,
+    't_start': 1.0
+})
+
+@spaces.GPU
+def generate_image(prompt, height, width, seed):
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+    batch_size = 1
+    height_lr, width_lr = get_target_lr_size(height / width, std_size=32)
+    stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)
+    stage_c_latent_shape_lr, stage_b_latent_shape_lr = calculate_latent_sizes(height_lr, width_lr, batch_size=batch_size)
+
+    batch = {'captions': [prompt] * batch_size}
+    conditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
+    unconditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)    
     
-    settings = provider.get_provider_default_settings()
-    settings.temperature = temperature
-    settings.top_k = top_k
-    settings.top_p = top_p
-    settings.max_tokens = max_tokens
-    settings.repeat_penalty = repeat_penalty
-    settings.stream = True
-
-    messages = BasicChatHistory()
-
-    for msn in history:
-        user = {
-            'role': Roles.user,
-            'content': msn[0]
-        }
-        assistant = {
-            'role': Roles.assistant,
-            'content': msn[1]
-        }
-        messages.add_message(user)
-        messages.add_message(assistant)
+    conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
+    unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
     
-    stream = agent.get_chat_response(
-        message,
-        llm_sampling_settings=settings,
-        chat_history=messages,
-        returns_streaming_generator=True,
-        print_output=False
-    )
-    
-    outputs = ""
-    for output in stream:
-        outputs += output
-        yield outputs
-
-description = """<p><center>
-<a href="https://huggingface.co/mistralai/Mistral-Nemo-Instruct-2407" target="_blank">[Instruct Model]</a>
-<a href="https://huggingface.co/mistralai/Mistral-Nemo-Base-2407" target="_blank">[Base Model]</a>
-<a href="https://huggingface.co/second-state/Mistral-Nemo-Instruct-2407-GGUF" target="_blank">[GGUF Version]</a>
-</center></p>
-"""
-
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a helpful assistant.", label="System message"),
-        gr.Slider(minimum=1, maximum=4096, value=2048, step=1, label="Max tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p",
-        ),
-        gr.Slider(
-            minimum=0,
-            maximum=100,
-            value=40,
-            step=1,
-            label="Top-k",
-        ),
-        gr.Slider(
-            minimum=0.0,
-            maximum=2.0,
-            value=1.1,
-            step=0.1,
-            label="Repetition penalty",
-        ),
+    with torch.no_grad():
+        models.generator.cuda()
+        with torch.cuda.amp.autocast(dtype=dtype):
+            sampled_c = generation_c(batch, models, extras, core, stage_c_latent_shape, stage_c_latent_shape_lr, device)
+        
+        models.generator.cpu()
+        torch.cuda.empty_cache()
+        
+        conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
+        unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
+        conditions_b['effnet'] = sampled_c
+        unconditions_b['effnet'] = torch.zeros_like(sampled_c)
+        
+        with torch.cuda.amp.autocast(dtype=dtype):
+            sampled = decode_b(conditions_b, unconditions_b, models_b, stage_b_latent_shape, extras_b, device, stage_a_tiled=True)
+        
+        torch.cuda.empty_cache()
+        imgs = show_images(sampled)
+        return imgs[0]
+
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Slider(minimum=256, maximum=2560, step=32, label="Height", value=1024),
+        gr.Slider(minimum=256, maximum=5120, step=32, label="Width", value=1024),
+        gr.Number(label="Seed", value=42)
     ],
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear",
-    submit_btn="Send",
-    title="Chat with Mistral-NeMo using llama.cpp", 
-    description=description,
-    chatbot=gr.Chatbot(
-        scale=1, 
-        likeable=False,
-        show_copy_button=True
-    )
+    outputs=gr.Image(type="pil"),
+    title="UltraPixel Image Generation",
+    description="Generate high-resolution images using UltraPixel model.",
+    theme='bethecloud/storj_theme'
 )
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+iface.launch()

configs/inference/controlnet_c_3b_canny.yaml

@@ -0,0 +1,14 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+# ControlNet specific
+controlnet_blocks: [0, 4, 8, 12, 51, 55, 59, 63]
+controlnet_filter: CannyFilter
+controlnet_filter_params: 
+  resize: 224
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+controlnet_checkpoint_path: models/canny.safetensors

configs/inference/controlnet_c_3b_identity.yaml

@@ -0,0 +1,17 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+# ControlNet specific
+controlnet_bottleneck_mode: 'simple'
+controlnet_blocks: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63]
+controlnet_filter: IdentityFilter
+controlnet_filter_params: 
+  max_faces: 4
+  p_drop: 0.00
+  p_full: 0.0
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+controlnet_checkpoint_path: 

configs/inference/controlnet_c_3b_inpainting.yaml

@@ -0,0 +1,15 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+# ControlNet specific
+controlnet_blocks: [0, 4, 8, 12, 51, 55, 59, 63]
+controlnet_filter: InpaintFilter
+controlnet_filter_params: 
+  thresold: [0.04, 0.4]
+  p_outpaint: 0.4
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+controlnet_checkpoint_path: models/inpainting.safetensors

configs/inference/controlnet_c_3b_sr.yaml

@@ -0,0 +1,15 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+# ControlNet specific
+controlnet_bottleneck_mode: 'large'
+controlnet_blocks: [0, 4, 8, 12, 51, 55, 59, 63]
+controlnet_filter: SREffnetFilter
+controlnet_filter_params: 
+  scale_factor: 0.5
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+controlnet_checkpoint_path: models/super_resolution.safetensors

configs/inference/lora_c_3b.yaml

@@ -0,0 +1,15 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+# LoRA specific
+module_filters: ['.attn']
+rank: 4
+train_tokens:
+  # - ['^snail', null] # token starts with "snail" -> "snail" & "snails", don't need to be reinitialized
+  - ['[fernando]', '^dog</w>'] # custom token [snail], initialize as avg of snail & snails
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+lora_checkpoint_path: models/lora_fernando_10k.safetensors

configs/inference/stage_b_1b.yaml

@@ -0,0 +1,13 @@
+# GLOBAL STUFF
+model_version: 700M
+dtype: bfloat16
+
+# For demonstration purposes in reconstruct_images.ipynb
+webdataset_path: path to your dataset
+batch_size: 1
+image_size: 2048
+grad_accum_steps: 1
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+stage_a_checkpoint_path: models/stage_a.safetensors
+generator_checkpoint_path: models/stage_b_lite_bf16.safetensors

configs/inference/stage_b_3b.yaml

@@ -0,0 +1,13 @@
+# GLOBAL STUFF
+model_version: 3B
+dtype: bfloat16
+
+# For demonstration purposes in reconstruct_images.ipynb
+webdataset_path: path to your dataset
+batch_size: 4
+image_size: 1024
+grad_accum_steps: 1
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+stage_a_checkpoint_path: models/stage_a.safetensors
+generator_checkpoint_path: models/stage_b_lite_bf16.safetensors

configs/inference/stage_c_1b.yaml

@@ -0,0 +1,7 @@
+# GLOBAL STUFF
+model_version: 1B
+dtype: bfloat16
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_lite_bf16.safetensors

configs/inference/stage_c_3b.yaml

@@ -0,0 +1,7 @@
+# GLOBAL STUFF
+model_version: 3.6B
+dtype: bfloat16
+
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors

configs/training/cfg_control_lr.yaml

@@ -0,0 +1,47 @@
+# GLOBAL STUFF
+experiment_id: Ultrapixel_controlnet
+
+checkpoint_path: checkpoint output path
+output_path:  visual results output path
+model_version: 3.6B
+dtype: float32
+# # WandB
+# wandb_project: StableCascade
+# wandb_entity: wandb_username
+#module_filters: ['.depthwise', '.mapper', '.attn', '.channelwise' ]
+#rank: 32
+# TRAINING PARAMS
+lr: 1.0e-4
+batch_size: 12
+#image_size: [1536, 2048, 2560, 3072, 4096]
+image_size: [1024, 2048, 2560, 3072, 3584, 3840, 4096, 4608]
+#image_size: [  1024, 1536, 2048, 2560,  3072, 3584, 3840, 4096, 4608]
+#image_size: [  1024, 1280]
+multi_aspect_ratio: [1/1, 1/2, 1/3, 2/3, 3/4, 1/5, 2/5, 3/5, 4/5, 1/6, 5/6, 9/16]
+grad_accum_steps: 2
+updates: 40000
+backup_every: 5000
+save_every: 256
+warmup_updates: 1
+use_fsdp: True
+
+# ControlNet specific
+controlnet_blocks: [0, 4, 8, 12, 51, 55, 59, 63]
+controlnet_filter: CannyFilter
+controlnet_filter_params: 
+  resize: 224
+# offset_noise: 0.1
+
+# GDF
+adaptive_loss_weight: True
+
+ema_start_iters: 10
+ema_iters: 50
+ema_beta: 0.9
+
+webdataset_path: path to your training dataset
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors
+controlnet_checkpoint_path: pretrained controlnet path
+

configs/training/lora_personalization.yaml

@@ -0,0 +1,37 @@
+# GLOBAL STUFF
+experiment_id: roubao_cat_personalized
+
+checkpoint_path: checkpoint output path
+output_path:  visual results output path
+model_version: 3.6B
+dtype: float32
+
+module_filters: [ '.attn']
+rank: 4
+train_tokens:
+  # - ['^snail', null] # token starts with "snail" -> "snail" & "snails", don't need to be reinitialized
+  - ['[roubaobao]', '^cat</w>'] # custom token [snail], initialize as avg of snail & snails
+# TRAINING PARAMS
+lr: 1.0e-4
+batch_size: 4
+
+image_size: [1024, 2048, 2560, 3072, 3584, 3840, 4096, 4608]
+multi_aspect_ratio: [1/1, 1/2, 1/3, 2/3, 3/4, 1/5, 2/5, 3/5, 4/5, 1/6, 5/6, 9/16]
+grad_accum_steps: 2
+updates: 40000
+backup_every: 5000
+save_every: 512
+warmup_updates: 1
+use_ddp: True
+
+# GDF
+adaptive_loss_weight: True
+
+
+tmp_prompt: a photo of a cat [roubaobao]
+webdataset_path: path to your personalized training dataset
+effnet_checkpoint_path:  models/effnet_encoder.safetensors
+previewer_checkpoint_path:  models/previewer.safetensors
+generator_checkpoint_path:  models/stage_c_bf16.safetensors
+ultrapixel_path: models/ultrapixel_t2i.safetensors
+

configs/training/t2i.yaml

@@ -0,0 +1,29 @@
+# GLOBAL STUFF
+experiment_id: ultrapixel_t2i
+#strc_fixlrt_norm3_lite_1024_hrft_newdata
+checkpoint_path: checkpoint output path   #output model directory
+output_path: visual results output path       #experiment output directory
+model_version: 3.6B    # finetune large  stage c model of stablecascade
+dtype: float32
+
+
+# TRAINING PARAMS
+lr: 1.0e-4
+batch_size: 4   # gpu_number * num_per_gpu * grad_accum_steps
+image_size: [1024, 2048, 2560, 3072, 3584, 3840, 4096, 4608]  # possible image resolution
+multi_aspect_ratio: [1/1, 1/2, 1/3, 2/3, 3/4, 1/5, 2/5, 3/5, 4/5, 1/6, 5/6, 9/16]
+grad_accum_steps: 2
+updates: 40000
+backup_every: 5000
+save_every: 256
+warmup_updates: 1
+use_ddp: True
+
+# GDF
+adaptive_loss_weight: True
+
+
+webdataset_path: path to your personalized training dataset
+effnet_checkpoint_path: models/effnet_encoder.safetensors
+previewer_checkpoint_path: models/previewer.safetensors
+generator_checkpoint_path: models/stage_c_bf16.safetensors

core/__init__.py

@@ -0,0 +1,372 @@
+import os
+import yaml
+import torch
+from torch import nn
+import wandb
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from torch.utils.data import Dataset, DataLoader
+
+from torch.distributed import init_process_group, destroy_process_group, barrier
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    FullStateDictConfig,
+    MixedPrecision,
+    ShardingStrategy,
+    StateDictType
+)
+
+from .utils import Base, EXPECTED, EXPECTED_TRAIN
+from .utils import create_folder_if_necessary, safe_save, load_or_fail
+
+# pylint: disable=unused-argument
+class WarpCore(ABC):
+    @dataclass(frozen=True)
+    class Config(Base):
+        experiment_id: str = EXPECTED_TRAIN
+        checkpoint_path: str = EXPECTED_TRAIN
+        output_path: str = EXPECTED_TRAIN
+        checkpoint_extension: str = "safetensors"
+        dist_file_subfolder: str = ""
+        allow_tf32: bool = True
+
+        wandb_project: str = None
+        wandb_entity: str = None
+
+    @dataclass() # not frozen, means that fields are mutable
+    class Info(): # not inheriting from Base, because we don't want to enforce the default fields
+        wandb_run_id: str = None
+        total_steps: int = 0
+        iter: int = 0
+
+    @dataclass(frozen=True)
+    class Data(Base):
+        dataset: Dataset = EXPECTED
+        dataloader: DataLoader  = EXPECTED
+        iterator: any = EXPECTED
+
+    @dataclass(frozen=True)
+    class Models(Base):
+        pass
+
+    @dataclass(frozen=True)
+    class Optimizers(Base):
+        pass
+
+    @dataclass(frozen=True)
+    class Schedulers(Base):
+        pass
+
+    @dataclass(frozen=True)
+    class Extras(Base):
+        pass
+    # ---------------------------------------
+    info: Info
+    config: Config
+
+    # FSDP stuff
+    fsdp_defaults = {
+        "sharding_strategy": ShardingStrategy.SHARD_GRAD_OP,
+        "cpu_offload": None,
+        "mixed_precision": MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.bfloat16,
+            buffer_dtype=torch.bfloat16,
+        ),
+        "limit_all_gathers": True,
+    }
+    fsdp_fullstate_save_policy = FullStateDictConfig(
+        offload_to_cpu=True, rank0_only=True
+    )
+    # ------------
+
+    # OVERRIDEABLE METHODS
+    
+    # [optionally] setup extra stuff, will be called BEFORE the models & optimizers are setup
+    def setup_extras_pre(self) -> Extras:
+        return self.Extras()
+
+    # setup dataset & dataloader, return a dict contained dataser, dataloader and/or iterator
+    @abstractmethod
+    def setup_data(self, extras: Extras) -> Data:
+        raise NotImplementedError("This method needs to be overriden")
+
+    # return a dict with all models that are going to be used in the training
+    @abstractmethod
+    def setup_models(self, extras: Extras) -> Models:
+        raise NotImplementedError("This method needs to be overriden")
+
+    # return a dict with all optimizers that are going to be used in the training
+    @abstractmethod
+    def setup_optimizers(self, extras: Extras, models: Models) -> Optimizers:
+        raise NotImplementedError("This method needs to be overriden")
+
+    # [optionally] return a dict with all schedulers that are going to be used in the training
+    def setup_schedulers(self, extras: Extras, models: Models, optimizers: Optimizers) -> Schedulers:
+        return self.Schedulers()
+
+    # [optionally] setup extra stuff, will be called AFTER the models & optimizers are setup
+    def setup_extras_post(self, extras: Extras, models: Models, optimizers: Optimizers, schedulers: Schedulers) -> Extras:
+        return self.Extras.from_dict(extras.to_dict())
+
+    # perform the training here
+    @abstractmethod
+    def train(self, data: Data, extras: Extras, models: Models, optimizers: Optimizers, schedulers: Schedulers):
+        raise NotImplementedError("This method needs to be overriden")
+    # ------------
+
+    def setup_info(self, full_path=None) -> Info:
+        if full_path is None:
+            full_path = (f"{self.config.checkpoint_path}/{self.config.experiment_id}/info.json")
+        info_dict = load_or_fail(full_path, wandb_run_id=None) or {}
+        info_dto = self.Info(**info_dict)
+        if info_dto.total_steps > 0 and self.is_main_node:
+            print(">>> RESUMING TRAINING FROM ITER ", info_dto.total_steps)
+        return info_dto
+
+    def setup_config(self, config_file_path=None, config_dict=None, training=True) -> Config:
+        if config_file_path is not None:
+            if config_file_path.endswith(".yml") or config_file_path.endswith(".yaml"):
+                with open(config_file_path, "r", encoding="utf-8") as file:
+                    loaded_config = yaml.safe_load(file)
+            elif config_file_path.endswith(".json"):
+                with open(config_file_path, "r", encoding="utf-8") as file:
+                    loaded_config = json.load(file)
+            else:
+                raise ValueError("Config file must be either a .yml|.yaml or .json file")
+            return self.Config.from_dict({**loaded_config, 'training': training})
+        if config_dict is not None:
+            return self.Config.from_dict({**config_dict, 'training': training})
+        return self.Config(training=training)
+
+    def setup_ddp(self, experiment_id, single_gpu=False):
+        if not single_gpu:
+            local_rank = int(os.environ.get("SLURM_LOCALID"))
+            process_id = int(os.environ.get("SLURM_PROCID"))
+            world_size = int(os.environ.get("SLURM_NNODES")) * torch.cuda.device_count()
+
+            self.process_id = process_id
+            self.is_main_node = process_id == 0
+            self.device = torch.device(local_rank)
+            self.world_size = world_size
+
+            dist_file_path = f"{os.getcwd()}/{self.config.dist_file_subfolder}dist_file_{experiment_id}"
+            # if os.path.exists(dist_file_path) and self.is_main_node:
+            #     os.remove(dist_file_path)
+
+            torch.cuda.set_device(local_rank)
+            init_process_group(
+                backend="nccl",
+                rank=process_id,
+                world_size=world_size,
+                init_method=f"file://{dist_file_path}",
+            )
+            print(f"[GPU {process_id}] READY")
+        else:
+            print("Running in single thread, DDP not enabled.")
+
+    def setup_wandb(self):
+        if self.is_main_node and self.config.wandb_project is not None:
+            self.info.wandb_run_id = self.info.wandb_run_id or wandb.util.generate_id()
+            wandb.init(project=self.config.wandb_project, entity=self.config.wandb_entity, name=self.config.experiment_id, id=self.info.wandb_run_id, resume="allow", config=self.config.to_dict())
+
+            if self.info.total_steps > 0:
+                wandb.alert(title=f"Training {self.info.wandb_run_id} resumed", text=f"Training {self.info.wandb_run_id} resumed from step {self.info.total_steps}")
+            else:
+                wandb.alert(title=f"Training {self.info.wandb_run_id} started", text=f"Training {self.info.wandb_run_id} started")
+
+    # LOAD UTILITIES ----------
+    def load_model(self, model, model_id=None, full_path=None, strict=True):
+        print('in line 181 load model', type(model), model_id, full_path, strict)
+        if model_id is not None and full_path is None:
+            full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{model_id}.{self.config.checkpoint_extension}"
+        elif full_path is None and model_id is None:
+            raise ValueError(
+                "This method expects either 'model_id' or 'full_path' to be defined"
+            )
+
+        checkpoint = load_or_fail(full_path, wandb_run_id=self.info.wandb_run_id if self.is_main_node else None)
+        if checkpoint is not None:
+            model.load_state_dict(checkpoint, strict=strict)
+            del checkpoint
+
+        return model
+
+    def load_optimizer(self, optim, optim_id=None, full_path=None, fsdp_model=None):
+        if optim_id is not None and full_path is None:
+            full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{optim_id}.pt"
+        elif full_path is None and optim_id is None:
+            raise ValueError(
+                "This method expects either 'optim_id' or 'full_path' to be defined"
+            )
+
+        checkpoint = load_or_fail(full_path, wandb_run_id=self.info.wandb_run_id if self.is_main_node else None)
+        if checkpoint is not None:
+            try:
+                if fsdp_model is not None:
+                    sharded_optimizer_state_dict = (
+                        FSDP.scatter_full_optim_state_dict(  # <---- FSDP
+                            checkpoint
+                            if (
+                                self.is_main_node
+                                or self.fsdp_defaults["sharding_strategy"]
+                                == ShardingStrategy.NO_SHARD
+                            )
+                            else None,
+                            fsdp_model,
+                        )
+                    )
+                    optim.load_state_dict(sharded_optimizer_state_dict)
+                    del checkpoint, sharded_optimizer_state_dict
+                else:
+                    optim.load_state_dict(checkpoint)
+            # pylint: disable=broad-except
+            except Exception as e:
+                print("!!! Failed loading optimizer, skipping... Exception:", e)
+
+        return optim
+
+    # SAVE UTILITIES ----------
+    def save_info(self, info, suffix=""):
+        full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/info{suffix}.json"
+        create_folder_if_necessary(full_path)
+        if self.is_main_node:
+            safe_save(vars(self.info), full_path)
+
+    def save_model(self, model, model_id=None, full_path=None, is_fsdp=False):
+        if model_id is not None and full_path is None:
+            full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{model_id}.{self.config.checkpoint_extension}"
+        elif full_path is None and model_id is None:
+            raise ValueError(
+                "This method expects either 'model_id' or 'full_path' to be defined"
+            )
+        create_folder_if_necessary(full_path)
+        if is_fsdp:
+            with FSDP.summon_full_params(model):
+                pass
+            with FSDP.state_dict_type(
+                model, StateDictType.FULL_STATE_DICT, self.fsdp_fullstate_save_policy
+            ):
+                checkpoint = model.state_dict()
+            if self.is_main_node:
+                safe_save(checkpoint, full_path)
+            del checkpoint
+        else:
+            if self.is_main_node:
+                checkpoint = model.state_dict()
+                safe_save(checkpoint, full_path)
+                del checkpoint
+
+    def save_optimizer(self, optim, optim_id=None, full_path=None, fsdp_model=None):
+        if optim_id is not None and full_path is None:
+            full_path = f"{self.config.checkpoint_path}/{self.config.experiment_id}/{optim_id}.pt"
+        elif full_path is None and optim_id is None:
+            raise ValueError(
+                "This method expects either 'optim_id' or 'full_path' to be defined"
+            )
+        create_folder_if_necessary(full_path)
+        if fsdp_model is not None:
+            optim_statedict = FSDP.full_optim_state_dict(fsdp_model, optim)
+            if self.is_main_node:
+                safe_save(optim_statedict, full_path)
+            del optim_statedict
+        else:
+            if self.is_main_node:
+                checkpoint = optim.state_dict()
+                safe_save(checkpoint, full_path)
+                del checkpoint
+    # -----
+
+    def __init__(self, config_file_path=None, config_dict=None, device="cpu", training=True):
+        # Temporary setup, will be overriden by setup_ddp if required
+        self.device = device
+        self.process_id = 0
+        self.is_main_node = True
+        self.world_size = 1
+        # ----
+
+        self.config: self.Config = self.setup_config(config_file_path, config_dict, training)
+        self.info: self.Info = self.setup_info()
+
+    def __call__(self, single_gpu=False):
+        self.setup_ddp(self.config.experiment_id, single_gpu=single_gpu)  # this will change the device to the CUDA rank
+        self.setup_wandb()
+        if self.config.allow_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
+        if self.is_main_node:
+            print()
+            print("**STARTIG JOB WITH CONFIG:**")
+            print(yaml.dump(self.config.to_dict(), default_flow_style=False))
+            print("------------------------------------")
+            print()
+            print("**INFO:**")
+            print(yaml.dump(vars(self.info), default_flow_style=False))
+            print("------------------------------------")
+            print()
+
+        # SETUP STUFF
+        extras = self.setup_extras_pre()
+        assert extras is not None, "setup_extras_pre() must return a DTO"
+
+        data = self.setup_data(extras)
+        assert data is not None, "setup_data() must return a DTO"
+        if self.is_main_node:
+            print("**DATA:**")
+            print(yaml.dump({k:type(v).__name__ for k, v in data.to_dict().items()}, default_flow_style=False))
+            print("------------------------------------")
+            print()
+
+        models = self.setup_models(extras)
+        assert models is not None, "setup_models() must return a DTO"
+        if self.is_main_node:
+            print("**MODELS:**")
+            print(yaml.dump({
+                k:f"{type(v).__name__} - {f'trainable params {sum(p.numel() for p in v.parameters() if p.requires_grad)}' if isinstance(v, nn.Module) else 'Not a nn.Module'}" for k, v in models.to_dict().items()
+            }, default_flow_style=False))
+            print("------------------------------------")
+            print()
+
+        optimizers = self.setup_optimizers(extras, models)
+        assert optimizers is not None, "setup_optimizers() must return a DTO"
+        if self.is_main_node:
+            print("**OPTIMIZERS:**")
+            print(yaml.dump({k:type(v).__name__ for k, v in optimizers.to_dict().items()}, default_flow_style=False))
+            print("------------------------------------")
+            print()
+
+        schedulers = self.setup_schedulers(extras, models, optimizers)
+        assert schedulers is not None, "setup_schedulers() must return a DTO"
+        if self.is_main_node:
+            print("**SCHEDULERS:**")
+            print(yaml.dump({k:type(v).__name__ for k, v in schedulers.to_dict().items()}, default_flow_style=False))
+            print("------------------------------------")
+            print()
+
+        post_extras =self.setup_extras_post(extras, models, optimizers, schedulers)
+        assert post_extras is not None, "setup_extras_post() must return a DTO"
+        extras = self.Extras.from_dict({ **extras.to_dict(),**post_extras.to_dict() })
+        if self.is_main_node:
+            print("**EXTRAS:**")
+            print(yaml.dump({k:f"{v}" for k, v in extras.to_dict().items()}, default_flow_style=False))
+            print("------------------------------------")
+            print()
+        # -------
+
+        # TRAIN
+        if self.is_main_node:
+            print("**TRAINING STARTING...**")
+        self.train(data, extras, models, optimizers, schedulers)
+
+        if single_gpu is False:
+            barrier()
+            destroy_process_group()
+        if self.is_main_node:
+            print()
+            print("------------------------------------")
+            print()
+            print("**TRAINING COMPLETE**")
+            if self.config.wandb_project is not None:
+                wandb.alert(title=f"Training {self.info.wandb_run_id} finished", text=f"Training {self.info.wandb_run_id} finished")

core/data/__init__.py

@@ -0,0 +1,69 @@
+import json
+import subprocess
+import yaml
+import os
+from .bucketeer import Bucketeer
+
+class MultiFilter():
+    def __init__(self, rules, default=False):
+        self.rules = rules
+        self.default = default
+
+    def __call__(self, x):
+        try:
+            x_json = x['json']
+            if isinstance(x_json, bytes):
+                x_json = json.loads(x_json) 
+            validations = []
+            for k, r in self.rules.items():
+                if isinstance(k, tuple):
+                    v = r(*[x_json[kv] for kv in k])
+                else:
+                    v = r(x_json[k])
+                validations.append(v)
+            return all(validations)
+        except Exception:
+            return False
+
+class MultiGetter():
+    def __init__(self, rules):
+        self.rules = rules
+
+    def __call__(self, x_json):
+        if isinstance(x_json, bytes):
+            x_json = json.loads(x_json) 
+        outputs = []
+        for k, r in self.rules.items():
+            if isinstance(k, tuple):
+                v = r(*[x_json[kv] for kv in k])
+            else:
+                v = r(x_json[k])
+            outputs.append(v)
+        if len(outputs) == 1:
+            outputs = outputs[0]
+        return outputs
+
+def setup_webdataset_path(paths, cache_path=None):
+    if cache_path is None or not os.path.exists(cache_path):
+        tar_paths = []
+        if isinstance(paths, str):
+            paths = [paths]
+        for path in paths:
+            if path.strip().endswith(".tar"):
+                # Avoid looking up s3 if we already have a tar file
+                tar_paths.append(path)
+                continue
+            bucket = "/".join(path.split("/")[:3])
+            result = subprocess.run([f"aws s3 ls {path} --recursive | awk '{{print $4}}'"], stdout=subprocess.PIPE, shell=True, check=True)
+            files = result.stdout.decode('utf-8').split()
+            files = [f"{bucket}/{f}" for f in files if f.endswith(".tar")]
+            tar_paths += files
+
+        with open(cache_path, 'w', encoding='utf-8') as outfile:
+            yaml.dump(tar_paths, outfile, default_flow_style=False)
+    else:
+        with open(cache_path, 'r', encoding='utf-8') as file:
+            tar_paths = yaml.safe_load(file)
+
+    tar_paths_str = ",".join([f"{p}" for p in tar_paths])
+    return f"pipe:aws s3 cp {{ {tar_paths_str} }} -"

core/data/bucketeer.py

@@ -0,0 +1,88 @@
+import torch
+import torchvision
+import numpy as np
+from torchtools.transforms import SmartCrop
+import math
+
+class Bucketeer():
+    def __init__(self, dataloader, density=256*256, factor=8, ratios=[1/1, 1/2, 3/4, 3/5, 4/5, 6/9, 9/16], reverse_list=True, randomize_p=0.3, randomize_q=0.2, crop_mode='random', p_random_ratio=0.0, interpolate_nearest=False):
+        assert crop_mode in ['center', 'random', 'smart']
+        self.crop_mode = crop_mode
+        self.ratios = ratios
+        if reverse_list:
+            for r in list(ratios):
+                if 1/r not in self.ratios:
+                    self.ratios.append(1/r)
+        self.sizes = {}
+        for dd in density:
+          self.sizes[dd]= [(int(((dd/r)**0.5//factor)*factor), int(((dd*r)**0.5//factor)*factor)) for r in ratios]   
+    
+        self.batch_size = dataloader.batch_size
+        self.iterator = iter(dataloader)
+        all_sizes =  []
+        for k, vs in self.sizes.items():
+          all_sizes += vs
+        self.buckets = {s: [] for s in all_sizes}
+        self.smartcrop = SmartCrop(int(density**0.5), randomize_p, randomize_q) if self.crop_mode=='smart' else None
+        self.p_random_ratio = p_random_ratio
+        self.interpolate_nearest = interpolate_nearest
+
+    def get_available_batch(self):
+        for b in self.buckets:
+            if len(self.buckets[b]) >= self.batch_size:
+                batch = self.buckets[b][:self.batch_size]
+                self.buckets[b] = self.buckets[b][self.batch_size:]
+                return batch
+        return None
+
+    def get_closest_size(self, x):
+        w, h = x.size(-1), x.size(-2)
+
+            
+        best_size_idx = np.argmin([abs(w/h-r) for r in self.ratios])
+        find_dict = {dd : abs(w*h - self.sizes[dd][best_size_idx][0]*self.sizes[dd][best_size_idx][1]) for dd, vv in self.sizes.items()}
+        min_ = find_dict[list(find_dict.keys())[0]]
+        find_size = self.sizes[list(find_dict.keys())[0]][best_size_idx]
+        for dd, val in find_dict.items():
+          if val < min_:
+            min_ = val
+            find_size = self.sizes[dd][best_size_idx]  
+            
+        return find_size
+
+    def get_resize_size(self, orig_size, tgt_size):
+        if (tgt_size[1]/tgt_size[0] - 1) * (orig_size[1]/orig_size[0] - 1) >= 0:
+            alt_min = int(math.ceil(max(tgt_size)*min(orig_size)/max(orig_size)))
+            resize_size = max(alt_min, min(tgt_size))
+        else:
+            alt_max = int(math.ceil(min(tgt_size)*max(orig_size)/min(orig_size)))
+            resize_size = max(alt_max, max(tgt_size))
+
+        return resize_size
+
+    def __next__(self):
+        batch = self.get_available_batch()
+        while batch is None:
+            elements = next(self.iterator)
+            for dct in elements:
+                img = dct['images']
+                size = self.get_closest_size(img)
+                resize_size = self.get_resize_size(img.shape[-2:], size)
+              
+                if self.interpolate_nearest:
+                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.NEAREST)
+                else:
+                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True)
+                if self.crop_mode == 'center':
+                    img = torchvision.transforms.functional.center_crop(img, size)
+                elif self.crop_mode == 'random':
+                    img = torchvision.transforms.RandomCrop(size)(img)
+                elif self.crop_mode == 'smart':
+                    self.smartcrop.output_size = size
+                    img = self.smartcrop(img)
+                
+                self.buckets[size].append({**{'images': img}, **{k:dct[k] for k in dct if k != 'images'}})
+            batch = self.get_available_batch()
+
+        out = {k:[batch[i][k] for i in range(len(batch))] for k in batch[0]}
+        return {k: torch.stack(o, dim=0) if isinstance(o[0], torch.Tensor) else o for k, o in out.items()}

core/data/bucketeer_deg.py

@@ -0,0 +1,91 @@
+import torch
+import torchvision
+import numpy as np
+from torchtools.transforms import SmartCrop
+import math
+
+class Bucketeer():
+    def __init__(self, dataloader, density=256*256, factor=8, ratios=[1/1, 1/2, 3/4, 3/5, 4/5, 6/9, 9/16], reverse_list=True, randomize_p=0.3, randomize_q=0.2, crop_mode='random', p_random_ratio=0.0, interpolate_nearest=False):
+        assert crop_mode in ['center', 'random', 'smart']
+        self.crop_mode = crop_mode
+        self.ratios = ratios
+        if reverse_list:
+            for r in list(ratios):
+                if 1/r not in self.ratios:
+                    self.ratios.append(1/r)
+        self.sizes = {}
+        for dd in density:
+          self.sizes[dd]= [(int(((dd/r)**0.5//factor)*factor), int(((dd*r)**0.5//factor)*factor)) for r in ratios]   
+        print('in line 17 buckteer', self.sizes)     
+        self.batch_size = dataloader.batch_size
+        self.iterator = iter(dataloader)
+        all_sizes =  []
+        for k, vs in self.sizes.items():
+          all_sizes += vs
+        self.buckets = {s: [] for s in all_sizes}
+        self.smartcrop = SmartCrop(int(density**0.5), randomize_p, randomize_q) if self.crop_mode=='smart' else None
+        self.p_random_ratio = p_random_ratio
+        self.interpolate_nearest = interpolate_nearest
+
+    def get_available_batch(self):
+        for b in self.buckets:
+            if len(self.buckets[b]) >= self.batch_size:
+                batch = self.buckets[b][:self.batch_size]
+                self.buckets[b] = self.buckets[b][self.batch_size:]
+                return batch
+        return None
+
+    def get_closest_size(self, x):
+        w, h = x.size(-1), x.size(-2)
+        #if self.p_random_ratio > 0 and np.random.rand() < self.p_random_ratio:
+        #    best_size_idx = np.random.randint(len(self.ratios))
+            #print('in line 41 get closes size', best_size_idx, x.shape, self.p_random_ratio)
+        #else:
+            
+        best_size_idx = np.argmin([abs(w/h-r) for r in self.ratios])
+        find_dict = {dd : abs(w*h - self.sizes[dd][best_size_idx][0]*self.sizes[dd][best_size_idx][1]) for dd, vv in self.sizes.items()}
+        min_ = find_dict[list(find_dict.keys())[0]]
+        find_size = self.sizes[list(find_dict.keys())[0]][best_size_idx]
+        for dd, val in find_dict.items():
+          if val < min_:
+            min_ = val
+            find_size = self.sizes[dd][best_size_idx]  
+            
+        return find_size
+
+    def get_resize_size(self, orig_size, tgt_size):
+        if (tgt_size[1]/tgt_size[0] - 1) * (orig_size[1]/orig_size[0] - 1) >= 0:
+            alt_min = int(math.ceil(max(tgt_size)*min(orig_size)/max(orig_size)))
+            resize_size = max(alt_min, min(tgt_size))
+        else:
+            alt_max = int(math.ceil(min(tgt_size)*max(orig_size)/min(orig_size)))
+            resize_size = max(alt_max, max(tgt_size))
+        #print('in line 50', orig_size, tgt_size, resize_size)
+        return resize_size
+
+    def __next__(self):
+        batch = self.get_available_batch()
+        while batch is None:
+            elements = next(self.iterator)
+            for dct in elements:
+                img = dct['images']
+                size = self.get_closest_size(img)
+                resize_size = self.get_resize_size(img.shape[-2:], size)
+                #print('in line 74', img.size(), resize_size)
+                if self.interpolate_nearest:
+                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.NEAREST)
+                else:
+                    img = torchvision.transforms.functional.resize(img, resize_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True)
+                if self.crop_mode == 'center':
+                    img = torchvision.transforms.functional.center_crop(img, size)
+                elif self.crop_mode == 'random':
+                    img = torchvision.transforms.RandomCrop(size)(img)
+                elif self.crop_mode == 'smart':
+                    self.smartcrop.output_size = size
+                    img = self.smartcrop(img)
+                print('in line 86 bucketeer', type(img), img.shape, torch.max(img), torch.min(img))
+                self.buckets[size].append({**{'images': img}, **{k:dct[k] for k in dct if k != 'images'}})
+            batch = self.get_available_batch()
+
+        out = {k:[batch[i][k] for i in range(len(batch))] for k in batch[0]}
+        return {k: torch.stack(o, dim=0) if isinstance(o[0], torch.Tensor) else o for k, o in out.items()}

core/data/deg_kair_utils/utils_alignfaces.py

@@ -0,0 +1,263 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Apr 24 15:43:29 2017
+@author: zhaoy
+"""
+import cv2
+import numpy as np
+from skimage import transform as trans
+
+# reference facial points, a list of coordinates (x,y)
+REFERENCE_FACIAL_POINTS = [
+    [30.29459953, 51.69630051],
+    [65.53179932, 51.50139999],
+    [48.02519989, 71.73660278],
+    [33.54930115, 92.3655014],
+    [62.72990036, 92.20410156]
+]
+
+DEFAULT_CROP_SIZE = (96, 112)
+
+
+def _umeyama(src, dst, estimate_scale=True, scale=1.0):
+    """Estimate N-D similarity transformation with or without scaling.
+    Parameters
+    ----------
+    src : (M, N) array
+        Source coordinates.
+    dst : (M, N) array
+        Destination coordinates.
+    estimate_scale : bool
+        Whether to estimate scaling factor.
+    Returns
+    -------
+    T : (N + 1, N + 1)
+        The homogeneous similarity transformation matrix. The matrix contains
+        NaN values only if the problem is not well-conditioned.
+    References
+    ----------
+    .. [1] "Least-squares estimation of transformation parameters between two
+            point patterns", Shinji Umeyama, PAMI 1991, :DOI:`10.1109/34.88573`
+    """
+
+    num = src.shape[0]
+    dim = src.shape[1]
+
+    # Compute mean of src and dst.
+    src_mean = src.mean(axis=0)
+    dst_mean = dst.mean(axis=0)
+
+    # Subtract mean from src and dst.
+    src_demean = src - src_mean
+    dst_demean = dst - dst_mean
+
+    # Eq. (38).
+    A = dst_demean.T @ src_demean / num
+
+    # Eq. (39).
+    d = np.ones((dim,), dtype=np.double)
+    if np.linalg.det(A) < 0:
+        d[dim - 1] = -1
+
+    T = np.eye(dim + 1, dtype=np.double)
+
+    U, S, V = np.linalg.svd(A)
+
+    # Eq. (40) and (43).
+    rank = np.linalg.matrix_rank(A)
+    if rank == 0:
+        return np.nan * T
+    elif rank == dim - 1:
+        if np.linalg.det(U) * np.linalg.det(V) > 0:
+            T[:dim, :dim] = U @ V
+        else:
+            s = d[dim - 1]
+            d[dim - 1] = -1
+            T[:dim, :dim] = U @ np.diag(d) @ V
+            d[dim - 1] = s
+    else:
+        T[:dim, :dim] = U @ np.diag(d) @ V
+
+    if estimate_scale:
+        # Eq. (41) and (42).
+        scale = 1.0 / src_demean.var(axis=0).sum() * (S @ d)
+    else:
+        scale = scale
+
+    T[:dim, dim] = dst_mean - scale * (T[:dim, :dim] @ src_mean.T)
+    T[:dim, :dim] *= scale
+
+    return T, scale
+
+
+class FaceWarpException(Exception):
+    def __str__(self):
+        return 'In File {}:{}'.format(
+            __file__, super.__str__(self))
+
+
+def get_reference_facial_points(output_size=None,
+                                inner_padding_factor=0.0,
+                                outer_padding=(0, 0),
+                                default_square=False):
+    tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)
+    tmp_crop_size = np.array(DEFAULT_CROP_SIZE)
+
+    # 0) make the inner region a square
+    if default_square:
+        size_diff = max(tmp_crop_size) - tmp_crop_size
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += size_diff
+
+    if (output_size and
+            output_size[0] == tmp_crop_size[0] and
+            output_size[1] == tmp_crop_size[1]):
+        print('output_size == DEFAULT_CROP_SIZE {}: return default reference points'.format(tmp_crop_size))
+        return tmp_5pts
+
+    if (inner_padding_factor == 0 and
+            outer_padding == (0, 0)):
+        if output_size is None:
+            print('No paddings to do: return default reference points')
+            return tmp_5pts
+        else:
+            raise FaceWarpException(
+                'No paddings to do, output_size must be None or {}'.format(tmp_crop_size))
+
+    # check output size
+    if not (0 <= inner_padding_factor <= 1.0):
+        raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')
+
+    if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0)
+            and output_size is None):
+        output_size = tmp_crop_size * \
+                      (1 + inner_padding_factor * 2).astype(np.int32)
+        output_size += np.array(outer_padding)
+        print('              deduced from paddings, output_size = ', output_size)
+
+    if not (outer_padding[0] < output_size[0]
+            and outer_padding[1] < output_size[1]):
+        raise FaceWarpException('Not (outer_padding[0] < output_size[0]'
+                                'and outer_padding[1] < output_size[1])')
+
+    # 1) pad the inner region according inner_padding_factor
+    # print('---> STEP1: pad the inner region according inner_padding_factor')
+    if inner_padding_factor > 0:
+        size_diff = tmp_crop_size * inner_padding_factor * 2
+        tmp_5pts += size_diff / 2
+        tmp_crop_size += np.round(size_diff).astype(np.int32)
+
+    # print('              crop_size = ', tmp_crop_size)
+    # print('              reference_5pts = ', tmp_5pts)
+
+    # 2) resize the padded inner region
+    # print('---> STEP2: resize the padded inner region')
+    size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 2
+    # print('              crop_size = ', tmp_crop_size)
+    # print('              size_bf_outer_pad = ', size_bf_outer_pad)
+
+    if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:
+        raise FaceWarpException('Must have (output_size - outer_padding)'
+                                '= some_scale * (crop_size * (1.0 + inner_padding_factor)')
+
+    scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]
+    # print('              resize scale_factor = ', scale_factor)
+    tmp_5pts = tmp_5pts * scale_factor
+    #    size_diff = tmp_crop_size * (scale_factor - min(scale_factor))
+    #    tmp_5pts = tmp_5pts + size_diff / 2
+    tmp_crop_size = size_bf_outer_pad
+    # print('              crop_size = ', tmp_crop_size)
+    # print('              reference_5pts = ', tmp_5pts)
+
+    # 3) add outer_padding to make output_size
+    reference_5point = tmp_5pts + np.array(outer_padding)
+    tmp_crop_size = output_size
+    # print('---> STEP3: add outer_padding to make output_size')
+    # print('              crop_size = ', tmp_crop_size)
+    # print('              reference_5pts = ', tmp_5pts)
+    #
+    # print('===> end get_reference_facial_points\n')
+
+    return reference_5point
+
+
+def get_affine_transform_matrix(src_pts, dst_pts):
+    tfm = np.float32([[1, 0, 0], [0, 1, 0]])
+    n_pts = src_pts.shape[0]
+    ones = np.ones((n_pts, 1), src_pts.dtype)
+    src_pts_ = np.hstack([src_pts, ones])
+    dst_pts_ = np.hstack([dst_pts, ones])
+
+    A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)
+
+    if rank == 3:
+        tfm = np.float32([
+            [A[0, 0], A[1, 0], A[2, 0]],
+            [A[0, 1], A[1, 1], A[2, 1]]
+        ])
+    elif rank == 2:
+        tfm = np.float32([
+            [A[0, 0], A[1, 0], 0],
+            [A[0, 1], A[1, 1], 0]
+        ])
+
+    return tfm
+
+
+def warp_and_crop_face(src_img,
+                       facial_pts,
+                       reference_pts=None,
+                       crop_size=(96, 112),
+                       align_type='smilarity'): #smilarity cv2_affine affine
+    if reference_pts is None:
+        if crop_size[0] == 96 and crop_size[1] == 112:
+            reference_pts = REFERENCE_FACIAL_POINTS
+        else:
+            default_square = False
+            inner_padding_factor = 0
+            outer_padding = (0, 0)
+            output_size = crop_size
+
+            reference_pts = get_reference_facial_points(output_size,
+                                                        inner_padding_factor,
+                                                        outer_padding,
+                                                        default_square)
+
+    ref_pts = np.float32(reference_pts)
+    ref_pts_shp = ref_pts.shape
+    if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:
+        raise FaceWarpException(
+            'reference_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if ref_pts_shp[0] == 2:
+        ref_pts = ref_pts.T
+
+    src_pts = np.float32(facial_pts)
+    src_pts_shp = src_pts.shape
+    if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:
+        raise FaceWarpException(
+            'facial_pts.shape must be (K,2) or (2,K) and K>2')
+
+    if src_pts_shp[0] == 2:
+        src_pts = src_pts.T
+
+    if src_pts.shape != ref_pts.shape:
+        raise FaceWarpException(
+            'facial_pts and reference_pts must have the same shape')
+
+    if align_type is 'cv2_affine':
+        tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])
+        tfm_inv = cv2.getAffineTransform(ref_pts[0:3], src_pts[0:3])
+    elif align_type is 'affine':
+        tfm = get_affine_transform_matrix(src_pts, ref_pts)
+        tfm_inv = get_affine_transform_matrix(ref_pts, src_pts)
+    else:
+        params, scale = _umeyama(src_pts, ref_pts)
+        tfm = params[:2, :]
+
+        params, _ = _umeyama(ref_pts, src_pts, False, scale=1.0/scale)
+        tfm_inv = params[:2, :]
+
+    face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]), flags=3)
+
+    return face_img, tfm_inv

core/data/deg_kair_utils/utils_blindsr.py

@@ -0,0 +1,631 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import cv2
+import torch
+
+from core.data.deg_kair_utils import utils_image as util
+
+import random
+from scipy import ndimage
+import scipy
+import scipy.stats as ss
+from scipy.interpolate import interp2d
+from scipy.linalg import orth
+
+
+
+
+"""
+# --------------------------------------------
+# Super-Resolution
+# --------------------------------------------
+#
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# From 2019/03--2021/08
+# --------------------------------------------
+"""
+
+def modcrop_np(img, sf):
+    '''
+    Args:
+        img: numpy image, WxH or WxHxC
+        sf: scale factor
+
+    Return:
+        cropped image
+    '''
+    w, h = img.shape[:2]
+    im = np.copy(img)
+    return im[:w - w % sf, :h - h % sf, ...]
+
+
+"""
+# --------------------------------------------
+# anisotropic Gaussian kernels
+# --------------------------------------------
+"""
+def analytic_kernel(k):
+    """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
+    k_size = k.shape[0]
+    # Calculate the big kernels size
+    big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
+    # Loop over the small kernel to fill the big one
+    for r in range(k_size):
+        for c in range(k_size):
+            big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
+    # Crop the edges of the big kernel to ignore very small values and increase run time of SR
+    crop = k_size // 2
+    cropped_big_k = big_k[crop:-crop, crop:-crop]
+    # Normalize to 1
+    return cropped_big_k / cropped_big_k.sum()
+
+
+def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
+    """ generate an anisotropic Gaussian kernel
+    Args:
+        ksize : e.g., 15, kernel size
+        theta : [0,  pi], rotation angle range
+        l1    : [0.1,50], scaling of eigenvalues
+        l2    : [0.1,l1], scaling of eigenvalues
+        If l1 = l2, will get an isotropic Gaussian kernel.
+
+    Returns:
+        k     : kernel
+    """
+
+    v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
+    V = np.array([[v[0], v[1]], [v[1], -v[0]]])
+    D = np.array([[l1, 0], [0, l2]])
+    Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
+    k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
+
+    return k
+
+
+def gm_blur_kernel(mean, cov, size=15):
+    center = size / 2.0 + 0.5
+    k = np.zeros([size, size])
+    for y in range(size):
+        for x in range(size):
+            cy = y - center + 1
+            cx = x - center + 1
+            k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
+
+    k = k / np.sum(k)
+    return k
+
+
+def shift_pixel(x, sf, upper_left=True):
+    """shift pixel for super-resolution with different scale factors
+    Args:
+        x: WxHxC or WxH
+        sf: scale factor
+        upper_left: shift direction
+    """
+    h, w = x.shape[:2]
+    shift = (sf-1)*0.5
+    xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
+    if upper_left:
+        x1 = xv + shift
+        y1 = yv + shift
+    else:
+        x1 = xv - shift
+        y1 = yv - shift
+
+    x1 = np.clip(x1, 0, w-1)
+    y1 = np.clip(y1, 0, h-1)
+
+    if x.ndim == 2:
+        x = interp2d(xv, yv, x)(x1, y1)
+    if x.ndim == 3:
+        for i in range(x.shape[-1]):
+            x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
+
+    return x
+
+
+def blur(x, k):
+    '''
+    x: image, NxcxHxW
+    k: kernel, Nx1xhxw
+    '''
+    n, c = x.shape[:2]
+    p1, p2 = (k.shape[-2]-1)//2, (k.shape[-1]-1)//2
+    x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
+    k = k.repeat(1,c,1,1)
+    k = k.view(-1, 1, k.shape[2], k.shape[3])
+    x = x.view(1, -1, x.shape[2], x.shape[3])
+    x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n*c)
+    x = x.view(n, c, x.shape[2], x.shape[3])
+
+    return x
+
+
+
+def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
+    """"
+    # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
+    # Kai Zhang
+    # min_var = 0.175 * sf  # variance of the gaussian kernel will be sampled between min_var and max_var
+    # max_var = 2.5 * sf
+    """
+    # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
+    lambda_1 = min_var + np.random.rand() * (max_var - min_var)
+    lambda_2 = min_var + np.random.rand() * (max_var - min_var)
+    theta = np.random.rand() * np.pi  # random theta
+    noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
+
+    # Set COV matrix using Lambdas and Theta
+    LAMBDA = np.diag([lambda_1, lambda_2])
+    Q = np.array([[np.cos(theta), -np.sin(theta)],
+                  [np.sin(theta), np.cos(theta)]])
+    SIGMA = Q @ LAMBDA @ Q.T
+    INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
+
+    # Set expectation position (shifting kernel for aligned image)
+    MU = k_size // 2 - 0.5*(scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
+    MU = MU[None, None, :, None]
+
+    # Create meshgrid for Gaussian
+    [X,Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
+    Z = np.stack([X, Y], 2)[:, :, :, None]
+
+    # Calcualte Gaussian for every pixel of the kernel
+    ZZ = Z-MU
+    ZZ_t = ZZ.transpose(0,1,3,2)
+    raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
+
+    # shift the kernel so it will be centered
+    #raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
+
+    # Normalize the kernel and return
+    #kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
+    kernel = raw_kernel / np.sum(raw_kernel)
+    return kernel
+
+
+def fspecial_gaussian(hsize, sigma):
+    hsize = [hsize, hsize]
+    siz = [(hsize[0]-1.0)/2.0, (hsize[1]-1.0)/2.0]
+    std = sigma
+    [x, y] = np.meshgrid(np.arange(-siz[1], siz[1]+1), np.arange(-siz[0], siz[0]+1))
+    arg = -(x*x + y*y)/(2*std*std)
+    h = np.exp(arg)
+    h[h < scipy.finfo(float).eps * h.max()] = 0
+    sumh = h.sum()
+    if sumh != 0:
+        h = h/sumh
+    return h
+
+
+def fspecial_laplacian(alpha):
+    alpha = max([0, min([alpha,1])])
+    h1 = alpha/(alpha+1)
+    h2 = (1-alpha)/(alpha+1)
+    h = [[h1, h2, h1], [h2, -4/(alpha+1), h2], [h1, h2, h1]]
+    h = np.array(h)
+    return h
+
+
+def fspecial(filter_type, *args, **kwargs):
+    '''
+    python code from:
+    https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
+    '''
+    if filter_type == 'gaussian':
+        return fspecial_gaussian(*args, **kwargs)
+    if filter_type == 'laplacian':
+        return fspecial_laplacian(*args, **kwargs)
+
+"""
+# --------------------------------------------
+# degradation models
+# --------------------------------------------
+"""
+
+
+def bicubic_degradation(x, sf=3):
+    '''
+    Args:
+        x: HxWxC image, [0, 1]
+        sf: down-scale factor
+
+    Return:
+        bicubicly downsampled LR image
+    '''
+    x = util.imresize_np(x, scale=1/sf)
+    return x
+
+
+def srmd_degradation(x, k, sf=3):
+    ''' blur + bicubic downsampling
+
+    Args:
+        x: HxWxC image, [0, 1]
+        k: hxw, double
+        sf: down-scale factor
+
+    Return:
+        downsampled LR image
+
+    Reference:
+        @inproceedings{zhang2018learning,
+          title={Learning a single convolutional super-resolution network for multiple degradations},
+          author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+          booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+          pages={3262--3271},
+          year={2018}
+        }
+    '''
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')  # 'nearest' | 'mirror'
+    x = bicubic_degradation(x, sf=sf)
+    return x
+
+
+def dpsr_degradation(x, k, sf=3):
+
+    ''' bicubic downsampling + blur
+
+    Args:
+        x: HxWxC image, [0, 1]
+        k: hxw, double
+        sf: down-scale factor
+
+    Return:
+        downsampled LR image
+
+    Reference:
+        @inproceedings{zhang2019deep,
+          title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
+          author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+          booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+          pages={1671--1681},
+          year={2019}
+        }
+    '''
+    x = bicubic_degradation(x, sf=sf)
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+    return x
+
+
+def classical_degradation(x, k, sf=3):
+    ''' blur + downsampling
+
+    Args:
+        x: HxWxC image, [0, 1]/[0, 255]
+        k: hxw, double
+        sf: down-scale factor
+
+    Return:
+        downsampled LR image
+    '''
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+    #x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
+    st = 0
+    return x[st::sf, st::sf, ...]
+
+
+def add_sharpening(img, weight=0.5, radius=50, threshold=10):
+    """USM sharpening. borrowed from real-ESRGAN
+    Input image: I; Blurry image: B.
+    1. K = I + weight * (I - B)
+    2. Mask = 1 if abs(I - B) > threshold, else: 0
+    3. Blur mask:
+    4. Out = Mask * K + (1 - Mask) * I
+    Args:
+        img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
+        weight (float): Sharp weight. Default: 1.
+        radius (float): Kernel size of Gaussian blur. Default: 50.
+        threshold (int):
+    """
+    if radius % 2 == 0:
+        radius += 1
+    blur = cv2.GaussianBlur(img, (radius, radius), 0)
+    residual = img - blur
+    mask = np.abs(residual) * 255 > threshold
+    mask = mask.astype('float32')
+    soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
+
+    K = img + weight * residual
+    K = np.clip(K, 0, 1)
+    return soft_mask * K + (1 - soft_mask) * img
+
+
+def add_blur(img, sf=4):
+    wd2 = 4.0 + sf
+    wd = 2.0 + 0.2*sf
+    if random.random() < 0.5:
+        l1 = wd2*random.random()
+        l2 = wd2*random.random()
+        k = anisotropic_Gaussian(ksize=2*random.randint(2,11)+3, theta=random.random()*np.pi, l1=l1, l2=l2)
+    else:
+        k = fspecial('gaussian', 2*random.randint(2,11)+3, wd*random.random())
+    img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
+
+    return img
+
+
+def add_resize(img, sf=4):
+    rnum = np.random.rand()
+    if rnum > 0.8:  # up
+        sf1 = random.uniform(1, 2)
+    elif rnum < 0.7:  # down
+        sf1 = random.uniform(0.5/sf, 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 = np.clip(img, 0.0, 1.0)
+
+    return img
+
+
+def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
+    noise_level = random.randint(noise_level1, noise_level2)
+    rnum = np.random.rand()
+    if rnum > 0.6:   # add color Gaussian noise
+        img += np.random.normal(0, noise_level/255.0, img.shape).astype(np.float32)
+    elif rnum < 0.4: # add grayscale Gaussian noise
+        img += np.random.normal(0, noise_level/255.0, (*img.shape[:2], 1)).astype(np.float32)
+    else:            # add  noise
+        L = noise_level2/255.
+        D = np.diag(np.random.rand(3))
+        U = orth(np.random.rand(3,3))
+        conv = np.dot(np.dot(np.transpose(U), D), U)
+        img += np.random.multivariate_normal([0,0,0], np.abs(L**2*conv), img.shape[:2]).astype(np.float32)
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_speckle_noise(img, noise_level1=2, noise_level2=25):
+    noise_level = random.randint(noise_level1, noise_level2)
+    img = np.clip(img, 0.0, 1.0)
+    rnum = random.random()
+    if rnum > 0.6:
+        img += img*np.random.normal(0, noise_level/255.0, img.shape).astype(np.float32)
+    elif rnum < 0.4:
+        img += img*np.random.normal(0, noise_level/255.0, (*img.shape[:2], 1)).astype(np.float32)
+    else:
+        L = noise_level2/255.
+        D = np.diag(np.random.rand(3))
+        U = orth(np.random.rand(3,3))
+        conv = np.dot(np.dot(np.transpose(U), D), U)
+        img += img*np.random.multivariate_normal([0,0,0], np.abs(L**2*conv), img.shape[:2]).astype(np.float32)
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_Poisson_noise(img):
+    img = np.clip((img * 255.0).round(), 0, 255) / 255.
+    vals = 10**(2*random.random()+2.0)  # [2, 4]
+    if random.random() < 0.5:
+        img = np.random.poisson(img * vals).astype(np.float32) / vals
+    else:
+        img_gray = np.dot(img[...,:3], [0.299, 0.587, 0.114])
+        img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
+        noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
+        img += noise_gray[:, :, np.newaxis]
+    img = np.clip(img, 0.0, 1.0)
+    return img
+
+
+def add_JPEG_noise(img):
+    quality_factor = random.randint(30, 95)
+    img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
+    result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
+    img = cv2.imdecode(encimg, 1)
+    img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
+    return img
+
+
+def random_crop(lq, hq, sf=4, lq_patchsize=64):
+    h, w = lq.shape[:2]
+    rnd_h = random.randint(0, h-lq_patchsize)
+    rnd_w = random.randint(0, w-lq_patchsize)
+    lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
+
+    rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
+    hq = hq[rnd_h_H:rnd_h_H + lq_patchsize*sf, rnd_w_H:rnd_w_H + lq_patchsize*sf, :]
+    return lq, hq
+
+
+def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
+    """
+    This is the degradation model of BSRGAN from the paper
+    "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
+    ----------
+    img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
+    sf: scale factor
+    isp_model: camera ISP model
+
+    Returns
+    -------
+    img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+    hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+    """
+    isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
+    sf_ori = sf
+
+    h1, w1 = img.shape[:2]
+    img = img.copy()[:h1 - h1 % sf, :w1 - w1 % sf, ...]  # mod crop
+    h, w = img.shape[:2]
+
+    if h < lq_patchsize*sf or w < lq_patchsize*sf:
+        raise ValueError(f'img size ({h1}X{w1}) is too small!')
+
+    hq = img.copy()
+
+    if sf == 4 and random.random() < scale2_prob:   # downsample1
+        if np.random.rand() < 0.5:
+            img = cv2.resize(img, (int(1/2*img.shape[1]), int(1/2*img.shape[0])), interpolation=random.choice([1,2,3]))
+        else:
+            img = util.imresize_np(img, 1/2, True)
+        img = np.clip(img, 0.0, 1.0)
+        sf = 2
+
+    shuffle_order = random.sample(range(7), 7)
+    idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
+    if idx1 > idx2:  # keep downsample3 last
+        shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
+
+    for i in shuffle_order:
+
+        if i == 0:
+            img = add_blur(img, sf=sf)
+
+        elif i == 1:
+            img = add_blur(img, sf=sf)
+
+        elif i == 2:
+            a, b = img.shape[1], img.shape[0]
+            # downsample2
+            if random.random() < 0.75:
+                sf1 = random.uniform(1,2*sf)
+                img = cv2.resize(img, (int(1/sf1*img.shape[1]), int(1/sf1*img.shape[0])), interpolation=random.choice([1,2,3]))
+            else:
+                k = fspecial('gaussian', 25, random.uniform(0.1, 0.6*sf))
+                k_shifted = shift_pixel(k, sf)
+                k_shifted = k_shifted/k_shifted.sum()  # blur with shifted kernel
+                img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
+                img = img[0::sf, 0::sf, ...]  # nearest downsampling
+            img = np.clip(img, 0.0, 1.0)
+
+        elif i == 3:
+            # downsample3
+            img = cv2.resize(img, (int(1/sf*a), int(1/sf*b)), interpolation=random.choice([1,2,3]))
+            img = np.clip(img, 0.0, 1.0)
+
+        elif i == 4:
+            # add Gaussian noise
+            img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+
+        elif i == 5:
+            # add JPEG noise
+            if random.random() < jpeg_prob:
+                img = add_JPEG_noise(img)
+
+        elif i == 6:
+            # add processed camera sensor noise
+            if random.random() < isp_prob and isp_model is not None:
+                with torch.no_grad():
+                    img, hq = isp_model.forward(img.copy(), hq)
+
+    # add final JPEG compression noise
+    img = add_JPEG_noise(img)
+
+    # random crop
+    img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
+
+    return img, hq
+
+
+
+
+def degradation_bsrgan_plus(img, sf=4, shuffle_prob=0.5, use_sharp=False, lq_patchsize=64, isp_model=None):
+    """
+    This is an extended degradation model by combining
+    the degradation models of BSRGAN and Real-ESRGAN
+    ----------
+    img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
+    sf: scale factor
+    use_shuffle: the degradation shuffle
+    use_sharp: sharpening the img
+
+    Returns
+    -------
+    img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+    hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+    """
+
+    h1, w1 = img.shape[:2]
+    img = img.copy()[:h1 - h1 % sf, :w1 - w1 % sf, ...]  # mod crop
+    h, w = img.shape[:2]
+
+    if h < lq_patchsize*sf or w < lq_patchsize*sf:
+        raise ValueError(f'img size ({h1}X{w1}) is too small!')
+
+    if use_sharp:
+        img = add_sharpening(img)
+    hq = img.copy()
+
+    if random.random() < shuffle_prob:
+        shuffle_order = random.sample(range(13), 13)
+    else:
+        shuffle_order = list(range(13))
+        # local shuffle for noise, JPEG is always the last one
+        shuffle_order[2:6] = random.sample(shuffle_order[2:6], len(range(2, 6)))
+        shuffle_order[9:13] = random.sample(shuffle_order[9:13], len(range(9, 13)))
+
+    poisson_prob, speckle_prob, isp_prob = 0.1, 0.1, 0.1
+
+    for i in shuffle_order:
+        if i == 0:
+            img = add_blur(img, sf=sf)
+        elif i == 1:
+            img = add_resize(img, sf=sf)
+        elif i == 2:
+            img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+        elif i == 3:
+            if random.random() < poisson_prob:
+                img = add_Poisson_noise(img)
+        elif i == 4:
+            if random.random() < speckle_prob:
+                img = add_speckle_noise(img)
+        elif i == 5:
+            if random.random() < isp_prob and isp_model is not None:
+                with torch.no_grad():
+                    img, hq = isp_model.forward(img.copy(), hq)
+        elif i == 6:
+            img = add_JPEG_noise(img)
+        elif i == 7:
+            img = add_blur(img, sf=sf)
+        elif i == 8:
+            img = add_resize(img, sf=sf)
+        elif i == 9:
+            img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+        elif i == 10:
+            if random.random() < poisson_prob:
+                img = add_Poisson_noise(img)
+        elif i == 11:
+            if random.random() < speckle_prob:
+                img = add_speckle_noise(img)
+        elif i == 12:
+            if random.random() < isp_prob and isp_model is not None:
+                with torch.no_grad():
+                    img, hq = isp_model.forward(img.copy(), hq)
+        else:
+            print('check the shuffle!')
+
+    # resize to desired size
+    img = cv2.resize(img, (int(1/sf*hq.shape[1]), int(1/sf*hq.shape[0])), interpolation=random.choice([1, 2, 3]))
+
+    # add final JPEG compression noise
+    img = add_JPEG_noise(img)
+
+    # random crop
+    img, hq = random_crop(img, hq, sf, lq_patchsize)
+
+    return img, hq
+
+
+
+if __name__ == '__main__':
+    img = util.imread_uint('utils/test.png', 3)
+    img = util.uint2single(img)
+    sf = 4
+    
+    for i in range(20):
+        img_lq, img_hq = degradation_bsrgan(img, sf=sf, lq_patchsize=72)
+        print(i)
+        lq_nearest =  cv2.resize(util.single2uint(img_lq), (int(sf*img_lq.shape[1]), int(sf*img_lq.shape[0])), interpolation=0)
+        img_concat = np.concatenate([lq_nearest, util.single2uint(img_hq)], axis=1)
+        util.imsave(img_concat, str(i)+'.png')
+
+#    for i in range(10):
+#        img_lq, img_hq = degradation_bsrgan_plus(img, sf=sf, shuffle_prob=0.1, use_sharp=True, lq_patchsize=64)
+#        print(i)
+#        lq_nearest =  cv2.resize(util.single2uint(img_lq), (int(sf*img_lq.shape[1]), int(sf*img_lq.shape[0])), interpolation=0)
+#        img_concat = np.concatenate([lq_nearest, util.single2uint(img_hq)], axis=1)
+#        util.imsave(img_concat, str(i)+'.png')
+
+#    run utils/utils_blindsr.py

core/data/deg_kair_utils/utils_bnorm.py

@@ -0,0 +1,91 @@
+import torch
+import torch.nn as nn
+
+
+"""
+# --------------------------------------------
+# Batch Normalization
+# --------------------------------------------
+
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# 01/Jan/2019
+# --------------------------------------------
+"""
+
+
+# --------------------------------------------
+# remove/delete specified layer
+# --------------------------------------------
+def deleteLayer(model, layer_type=nn.BatchNorm2d):
+    ''' Kai Zhang, 11/Jan/2019.
+    '''
+    for k, m in list(model.named_children()):
+        if isinstance(m, layer_type):
+            del model._modules[k]
+        deleteLayer(m, layer_type)
+
+
+# --------------------------------------------
+# merge bn, "conv+bn" --> "conv"
+# --------------------------------------------
+def merge_bn(model):
+    ''' Kai Zhang, 11/Jan/2019.
+    merge all 'Conv+BN' (or 'TConv+BN') into 'Conv' (or 'TConv')
+    based on https://github.com/pytorch/pytorch/pull/901
+    '''
+    prev_m = None
+    for k, m in list(model.named_children()):
+        if (isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d)) and (isinstance(prev_m, nn.Conv2d) or isinstance(prev_m, nn.Linear) or isinstance(prev_m, nn.ConvTranspose2d)):
+
+            w = prev_m.weight.data
+
+            if prev_m.bias is None:
+                zeros = torch.Tensor(prev_m.out_channels).zero_().type(w.type())
+                prev_m.bias = nn.Parameter(zeros)
+            b = prev_m.bias.data
+
+            invstd = m.running_var.clone().add_(m.eps).pow_(-0.5)
+            if isinstance(prev_m, nn.ConvTranspose2d):
+                w.mul_(invstd.view(1, w.size(1), 1, 1).expand_as(w))
+            else:
+                w.mul_(invstd.view(w.size(0), 1, 1, 1).expand_as(w))
+            b.add_(-m.running_mean).mul_(invstd)
+            if m.affine:
+                if isinstance(prev_m, nn.ConvTranspose2d):
+                    w.mul_(m.weight.data.view(1, w.size(1), 1, 1).expand_as(w))
+                else:
+                    w.mul_(m.weight.data.view(w.size(0), 1, 1, 1).expand_as(w))
+                b.mul_(m.weight.data).add_(m.bias.data)
+
+            del model._modules[k]
+        prev_m = m
+        merge_bn(m)
+
+
+# --------------------------------------------
+# add bn, "conv" --> "conv+bn"
+# --------------------------------------------
+def add_bn(model):
+    ''' Kai Zhang, 11/Jan/2019.
+    '''
+    for k, m in list(model.named_children()):
+        if (isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear) or isinstance(m, nn.ConvTranspose2d)):
+            b = nn.BatchNorm2d(m.out_channels, momentum=0.1, affine=True)
+            b.weight.data.fill_(1)
+            new_m = nn.Sequential(model._modules[k], b)
+            model._modules[k] = new_m
+        add_bn(m)
+
+
+# --------------------------------------------
+# tidy model after removing bn
+# --------------------------------------------
+def tidy_sequential(model):
+    ''' Kai Zhang, 11/Jan/2019.
+    '''
+    for k, m in list(model.named_children()):
+        if isinstance(m, nn.Sequential):
+            if m.__len__() == 1:
+                model._modules[k] = m.__getitem__(0)
+        tidy_sequential(m)

core/data/deg_kair_utils/utils_deblur.py

@@ -0,0 +1,655 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import scipy
+from scipy import fftpack
+import torch
+
+from math import cos, sin
+from numpy import zeros, ones, prod, array, pi, log, min, mod, arange, sum, mgrid, exp, pad, round
+from numpy.random import randn, rand
+from scipy.signal import convolve2d
+import cv2
+import random
+# import utils_image as util
+
+'''
+modified by Kai Zhang (github: https://github.com/cszn)
+03/03/2019
+'''
+
+
+def get_uperleft_denominator(img, kernel):
+    '''
+    img: HxWxC
+    kernel: hxw
+    denominator: HxWx1
+    upperleft: HxWxC
+    '''
+    V = psf2otf(kernel, img.shape[:2])
+    denominator = np.expand_dims(np.abs(V)**2, axis=2)
+    upperleft = np.expand_dims(np.conj(V), axis=2) * np.fft.fft2(img, axes=[0, 1])
+    return upperleft, denominator
+
+
+def get_uperleft_denominator_pytorch(img, kernel):
+    '''
+    img: NxCxHxW
+    kernel: Nx1xhxw
+    denominator: Nx1xHxW
+    upperleft: NxCxHxWx2
+    '''
+    V = p2o(kernel, img.shape[-2:])  # Nx1xHxWx2
+    denominator = V[..., 0]**2+V[..., 1]**2  # Nx1xHxW
+    upperleft = cmul(cconj(V), rfft(img))  # Nx1xHxWx2 * NxCxHxWx2
+    return upperleft, denominator
+
+
+def c2c(x):
+    return torch.from_numpy(np.stack([np.float32(x.real), np.float32(x.imag)], axis=-1))
+
+
+def r2c(x):
+    return torch.stack([x, torch.zeros_like(x)], -1)
+
+
+def cdiv(x, y):
+    a, b = x[..., 0], x[..., 1]
+    c, d = y[..., 0], y[..., 1]
+    cd2 = c**2 + d**2
+    return torch.stack([(a*c+b*d)/cd2, (b*c-a*d)/cd2], -1)
+
+
+def cabs(x):
+    return torch.pow(x[..., 0]**2+x[..., 1]**2, 0.5)
+
+
+def cmul(t1, t2):
+    '''
+    complex multiplication
+    t1: NxCxHxWx2
+    output: NxCxHxWx2
+    '''
+    real1, imag1 = t1[..., 0], t1[..., 1]
+    real2, imag2 = t2[..., 0], t2[..., 1]
+    return torch.stack([real1 * real2 - imag1 * imag2, real1 * imag2 + imag1 * real2], dim=-1)
+
+
+def cconj(t, inplace=False):
+    '''
+    # complex's conjugation
+    t: NxCxHxWx2
+    output: NxCxHxWx2
+    '''
+    c = t.clone() if not inplace else t
+    c[..., 1] *= -1
+    return c
+
+
+def rfft(t):
+    return torch.rfft(t, 2, onesided=False)
+
+
+def irfft(t):
+    return torch.irfft(t, 2, onesided=False)
+
+
+def fft(t):
+    return torch.fft(t, 2)
+
+
+def ifft(t):
+    return torch.ifft(t, 2)
+
+
+def p2o(psf, shape):
+    '''
+    # psf: NxCxhxw
+    # shape: [H,W]
+    # otf: NxCxHxWx2
+    '''
+    otf = torch.zeros(psf.shape[:-2] + shape).type_as(psf)
+    otf[...,:psf.shape[2],:psf.shape[3]].copy_(psf)
+    for axis, axis_size in enumerate(psf.shape[2:]):
+        otf = torch.roll(otf, -int(axis_size / 2), dims=axis+2)
+    otf = torch.rfft(otf, 2, onesided=False)
+    n_ops = torch.sum(torch.tensor(psf.shape).type_as(psf) * torch.log2(torch.tensor(psf.shape).type_as(psf)))
+    otf[...,1][torch.abs(otf[...,1])<n_ops*2.22e-16] = torch.tensor(0).type_as(psf)
+    return otf
+
+
+
+# otf2psf: not sure where I got this one from. Maybe translated from Octave source code or whatever. It's just math.
+def otf2psf(otf, outsize=None):
+    insize = np.array(otf.shape)
+    psf = np.fft.ifftn(otf, axes=(0, 1))
+    for axis, axis_size in enumerate(insize):
+        psf = np.roll(psf, np.floor(axis_size / 2).astype(int), axis=axis)
+    if type(outsize) != type(None):
+        insize = np.array(otf.shape)
+        outsize = np.array(outsize)
+        n = max(np.size(outsize), np.size(insize))
+        # outsize = postpad(outsize(:), n, 1);
+        # insize = postpad(insize(:) , n, 1);
+        colvec_out = outsize.flatten().reshape((np.size(outsize), 1))
+        colvec_in = insize.flatten().reshape((np.size(insize), 1))
+        outsize = np.pad(colvec_out, ((0, max(0, n - np.size(colvec_out))), (0, 0)), mode="constant")
+        insize = np.pad(colvec_in, ((0, max(0, n - np.size(colvec_in))), (0, 0)), mode="constant")
+
+        pad = (insize - outsize) / 2
+        if np.any(pad < 0):
+            print("otf2psf error: OUTSIZE must be smaller than or equal than OTF size")
+        prepad = np.floor(pad)
+        postpad = np.ceil(pad)
+        dims_start = prepad.astype(int)
+        dims_end = (insize - postpad).astype(int)
+        for i in range(len(dims_start.shape)):
+            psf = np.take(psf, range(dims_start[i][0], dims_end[i][0]), axis=i)
+    n_ops = np.sum(otf.size * np.log2(otf.shape))
+    psf = np.real_if_close(psf, tol=n_ops)
+    return psf
+
+
+# psf2otf copied/modified from https://github.com/aboucaud/pypher/blob/master/pypher/pypher.py
+def psf2otf(psf, shape=None):
+    """
+    Convert point-spread function to optical transfer function.
+    Compute the Fast Fourier Transform (FFT) of the point-spread
+    function (PSF) array and creates the optical transfer function (OTF)
+    array that is not influenced by the PSF off-centering.
+    By default, the OTF array is the same size as the PSF array.
+    To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
+    post-pads the PSF array (down or to the right) with zeros to match
+    dimensions specified in OUTSIZE, then circularly shifts the values of
+    the PSF array up (or to the left) until the central pixel reaches (1,1)
+    position.
+    Parameters
+    ----------
+    psf : `numpy.ndarray`
+        PSF array
+    shape : int
+        Output shape of the OTF array
+    Returns
+    -------
+    otf : `numpy.ndarray`
+        OTF array
+    Notes
+    -----
+    Adapted from MATLAB psf2otf function
+    """
+    if type(shape) == type(None):
+        shape = psf.shape
+    shape = np.array(shape)
+    if np.all(psf == 0):
+        # return np.zeros_like(psf)
+        return np.zeros(shape)
+    if len(psf.shape) == 1:
+        psf = psf.reshape((1, psf.shape[0]))
+    inshape = psf.shape
+    psf = zero_pad(psf, shape, position='corner')
+    for axis, axis_size in enumerate(inshape):
+        psf = np.roll(psf, -int(axis_size / 2), axis=axis)
+    # Compute the OTF
+    otf = np.fft.fft2(psf, axes=(0, 1))
+    # Estimate the rough number of operations involved in the FFT
+    # and discard the PSF imaginary part if within roundoff error
+    # roundoff error  = machine epsilon = sys.float_info.epsilon
+    # or np.finfo().eps
+    n_ops = np.sum(psf.size * np.log2(psf.shape))
+    otf = np.real_if_close(otf, tol=n_ops)
+    return otf
+
+
+def zero_pad(image, shape, position='corner'):
+    """
+    Extends image to a certain size with zeros
+    Parameters
+    ----------
+    image: real 2d `numpy.ndarray`
+        Input image
+    shape: tuple of int
+        Desired output shape of the image
+    position : str, optional
+        The position of the input image in the output one:
+            * 'corner'
+                top-left corner (default)
+            * 'center'
+                centered
+    Returns
+    -------
+    padded_img: real `numpy.ndarray`
+        The zero-padded image
+    """
+    shape = np.asarray(shape, dtype=int)
+    imshape = np.asarray(image.shape, dtype=int)
+    if np.alltrue(imshape == shape):
+        return image
+    if np.any(shape <= 0):
+        raise ValueError("ZERO_PAD: null or negative shape given")
+    dshape = shape - imshape
+    if np.any(dshape < 0):
+        raise ValueError("ZERO_PAD: target size smaller than source one")
+    pad_img = np.zeros(shape, dtype=image.dtype)
+    idx, idy = np.indices(imshape)
+    if position == 'center':
+        if np.any(dshape % 2 != 0):
+            raise ValueError("ZERO_PAD: source and target shapes "
+                             "have different parity.")
+        offx, offy = dshape // 2
+    else:
+        offx, offy = (0, 0)
+    pad_img[idx + offx, idy + offy] = image
+    return pad_img
+
+
+'''
+Reducing boundary artifacts
+'''
+
+
+def opt_fft_size(n):
+    '''
+    Kai Zhang (github: https://github.com/cszn)
+    03/03/2019
+    #  opt_fft_size.m
+    # compute an optimal data length for Fourier transforms
+    # written by Sunghyun Cho (sodomau@postech.ac.kr)
+    # persistent opt_fft_size_LUT;
+    '''
+
+    LUT_size = 2048
+    # print("generate opt_fft_size_LUT")
+    opt_fft_size_LUT = np.zeros(LUT_size)
+
+    e2 = 1
+    while e2 <= LUT_size:
+        e3 = e2
+        while e3 <= LUT_size:
+            e5 = e3
+            while e5 <= LUT_size:
+                e7 = e5
+                while e7 <= LUT_size:
+                    if e7 <= LUT_size:
+                        opt_fft_size_LUT[e7-1] = e7
+                    if e7*11 <= LUT_size:
+                        opt_fft_size_LUT[e7*11-1] = e7*11
+                    if e7*13 <= LUT_size:
+                        opt_fft_size_LUT[e7*13-1] = e7*13
+                    e7 = e7 * 7
+                e5 = e5 * 5
+            e3 = e3 * 3
+        e2 = e2 * 2
+
+    nn = 0
+    for i in range(LUT_size, 0, -1):
+        if opt_fft_size_LUT[i-1] != 0:
+            nn = i-1
+        else:
+            opt_fft_size_LUT[i-1] = nn+1
+
+    m = np.zeros(len(n))
+    for c in range(len(n)):
+        nn = n[c]
+        if nn <= LUT_size:
+            m[c] = opt_fft_size_LUT[nn-1]
+        else:
+            m[c] = -1
+    return m
+
+
+def wrap_boundary_liu(img, img_size):
+
+    """
+    Reducing boundary artifacts in image deconvolution
+    Renting Liu, Jiaya Jia
+    ICIP 2008
+    """
+    if img.ndim == 2:
+        ret = wrap_boundary(img, img_size)
+    elif img.ndim == 3:
+        ret = [wrap_boundary(img[:, :, i], img_size) for i in range(3)]
+        ret = np.stack(ret, 2)
+    return ret
+
+
+def wrap_boundary(img, img_size):
+
+    """
+    python code from:
+    https://github.com/ys-koshelev/nla_deblur/blob/90fe0ab98c26c791dcbdf231fe6f938fca80e2a0/boundaries.py
+    Reducing boundary artifacts in image deconvolution
+    Renting Liu, Jiaya Jia
+    ICIP 2008
+    """
+    (H, W) = np.shape(img)
+    H_w = int(img_size[0]) - H
+    W_w = int(img_size[1]) - W
+
+    # ret = np.zeros((img_size[0], img_size[1]));
+    alpha = 1
+    HG = img[:, :]
+
+    r_A = np.zeros((alpha*2+H_w, W))
+    r_A[:alpha, :] = HG[-alpha:, :]
+    r_A[-alpha:, :] = HG[:alpha, :]
+    a = np.arange(H_w)/(H_w-1)
+    # r_A(alpha+1:end-alpha, 1) = (1-a)*r_A(alpha,1) + a*r_A(end-alpha+1,1)
+    r_A[alpha:-alpha, 0] = (1-a)*r_A[alpha-1, 0] + a*r_A[-alpha, 0]
+    # r_A(alpha+1:end-alpha, end) = (1-a)*r_A(alpha,end) + a*r_A(end-alpha+1,end)
+    r_A[alpha:-alpha, -1] = (1-a)*r_A[alpha-1, -1] + a*r_A[-alpha, -1]
+
+    r_B = np.zeros((H, alpha*2+W_w))
+    r_B[:, :alpha] = HG[:, -alpha:]
+    r_B[:, -alpha:] = HG[:, :alpha]
+    a = np.arange(W_w)/(W_w-1)
+    r_B[0, alpha:-alpha] = (1-a)*r_B[0, alpha-1] + a*r_B[0, -alpha]
+    r_B[-1, alpha:-alpha] = (1-a)*r_B[-1, alpha-1] + a*r_B[-1, -alpha]
+
+    if alpha == 1:
+        A2 = solve_min_laplacian(r_A[alpha-1:, :])
+        B2 = solve_min_laplacian(r_B[:, alpha-1:])
+        r_A[alpha-1:, :] = A2
+        r_B[:, alpha-1:] = B2
+    else:
+        A2 = solve_min_laplacian(r_A[alpha-1:-alpha+1, :])
+        r_A[alpha-1:-alpha+1, :] = A2
+        B2 = solve_min_laplacian(r_B[:, alpha-1:-alpha+1])
+        r_B[:, alpha-1:-alpha+1] = B2
+    A = r_A
+    B = r_B
+
+    r_C = np.zeros((alpha*2+H_w, alpha*2+W_w))
+    r_C[:alpha, :] = B[-alpha:, :]
+    r_C[-alpha:, :] = B[:alpha, :]
+    r_C[:, :alpha] = A[:, -alpha:]
+    r_C[:, -alpha:] = A[:, :alpha]
+
+    if alpha == 1:
+        C2 = C2 = solve_min_laplacian(r_C[alpha-1:, alpha-1:])
+        r_C[alpha-1:, alpha-1:] = C2
+    else:
+        C2 = solve_min_laplacian(r_C[alpha-1:-alpha+1, alpha-1:-alpha+1])
+        r_C[alpha-1:-alpha+1, alpha-1:-alpha+1] = C2
+    C = r_C
+    # return C
+    A = A[alpha-1:-alpha-1, :]
+    B = B[:, alpha:-alpha]
+    C = C[alpha:-alpha, alpha:-alpha]
+    ret = np.vstack((np.hstack((img, B)), np.hstack((A, C))))
+    return ret
+
+
+def solve_min_laplacian(boundary_image):
+    (H, W) = np.shape(boundary_image)
+
+    # Laplacian
+    f = np.zeros((H, W))
+    # boundary image contains image intensities at boundaries
+    boundary_image[1:-1, 1:-1] = 0
+    j = np.arange(2, H)-1
+    k = np.arange(2, W)-1
+    f_bp = np.zeros((H, W))
+    f_bp[np.ix_(j, k)] = -4*boundary_image[np.ix_(j, k)] + boundary_image[np.ix_(j, k+1)] + boundary_image[np.ix_(j, k-1)] + boundary_image[np.ix_(j-1, k)] + boundary_image[np.ix_(j+1, k)]
+    
+    del(j, k)
+    f1 = f - f_bp  # subtract boundary points contribution
+    del(f_bp, f)
+
+    # DST Sine Transform algo starts here
+    f2 = f1[1:-1,1:-1]
+    del(f1)
+
+    # compute sine tranform
+    if f2.shape[1] == 1:
+        tt = fftpack.dst(f2, type=1, axis=0)/2
+    else:
+        tt = fftpack.dst(f2, type=1)/2
+
+    if tt.shape[0] == 1:
+        f2sin = np.transpose(fftpack.dst(np.transpose(tt), type=1, axis=0)/2)
+    else:
+        f2sin = np.transpose(fftpack.dst(np.transpose(tt), type=1)/2) 
+    del(f2)
+
+    # compute Eigen Values
+    [x, y] = np.meshgrid(np.arange(1, W-1), np.arange(1, H-1))
+    denom = (2*np.cos(np.pi*x/(W-1))-2) + (2*np.cos(np.pi*y/(H-1)) - 2)
+
+    # divide
+    f3 = f2sin/denom
+    del(f2sin, x, y)
+
+    # compute Inverse Sine Transform
+    if f3.shape[0] == 1:
+        tt = fftpack.idst(f3*2, type=1, axis=1)/(2*(f3.shape[1]+1))
+    else:
+        tt = fftpack.idst(f3*2, type=1, axis=0)/(2*(f3.shape[0]+1))
+    del(f3)
+    if tt.shape[1] == 1:
+        img_tt = np.transpose(fftpack.idst(np.transpose(tt)*2, type=1)/(2*(tt.shape[0]+1)))
+    else:
+        img_tt = np.transpose(fftpack.idst(np.transpose(tt)*2, type=1, axis=0)/(2*(tt.shape[1]+1)))
+    del(tt)
+
+    # put solution in inner points; outer points obtained from boundary image
+    img_direct = boundary_image
+    img_direct[1:-1, 1:-1] = 0
+    img_direct[1:-1, 1:-1] = img_tt
+    return img_direct
+
+
+"""
+Created on Thu Jan 18 15:36:32 2018
+@author: italo
+https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
+"""
+
+"""
+Syntax
+h = fspecial(type)
+h = fspecial('average',hsize)
+h = fspecial('disk',radius)
+h = fspecial('gaussian',hsize,sigma)
+h = fspecial('laplacian',alpha)
+h = fspecial('log',hsize,sigma)
+h = fspecial('motion',len,theta)
+h = fspecial('prewitt')
+h = fspecial('sobel')
+"""
+
+
+def fspecial_average(hsize=3):
+    """Smoothing filter"""
+    return np.ones((hsize, hsize))/hsize**2
+
+
+def fspecial_disk(radius):
+    """Disk filter"""
+    raise(NotImplemented)
+    rad = 0.6
+    crad = np.ceil(rad-0.5)
+    [x, y] = np.meshgrid(np.arange(-crad, crad+1), np.arange(-crad, crad+1))
+    maxxy = np.zeros(x.shape)
+    maxxy[abs(x) >= abs(y)] = abs(x)[abs(x) >= abs(y)]
+    maxxy[abs(y) >= abs(x)] = abs(y)[abs(y) >= abs(x)]
+    minxy = np.zeros(x.shape)
+    minxy[abs(x) <= abs(y)] = abs(x)[abs(x) <= abs(y)]
+    minxy[abs(y) <= abs(x)] = abs(y)[abs(y) <= abs(x)]
+    m1 = (rad**2 <  (maxxy+0.5)**2 + (minxy-0.5)**2)*(minxy-0.5) +\
+         (rad**2 >= (maxxy+0.5)**2 + (minxy-0.5)**2)*\
+         np.sqrt((rad**2 + 0j) - (maxxy + 0.5)**2)
+    m2 = (rad**2 >  (maxxy-0.5)**2 + (minxy+0.5)**2)*(minxy+0.5) +\
+         (rad**2 <= (maxxy-0.5)**2 + (minxy+0.5)**2)*\
+         np.sqrt((rad**2 + 0j) - (maxxy - 0.5)**2)
+    h = None
+    return h
+
+
+def fspecial_gaussian(hsize, sigma):
+    hsize = [hsize, hsize]
+    siz = [(hsize[0]-1.0)/2.0, (hsize[1]-1.0)/2.0]
+    std = sigma
+    [x, y] = np.meshgrid(np.arange(-siz[1], siz[1]+1), np.arange(-siz[0], siz[0]+1))
+    arg = -(x*x + y*y)/(2*std*std)
+    h = np.exp(arg)
+    h[h < scipy.finfo(float).eps * h.max()] = 0
+    sumh = h.sum()
+    if sumh != 0:
+        h = h/sumh
+    return h
+
+
+def fspecial_laplacian(alpha):
+    alpha = max([0, min([alpha,1])])
+    h1 = alpha/(alpha+1)
+    h2 = (1-alpha)/(alpha+1)
+    h = [[h1, h2, h1], [h2, -4/(alpha+1), h2], [h1, h2, h1]]
+    h = np.array(h)
+    return h
+
+
+def fspecial_log(hsize, sigma):
+    raise(NotImplemented)
+
+
+def fspecial_motion(motion_len, theta):
+    raise(NotImplemented)
+
+
+def fspecial_prewitt():
+    return np.array([[1, 1, 1], [0, 0, 0], [-1, -1, -1]])
+
+
+def fspecial_sobel():
+    return np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
+
+
+def fspecial(filter_type, *args, **kwargs):
+    '''
+    python code from:
+    https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
+    '''
+    if filter_type == 'average':
+        return fspecial_average(*args, **kwargs)
+    if filter_type == 'disk':
+        return fspecial_disk(*args, **kwargs)
+    if filter_type == 'gaussian':
+        return fspecial_gaussian(*args, **kwargs)
+    if filter_type == 'laplacian':
+        return fspecial_laplacian(*args, **kwargs)
+    if filter_type == 'log':
+        return fspecial_log(*args, **kwargs)
+    if filter_type == 'motion':
+        return fspecial_motion(*args, **kwargs)
+    if filter_type == 'prewitt':
+        return fspecial_prewitt(*args, **kwargs)
+    if filter_type == 'sobel':
+        return fspecial_sobel(*args, **kwargs)
+
+
+def fspecial_gauss(size, sigma):
+    x, y = mgrid[-size // 2 + 1 : size // 2 + 1, -size // 2 + 1 : size // 2 + 1]
+    g = exp(-((x ** 2 + y ** 2) / (2.0 * sigma ** 2)))
+    return g / g.sum()
+
+
+def blurkernel_synthesis(h=37, w=None):
+    # https://github.com/tkkcc/prior/blob/879a0b6c117c810776d8cc6b63720bf29f7d0cc4/util/gen_kernel.py
+    w = h if w is None else w
+    kdims = [h, w]
+    x = randomTrajectory(250)
+    k = None
+    while k is None:
+        k = kernelFromTrajectory(x)
+
+    # center pad to kdims
+    pad_width = ((kdims[0] - k.shape[0]) // 2, (kdims[1] - k.shape[1]) // 2)
+    pad_width = [(pad_width[0],), (pad_width[1],)]
+    
+    if pad_width[0][0]<0 or pad_width[1][0]<0:
+        k = k[0:h, 0:h]
+    else:
+        k = pad(k, pad_width, "constant")
+    x1,x2 = k.shape
+    if np.random.randint(0, 4) == 1:
+        k = cv2.resize(k, (random.randint(x1, 5*x1), random.randint(x2, 5*x2)), interpolation=cv2.INTER_LINEAR)
+        y1, y2 = k.shape
+        k = k[(y1-x1)//2: (y1-x1)//2+x1, (y2-x2)//2: (y2-x2)//2+x2]
+        
+    if sum(k)<0.1:
+        k = fspecial_gaussian(h, 0.1+6*np.random.rand(1))
+    k = k / sum(k)
+    # import matplotlib.pyplot as plt
+    # plt.imshow(k, interpolation="nearest", cmap="gray")
+    # plt.show()
+    return k
+
+
+def kernelFromTrajectory(x):
+    h = 5 - log(rand()) / 0.15
+    h = round(min([h, 27])).astype(int)
+    h = h + 1 - h % 2
+    w = h
+    k = zeros((h, w))
+
+    xmin = min(x[0])
+    xmax = max(x[0])
+    ymin = min(x[1])
+    ymax = max(x[1])
+    xthr = arange(xmin, xmax, (xmax - xmin) / w)
+    ythr = arange(ymin, ymax, (ymax - ymin) / h)
+
+    for i in range(1, xthr.size):
+        for j in range(1, ythr.size):
+            idx = (
+                (x[0, :] >= xthr[i - 1])
+                & (x[0, :] < xthr[i])
+                & (x[1, :] >= ythr[j - 1])
+                & (x[1, :] < ythr[j])
+            )
+            k[i - 1, j - 1] = sum(idx)
+    if sum(k) == 0:
+        return
+    k = k / sum(k)
+    k = convolve2d(k, fspecial_gauss(3, 1), "same")
+    k = k / sum(k)
+    return k
+
+
+def randomTrajectory(T):
+    x = zeros((3, T))
+    v = randn(3, T)
+    r = zeros((3, T))
+    trv = 1 / 1
+    trr = 2 * pi / T
+    for t in range(1, T):
+        F_rot = randn(3) / (t + 1) + r[:, t - 1]
+        F_trans = randn(3) / (t + 1)
+        r[:, t] = r[:, t - 1] + trr * F_rot
+        v[:, t] = v[:, t - 1] + trv * F_trans
+        st = v[:, t]
+        st = rot3D(st, r[:, t])
+        x[:, t] = x[:, t - 1] + st
+    return x
+
+
+def rot3D(x, r):
+    Rx = array([[1, 0, 0], [0, cos(r[0]), -sin(r[0])], [0, sin(r[0]), cos(r[0])]])
+    Ry = array([[cos(r[1]), 0, sin(r[1])], [0, 1, 0], [-sin(r[1]), 0, cos(r[1])]])
+    Rz = array([[cos(r[2]), -sin(r[2]), 0], [sin(r[2]), cos(r[2]), 0], [0, 0, 1]])
+    R = Rz @ Ry @ Rx
+    x = R @ x
+    return x
+
+
+if __name__ == '__main__':
+    a = opt_fft_size([111])
+    print(a)
+
+    print(fspecial('gaussian', 5, 1))
+    
+    print(p2o(torch.zeros(1,1,4,4).float(),(14,14)).shape)
+
+    k = blurkernel_synthesis(11)
+    import matplotlib.pyplot as plt
+    plt.imshow(k, interpolation="nearest", cmap="gray")
+    plt.show()

core/data/deg_kair_utils/utils_dist.py

@@ -0,0 +1,201 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E501
+import functools
+import os
+import subprocess
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+# ----------------------------------
+# init
+# ----------------------------------
+def init_dist(launcher, backend='nccl', **kwargs):
+    if mp.get_start_method(allow_none=True) is None:
+        mp.set_start_method('spawn')
+    if launcher == 'pytorch':
+        _init_dist_pytorch(backend, **kwargs)
+    elif launcher == 'slurm':
+        _init_dist_slurm(backend, **kwargs)
+    else:
+        raise ValueError(f'Invalid launcher type: {launcher}')
+
+
+def _init_dist_pytorch(backend, **kwargs):
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def _init_dist_slurm(backend, port=None):
+    """Initialize slurm distributed training environment.
+    If argument ``port`` is not specified, then the master port will be system
+    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
+    environment variable, then a default port ``29500`` will be used.
+    Args:
+        backend (str): Backend of torch.distributed.
+        port (int, optional): Master port. Defaults to None.
+    """
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    addr = subprocess.getoutput(
+        f'scontrol show hostname {node_list} | head -n1')
+    # specify master port
+    if port is not None:
+        os.environ['MASTER_PORT'] = str(port)
+    elif 'MASTER_PORT' in os.environ:
+        pass  # use MASTER_PORT in the environment variable
+    else:
+        # 29500 is torch.distributed default port
+        os.environ['MASTER_PORT'] = '29500'
+    os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+    os.environ['RANK'] = str(proc_id)
+    dist.init_process_group(backend=backend)
+
+
+
+# ----------------------------------
+# get rank and world_size
+# ----------------------------------
+def get_dist_info():
+    if dist.is_available():
+        initialized = dist.is_initialized()
+    else:
+        initialized = False
+    if initialized:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    else:
+        rank = 0
+        world_size = 1
+    return rank, world_size
+
+
+def get_rank():
+    if not dist.is_available():
+        return 0
+
+    if not dist.is_initialized():
+        return 0
+
+    return dist.get_rank()
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+
+    if not dist.is_initialized():
+        return 1
+
+    return dist.get_world_size()
+
+
+def master_only(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        rank, _ = get_dist_info()
+        if rank == 0:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+
+
+
+
+# ----------------------------------
+# operation across ranks
+# ----------------------------------
+def reduce_sum(tensor):
+    if not dist.is_available():
+        return tensor
+
+    if not dist.is_initialized():
+        return tensor
+
+    tensor = tensor.clone()
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+
+    return tensor
+
+
+def gather_grad(params):
+    world_size = get_world_size()
+    
+    if world_size == 1:
+        return
+
+    for param in params:
+        if param.grad is not None:
+            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
+            param.grad.data.div_(world_size)
+
+
+def all_gather(data):
+    world_size = get_world_size()
+
+    if world_size == 1:
+        return [data]
+
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to('cuda')
+
+    local_size = torch.IntTensor([tensor.numel()]).to('cuda')
+    size_list = [torch.IntTensor([0]).to('cuda') for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.ByteTensor(size=(max_size,)).to('cuda'))
+
+    if local_size != max_size:
+        padding = torch.ByteTensor(size=(max_size - local_size,)).to('cuda')
+        tensor = torch.cat((tensor, padding), 0)
+
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def reduce_loss_dict(loss_dict):
+    world_size = get_world_size()
+
+    if world_size < 2:
+        return loss_dict
+
+    with torch.no_grad():
+        keys = []
+        losses = []
+
+        for k in sorted(loss_dict.keys()):
+            keys.append(k)
+            losses.append(loss_dict[k])
+
+        losses = torch.stack(losses, 0)
+        dist.reduce(losses, dst=0)
+
+        if dist.get_rank() == 0:
+            losses /= world_size
+
+        reduced_losses = {k: v for k, v in zip(keys, losses)}
+
+    return reduced_losses
+

core/data/deg_kair_utils/utils_googledownload.py

@@ -0,0 +1,93 @@
+import math
+import requests
+from tqdm import tqdm
+
+
+'''
+borrowed from 
+https://github.com/xinntao/BasicSR/blob/28883e15eedc3381d23235ff3cf7c454c4be87e6/basicsr/utils/download_util.py
+'''
+
+
+def sizeof_fmt(size, suffix='B'):
+    """Get human readable file size.
+    Args:
+        size (int): File size.
+        suffix (str): Suffix. Default: 'B'.
+    Return:
+        str: Formated file siz.
+    """
+    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
+        if abs(size) < 1024.0:
+            return f'{size:3.1f} {unit}{suffix}'
+        size /= 1024.0
+    return f'{size:3.1f} Y{suffix}'
+
+
+def download_file_from_google_drive(file_id, save_path):
+    """Download files from google drive.
+    Ref:
+    https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive  # noqa E501
+    Args:
+        file_id (str): File id.
+        save_path (str): Save path.
+    """
+
+    session = requests.Session()
+    URL = 'https://docs.google.com/uc?export=download'
+    params = {'id': file_id}
+
+    response = session.get(URL, params=params, stream=True)
+    token = get_confirm_token(response)
+    if token:
+        params['confirm'] = token
+        response = session.get(URL, params=params, stream=True)
+
+    # get file size
+    response_file_size = session.get(
+        URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
+    if 'Content-Range' in response_file_size.headers:
+        file_size = int(
+            response_file_size.headers['Content-Range'].split('/')[1])
+    else:
+        file_size = None
+
+    save_response_content(response, save_path, file_size)
+
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            return value
+    return None
+
+
+def save_response_content(response,
+                          destination,
+                          file_size=None,
+                          chunk_size=32768):
+    if file_size is not None:
+        pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
+
+        readable_file_size = sizeof_fmt(file_size)
+    else:
+        pbar = None
+
+    with open(destination, 'wb') as f:
+        downloaded_size = 0
+        for chunk in response.iter_content(chunk_size):
+            downloaded_size += chunk_size
+            if pbar is not None:
+                pbar.update(1)
+                pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} '
+                                     f'/ {readable_file_size}')
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)
+        if pbar is not None:
+            pbar.close()
+
+
+if __name__ == "__main__":
+    file_id = '1WNULM1e8gRNvsngVscsQ8tpaOqJ4mYtv'
+    save_path = 'BSRGAN.pth'
+    download_file_from_google_drive(file_id, save_path)

core/data/deg_kair_utils/utils_image.py

@@ -0,0 +1,1016 @@
+import os
+import math
+import random
+import numpy as np
+import torch
+import cv2
+from torchvision.utils import make_grid
+from datetime import datetime
+# import torchvision.transforms as transforms
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+# https://github.com/twhui/SRGAN-pyTorch
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def get_timestamp():
+    return datetime.now().strftime('%y%m%d-%H%M%S')
+
+
+def imshow(x, title=None, cbar=False, figsize=None):
+    plt.figure(figsize=figsize)
+    plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
+    if title:
+        plt.title(title)
+    if cbar:
+        plt.colorbar()
+    plt.show()
+
+
+def surf(Z, cmap='rainbow', figsize=None):
+    plt.figure(figsize=figsize)
+    ax3 = plt.axes(projection='3d')
+
+    w, h = Z.shape[:2]
+    xx = np.arange(0,w,1)
+    yy = np.arange(0,h,1)
+    X, Y = np.meshgrid(xx, yy)
+    ax3.plot_surface(X,Y,Z,cmap=cmap)
+    #ax3.contour(X,Y,Z, zdim='z',offset=-2，cmap=cmap)
+    plt.show()
+
+
+'''
+# --------------------------------------------
+# get image pathes
+# --------------------------------------------
+'''
+
+
+def get_image_paths(dataroot):
+    paths = None  # return None if dataroot is None
+    if isinstance(dataroot, str):
+        paths = sorted(_get_paths_from_images(dataroot))
+    elif isinstance(dataroot, list):
+        paths = []
+        for i in dataroot:
+            paths += sorted(_get_paths_from_images(i))
+    return paths
+
+
+def _get_paths_from_images(path):
+    assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+    images = []
+    for dirpath, _, fnames in sorted(os.walk(path)):
+        for fname in sorted(fnames):
+            if is_image_file(fname):
+                img_path = os.path.join(dirpath, fname)
+                images.append(img_path)
+    assert images, '{:s} has no valid image file'.format(path)
+    return images
+
+
+'''
+# --------------------------------------------
+# split large images into small images 
+# --------------------------------------------
+'''
+
+
+def patches_from_image(img, p_size=512, p_overlap=64, p_max=800):
+    w, h = img.shape[:2]
+    patches = []
+    if w > p_max and h > p_max:
+        w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int))
+        h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int))
+        w1.append(w-p_size)
+        h1.append(h-p_size)
+        # print(w1)
+        # print(h1)
+        for i in w1:
+            for j in h1:
+                patches.append(img[i:i+p_size, j:j+p_size,:])
+    else:
+        patches.append(img)
+
+    return patches
+
+
+def imssave(imgs, img_path):
+    """
+    imgs: list, N images of size WxHxC
+    """
+    img_name, ext = os.path.splitext(os.path.basename(img_path))
+    for i, img in enumerate(imgs):
+        if img.ndim == 3:
+            img = img[:, :, [2, 1, 0]]
+        new_path = os.path.join(os.path.dirname(img_path), img_name+str('_{:04d}'.format(i))+'.png')
+        cv2.imwrite(new_path, img)
+
+
+def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=512, p_overlap=96, p_max=800):
+    """
+    split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size), 
+    and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max)
+    will be splitted.
+
+    Args:
+        original_dataroot:
+        taget_dataroot:
+        p_size: size of small images
+        p_overlap: patch size in training is a good choice
+        p_max: images with smaller size than (p_max)x(p_max) keep unchanged.
+    """
+    paths = get_image_paths(original_dataroot)
+    for img_path in paths:
+        # img_name, ext = os.path.splitext(os.path.basename(img_path))
+        img = imread_uint(img_path, n_channels=n_channels)
+        patches = patches_from_image(img, p_size, p_overlap, p_max)
+        imssave(patches, os.path.join(taget_dataroot, os.path.basename(img_path)))
+        #if original_dataroot == taget_dataroot:
+        #del img_path
+
+'''
+# --------------------------------------------
+# makedir
+# --------------------------------------------
+'''
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def mkdirs(paths):
+    if isinstance(paths, str):
+        mkdir(paths)
+    else:
+        for path in paths:
+            mkdir(path)
+
+
+def mkdir_and_rename(path):
+    if os.path.exists(path):
+        new_name = path + '_archived_' + get_timestamp()
+        print('Path already exists. Rename it to [{:s}]'.format(new_name))
+        os.rename(path, new_name)
+    os.makedirs(path)
+
+
+'''
+# --------------------------------------------
+# read image from path
+# opencv is fast, but read BGR numpy image
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# get uint8 image of size HxWxn_channles (RGB)
+# --------------------------------------------
+def imread_uint(path, n_channels=3):
+    #  input: path
+    # output: HxWx3(RGB or GGG), or HxWx1 (G)
+    if n_channels == 1:
+        img = cv2.imread(path, 0)  # cv2.IMREAD_GRAYSCALE
+        img = np.expand_dims(img, axis=2)  # HxWx1
+    elif n_channels == 3:
+        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # BGR or G
+        if img.ndim == 2:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)  # GGG
+        else:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # RGB
+    return img
+
+
+# --------------------------------------------
+# matlab's imwrite
+# --------------------------------------------
+def imsave(img, img_path):
+    img = np.squeeze(img)
+    if img.ndim == 3:
+        img = img[:, :, [2, 1, 0]]
+    cv2.imwrite(img_path, img)
+
+def imwrite(img, img_path):
+    img = np.squeeze(img)
+    if img.ndim == 3:
+        img = img[:, :, [2, 1, 0]]
+    cv2.imwrite(img_path, img)
+
+
+
+# --------------------------------------------
+# get single image of size HxWxn_channles (BGR)
+# --------------------------------------------
+def read_img(path):
+    # read image by cv2
+    # return: Numpy float32, HWC, BGR, [0,1]
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)  # cv2.IMREAD_GRAYSCALE
+    img = img.astype(np.float32) / 255.
+    if img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    # some images have 4 channels
+    if img.shape[2] > 3:
+        img = img[:, :, :3]
+    return img
+
+
+'''
+# --------------------------------------------
+# image format conversion
+# --------------------------------------------
+# numpy(single) <--->  numpy(uint)
+# numpy(single) <--->  tensor
+# numpy(uint)   <--->  tensor
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# numpy(single) [0, 1] <--->  numpy(uint)
+# --------------------------------------------
+
+
+def uint2single(img):
+
+    return np.float32(img/255.)
+
+
+def single2uint(img):
+
+    return np.uint8((img.clip(0, 1)*255.).round())
+
+
+def uint162single(img):
+
+    return np.float32(img/65535.)
+
+
+def single2uint16(img):
+
+    return np.uint16((img.clip(0, 1)*65535.).round())
+
+
+# --------------------------------------------
+# numpy(uint) (HxWxC or HxW) <--->  tensor
+# --------------------------------------------
+
+
+# convert uint to 4-dimensional torch tensor
+def uint2tensor4(img):
+    if img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
+
+
+# convert uint to 3-dimensional torch tensor
+def uint2tensor3(img):
+    if img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
+
+
+# convert 2/3/4-dimensional torch tensor to uint
+def tensor2uint(img):
+    img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
+    if img.ndim == 3:
+        img = np.transpose(img, (1, 2, 0))
+    return np.uint8((img*255.0).round())
+
+
+# --------------------------------------------
+# numpy(single) (HxWxC) <--->  tensor
+# --------------------------------------------
+
+
+# convert single (HxWxC) to 3-dimensional torch tensor
+def single2tensor3(img):
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
+
+
+# convert single (HxWxC) to 4-dimensional torch tensor
+def single2tensor4(img):
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
+
+
+# convert torch tensor to single
+def tensor2single(img):
+    img = img.data.squeeze().float().cpu().numpy()
+    if img.ndim == 3:
+        img = np.transpose(img, (1, 2, 0))
+
+    return img
+
+# convert torch tensor to single
+def tensor2single3(img):
+    img = img.data.squeeze().float().cpu().numpy()
+    if img.ndim == 3:
+        img = np.transpose(img, (1, 2, 0))
+    elif img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    return img
+
+
+def single2tensor5(img):
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0)
+
+
+def single32tensor5(img):
+    return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0)
+
+
+def single42tensor4(img):
+    return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float()
+
+
+# from skimage.io import imread, imsave
+def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
+    '''
+    Converts a torch Tensor into an image Numpy array of BGR channel order
+    Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
+    Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
+    '''
+    tensor = tensor.squeeze().float().cpu().clamp_(*min_max)  # squeeze first, then clamp
+    tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])  # to range [0,1]
+    n_dim = tensor.dim()
+    if n_dim == 4:
+        n_img = len(tensor)
+        img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
+        img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
+    elif n_dim == 3:
+        img_np = tensor.numpy()
+        img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
+    elif n_dim == 2:
+        img_np = tensor.numpy()
+    else:
+        raise TypeError(
+            'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
+    if out_type == np.uint8:
+        img_np = (img_np * 255.0).round()
+        # Important. Unlike matlab, numpy.uint8() WILL NOT round by default.
+    return img_np.astype(out_type)
+
+
+'''
+# --------------------------------------------
+# Augmentation, flipe and/or rotate
+# --------------------------------------------
+# The following two are enough.
+# (1) augmet_img: numpy image of WxHxC or WxH
+# (2) augment_img_tensor4: tensor image 1xCxWxH
+# --------------------------------------------
+'''
+
+
+def augment_img(img, mode=0):
+    '''Kai Zhang (github: https://github.com/cszn)
+    '''
+    if mode == 0:
+        return img
+    elif mode == 1:
+        return np.flipud(np.rot90(img))
+    elif mode == 2:
+        return np.flipud(img)
+    elif mode == 3:
+        return np.rot90(img, k=3)
+    elif mode == 4:
+        return np.flipud(np.rot90(img, k=2))
+    elif mode == 5:
+        return np.rot90(img)
+    elif mode == 6:
+        return np.rot90(img, k=2)
+    elif mode == 7:
+        return np.flipud(np.rot90(img, k=3))
+
+
+def augment_img_tensor4(img, mode=0):
+    '''Kai Zhang (github: https://github.com/cszn)
+    '''
+    if mode == 0:
+        return img
+    elif mode == 1:
+        return img.rot90(1, [2, 3]).flip([2])
+    elif mode == 2:
+        return img.flip([2])
+    elif mode == 3:
+        return img.rot90(3, [2, 3])
+    elif mode == 4:
+        return img.rot90(2, [2, 3]).flip([2])
+    elif mode == 5:
+        return img.rot90(1, [2, 3])
+    elif mode == 6:
+        return img.rot90(2, [2, 3])
+    elif mode == 7:
+        return img.rot90(3, [2, 3]).flip([2])
+
+
+def augment_img_tensor(img, mode=0):
+    '''Kai Zhang (github: https://github.com/cszn)
+    '''
+    img_size = img.size()
+    img_np = img.data.cpu().numpy()
+    if len(img_size) == 3:
+        img_np = np.transpose(img_np, (1, 2, 0))
+    elif len(img_size) == 4:
+        img_np = np.transpose(img_np, (2, 3, 1, 0))
+    img_np = augment_img(img_np, mode=mode)
+    img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
+    if len(img_size) == 3:
+        img_tensor = img_tensor.permute(2, 0, 1)
+    elif len(img_size) == 4:
+        img_tensor = img_tensor.permute(3, 2, 0, 1)
+
+    return img_tensor.type_as(img)
+
+
+def augment_img_np3(img, mode=0):
+    if mode == 0:
+        return img
+    elif mode == 1:
+        return img.transpose(1, 0, 2)
+    elif mode == 2:
+        return img[::-1, :, :]
+    elif mode == 3:
+        img = img[::-1, :, :]
+        img = img.transpose(1, 0, 2)
+        return img
+    elif mode == 4:
+        return img[:, ::-1, :]
+    elif mode == 5:
+        img = img[:, ::-1, :]
+        img = img.transpose(1, 0, 2)
+        return img
+    elif mode == 6:
+        img = img[:, ::-1, :]
+        img = img[::-1, :, :]
+        return img
+    elif mode == 7:
+        img = img[:, ::-1, :]
+        img = img[::-1, :, :]
+        img = img.transpose(1, 0, 2)
+        return img
+
+
+def augment_imgs(img_list, hflip=True, rot=True):
+    # horizontal flip OR rotate
+    hflip = hflip and random.random() < 0.5
+    vflip = rot and random.random() < 0.5
+    rot90 = rot and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:
+            img = img[:, ::-1, :]
+        if vflip:
+            img = img[::-1, :, :]
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    return [_augment(img) for img in img_list]
+
+
+'''
+# --------------------------------------------
+# modcrop and shave
+# --------------------------------------------
+'''
+
+
+def modcrop(img_in, scale):
+    # img_in: Numpy, HWC or HW
+    img = np.copy(img_in)
+    if img.ndim == 2:
+        H, W = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r]
+    elif img.ndim == 3:
+        H, W, C = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r, :]
+    else:
+        raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
+    return img
+
+
+def shave(img_in, border=0):
+    # img_in: Numpy, HWC or HW
+    img = np.copy(img_in)
+    h, w = img.shape[:2]
+    img = img[border:h-border, border:w-border]
+    return img
+
+
+'''
+# --------------------------------------------
+# image processing process on numpy image
+# channel_convert(in_c, tar_type, img_list):
+# rgb2ycbcr(img, only_y=True):
+# bgr2ycbcr(img, only_y=True):
+# ycbcr2rgb(img):
+# --------------------------------------------
+'''
+
+
+def rgb2ycbcr(img, only_y=True):
+    '''same as matlab rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+                              [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def ycbcr2rgb(img):
+    '''same as matlab ycbcr2rgb
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                          [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
+    rlt = np.clip(rlt, 0, 255)
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def bgr2ycbcr(img, only_y=True):
+    '''bgr version of rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+                              [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def channel_convert(in_c, tar_type, img_list):
+    # conversion among BGR, gray and y
+    if in_c == 3 and tar_type == 'gray':  # BGR to gray
+        gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in gray_list]
+    elif in_c == 3 and tar_type == 'y':  # BGR to y
+        y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in y_list]
+    elif in_c == 1 and tar_type == 'RGB':  # gray/y to BGR
+        return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
+    else:
+        return img_list
+
+
+'''
+# --------------------------------------------
+# metric, PSNR, SSIM and PSNRB
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# PSNR
+# --------------------------------------------
+def calculate_psnr(img1, img2, border=0):
+    # img1 and img2 have range [0, 255]
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))
+
+
+# --------------------------------------------
+# SSIM
+# --------------------------------------------
+def calculate_ssim(img1, img2, border=0):
+    '''calculate SSIM
+    the same outputs as MATLAB's
+    img1, img2: [0, 255]
+    '''
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[2] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
+            return np.array(ssims).mean()
+        elif img1.shape[2] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+
+def ssim(img1, img2):
+    C1 = (0.01 * 255)**2
+    C2 = (0.03 * 255)**2
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()
+
+
+def _blocking_effect_factor(im):
+    block_size = 8
+
+    block_horizontal_positions = torch.arange(7, im.shape[3] - 1, 8)
+    block_vertical_positions = torch.arange(7, im.shape[2] - 1, 8)
+
+    horizontal_block_difference = (
+                (im[:, :, :, block_horizontal_positions] - im[:, :, :, block_horizontal_positions + 1]) ** 2).sum(
+        3).sum(2).sum(1)
+    vertical_block_difference = (
+                (im[:, :, block_vertical_positions, :] - im[:, :, block_vertical_positions + 1, :]) ** 2).sum(3).sum(
+        2).sum(1)
+
+    nonblock_horizontal_positions = np.setdiff1d(torch.arange(0, im.shape[3] - 1), block_horizontal_positions)
+    nonblock_vertical_positions = np.setdiff1d(torch.arange(0, im.shape[2] - 1), block_vertical_positions)
+
+    horizontal_nonblock_difference = (
+                (im[:, :, :, nonblock_horizontal_positions] - im[:, :, :, nonblock_horizontal_positions + 1]) ** 2).sum(
+        3).sum(2).sum(1)
+    vertical_nonblock_difference = (
+                (im[:, :, nonblock_vertical_positions, :] - im[:, :, nonblock_vertical_positions + 1, :]) ** 2).sum(
+        3).sum(2).sum(1)
+
+    n_boundary_horiz = im.shape[2] * (im.shape[3] // block_size - 1)
+    n_boundary_vert = im.shape[3] * (im.shape[2] // block_size - 1)
+    boundary_difference = (horizontal_block_difference + vertical_block_difference) / (
+                n_boundary_horiz + n_boundary_vert)
+
+    n_nonboundary_horiz = im.shape[2] * (im.shape[3] - 1) - n_boundary_horiz
+    n_nonboundary_vert = im.shape[3] * (im.shape[2] - 1) - n_boundary_vert
+    nonboundary_difference = (horizontal_nonblock_difference + vertical_nonblock_difference) / (
+                n_nonboundary_horiz + n_nonboundary_vert)
+
+    scaler = np.log2(block_size) / np.log2(min([im.shape[2], im.shape[3]]))
+    bef = scaler * (boundary_difference - nonboundary_difference)
+
+    bef[boundary_difference <= nonboundary_difference] = 0
+    return bef
+
+
+def calculate_psnrb(img1, img2, border=0):
+    """Calculate PSNR-B (Peak Signal-to-Noise Ratio).
+    Ref: Quality assessment of deblocked images, for JPEG image deblocking evaluation
+    # https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
+    Args:
+        img1 (ndarray): Images with range [0, 255].
+        img2 (ndarray): Images with range [0, 255].
+        border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the PSNR calculation.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+    Returns:
+        float: psnr result.
+    """
+
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+
+    if img1.ndim == 2:
+        img1, img2 = np.expand_dims(img1, 2), np.expand_dims(img2, 2)
+
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    # follow https://gitlab.com/Queuecumber/quantization-guided-ac/-/blob/master/metrics/psnrb.py
+    img1 = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0) / 255.
+    img2 = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0) / 255.
+
+    total = 0
+    for c in range(img1.shape[1]):
+        mse = torch.nn.functional.mse_loss(img1[:, c:c + 1, :, :], img2[:, c:c + 1, :, :], reduction='none')
+        bef = _blocking_effect_factor(img1[:, c:c + 1, :, :])
+
+        mse = mse.view(mse.shape[0], -1).mean(1)
+        total += 10 * torch.log10(1 / (mse + bef))
+
+    return float(total) / img1.shape[1]
+
+'''
+# --------------------------------------------
+# matlab's bicubic imresize (numpy and torch) [0, 1]
+# --------------------------------------------
+'''
+
+
+# matlab 'imresize' function, now only support 'bicubic'
+def cubic(x):
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
+        (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    if (scale < 1) and (antialiasing):
+        # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5+scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    P = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
+        1, P).expand(out_length, P)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
+    # apply cubic kernel
+    if (scale < 1) and (antialiasing):
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, P)
+
+    # If a column in weights is all zero, get rid of it. only consider the first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, P - 2)
+        weights = weights.narrow(1, 1, P - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, P - 2)
+        weights = weights.narrow(1, 0, P - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+# --------------------------------------------
+# imresize for tensor image [0, 1]
+# --------------------------------------------
+def imresize(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: pytorch tensor, CHW or HW [0,1]
+    # output: CHW or HW [0,1] w/o round
+    need_squeeze = True if img.dim() == 2 else False
+    if need_squeeze:
+        img.unsqueeze_(0)
+    in_C, in_H, in_W = img.size()
+    out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
+    img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:, :sym_len_Hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_He:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_C, out_H, in_W)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        for j in range(out_C):
+            out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
+    out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_Ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_We:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_C, out_H, out_W)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        for j in range(out_C):
+            out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])
+    if need_squeeze:
+        out_2.squeeze_()
+    return out_2
+
+
+# --------------------------------------------
+# imresize for numpy image [0, 1]
+# --------------------------------------------
+def imresize_np(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: Numpy, HWC or HW [0,1]
+    # output: HWC or HW [0,1] w/o round
+    img = torch.from_numpy(img)
+    need_squeeze = True if img.dim() == 2 else False
+    if need_squeeze:
+        img.unsqueeze_(2)
+
+    in_H, in_W, in_C = img.size()
+    out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
+    img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:sym_len_Hs, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[-sym_len_He:, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(out_H, in_W, in_C)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        for j in range(out_C):
+            out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
+    out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :sym_len_Ws, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, -sym_len_We:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(out_H, out_W, in_C)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        for j in range(out_C):
+            out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
+    if need_squeeze:
+        out_2.squeeze_()
+
+    return out_2.numpy()
+
+
+if __name__ == '__main__':
+    img = imread_uint('test.bmp', 3)
+#    img = uint2single(img)
+#    img_bicubic = imresize_np(img, 1/4)
+#    imshow(single2uint(img_bicubic))
+#
+#    img_tensor = single2tensor4(img)
+#    for i in range(8):
+#        imshow(np.concatenate((augment_img(img, i), tensor2single(augment_img_tensor4(img_tensor, i))), 1))
+    
+#    patches = patches_from_image(img, p_size=128, p_overlap=0, p_max=200)
+#    imssave(patches,'a.png')
+
+
+    
+    
+    
+    
+    

core/data/deg_kair_utils/utils_lmdb.py

@@ -0,0 +1,205 @@
+import cv2
+import lmdb
+import sys
+from multiprocessing import Pool
+from os import path as osp
+from tqdm import tqdm
+
+
+def make_lmdb_from_imgs(data_path,
+                        lmdb_path,
+                        img_path_list,
+                        keys,
+                        batch=5000,
+                        compress_level=1,
+                        multiprocessing_read=False,
+                        n_thread=40,
+                        map_size=None):
+    """Make lmdb from images.
+
+    Contents of lmdb. The file structure is:
+    example.lmdb
+    ├── data.mdb
+    ├── lock.mdb
+    ├── meta_info.txt
+
+    The data.mdb and lock.mdb are standard lmdb files and you can refer to
+    https://lmdb.readthedocs.io/en/release/ for more details.
+
+    The meta_info.txt is a specified txt file to record the meta information
+    of our datasets. It will be automatically created when preparing
+    datasets by our provided dataset tools.
+    Each line in the txt file records 1)image name (with extension),
+    2)image shape, and 3)compression level, separated by a white space.
+
+    For example, the meta information could be:
+    `000_00000000.png (720,1280,3) 1`, which means:
+    1) image name (with extension): 000_00000000.png;
+    2) image shape: (720,1280,3);
+    3) compression level: 1
+
+    We use the image name without extension as the lmdb key.
+
+    If `multiprocessing_read` is True, it will read all the images to memory
+    using multiprocessing. Thus, your server needs to have enough memory.
+
+    Args:
+        data_path (str): Data path for reading images.
+        lmdb_path (str): Lmdb save path.
+        img_path_list (str): Image path list.
+        keys (str): Used for lmdb keys.
+        batch (int): After processing batch images, lmdb commits.
+            Default: 5000.
+        compress_level (int): Compress level when encoding images. Default: 1.
+        multiprocessing_read (bool): Whether use multiprocessing to read all
+            the images to memory. Default: False.
+        n_thread (int): For multiprocessing.
+        map_size (int | None): Map size for lmdb env. If None, use the
+            estimated size from images. Default: None
+    """
+
+    assert len(img_path_list) == len(keys), ('img_path_list and keys should have the same length, '
+                                             f'but got {len(img_path_list)} and {len(keys)}')
+    print(f'Create lmdb for {data_path}, save to {lmdb_path}...')
+    print(f'Totoal images: {len(img_path_list)}')
+    if not lmdb_path.endswith('.lmdb'):
+        raise ValueError("lmdb_path must end with '.lmdb'.")
+    if osp.exists(lmdb_path):
+        print(f'Folder {lmdb_path} already exists. Exit.')
+        sys.exit(1)
+
+    if multiprocessing_read:
+        # read all the images to memory (multiprocessing)
+        dataset = {}  # use dict to keep the order for multiprocessing
+        shapes = {}
+        print(f'Read images with multiprocessing, #thread: {n_thread} ...')
+        pbar = tqdm(total=len(img_path_list), unit='image')
+
+        def callback(arg):
+            """get the image data and update pbar."""
+            key, dataset[key], shapes[key] = arg
+            pbar.update(1)
+            pbar.set_description(f'Read {key}')
+
+        pool = Pool(n_thread)
+        for path, key in zip(img_path_list, keys):
+            pool.apply_async(read_img_worker, args=(osp.join(data_path, path), key, compress_level), callback=callback)
+        pool.close()
+        pool.join()
+        pbar.close()
+        print(f'Finish reading {len(img_path_list)} images.')
+
+    # create lmdb environment
+    if map_size is None:
+        # obtain data size for one image
+        img = cv2.imread(osp.join(data_path, img_path_list[0]), cv2.IMREAD_UNCHANGED)
+        _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
+        data_size_per_img = img_byte.nbytes
+        print('Data size per image is: ', data_size_per_img)
+        data_size = data_size_per_img * len(img_path_list)
+        map_size = data_size * 10
+
+    env = lmdb.open(lmdb_path, map_size=map_size)
+
+    # write data to lmdb
+    pbar = tqdm(total=len(img_path_list), unit='chunk')
+    txn = env.begin(write=True)
+    txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
+    for idx, (path, key) in enumerate(zip(img_path_list, keys)):
+        pbar.update(1)
+        pbar.set_description(f'Write {key}')
+        key_byte = key.encode('ascii')
+        if multiprocessing_read:
+            img_byte = dataset[key]
+            h, w, c = shapes[key]
+        else:
+            _, img_byte, img_shape = read_img_worker(osp.join(data_path, path), key, compress_level)
+            h, w, c = img_shape
+
+        txn.put(key_byte, img_byte)
+        # write meta information
+        txt_file.write(f'{key}.png ({h},{w},{c}) {compress_level}\n')
+        if idx % batch == 0:
+            txn.commit()
+            txn = env.begin(write=True)
+    pbar.close()
+    txn.commit()
+    env.close()
+    txt_file.close()
+    print('\nFinish writing lmdb.')
+
+
+def read_img_worker(path, key, compress_level):
+    """Read image worker.
+
+    Args:
+        path (str): Image path.
+        key (str): Image key.
+        compress_level (int): Compress level when encoding images.
+
+    Returns:
+        str: Image key.
+        byte: Image byte.
+        tuple[int]: Image shape.
+    """
+
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    # deal with `libpng error: Read Error`
+    if img is None:
+        print(f'To deal with `libpng error: Read Error`, use PIL to load {path}')
+        from PIL import Image
+        import numpy as np
+        img = Image.open(path)
+        img = np.asanyarray(img)
+        img = img[:, :, [2, 1, 0]]
+
+    if img.ndim == 2:
+        h, w = img.shape
+        c = 1
+    else:
+        h, w, c = img.shape
+    _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
+    return (key, img_byte, (h, w, c))
+
+
+class LmdbMaker():
+    """LMDB Maker.
+
+    Args:
+        lmdb_path (str): Lmdb save path.
+        map_size (int): Map size for lmdb env. Default: 1024 ** 4, 1TB.
+        batch (int): After processing batch images, lmdb commits.
+            Default: 5000.
+        compress_level (int): Compress level when encoding images. Default: 1.
+    """
+
+    def __init__(self, lmdb_path, map_size=1024**4, batch=5000, compress_level=1):
+        if not lmdb_path.endswith('.lmdb'):
+            raise ValueError("lmdb_path must end with '.lmdb'.")
+        if osp.exists(lmdb_path):
+            print(f'Folder {lmdb_path} already exists. Exit.')
+            sys.exit(1)
+
+        self.lmdb_path = lmdb_path
+        self.batch = batch
+        self.compress_level = compress_level
+        self.env = lmdb.open(lmdb_path, map_size=map_size)
+        self.txn = self.env.begin(write=True)
+        self.txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
+        self.counter = 0
+
+    def put(self, img_byte, key, img_shape):
+        self.counter += 1
+        key_byte = key.encode('ascii')
+        self.txn.put(key_byte, img_byte)
+        # write meta information
+        h, w, c = img_shape
+        self.txt_file.write(f'{key}.png ({h},{w},{c}) {self.compress_level}\n')
+        if self.counter % self.batch == 0:
+            self.txn.commit()
+            self.txn = self.env.begin(write=True)
+
+    def close(self):
+        self.txn.commit()
+        self.env.close()
+        self.txt_file.close()

core/data/deg_kair_utils/utils_logger.py

@@ -0,0 +1,66 @@
+import sys
+import datetime
+import logging
+
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+def log(*args, **kwargs):
+    print(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), *args, **kwargs)
+
+
+'''
+# --------------------------------------------
+# logger
+# --------------------------------------------
+'''
+
+
+def logger_info(logger_name, log_path='default_logger.log'):
+    ''' set up logger
+    modified by Kai Zhang (github: https://github.com/cszn)
+    '''
+    log = logging.getLogger(logger_name)
+    if log.hasHandlers():
+        print('LogHandlers exist!')
+    else:
+        print('LogHandlers setup!')
+        level = logging.INFO
+        formatter = logging.Formatter('%(asctime)s.%(msecs)03d : %(message)s', datefmt='%y-%m-%d %H:%M:%S')
+        fh = logging.FileHandler(log_path, mode='a')
+        fh.setFormatter(formatter)
+        log.setLevel(level)
+        log.addHandler(fh)
+        # print(len(log.handlers))
+
+        sh = logging.StreamHandler()
+        sh.setFormatter(formatter)
+        log.addHandler(sh)
+
+
+'''
+# --------------------------------------------
+# print to file and std_out simultaneously
+# --------------------------------------------
+'''
+
+
+class logger_print(object):
+    def __init__(self, log_path="default.log"):
+        self.terminal = sys.stdout
+        self.log = open(log_path, 'a')
+
+    def write(self, message):
+        self.terminal.write(message)
+        self.log.write(message)  # write the message
+
+    def flush(self):
+        pass

core/data/deg_kair_utils/utils_mat.py

@@ -0,0 +1,88 @@
+import os
+import json
+import scipy.io as spio
+import pandas as pd
+
+
+def loadmat(filename):
+    '''
+    this function should be called instead of direct spio.loadmat
+    as it cures the problem of not properly recovering python dictionaries
+    from mat files. It calls the function check keys to cure all entries
+    which are still mat-objects
+    '''
+    data = spio.loadmat(filename, struct_as_record=False, squeeze_me=True)
+    return dict_to_nonedict(_check_keys(data))
+
+def _check_keys(dict):
+    '''
+    checks if entries in dictionary are mat-objects. If yes
+    todict is called to change them to nested dictionaries
+    '''
+    for key in dict:
+        if isinstance(dict[key], spio.matlab.mio5_params.mat_struct):
+            dict[key] = _todict(dict[key])
+    return dict
+
+def _todict(matobj):
+    '''
+    A recursive function which constructs from matobjects nested dictionaries
+    '''
+    dict = {}
+    for strg in matobj._fieldnames:
+        elem = matobj.__dict__[strg]
+        if isinstance(elem, spio.matlab.mio5_params.mat_struct):
+            dict[strg] = _todict(elem)
+        else:
+            dict[strg] = elem
+    return dict
+
+
+def dict_to_nonedict(opt):
+    if isinstance(opt, dict):
+        new_opt = dict()
+        for key, sub_opt in opt.items():
+            new_opt[key] = dict_to_nonedict(sub_opt)
+        return NoneDict(**new_opt)
+    elif isinstance(opt, list):
+        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+    else:
+        return opt
+
+
+class NoneDict(dict):
+    def __missing__(self, key):
+        return None
+
+
+def mat2json(mat_path=None, filepath = None):
+    """
+    Converts .mat file to .json and writes new file
+    Parameters
+    ----------
+    mat_path: Str
+        path/filename .mat存放路径
+    filepath: Str
+        如果需要保存成json, 添加这一路径. 否则不保存
+    Returns
+        返回转化的字典
+    -------
+    None
+    Examples
+    --------
+    >>> mat2json(blah blah)
+    """
+
+    matlabFile = loadmat(mat_path)
+    #pop all those dumb fields that don't let you jsonize file
+    matlabFile.pop('__header__')
+    matlabFile.pop('__version__')
+    matlabFile.pop('__globals__')
+    #jsonize the file - orientation is 'index'
+    matlabFile = pd.Series(matlabFile).to_json()
+
+    if filepath:
+        json_path = os.path.splitext(os.path.split(mat_path)[1])[0] + '.json'
+        with open(json_path, 'w') as f:
+                f.write(matlabFile)
+    return matlabFile

core/data/deg_kair_utils/utils_matconvnet.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import torch
+from collections import OrderedDict
+
+# import scipy.io as io
+import hdf5storage
+
+"""
+# --------------------------------------------
+# Convert matconvnet SimpleNN model into pytorch model
+# --------------------------------------------
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# 28/Nov/2019
+# --------------------------------------------
+"""
+
+
+def weights2tensor(x, squeeze=False, in_features=None, out_features=None):
+    """Modified version of https://github.com/albanie/pytorch-mcn
+    Adjust memory layout and load weights as torch tensor
+    Args:
+        x (ndaray): a numpy array, corresponding to a set of network weights
+           stored in column major order
+        squeeze (bool) [False]: whether to squeeze the tensor (i.e. remove
+           singletons from the trailing dimensions. So after converting to
+           pytorch layout (C_out, C_in, H, W), if the shape is (A, B, 1, 1)
+           it will be reshaped to a matrix with shape (A,B).
+        in_features (int :: None): used to reshape weights for a linear block.
+        out_features (int :: None): used to reshape weights for a linear block.
+    Returns:
+        torch.tensor: a permuted sets of weights, matching the pytorch layout
+        convention
+    """
+    if x.ndim == 4:
+        x = x.transpose((3, 2, 0, 1))
+# for FFDNet, pixel-shuffle layer
+#        if x.shape[1]==13:
+#            x=x[:,[0,2,1,3,  4,6,5,7, 8,10,9,11, 12],:,:]
+#        if x.shape[0]==12:   
+#            x=x[[0,2,1,3,  4,6,5,7, 8,10,9,11],:,:,:]
+#        if x.shape[1]==5:
+#            x=x[:,[0,2,1,3,  4],:,:]
+#        if x.shape[0]==4:   
+#            x=x[[0,2,1,3],:,:,:]
+## for SRMD, pixel-shuffle layer
+#        if x.shape[0]==12:   
+#            x=x[[0,2,1,3,  4,6,5,7, 8,10,9,11],:,:,:]
+#        if x.shape[0]==27:
+#            x=x[[0,3,6,1,4,7,2,5,8, 0+9,3+9,6+9,1+9,4+9,7+9,2+9,5+9,8+9, 0+18,3+18,6+18,1+18,4+18,7+18,2+18,5+18,8+18],:,:,:]
+#        if x.shape[0]==48:   
+#            x=x[[0,4,8,12,1,5,9,13,2,6,10,14,3,7,11,15,  0+16,4+16,8+16,12+16,1+16,5+16,9+16,13+16,2+16,6+16,10+16,14+16,3+16,7+16,11+16,15+16,  0+32,4+32,8+32,12+32,1+32,5+32,9+32,13+32,2+32,6+32,10+32,14+32,3+32,7+32,11+32,15+32],:,:,:]
+
+    elif x.ndim == 3:  # add by Kai
+        x = x[:,:,:,None]
+        x = x.transpose((3, 2, 0, 1))
+    elif x.ndim == 2:
+        if x.shape[1] == 1:
+            x = x.flatten()
+    if squeeze:
+        if in_features and out_features:
+            x = x.reshape((out_features, in_features))
+        x = np.squeeze(x)
+    return torch.from_numpy(np.ascontiguousarray(x))
+
+
+def save_model(network, save_path):
+    state_dict = network.state_dict()
+    for key, param in state_dict.items():
+        state_dict[key] = param.cpu()
+    torch.save(state_dict, save_path)
+
+
+if __name__ == '__main__':
+    
+    
+#    from utils import utils_logger
+#    import logging
+#    utils_logger.logger_info('a', 'a.log')
+#    logger = logging.getLogger('a')
+#    
+    # mcn = hdf5storage.loadmat('/model_zoo/matfile/FFDNet_Clip_gray.mat')
+    mcn = hdf5storage.loadmat('models/modelcolor.mat')
+    
+    
+    #logger.info(mcn['CNNdenoiser'][0][0][0][1][0][0][0][0])
+
+    mat_net = OrderedDict()
+    for idx in range(25):
+        mat_net[str(idx)] = OrderedDict()
+        count = -1
+        
+        print(idx)
+        for i in range(13):
+            
+            if mcn['CNNdenoiser'][0][idx][0][i][0][0][0][0] == 'conv':
+                
+                count += 1
+                w = mcn['CNNdenoiser'][0][idx][0][i][0][1][0][0]
+               # print(w.shape)
+                w = weights2tensor(w)
+               # print(w.shape)
+                
+                b = mcn['CNNdenoiser'][0][idx][0][i][0][1][0][1]
+                b = weights2tensor(b)
+                print(b.shape)
+
+                mat_net[str(idx)]['model.{:d}.weight'.format(count*2)] = w
+                mat_net[str(idx)]['model.{:d}.bias'.format(count*2)] = b
+
+    torch.save(mat_net, 'model_zoo/modelcolor.pth')
+   
+
+
+#    from models.network_dncnn import IRCNN as net
+#    network = net(in_nc=3, out_nc=3, nc=64)
+#    state_dict = network.state_dict()
+#
+#    #show_kv(state_dict)
+#
+#    for i in range(len(mcn['net'][0][0][0])):
+#        print(mcn['net'][0][0][0][i][0][0][0][0])
+#
+#    count = -1
+#    mat_net = OrderedDict()
+#    for i in range(len(mcn['net'][0][0][0])):
+#        if mcn['net'][0][0][0][i][0][0][0][0] == 'conv':
+#            
+#            count += 1
+#            w = mcn['net'][0][0][0][i][0][1][0][0]
+#            print(w.shape)
+#            w = weights2tensor(w)
+#            print(w.shape)
+#            
+#            b = mcn['net'][0][0][0][i][0][1][0][1]
+#            b = weights2tensor(b)
+#            print(b.shape)
+#            
+#            mat_net['model.{:d}.weight'.format(count*2)] = w
+#            mat_net['model.{:d}.bias'.format(count*2)] = b
+#
+#    torch.save(mat_net, 'E:/pytorch/KAIR_ongoing/model_zoo/ffdnet_gray_clip.pth')
+#    
+#    
+#
+#    crt_net = torch.load('E:/pytorch/KAIR_ongoing/model_zoo/imdn_x4.pth')
+#    def show_kv(net):
+#        for k, v in net.items():
+#            print(k)
+#
+#    show_kv(crt_net)
+
+
+#    from models.network_dncnn import DnCNN as net
+#    network = net(in_nc=2, out_nc=1, nc=64, nb=20, act_mode='R')
+
+#    from models.network_srmd import SRMD as net
+#    #network = net(in_nc=1, out_nc=1, nc=64, nb=15, act_mode='R')
+#    network = net(in_nc=19, out_nc=3, nc=128, nb=12, upscale=4, act_mode='R', upsample_mode='pixelshuffle')
+#    
+#    from models.network_rrdb import RRDB as net
+#    network = net(in_nc=3, out_nc=3, nc=64, nb=23, gc=32, upscale=4, act_mode='L', upsample_mode='upconv')
+#    
+#    state_dict = network.state_dict()
+#    for key, param in state_dict.items():
+#        print(key)
+#    from models.network_imdn import IMDN as net
+#    network = net(in_nc=3, out_nc=3, nc=64, nb=8, upscale=4, act_mode='L', upsample_mode='pixelshuffle')
+#    state_dict = network.state_dict()
+#    mat_net = OrderedDict()
+#    for ((key, param),(key2, param2)) in zip(state_dict.items(), crt_net.items()):
+#        mat_net[key] = param2
+#    torch.save(mat_net, 'model_zoo/imdn_x4_1.pth') 
+#        
+
+#    net_old = torch.load('net_old.pth')
+#    def show_kv(net):
+#        for k, v in net.items():
+#            print(k)
+#
+#    show_kv(net_old)
+#    from models.network_dpsr import MSRResNet_prior as net
+#    model = net(in_nc=4, out_nc=3, nc=96, nb=16, upscale=4, act_mode='R', upsample_mode='pixelshuffle')
+#    state_dict = network.state_dict()
+#    net_new = OrderedDict()
+#    for ((key, param),(key_old, param_old)) in zip(state_dict.items(), net_old.items()):
+#        net_new[key] = param_old
+#    torch.save(net_new, 'net_new.pth') 
+
+
+   # print(key)
+      #  print(param.size())
+
+
+
+    # run utils/utils_matconvnet.py

core/data/deg_kair_utils/utils_model.py

@@ -0,0 +1,330 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import torch
+from utils import utils_image as util
+import re
+import glob
+import os
+
+
+'''
+# --------------------------------------------
+# Model
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+'''
+
+
+def find_last_checkpoint(save_dir, net_type='G', pretrained_path=None):
+    """
+    # ---------------------------------------
+    # Kai Zhang (github: https://github.com/cszn)
+    # 03/Mar/2019
+    # ---------------------------------------
+    Args:
+        save_dir: model folder
+        net_type: 'G' or 'D' or 'optimizerG' or 'optimizerD'
+        pretrained_path: pretrained model path. If save_dir does not have any model, load from pretrained_path
+
+    Return:
+        init_iter: iteration number
+        init_path: model path
+    # ---------------------------------------
+    """
+
+    file_list = glob.glob(os.path.join(save_dir, '*_{}.pth'.format(net_type)))
+    if file_list:
+        iter_exist = []
+        for file_ in file_list:
+            iter_current = re.findall(r"(\d+)_{}.pth".format(net_type), file_)
+            iter_exist.append(int(iter_current[0]))
+        init_iter = max(iter_exist)
+        init_path = os.path.join(save_dir, '{}_{}.pth'.format(init_iter, net_type))
+    else:
+        init_iter = 0
+        init_path = pretrained_path
+    return init_iter, init_path
+
+
+def test_mode(model, L, mode=0, refield=32, min_size=256, sf=1, modulo=1):
+    '''
+    # ---------------------------------------
+    # Kai Zhang (github: https://github.com/cszn)
+    # 03/Mar/2019
+    # ---------------------------------------
+    Args:
+        model: trained model
+        L: input Low-quality image
+        mode:
+            (0) normal: test(model, L)
+            (1) pad: test_pad(model, L, modulo=16)
+            (2) split: test_split(model, L, refield=32, min_size=256, sf=1, modulo=1)
+            (3) x8: test_x8(model, L, modulo=1) ^_^
+            (4) split and x8: test_split_x8(model, L, refield=32, min_size=256, sf=1, modulo=1)
+        refield: effective receptive filed of the network, 32 is enough
+            useful when split, i.e., mode=2, 4
+        min_size: min_sizeXmin_size image, e.g., 256X256 image
+            useful when split, i.e., mode=2, 4
+        sf: scale factor for super-resolution, otherwise 1
+        modulo: 1 if split
+            useful when pad, i.e., mode=1
+
+    Returns:
+        E: estimated image
+    # ---------------------------------------
+    '''
+    if mode == 0:
+        E = test(model, L)
+    elif mode == 1:
+        E = test_pad(model, L, modulo, sf)
+    elif mode == 2:
+        E = test_split(model, L, refield, min_size, sf, modulo)
+    elif mode == 3:
+        E = test_x8(model, L, modulo, sf)
+    elif mode == 4:
+        E = test_split_x8(model, L, refield, min_size, sf, modulo)
+    return E
+
+
+'''
+# --------------------------------------------
+# normal (0)
+# --------------------------------------------
+'''
+
+
+def test(model, L):
+    E = model(L)
+    return E
+
+
+'''
+# --------------------------------------------
+# pad (1)
+# --------------------------------------------
+'''
+
+
+def test_pad(model, L, modulo=16, sf=1):
+    h, w = L.size()[-2:]
+    paddingBottom = int(np.ceil(h/modulo)*modulo-h)
+    paddingRight = int(np.ceil(w/modulo)*modulo-w)
+    L = torch.nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(L)
+    E = model(L)
+    E = E[..., :h*sf, :w*sf]
+    return E
+
+
+'''
+# --------------------------------------------
+# split (function)
+# --------------------------------------------
+'''
+
+
+def test_split_fn(model, L, refield=32, min_size=256, sf=1, modulo=1):
+    """
+    Args:
+        model: trained model
+        L: input Low-quality image
+        refield: effective receptive filed of the network, 32 is enough
+        min_size: min_sizeXmin_size image, e.g., 256X256 image
+        sf: scale factor for super-resolution, otherwise 1
+        modulo: 1 if split
+
+    Returns:
+        E: estimated result
+    """
+    h, w = L.size()[-2:]
+    if h*w <= min_size**2:
+        L = torch.nn.ReplicationPad2d((0, int(np.ceil(w/modulo)*modulo-w), 0, int(np.ceil(h/modulo)*modulo-h)))(L)
+        E = model(L)
+        E = E[..., :h*sf, :w*sf]
+    else:
+        top = slice(0, (h//2//refield+1)*refield)
+        bottom = slice(h - (h//2//refield+1)*refield, h)
+        left = slice(0, (w//2//refield+1)*refield)
+        right = slice(w - (w//2//refield+1)*refield, w)
+        Ls = [L[..., top, left], L[..., top, right], L[..., bottom, left], L[..., bottom, right]]
+
+        if h * w <= 4*(min_size**2):
+            Es = [model(Ls[i]) for i in range(4)]
+        else:
+            Es = [test_split_fn(model, Ls[i], refield=refield, min_size=min_size, sf=sf, modulo=modulo) for i in range(4)]
+
+        b, c = Es[0].size()[:2]
+        E = torch.zeros(b, c, sf * h, sf * w).type_as(L)
+
+        E[..., :h//2*sf, :w//2*sf] = Es[0][..., :h//2*sf, :w//2*sf]
+        E[..., :h//2*sf, w//2*sf:w*sf] = Es[1][..., :h//2*sf, (-w + w//2)*sf:]
+        E[..., h//2*sf:h*sf, :w//2*sf] = Es[2][..., (-h + h//2)*sf:, :w//2*sf]
+        E[..., h//2*sf:h*sf, w//2*sf:w*sf] = Es[3][..., (-h + h//2)*sf:, (-w + w//2)*sf:]
+    return E
+
+
+'''
+# --------------------------------------------
+# split (2)
+# --------------------------------------------
+'''
+
+
+def test_split(model, L, refield=32, min_size=256, sf=1, modulo=1):
+    E = test_split_fn(model, L, refield=refield, min_size=min_size, sf=sf, modulo=modulo)
+    return E
+
+
+'''
+# --------------------------------------------
+# x8 (3)
+# --------------------------------------------
+'''
+
+
+def test_x8(model, L, modulo=1, sf=1):
+    E_list = [test_pad(model, util.augment_img_tensor4(L, mode=i), modulo=modulo, sf=sf) for i in range(8)]
+    for i in range(len(E_list)):
+        if i == 3 or i == 5:
+            E_list[i] = util.augment_img_tensor4(E_list[i], mode=8 - i)
+        else:
+            E_list[i] = util.augment_img_tensor4(E_list[i], mode=i)
+    output_cat = torch.stack(E_list, dim=0)
+    E = output_cat.mean(dim=0, keepdim=False)
+    return E
+
+
+'''
+# --------------------------------------------
+# split and x8 (4)
+# --------------------------------------------
+'''
+
+
+def test_split_x8(model, L, refield=32, min_size=256, sf=1, modulo=1):
+    E_list = [test_split_fn(model, util.augment_img_tensor4(L, mode=i), refield=refield, min_size=min_size, sf=sf, modulo=modulo) for i in range(8)]
+    for k, i in enumerate(range(len(E_list))):
+        if i==3 or i==5:
+            E_list[k] = util.augment_img_tensor4(E_list[k], mode=8-i)
+        else:
+            E_list[k] = util.augment_img_tensor4(E_list[k], mode=i)
+    output_cat = torch.stack(E_list, dim=0)
+    E = output_cat.mean(dim=0, keepdim=False)
+    return E
+
+
+'''
+# ^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-
+# _^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^
+# ^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-^_^-
+'''
+
+
+'''
+# --------------------------------------------
+# print
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# print model
+# --------------------------------------------
+def print_model(model):
+    msg = describe_model(model)
+    print(msg)
+
+
+# --------------------------------------------
+# print params
+# --------------------------------------------
+def print_params(model):
+    msg = describe_params(model)
+    print(msg)
+
+
+'''
+# --------------------------------------------
+# information
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# model inforation
+# --------------------------------------------
+def info_model(model):
+    msg = describe_model(model)
+    return msg
+
+
+# --------------------------------------------
+# params inforation
+# --------------------------------------------
+def info_params(model):
+    msg = describe_params(model)
+    return msg
+
+
+'''
+# --------------------------------------------
+# description
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# model name and total number of parameters
+# --------------------------------------------
+def describe_model(model):
+    if isinstance(model, torch.nn.DataParallel):
+        model = model.module
+    msg = '\n'
+    msg += 'models name: {}'.format(model.__class__.__name__) + '\n'
+    msg += 'Params number: {}'.format(sum(map(lambda x: x.numel(), model.parameters()))) + '\n'
+    msg += 'Net structure:\n{}'.format(str(model)) + '\n'
+    return msg
+
+
+# --------------------------------------------
+# parameters description
+# --------------------------------------------
+def describe_params(model):
+    if isinstance(model, torch.nn.DataParallel):
+        model = model.module
+    msg = '\n'
+    msg += ' | {:^6s} | {:^6s} | {:^6s} | {:^6s} || {:<20s}'.format('mean', 'min', 'max', 'std', 'shape', 'param_name') + '\n'
+    for name, param in model.state_dict().items():
+        if not 'num_batches_tracked' in name:
+            v = param.data.clone().float()
+            msg += ' | {:>6.3f} | {:>6.3f} | {:>6.3f} | {:>6.3f} | {} || {:s}'.format(v.mean(), v.min(), v.max(), v.std(), v.shape, name) + '\n'
+    return msg
+
+
+if __name__ == '__main__':
+
+    class Net(torch.nn.Module):
+        def __init__(self, in_channels=3, out_channels=3):
+            super(Net, self).__init__()
+            self.conv = torch.nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1)
+
+        def forward(self, x):
+            x = self.conv(x)
+            return x
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+
+    model = Net()
+    model = model.eval()
+    print_model(model)
+    print_params(model)
+    x = torch.randn((2,3,401,401))
+    torch.cuda.empty_cache()
+    with torch.no_grad():
+        for mode in range(5):
+            y = test_mode(model, x, mode, refield=32, min_size=256, sf=1, modulo=1)
+            print(y.shape)
+
+    # run utils/utils_model.py

core/data/deg_kair_utils/utils_modelsummary.py

@@ -0,0 +1,485 @@
+import torch.nn as nn
+import torch
+import numpy as np
+
+'''
+---- 1) FLOPs: floating point operations
+---- 2) #Activations: the number of elements of all ‘Conv2d’ outputs
+---- 3) #Conv2d: the number of ‘Conv2d’ layers
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 21/July/2020
+# --------------------------------------------
+# Reference
+https://github.com/sovrasov/flops-counter.pytorch.git
+
+# If you use this code, please consider the following citation:
+
+@inproceedings{zhang2020aim, % 
+  title={AIM 2020 Challenge on Efficient Super-Resolution: Methods and Results},
+  author={Kai Zhang and Martin Danelljan and Yawei Li and Radu Timofte and others},
+  booktitle={European Conference on Computer Vision Workshops},
+  year={2020}
+}
+# --------------------------------------------
+'''
+
+def get_model_flops(model, input_res, print_per_layer_stat=True,
+                              input_constructor=None):
+    assert type(input_res) is tuple, 'Please provide the size of the input image.'
+    assert len(input_res) >= 3, 'Input image should have 3 dimensions.'
+    flops_model = add_flops_counting_methods(model)
+    flops_model.eval().start_flops_count()
+    if input_constructor:
+        input = input_constructor(input_res)
+        _ = flops_model(**input)
+    else:
+        device = list(flops_model.parameters())[-1].device
+        batch = torch.FloatTensor(1, *input_res).to(device)
+        _ = flops_model(batch)
+
+    if print_per_layer_stat:
+        print_model_with_flops(flops_model)
+    flops_count = flops_model.compute_average_flops_cost()
+    flops_model.stop_flops_count()
+
+    return flops_count
+
+def get_model_activation(model, input_res, input_constructor=None):
+    assert type(input_res) is tuple, 'Please provide the size of the input image.'
+    assert len(input_res) >= 3, 'Input image should have 3 dimensions.'
+    activation_model = add_activation_counting_methods(model)
+    activation_model.eval().start_activation_count()
+    if input_constructor:
+        input = input_constructor(input_res)
+        _ = activation_model(**input)
+    else:
+        device = list(activation_model.parameters())[-1].device
+        batch = torch.FloatTensor(1, *input_res).to(device)
+        _ = activation_model(batch)
+
+    activation_count, num_conv = activation_model.compute_average_activation_cost()
+    activation_model.stop_activation_count()
+
+    return activation_count, num_conv
+
+
+def get_model_complexity_info(model, input_res, print_per_layer_stat=True, as_strings=True,
+                              input_constructor=None):
+    assert type(input_res) is tuple
+    assert len(input_res) >= 3
+    flops_model = add_flops_counting_methods(model)
+    flops_model.eval().start_flops_count()
+    if input_constructor:
+        input = input_constructor(input_res)
+        _ = flops_model(**input)
+    else:
+        batch = torch.FloatTensor(1, *input_res)
+        _ = flops_model(batch)
+
+    if print_per_layer_stat:
+        print_model_with_flops(flops_model)
+    flops_count = flops_model.compute_average_flops_cost()
+    params_count = get_model_parameters_number(flops_model)
+    flops_model.stop_flops_count()
+
+    if as_strings:
+        return flops_to_string(flops_count), params_to_string(params_count)
+
+    return flops_count, params_count
+
+
+def flops_to_string(flops, units='GMac', precision=2):
+    if units is None:
+        if flops // 10**9 > 0:
+            return str(round(flops / 10.**9, precision)) + ' GMac'
+        elif flops // 10**6 > 0:
+            return str(round(flops / 10.**6, precision)) + ' MMac'
+        elif flops // 10**3 > 0:
+            return str(round(flops / 10.**3, precision)) + ' KMac'
+        else:
+            return str(flops) + ' Mac'
+    else:
+        if units == 'GMac':
+            return str(round(flops / 10.**9, precision)) + ' ' + units
+        elif units == 'MMac':
+            return str(round(flops / 10.**6, precision)) + ' ' + units
+        elif units == 'KMac':
+            return str(round(flops / 10.**3, precision)) + ' ' + units
+        else:
+            return str(flops) + ' Mac'
+
+
+def params_to_string(params_num):
+    if params_num // 10 ** 6 > 0:
+        return str(round(params_num / 10 ** 6, 2)) + ' M'
+    elif params_num // 10 ** 3:
+        return str(round(params_num / 10 ** 3, 2)) + ' k'
+    else:
+        return str(params_num)
+
+
+def print_model_with_flops(model, units='GMac', precision=3):
+    total_flops = model.compute_average_flops_cost()
+
+    def accumulate_flops(self):
+        if is_supported_instance(self):
+            return self.__flops__ / model.__batch_counter__
+        else:
+            sum = 0
+            for m in self.children():
+                sum += m.accumulate_flops()
+            return sum
+
+    def flops_repr(self):
+        accumulated_flops_cost = self.accumulate_flops()
+        return ', '.join([flops_to_string(accumulated_flops_cost, units=units, precision=precision),
+                          '{:.3%} MACs'.format(accumulated_flops_cost / total_flops),
+                          self.original_extra_repr()])
+
+    def add_extra_repr(m):
+        m.accumulate_flops = accumulate_flops.__get__(m)
+        flops_extra_repr = flops_repr.__get__(m)
+        if m.extra_repr != flops_extra_repr:
+            m.original_extra_repr = m.extra_repr
+            m.extra_repr = flops_extra_repr
+            assert m.extra_repr != m.original_extra_repr
+
+    def del_extra_repr(m):
+        if hasattr(m, 'original_extra_repr'):
+            m.extra_repr = m.original_extra_repr
+            del m.original_extra_repr
+        if hasattr(m, 'accumulate_flops'):
+            del m.accumulate_flops
+
+    model.apply(add_extra_repr)
+    print(model)
+    model.apply(del_extra_repr)
+
+
+def get_model_parameters_number(model):
+    params_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return params_num
+
+
+def add_flops_counting_methods(net_main_module):
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    # embed()
+    net_main_module.start_flops_count = start_flops_count.__get__(net_main_module)
+    net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module)
+    net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module)
+    net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(net_main_module)
+
+    net_main_module.reset_flops_count()
+    return net_main_module
+
+
+def compute_average_flops_cost(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Returns current mean flops consumption per image.
+
+    """
+
+    flops_sum = 0
+    for module in self.modules():
+        if is_supported_instance(module):
+            flops_sum += module.__flops__
+
+    return flops_sum
+
+
+def start_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Activates the computation of mean flops consumption per image.
+    Call it before you run the network.
+
+    """
+    self.apply(add_flops_counter_hook_function)
+
+
+def stop_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Stops computing the mean flops consumption per image.
+    Call whenever you want to pause the computation.
+
+    """
+    self.apply(remove_flops_counter_hook_function)
+
+
+def reset_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+
+    Resets statistics computed so far.
+
+    """
+    self.apply(add_flops_counter_variable_or_reset)
+
+
+def add_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            return
+
+        if isinstance(module, (nn.Conv2d, nn.Conv3d, nn.ConvTranspose2d)):
+            handle = module.register_forward_hook(conv_flops_counter_hook)
+        elif isinstance(module, (nn.ReLU, nn.PReLU, nn.ELU, nn.LeakyReLU, nn.ReLU6)):
+            handle = module.register_forward_hook(relu_flops_counter_hook)
+        elif isinstance(module, nn.Linear):
+            handle = module.register_forward_hook(linear_flops_counter_hook)
+        elif isinstance(module, (nn.BatchNorm2d)):
+            handle = module.register_forward_hook(bn_flops_counter_hook)
+        else:
+            handle = module.register_forward_hook(empty_flops_counter_hook)
+        module.__flops_handle__ = handle
+
+
+def remove_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            module.__flops_handle__.remove()
+            del module.__flops_handle__
+
+
+def add_flops_counter_variable_or_reset(module):
+    if is_supported_instance(module):
+        module.__flops__ = 0
+
+
+# ---- Internal functions
+def is_supported_instance(module):
+    if isinstance(module,
+                  (
+                          nn.Conv2d, nn.ConvTranspose2d,
+                          nn.BatchNorm2d,
+                          nn.Linear,
+                          nn.ReLU, nn.PReLU, nn.ELU, nn.LeakyReLU, nn.ReLU6,
+                  )):
+        return True
+
+    return False
+
+
+def conv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    # input = input[0]
+
+    batch_size = output.shape[0]
+    output_dims = list(output.shape[2:])
+
+    kernel_dims = list(conv_module.kernel_size)
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    filters_per_channel = out_channels // groups
+    conv_per_position_flops = np.prod(kernel_dims) * in_channels * filters_per_channel
+
+    active_elements_count = batch_size * np.prod(output_dims)
+    overall_conv_flops = int(conv_per_position_flops) * int(active_elements_count)
+
+    # overall_flops = overall_conv_flops
+
+    conv_module.__flops__ += int(overall_conv_flops)
+    # conv_module.__output_dims__ = output_dims
+
+
+def relu_flops_counter_hook(module, input, output):
+    active_elements_count = output.numel()
+    module.__flops__ += int(active_elements_count)
+    # print(module.__flops__, id(module))
+    # print(module)
+
+
+def linear_flops_counter_hook(module, input, output):
+    input = input[0]
+    if len(input.shape) == 1:
+        batch_size = 1
+        module.__flops__ += int(batch_size * input.shape[0] * output.shape[0])
+    else:
+        batch_size = input.shape[0]
+        module.__flops__ += int(batch_size * input.shape[1] * output.shape[1])
+
+
+def bn_flops_counter_hook(module, input, output):
+    # input = input[0]
+    # TODO: need to check here
+    # batch_flops = np.prod(input.shape)
+    # if module.affine:
+    #     batch_flops *= 2
+    # module.__flops__ += int(batch_flops)
+    batch = output.shape[0]
+    output_dims = output.shape[2:]
+    channels = module.num_features
+    batch_flops = batch * channels * np.prod(output_dims)
+    if module.affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+
+# ---- Count the number of convolutional layers and the activation
+def add_activation_counting_methods(net_main_module):
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    # embed()
+    net_main_module.start_activation_count = start_activation_count.__get__(net_main_module)
+    net_main_module.stop_activation_count = stop_activation_count.__get__(net_main_module)
+    net_main_module.reset_activation_count = reset_activation_count.__get__(net_main_module)
+    net_main_module.compute_average_activation_cost = compute_average_activation_cost.__get__(net_main_module)
+
+    net_main_module.reset_activation_count()
+    return net_main_module
+
+
+def compute_average_activation_cost(self):
+    """
+    A method that will be available after add_activation_counting_methods() is called
+    on a desired net object.
+
+    Returns current mean activation consumption per image.
+
+    """
+
+    activation_sum = 0
+    num_conv = 0
+    for module in self.modules():
+        if is_supported_instance_for_activation(module):
+            activation_sum += module.__activation__
+            num_conv += module.__num_conv__
+    return activation_sum, num_conv
+
+
+def start_activation_count(self):
+    """
+    A method that will be available after add_activation_counting_methods() is called
+    on a desired net object.
+
+    Activates the computation of mean activation consumption per image.
+    Call it before you run the network.
+
+    """
+    self.apply(add_activation_counter_hook_function)
+
+
+def stop_activation_count(self):
+    """
+    A method that will be available after add_activation_counting_methods() is called
+    on a desired net object.
+
+    Stops computing the mean activation consumption per image.
+    Call whenever you want to pause the computation.
+
+    """
+    self.apply(remove_activation_counter_hook_function)
+
+
+def reset_activation_count(self):
+    """
+    A method that will be available after add_activation_counting_methods() is called
+    on a desired net object.
+
+    Resets statistics computed so far.
+
+    """
+    self.apply(add_activation_counter_variable_or_reset)
+
+
+def add_activation_counter_hook_function(module):
+    if is_supported_instance_for_activation(module):
+        if hasattr(module, '__activation_handle__'):
+            return
+
+        if isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+            handle = module.register_forward_hook(conv_activation_counter_hook)
+            module.__activation_handle__ = handle
+
+
+def remove_activation_counter_hook_function(module):
+    if is_supported_instance_for_activation(module):
+        if hasattr(module, '__activation_handle__'):
+            module.__activation_handle__.remove()
+            del module.__activation_handle__
+
+
+def add_activation_counter_variable_or_reset(module):
+    if is_supported_instance_for_activation(module):
+        module.__activation__ = 0
+        module.__num_conv__ = 0
+
+
+def is_supported_instance_for_activation(module):
+    if isinstance(module,
+                  (
+                          nn.Conv2d, nn.ConvTranspose2d,
+                  )):
+        return True
+
+    return False
+
+def conv_activation_counter_hook(module, input, output):
+    """
+    Calculate the activations in the convolutional operation.
+    Reference: Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, Piotr Dollár, Designing Network Design Spaces.
+    :param module:
+    :param input:
+    :param output:
+    :return:
+    """
+    module.__activation__ += output.numel()
+    module.__num_conv__ += 1
+
+
+def empty_flops_counter_hook(module, input, output):
+    module.__flops__ += 0
+
+
+def upsample_flops_counter_hook(module, input, output):
+    output_size = output[0]
+    batch_size = output_size.shape[0]
+    output_elements_count = batch_size
+    for val in output_size.shape[1:]:
+        output_elements_count *= val
+    module.__flops__ += int(output_elements_count)
+
+
+def pool_flops_counter_hook(module, input, output):
+    input = input[0]
+    module.__flops__ += int(np.prod(input.shape))
+
+
+def dconv_flops_counter_hook(dconv_module, input, output):
+    input = input[0]
+
+    batch_size = input.shape[0]
+    output_dims = list(output.shape[2:])
+
+    m_channels, in_channels, kernel_dim1, _, = dconv_module.weight.shape
+    out_channels, _, kernel_dim2, _, = dconv_module.projection.shape
+    # groups = dconv_module.groups
+
+    # filters_per_channel = out_channels // groups
+    conv_per_position_flops1 = kernel_dim1 ** 2 * in_channels * m_channels
+    conv_per_position_flops2 = kernel_dim2 ** 2 * out_channels * m_channels
+    active_elements_count = batch_size * np.prod(output_dims)
+
+    overall_conv_flops = (conv_per_position_flops1 + conv_per_position_flops2) * active_elements_count
+    overall_flops = overall_conv_flops
+
+    dconv_module.__flops__ += int(overall_flops)
+    # dconv_module.__output_dims__ = output_dims
+
+
+
+
+

core/data/deg_kair_utils/utils_option.py

@@ -0,0 +1,255 @@
+import os
+from collections import OrderedDict
+from datetime import datetime
+import json
+import re
+import glob
+
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+def get_timestamp():
+    return datetime.now().strftime('_%y%m%d_%H%M%S')
+
+
+def parse(opt_path, is_train=True):
+
+    # ----------------------------------------
+    # remove comments starting with '//'
+    # ----------------------------------------
+    json_str = ''
+    with open(opt_path, 'r') as f:
+        for line in f:
+            line = line.split('//')[0] + '\n'
+            json_str += line
+
+    # ----------------------------------------
+    # initialize opt
+    # ----------------------------------------
+    opt = json.loads(json_str, object_pairs_hook=OrderedDict)
+
+    opt['opt_path'] = opt_path
+    opt['is_train'] = is_train
+
+    # ----------------------------------------
+    # set default
+    # ----------------------------------------
+    if 'merge_bn' not in opt:
+        opt['merge_bn'] = False
+        opt['merge_bn_startpoint'] = -1
+
+    if 'scale' not in opt:
+        opt['scale'] = 1
+
+    # ----------------------------------------
+    # datasets
+    # ----------------------------------------
+    for phase, dataset in opt['datasets'].items():
+        phase = phase.split('_')[0]
+        dataset['phase'] = phase
+        dataset['scale'] = opt['scale']  # broadcast
+        dataset['n_channels'] = opt['n_channels']  # broadcast
+        if 'dataroot_H' in dataset and dataset['dataroot_H'] is not None:
+            dataset['dataroot_H'] = os.path.expanduser(dataset['dataroot_H'])
+        if 'dataroot_L' in dataset and dataset['dataroot_L'] is not None:
+            dataset['dataroot_L'] = os.path.expanduser(dataset['dataroot_L'])
+
+    # ----------------------------------------
+    # path
+    # ----------------------------------------
+    for key, path in opt['path'].items():
+        if path and key in opt['path']:
+            opt['path'][key] = os.path.expanduser(path)
+
+    path_task = os.path.join(opt['path']['root'], opt['task'])
+    opt['path']['task'] = path_task
+    opt['path']['log'] = path_task
+    opt['path']['options'] = os.path.join(path_task, 'options')
+
+    if is_train:
+        opt['path']['models'] = os.path.join(path_task, 'models')
+        opt['path']['images'] = os.path.join(path_task, 'images')
+    else:  # test
+        opt['path']['images'] = os.path.join(path_task, 'test_images')
+
+    # ----------------------------------------
+    # network
+    # ----------------------------------------
+    opt['netG']['scale'] = opt['scale'] if 'scale' in opt else 1
+
+    # ----------------------------------------
+    # GPU devices
+    # ----------------------------------------
+    gpu_list = ','.join(str(x) for x in opt['gpu_ids'])
+    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+    print('export CUDA_VISIBLE_DEVICES=' + gpu_list)
+
+    # ----------------------------------------
+    # default setting for distributeddataparallel
+    # ----------------------------------------
+    if 'find_unused_parameters' not in opt:
+        opt['find_unused_parameters'] = True
+    if 'use_static_graph' not in opt:
+        opt['use_static_graph'] = False
+    if 'dist' not in opt:
+        opt['dist'] = False
+    opt['num_gpu'] = len(opt['gpu_ids'])
+    print('number of GPUs is: ' + str(opt['num_gpu']))
+
+    # ----------------------------------------
+    # default setting for perceptual loss
+    # ----------------------------------------
+    if 'F_feature_layer' not in opt['train']:
+        opt['train']['F_feature_layer'] = 34  # 25; [2,7,16,25,34]
+    if 'F_weights' not in opt['train']:
+        opt['train']['F_weights'] = 1.0  # 1.0; [0.1,0.1,1.0,1.0,1.0]
+    if 'F_lossfn_type' not in opt['train']:
+        opt['train']['F_lossfn_type'] = 'l1'
+    if 'F_use_input_norm' not in opt['train']:
+        opt['train']['F_use_input_norm'] = True
+    if 'F_use_range_norm' not in opt['train']:
+        opt['train']['F_use_range_norm'] = False
+
+    # ----------------------------------------
+    # default setting for optimizer
+    # ----------------------------------------
+    if 'G_optimizer_type' not in opt['train']:
+        opt['train']['G_optimizer_type'] = "adam"
+    if 'G_optimizer_betas' not in opt['train']:
+        opt['train']['G_optimizer_betas'] = [0.9,0.999]
+    if 'G_scheduler_restart_weights' not in opt['train']:
+        opt['train']['G_scheduler_restart_weights'] = 1
+    if 'G_optimizer_wd' not in opt['train']:
+        opt['train']['G_optimizer_wd'] = 0
+    if 'G_optimizer_reuse' not in opt['train']:
+        opt['train']['G_optimizer_reuse'] = False
+    if 'netD' in opt and 'D_optimizer_reuse' not in opt['train']:
+        opt['train']['D_optimizer_reuse'] = False
+
+    # ----------------------------------------
+    # default setting of strict for model loading
+    # ----------------------------------------
+    if 'G_param_strict' not in opt['train']:
+        opt['train']['G_param_strict'] = True
+    if 'netD' in opt and 'D_param_strict' not in opt['path']:
+        opt['train']['D_param_strict'] = True
+    if 'E_param_strict' not in opt['path']:
+        opt['train']['E_param_strict'] = True
+
+    # ----------------------------------------
+    # Exponential Moving Average
+    # ----------------------------------------
+    if 'E_decay' not in opt['train']:
+        opt['train']['E_decay'] = 0
+
+    # ----------------------------------------
+    # default setting for discriminator
+    # ----------------------------------------
+    if 'netD' in opt:
+        if 'net_type' not in opt['netD']:
+            opt['netD']['net_type'] = 'discriminator_patchgan'  # discriminator_unet
+        if 'in_nc' not in opt['netD']:
+            opt['netD']['in_nc'] = 3
+        if 'base_nc' not in opt['netD']:
+            opt['netD']['base_nc'] = 64
+        if 'n_layers' not in opt['netD']:
+            opt['netD']['n_layers'] = 3
+        if 'norm_type' not in opt['netD']:
+            opt['netD']['norm_type'] = 'spectral'
+
+
+    return opt
+
+
+def find_last_checkpoint(save_dir, net_type='G', pretrained_path=None):
+    """
+    Args: 
+        save_dir: model folder
+        net_type: 'G' or 'D' or 'optimizerG' or 'optimizerD'
+        pretrained_path: pretrained model path. If save_dir does not have any model, load from pretrained_path
+
+    Return:
+        init_iter: iteration number
+        init_path: model path
+    """
+    file_list = glob.glob(os.path.join(save_dir, '*_{}.pth'.format(net_type)))
+    if file_list:
+        iter_exist = []
+        for file_ in file_list:
+            iter_current = re.findall(r"(\d+)_{}.pth".format(net_type), file_)
+            iter_exist.append(int(iter_current[0]))
+        init_iter = max(iter_exist)
+        init_path = os.path.join(save_dir, '{}_{}.pth'.format(init_iter, net_type))
+    else:
+        init_iter = 0
+        init_path = pretrained_path
+    return init_iter, init_path
+
+
+'''
+# --------------------------------------------
+# convert the opt into json file
+# --------------------------------------------
+'''
+
+
+def save(opt):
+    opt_path = opt['opt_path']
+    opt_path_copy = opt['path']['options']
+    dirname, filename_ext = os.path.split(opt_path)
+    filename, ext = os.path.splitext(filename_ext)
+    dump_path = os.path.join(opt_path_copy, filename+get_timestamp()+ext)
+    with open(dump_path, 'w') as dump_file:
+        json.dump(opt, dump_file, indent=2)
+
+
+'''
+# --------------------------------------------
+# dict to string for logger
+# --------------------------------------------
+'''
+
+
+def dict2str(opt, indent_l=1):
+    msg = ''
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_l * 2) + k + ':[\n'
+            msg += dict2str(v, indent_l + 1)
+            msg += ' ' * (indent_l * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_l * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+
+'''
+# --------------------------------------------
+# convert OrderedDict to NoneDict,
+# return None for missing key
+# --------------------------------------------
+'''
+
+
+def dict_to_nonedict(opt):
+    if isinstance(opt, dict):
+        new_opt = dict()
+        for key, sub_opt in opt.items():
+            new_opt[key] = dict_to_nonedict(sub_opt)
+        return NoneDict(**new_opt)
+    elif isinstance(opt, list):
+        return [dict_to_nonedict(sub_opt) for sub_opt in opt]
+    else:
+        return opt
+
+
+class NoneDict(dict):
+    def __missing__(self, key):
+        return None

core/data/deg_kair_utils/utils_params.py

@@ -0,0 +1,135 @@
+import torch
+
+import torchvision
+
+from models import basicblock as B
+
+def show_kv(net):
+    for k, v in net.items():
+        print(k)
+
+# should run train debug mode first to get an initial model
+#crt_net = torch.load('../../experiments/debug_SRResNet_bicx4_in3nf64nb16/models/8_G.pth')
+#
+#for k, v in crt_net.items():
+#    print(k)
+#for k, v in crt_net.items():
+#    if k in pretrained_net:
+#        crt_net[k] = pretrained_net[k]
+#        print('replace ... ', k)
+
+# x2 -> x4
+#crt_net['model.5.weight'] = pretrained_net['model.2.weight']
+#crt_net['model.5.bias'] = pretrained_net['model.2.bias']
+#crt_net['model.8.weight'] = pretrained_net['model.5.weight']
+#crt_net['model.8.bias'] = pretrained_net['model.5.bias']
+#crt_net['model.10.weight'] = pretrained_net['model.7.weight']
+#crt_net['model.10.bias'] = pretrained_net['model.7.bias']
+#torch.save(crt_net, '../pretrained_tmp.pth')
+
+# x2 -> x3
+'''
+in_filter = pretrained_net['model.2.weight'] # 256, 64, 3, 3
+new_filter = torch.Tensor(576, 64, 3, 3)
+new_filter[0:256, :, :, :] = in_filter
+new_filter[256:512, :, :, :] = in_filter
+new_filter[512:, :, :, :] = in_filter[0:576-512, :, :, :]
+crt_net['model.2.weight'] = new_filter
+
+in_bias = pretrained_net['model.2.bias']  # 256, 64, 3, 3
+new_bias = torch.Tensor(576)
+new_bias[0:256] = in_bias
+new_bias[256:512] = in_bias
+new_bias[512:] = in_bias[0:576 - 512]
+crt_net['model.2.bias'] = new_bias
+
+torch.save(crt_net, '../pretrained_tmp.pth')
+'''
+
+# x2 -> x8
+'''
+crt_net['model.5.weight'] = pretrained_net['model.2.weight']
+crt_net['model.5.bias'] = pretrained_net['model.2.bias']
+crt_net['model.8.weight'] = pretrained_net['model.2.weight']
+crt_net['model.8.bias'] = pretrained_net['model.2.bias']
+crt_net['model.11.weight'] = pretrained_net['model.5.weight']
+crt_net['model.11.bias'] = pretrained_net['model.5.bias']
+crt_net['model.13.weight'] = pretrained_net['model.7.weight']
+crt_net['model.13.bias'] = pretrained_net['model.7.bias']
+torch.save(crt_net, '../pretrained_tmp.pth')
+'''
+
+# x3/4/8 RGB -> Y
+
+def rgb2gray_net(net, only_input=True):
+
+    if only_input:
+        in_filter = net['0.weight']
+        in_new_filter = in_filter[:,0,:,:]*0.2989 + in_filter[:,1,:,:]*0.587 + in_filter[:,2,:,:]*0.114
+        in_new_filter.unsqueeze_(1)
+        net['0.weight'] = in_new_filter
+
+#    out_filter = pretrained_net['model.13.weight']
+#    out_new_filter = out_filter[0, :, :, :] * 0.2989 + out_filter[1, :, :, :] * 0.587 + \
+#        out_filter[2, :, :, :] * 0.114
+#    out_new_filter.unsqueeze_(0)
+#    crt_net['model.13.weight'] = out_new_filter
+#    out_bias = pretrained_net['model.13.bias']
+#    out_new_bias = out_bias[0] * 0.2989 + out_bias[1] * 0.587 + out_bias[2] * 0.114
+#    out_new_bias = torch.Tensor(1).fill_(out_new_bias)
+#    crt_net['model.13.bias'] = out_new_bias
+
+#    torch.save(crt_net, '../pretrained_tmp.pth')
+
+    return net
+
+
+
+if __name__ == '__main__':
+    
+    net = torchvision.models.vgg19(pretrained=True)
+    for k,v in net.features.named_parameters():
+        if k=='0.weight':
+            in_new_filter = v[:,0,:,:]*0.2989 + v[:,1,:,:]*0.587 + v[:,2,:,:]*0.114
+            in_new_filter.unsqueeze_(1)
+            v = in_new_filter
+            print(v.shape)
+            print(v[0,0,0,0])
+        if k=='0.bias':
+            in_new_bias = v
+            print(v[0])
+
+    print(net.features[0])
+
+    net.features[0] = B.conv(1, 64, mode='C') 
+
+    print(net.features[0])
+    net.features[0].weight.data=in_new_filter
+    net.features[0].bias.data=in_new_bias
+
+    for k,v in net.features.named_parameters():
+        if k=='0.weight':
+            print(v[0,0,0,0])
+        if k=='0.bias':
+            print(v[0])
+
+    # transfer parameters of old model to new one
+    model_old = torch.load(model_path)
+    state_dict = model.state_dict()
+    for ((key, param),(key2, param2)) in zip(model_old.items(), state_dict.items()):
+        state_dict[key2] = param
+        print([key, key2])
+       # print([param.size(), param2.size()])
+    torch.save(state_dict, 'model_new.pth') 
+
+
+   # rgb2gray_net(net)
+
+
+
+
+
+
+
+
+

core/data/deg_kair_utils/utils_receptivefield.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+
+# online calculation: https://fomoro.com/research/article/receptive-field-calculator#
+
+# [filter size, stride, padding]
+#Assume the two dimensions are the same
+#Each kernel requires the following parameters:
+# - k_i: kernel size
+# - s_i: stride
+# - p_i: padding (if padding is uneven, right padding will higher than left padding; "SAME" option in tensorflow)
+# 
+#Each layer i requires the following parameters to be fully represented: 
+# - n_i: number of feature (data layer has n_1 = imagesize )
+# - j_i: distance (projected to image pixel distance) between center of two adjacent features
+# - r_i: receptive field of a feature in layer i
+# - start_i: position of the first feature's receptive field in layer i (idx start from 0, negative means the center fall into padding)
+
+import math
+
+def outFromIn(conv, layerIn):
+    n_in = layerIn[0]
+    j_in = layerIn[1]
+    r_in = layerIn[2]
+    start_in = layerIn[3]
+    k = conv[0]
+    s = conv[1]
+    p = conv[2]
+  
+    n_out = math.floor((n_in - k + 2*p)/s) + 1
+    actualP = (n_out-1)*s - n_in + k 
+    pR = math.ceil(actualP/2)
+    pL = math.floor(actualP/2)
+  
+    j_out = j_in * s
+    r_out = r_in + (k - 1)*j_in
+    start_out = start_in + ((k-1)/2 - pL)*j_in
+    return n_out, j_out, r_out, start_out
+
+def printLayer(layer, layer_name):
+    print(layer_name + ":")
+    print(" n features: %s  jump: %s  receptive size: %s  start: %s " % (layer[0], layer[1], layer[2], layer[3]))
+
+
+
+layerInfos = []
+if __name__ == '__main__':
+
+    convnet =   [[3,1,1],[3,1,1],[3,1,1],[4,2,1],[2,2,0],[3,1,1]]
+    layer_names = ['conv1','conv2','conv3','conv4','conv5','conv6','conv7','conv8','conv9','conv10','conv11','conv12']
+    imsize = 128
+
+    print ("-------Net summary------")
+    currentLayer = [imsize, 1, 1, 0.5]
+    printLayer(currentLayer, "input image")
+    for i in range(len(convnet)):
+        currentLayer = outFromIn(convnet[i], currentLayer)
+        layerInfos.append(currentLayer)
+        printLayer(currentLayer, layer_names[i])
+
+
+# run utils/utils_receptivefield.py
+    

core/data/deg_kair_utils/utils_regularizers.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# SVD Orthogonal Regularization
+# --------------------------------------------
+def regularizer_orth(m):
+    """
+    # ----------------------------------------
+    # SVD Orthogonal Regularization
+    # ----------------------------------------
+    # Applies regularization to the training by performing the
+    # orthogonalization technique described in the paper
+    # This function is to be called by the torch.nn.Module.apply() method,
+    # which applies svd_orthogonalization() to every layer of the model.
+    # usage: net.apply(regularizer_orth)
+    # ----------------------------------------
+    """
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        w = m.weight.data.clone()
+        c_out, c_in, f1, f2 = w.size()
+        # dtype = m.weight.data.type()
+        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
+        # self.netG.apply(svd_orthogonalization)
+        u, s, v = torch.svd(w)
+        s[s > 1.5] = s[s > 1.5] - 1e-4
+        s[s < 0.5] = s[s < 0.5] + 1e-4
+        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
+        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
+    else:
+        pass
+
+
+# --------------------------------------------
+# SVD Orthogonal Regularization
+# --------------------------------------------
+def regularizer_orth2(m):
+    """
+    # ----------------------------------------
+    # Applies regularization to the training by performing the
+    # orthogonalization technique described in the paper
+    # This function is to be called by the torch.nn.Module.apply() method,
+    # which applies svd_orthogonalization() to every layer of the model.
+    # usage: net.apply(regularizer_orth2)
+    # ----------------------------------------
+    """
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        w = m.weight.data.clone()
+        c_out, c_in, f1, f2 = w.size()
+        # dtype = m.weight.data.type()
+        w = w.permute(2, 3, 1, 0).contiguous().view(f1*f2*c_in, c_out)
+        u, s, v = torch.svd(w)
+        s_mean = s.mean()
+        s[s > 1.5*s_mean] = s[s > 1.5*s_mean] - 1e-4
+        s[s < 0.5*s_mean] = s[s < 0.5*s_mean] + 1e-4
+        w = torch.mm(torch.mm(u, torch.diag(s)), v.t())
+        m.weight.data = w.view(f1, f2, c_in, c_out).permute(3, 2, 0, 1)  # .type(dtype)
+    else:
+        pass
+
+
+
+def regularizer_clip(m):
+    """
+    # ----------------------------------------
+    # usage: net.apply(regularizer_clip)
+    # ----------------------------------------
+    """
+    eps = 1e-4
+    c_min = -1.5
+    c_max = 1.5
+
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1 or classname.find('Linear') != -1:
+        w = m.weight.data.clone()
+        w[w > c_max] -= eps
+        w[w < c_min] += eps
+        m.weight.data = w
+
+        if m.bias is not None:
+            b = m.bias.data.clone()
+            b[b > c_max] -= eps
+            b[b < c_min] += eps
+            m.bias.data = b
+
+#    elif classname.find('BatchNorm2d') != -1:
+#
+#       rv = m.running_var.data.clone()
+#       rm = m.running_mean.data.clone()
+#
+#        if m.affine:
+#            m.weight.data
+#            m.bias.data

core/data/deg_kair_utils/utils_sisr.py

@@ -0,0 +1,848 @@
+# -*- coding: utf-8 -*-
+from utils import utils_image as util
+import random
+
+import scipy
+import scipy.stats as ss
+import scipy.io as io
+from scipy import ndimage
+from scipy.interpolate import interp2d
+
+import numpy as np
+import torch
+
+
+"""
+# --------------------------------------------
+# Super-Resolution
+# --------------------------------------------
+#
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# modified by Kai Zhang (github: https://github.com/cszn)
+# 03/03/2020
+# --------------------------------------------
+"""
+
+
+"""
+# --------------------------------------------
+# anisotropic Gaussian kernels
+# --------------------------------------------
+"""
+
+
+def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
+    """ generate an anisotropic Gaussian kernel
+    Args:
+        ksize : e.g., 15, kernel size
+        theta : [0,  pi], rotation angle range
+        l1    : [0.1,50], scaling of eigenvalues
+        l2    : [0.1,l1], scaling of eigenvalues
+        If l1 = l2, will get an isotropic Gaussian kernel.
+    Returns:
+        k     : kernel
+    """
+
+    v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
+    V = np.array([[v[0], v[1]], [v[1], -v[0]]])
+    D = np.array([[l1, 0], [0, l2]])
+    Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
+    k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
+
+    return k
+
+
+def gm_blur_kernel(mean, cov, size=15):
+    center = size / 2.0 + 0.5
+    k = np.zeros([size, size])
+    for y in range(size):
+        for x in range(size):
+            cy = y - center + 1
+            cx = x - center + 1
+            k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
+
+    k = k / np.sum(k)
+    return k
+
+
+"""
+# --------------------------------------------
+# calculate PCA projection matrix
+# --------------------------------------------
+"""
+
+
+def get_pca_matrix(x, dim_pca=15):
+    """
+    Args:
+        x: 225x10000 matrix
+        dim_pca: 15
+    Returns:
+        pca_matrix: 15x225
+    """
+    C = np.dot(x, x.T)
+    w, v = scipy.linalg.eigh(C)
+    pca_matrix = v[:, -dim_pca:].T
+
+    return pca_matrix
+
+
+def show_pca(x):
+    """
+    x: PCA projection matrix, e.g., 15x225
+    """
+    for i in range(x.shape[0]):
+        xc = np.reshape(x[i, :], (int(np.sqrt(x.shape[1])), -1), order="F")
+        util.surf(xc)
+
+
+def cal_pca_matrix(path='PCA_matrix.mat', ksize=15, l_max=12.0, dim_pca=15, num_samples=500):
+    kernels = np.zeros([ksize*ksize, num_samples], dtype=np.float32)
+    for i in range(num_samples):
+
+        theta = np.pi*np.random.rand(1)
+        l1    = 0.1+l_max*np.random.rand(1)
+        l2    = 0.1+(l1-0.1)*np.random.rand(1)
+
+        k = anisotropic_Gaussian(ksize=ksize, theta=theta[0], l1=l1[0], l2=l2[0])
+
+        # util.imshow(k)
+
+        kernels[:, i] = np.reshape(k, (-1), order="F")  # k.flatten(order='F')
+
+    # io.savemat('k.mat', {'k': kernels})
+
+    pca_matrix = get_pca_matrix(kernels, dim_pca=dim_pca)
+
+    io.savemat(path, {'p': pca_matrix})
+
+    return pca_matrix
+
+
+"""
+# --------------------------------------------
+# shifted anisotropic Gaussian kernels
+# --------------------------------------------
+"""
+
+
+def shifted_anisotropic_Gaussian(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
+    """"
+    # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
+    # Kai Zhang
+    # min_var = 0.175 * sf  # variance of the gaussian kernel will be sampled between min_var and max_var
+    # max_var = 2.5 * sf
+    """
+    # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
+    lambda_1 = min_var + np.random.rand() * (max_var - min_var)
+    lambda_2 = min_var + np.random.rand() * (max_var - min_var)
+    theta = np.random.rand() * np.pi  # random theta
+    noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
+
+    # Set COV matrix using Lambdas and Theta
+    LAMBDA = np.diag([lambda_1, lambda_2])
+    Q = np.array([[np.cos(theta), -np.sin(theta)],
+                  [np.sin(theta), np.cos(theta)]])
+    SIGMA = Q @ LAMBDA @ Q.T
+    INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
+
+    # Set expectation position (shifting kernel for aligned image)
+    MU = k_size // 2 - 0.5*(scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
+    MU = MU[None, None, :, None]
+
+    # Create meshgrid for Gaussian
+    [X,Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
+    Z = np.stack([X, Y], 2)[:, :, :, None]
+
+    # Calcualte Gaussian for every pixel of the kernel
+    ZZ = Z-MU
+    ZZ_t = ZZ.transpose(0,1,3,2)
+    raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
+
+    # shift the kernel so it will be centered
+    #raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
+
+    # Normalize the kernel and return
+    #kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
+    kernel = raw_kernel / np.sum(raw_kernel)
+    return kernel
+
+
+def gen_kernel(k_size=np.array([25, 25]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=12., noise_level=0):
+    """"
+    # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
+    # Kai Zhang
+    # min_var = 0.175 * sf  # variance of the gaussian kernel will be sampled between min_var and max_var
+    # max_var = 2.5 * sf
+    """
+    sf = random.choice([1, 2, 3, 4])
+    scale_factor = np.array([sf, sf])
+    # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
+    lambda_1 = min_var + np.random.rand() * (max_var - min_var)
+    lambda_2 = min_var + np.random.rand() * (max_var - min_var)
+    theta = np.random.rand() * np.pi  # random theta
+    noise = 0#-noise_level + np.random.rand(*k_size) * noise_level * 2
+
+    # Set COV matrix using Lambdas and Theta
+    LAMBDA = np.diag([lambda_1, lambda_2])
+    Q = np.array([[np.cos(theta), -np.sin(theta)],
+                  [np.sin(theta), np.cos(theta)]])
+    SIGMA = Q @ LAMBDA @ Q.T
+    INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
+
+    # Set expectation position (shifting kernel for aligned image)
+    MU = k_size // 2 - 0.5*(scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
+    MU = MU[None, None, :, None]
+
+    # Create meshgrid for Gaussian
+    [X,Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
+    Z = np.stack([X, Y], 2)[:, :, :, None]
+
+    # Calcualte Gaussian for every pixel of the kernel
+    ZZ = Z-MU
+    ZZ_t = ZZ.transpose(0,1,3,2)
+    raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
+
+    # shift the kernel so it will be centered
+    #raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
+
+    # Normalize the kernel and return
+    #kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
+    kernel = raw_kernel / np.sum(raw_kernel)
+    return kernel
+
+
+"""
+# --------------------------------------------
+# degradation models
+# --------------------------------------------
+"""
+
+
+def bicubic_degradation(x, sf=3):
+    '''
+    Args:
+        x: HxWxC image, [0, 1]
+        sf: down-scale factor
+    Return:
+        bicubicly downsampled LR image
+    '''
+    x = util.imresize_np(x, scale=1/sf)
+    return x
+
+
+def srmd_degradation(x, k, sf=3):
+    ''' blur + bicubic downsampling
+    Args:
+        x: HxWxC image, [0, 1]
+        k: hxw, double
+        sf: down-scale factor
+    Return:
+        downsampled LR image
+    Reference:
+        @inproceedings{zhang2018learning,
+          title={Learning a single convolutional super-resolution network for multiple degradations},
+          author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+          booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+          pages={3262--3271},
+          year={2018}
+        }
+    '''
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')  # 'nearest' | 'mirror'
+    x = bicubic_degradation(x, sf=sf)
+    return x
+
+
+def dpsr_degradation(x, k, sf=3):
+
+    ''' bicubic downsampling + blur
+    Args:
+        x: HxWxC image, [0, 1]
+        k: hxw, double
+        sf: down-scale factor
+    Return:
+        downsampled LR image
+    Reference:
+        @inproceedings{zhang2019deep,
+          title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
+          author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+          booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+          pages={1671--1681},
+          year={2019}
+        }
+    '''
+    x = bicubic_degradation(x, sf=sf)
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+    return x
+
+
+def classical_degradation(x, k, sf=3):
+    ''' blur + downsampling
+
+    Args:
+        x: HxWxC image, [0, 1]/[0, 255]
+        k: hxw, double
+        sf: down-scale factor
+
+    Return:
+        downsampled LR image
+    '''
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+    #x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
+    st = 0
+    return x[st::sf, st::sf, ...]
+
+
+def modcrop_np(img, sf):
+    '''
+    Args:
+        img: numpy image, WxH or WxHxC
+        sf: scale factor
+    Return:
+        cropped image
+    '''
+    w, h = img.shape[:2]
+    im = np.copy(img)
+    return im[:w - w % sf, :h - h % sf, ...]
+
+
+'''
+# =================
+# Numpy
+# =================
+'''
+
+
+def shift_pixel(x, sf, upper_left=True):
+    """shift pixel for super-resolution with different scale factors
+    Args:
+        x: WxHxC or WxH, image or kernel
+        sf: scale factor
+        upper_left: shift direction
+    """
+    h, w = x.shape[:2]
+    shift = (sf-1)*0.5
+    xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
+    if upper_left:
+        x1 = xv + shift
+        y1 = yv + shift
+    else:
+        x1 = xv - shift
+        y1 = yv - shift
+
+    x1 = np.clip(x1, 0, w-1)
+    y1 = np.clip(y1, 0, h-1)
+
+    if x.ndim == 2:
+        x = interp2d(xv, yv, x)(x1, y1)
+    if x.ndim == 3:
+        for i in range(x.shape[-1]):
+            x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
+
+    return x
+
+
+'''
+# =================
+# pytorch
+# =================
+'''
+
+
+def splits(a, sf):
+    '''
+    a: tensor NxCxWxHx2
+    sf: scale factor
+    out: tensor NxCx(W/sf)x(H/sf)x2x(sf^2)
+    '''
+    b = torch.stack(torch.chunk(a, sf, dim=2), dim=5)
+    b = torch.cat(torch.chunk(b, sf, dim=3), dim=5)
+    return b
+
+
+def c2c(x):
+    return torch.from_numpy(np.stack([np.float32(x.real), np.float32(x.imag)], axis=-1))
+
+
+def r2c(x):
+    return torch.stack([x, torch.zeros_like(x)], -1)
+
+
+def cdiv(x, y):
+    a, b = x[..., 0], x[..., 1]
+    c, d = y[..., 0], y[..., 1]
+    cd2 = c**2 + d**2
+    return torch.stack([(a*c+b*d)/cd2, (b*c-a*d)/cd2], -1)
+
+
+def csum(x, y):
+    return torch.stack([x[..., 0] + y, x[..., 1]], -1)
+
+
+def cabs(x):
+    return torch.pow(x[..., 0]**2+x[..., 1]**2, 0.5)
+
+
+def cmul(t1, t2):
+    '''
+    complex multiplication
+    t1: NxCxHxWx2
+    output: NxCxHxWx2
+    '''
+    real1, imag1 = t1[..., 0], t1[..., 1]
+    real2, imag2 = t2[..., 0], t2[..., 1]
+    return torch.stack([real1 * real2 - imag1 * imag2, real1 * imag2 + imag1 * real2], dim=-1)
+
+
+def cconj(t, inplace=False):
+    '''
+    # complex's conjugation
+    t: NxCxHxWx2
+    output: NxCxHxWx2
+    '''
+    c = t.clone() if not inplace else t
+    c[..., 1] *= -1
+    return c
+
+
+def rfft(t):
+    return torch.rfft(t, 2, onesided=False)
+
+
+def irfft(t):
+    return torch.irfft(t, 2, onesided=False)
+
+
+def fft(t):
+    return torch.fft(t, 2)
+
+
+def ifft(t):
+    return torch.ifft(t, 2)
+
+
+def p2o(psf, shape):
+    '''
+    Args:
+        psf: NxCxhxw
+        shape: [H,W]
+
+    Returns:
+        otf: NxCxHxWx2
+    '''
+    otf = torch.zeros(psf.shape[:-2] + shape).type_as(psf)
+    otf[...,:psf.shape[2],:psf.shape[3]].copy_(psf)
+    for axis, axis_size in enumerate(psf.shape[2:]):
+        otf = torch.roll(otf, -int(axis_size / 2), dims=axis+2)
+    otf = torch.rfft(otf, 2, onesided=False)
+    n_ops = torch.sum(torch.tensor(psf.shape).type_as(psf) * torch.log2(torch.tensor(psf.shape).type_as(psf)))
+    otf[...,1][torch.abs(otf[...,1])<n_ops*2.22e-16] = torch.tensor(0).type_as(psf)
+    return otf
+
+
+'''
+# =================
+PyTorch
+# =================
+'''
+
+def INVLS_pytorch(FB, FBC, F2B, FR, tau, sf=2):
+    '''
+    FB: NxCxWxHx2
+    F2B: NxCxWxHx2
+
+    x1 = FB.*FR;
+    FBR = BlockMM(nr,nc,Nb,m,x1);
+    invW = BlockMM(nr,nc,Nb,m,F2B);
+    invWBR = FBR./(invW + tau*Nb);
+    fun = @(block_struct) block_struct.data.*invWBR;
+    FCBinvWBR = blockproc(FBC,[nr,nc],fun);
+    FX = (FR-FCBinvWBR)/tau;
+    Xest = real(ifft2(FX));
+    '''
+    x1 = cmul(FB, FR)
+    FBR = torch.mean(splits(x1, sf), dim=-1, keepdim=False)
+    invW = torch.mean(splits(F2B, sf), dim=-1, keepdim=False)
+    invWBR = cdiv(FBR, csum(invW, tau))
+    FCBinvWBR = cmul(FBC, invWBR.repeat(1,1,sf,sf,1))
+    FX = (FR-FCBinvWBR)/tau
+    Xest = torch.irfft(FX, 2, onesided=False)
+    return Xest
+
+
+def real2complex(x):
+    return torch.stack([x, torch.zeros_like(x)], -1)
+
+
+def modcrop(img, sf):
+    '''
+    img: tensor image, NxCxWxH or CxWxH or WxH
+    sf: scale factor
+    '''
+    w, h = img.shape[-2:]
+    im = img.clone()
+    return im[..., :w - w % sf, :h - h % sf]
+
+
+def upsample(x, sf=3, center=False):
+    '''
+    x: tensor image, NxCxWxH
+    '''
+    st = (sf-1)//2 if center else 0
+    z = torch.zeros((x.shape[0], x.shape[1], x.shape[2]*sf, x.shape[3]*sf)).type_as(x)
+    z[..., st::sf, st::sf].copy_(x)
+    return z
+
+
+def downsample(x, sf=3, center=False):
+    st = (sf-1)//2 if center else 0
+    return x[..., st::sf, st::sf]
+
+
+def circular_pad(x, pad):
+    '''
+    # x[N, 1, W, H] -> x[N, 1, W + 2 pad, H + 2 pad] (pariodic padding)
+    '''
+    x = torch.cat([x, x[:, :, 0:pad, :]], dim=2)
+    x = torch.cat([x, x[:, :, :, 0:pad]], dim=3)
+    x = torch.cat([x[:, :, -2 * pad:-pad, :], x], dim=2)
+    x = torch.cat([x[:, :, :, -2 * pad:-pad], x], dim=3)
+    return x
+
+
+def pad_circular(input, padding):
+    # type: (Tensor, List[int]) -> Tensor
+    """
+    Arguments
+    :param input: tensor of shape :math:`(N, C_{\text{in}}, H, [W, D]))`
+    :param padding: (tuple): m-elem tuple where m is the degree of convolution
+    Returns
+    :return: tensor of shape :math:`(N, C_{\text{in}}, [D + 2 * padding[0],
+                                     H + 2 * padding[1]], W + 2 * padding[2]))`
+    """
+    offset = 3
+    for dimension in range(input.dim() - offset + 1):
+        input = dim_pad_circular(input, padding[dimension], dimension + offset)
+    return input
+
+
+def dim_pad_circular(input, padding, dimension):
+    # type: (Tensor, int, int) -> Tensor
+    input = torch.cat([input, input[[slice(None)] * (dimension - 1) +
+                      [slice(0, padding)]]], dim=dimension - 1)
+    input = torch.cat([input[[slice(None)] * (dimension - 1) +
+                      [slice(-2 * padding, -padding)]], input], dim=dimension - 1)
+    return input
+
+
+def imfilter(x, k):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+    '''
+    x = pad_circular(x, padding=((k.shape[-2]-1)//2, (k.shape[-1]-1)//2))
+    x = torch.nn.functional.conv2d(x, k, groups=x.shape[1])
+    return x
+
+
+def G(x, k, sf=3, center=False):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+    sf: scale factor
+    center: the first one or the moddle one
+
+    Matlab function:
+    tmp = imfilter(x,h,'circular');
+    y = downsample2(tmp,K);
+    '''
+    x = downsample(imfilter(x, k), sf=sf, center=center)
+    return x
+
+
+def Gt(x, k, sf=3, center=False):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+    sf: scale factor
+    center: the first one or the moddle one
+
+    Matlab function:
+    tmp = upsample2(x,K);
+    y = imfilter(tmp,h,'circular');
+    '''
+    x = imfilter(upsample(x, sf=sf, center=center), k)
+    return x
+
+
+def interpolation_down(x, sf, center=False):
+    mask = torch.zeros_like(x)
+    if center:
+        start = torch.tensor((sf-1)//2)
+        mask[..., start::sf, start::sf] = torch.tensor(1).type_as(x)
+        LR = x[..., start::sf, start::sf]
+    else:
+        mask[..., ::sf, ::sf] = torch.tensor(1).type_as(x)
+        LR = x[..., ::sf, ::sf]
+    y = x.mul(mask)
+
+    return LR, y, mask
+
+
+'''
+# =================
+Numpy
+# =================
+'''
+
+
+def blockproc(im, blocksize, fun):
+    xblocks = np.split(im, range(blocksize[0], im.shape[0], blocksize[0]), axis=0)
+    xblocks_proc = []
+    for xb in xblocks:
+        yblocks = np.split(xb, range(blocksize[1], im.shape[1], blocksize[1]), axis=1)
+        yblocks_proc = []
+        for yb in yblocks:
+            yb_proc = fun(yb)
+            yblocks_proc.append(yb_proc)
+        xblocks_proc.append(np.concatenate(yblocks_proc, axis=1))
+
+    proc = np.concatenate(xblocks_proc, axis=0)
+
+    return proc
+
+
+def fun_reshape(a):
+    return np.reshape(a, (-1,1,a.shape[-1]), order='F')
+
+
+def fun_mul(a, b):
+    return a*b
+
+
+def BlockMM(nr, nc, Nb, m, x1):
+    '''
+    myfun = @(block_struct) reshape(block_struct.data,m,1);
+    x1 = blockproc(x1,[nr nc],myfun);
+    x1 = reshape(x1,m,Nb);
+    x1 = sum(x1,2);
+    x = reshape(x1,nr,nc);
+    '''
+    fun = fun_reshape
+    x1 = blockproc(x1, blocksize=(nr, nc), fun=fun)
+    x1 = np.reshape(x1, (m, Nb, x1.shape[-1]), order='F')
+    x1 = np.sum(x1, 1)
+    x = np.reshape(x1, (nr, nc, x1.shape[-1]), order='F')
+    return x
+
+
+def INVLS(FB, FBC, F2B, FR, tau, Nb, nr, nc, m):
+    '''
+    x1 = FB.*FR;
+    FBR = BlockMM(nr,nc,Nb,m,x1);
+    invW = BlockMM(nr,nc,Nb,m,F2B);
+    invWBR = FBR./(invW + tau*Nb);
+    fun = @(block_struct) block_struct.data.*invWBR;
+    FCBinvWBR = blockproc(FBC,[nr,nc],fun);
+    FX = (FR-FCBinvWBR)/tau;
+    Xest = real(ifft2(FX));
+    '''
+    x1 = FB*FR
+    FBR = BlockMM(nr, nc, Nb, m, x1)
+    invW = BlockMM(nr, nc, Nb, m, F2B)
+    invWBR = FBR/(invW + tau*Nb)
+    FCBinvWBR = blockproc(FBC, [nr, nc], lambda im: fun_mul(im, invWBR))
+    FX = (FR-FCBinvWBR)/tau
+    Xest = np.real(np.fft.ifft2(FX, axes=(0, 1)))
+    return Xest
+
+
+def psf2otf(psf, shape=None):
+    """
+    Convert point-spread function to optical transfer function.
+    Compute the Fast Fourier Transform (FFT) of the point-spread
+    function (PSF) array and creates the optical transfer function (OTF)
+    array that is not influenced by the PSF off-centering.
+    By default, the OTF array is the same size as the PSF array.
+    To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
+    post-pads the PSF array (down or to the right) with zeros to match
+    dimensions specified in OUTSIZE, then circularly shifts the values of
+    the PSF array up (or to the left) until the central pixel reaches (1,1)
+    position.
+    Parameters
+    ----------
+    psf : `numpy.ndarray`
+        PSF array
+    shape : int
+        Output shape of the OTF array
+    Returns
+    -------
+    otf : `numpy.ndarray`
+        OTF array
+    Notes
+    -----
+    Adapted from MATLAB psf2otf function
+    """
+    if type(shape) == type(None):
+        shape = psf.shape
+    shape = np.array(shape)
+    if np.all(psf == 0):
+        # return np.zeros_like(psf)
+        return np.zeros(shape)
+    if len(psf.shape) == 1:
+        psf = psf.reshape((1, psf.shape[0]))
+    inshape = psf.shape
+    psf = zero_pad(psf, shape, position='corner')
+    for axis, axis_size in enumerate(inshape):
+        psf = np.roll(psf, -int(axis_size / 2), axis=axis)
+    # Compute the OTF
+    otf = np.fft.fft2(psf, axes=(0, 1))
+    # Estimate the rough number of operations involved in the FFT
+    # and discard the PSF imaginary part if within roundoff error
+    # roundoff error  = machine epsilon = sys.float_info.epsilon
+    # or np.finfo().eps
+    n_ops = np.sum(psf.size * np.log2(psf.shape))
+    otf = np.real_if_close(otf, tol=n_ops)
+    return otf
+
+
+def zero_pad(image, shape, position='corner'):
+    """
+    Extends image to a certain size with zeros
+    Parameters
+    ----------
+    image: real 2d `numpy.ndarray`
+        Input image
+    shape: tuple of int
+        Desired output shape of the image
+    position : str, optional
+        The position of the input image in the output one:
+            * 'corner'
+                top-left corner (default)
+            * 'center'
+                centered
+    Returns
+    -------
+    padded_img: real `numpy.ndarray`
+        The zero-padded image
+    """
+    shape = np.asarray(shape, dtype=int)
+    imshape = np.asarray(image.shape, dtype=int)
+    if np.alltrue(imshape == shape):
+        return image
+    if np.any(shape <= 0):
+        raise ValueError("ZERO_PAD: null or negative shape given")
+    dshape = shape - imshape
+    if np.any(dshape < 0):
+        raise ValueError("ZERO_PAD: target size smaller than source one")
+    pad_img = np.zeros(shape, dtype=image.dtype)
+    idx, idy = np.indices(imshape)
+    if position == 'center':
+        if np.any(dshape % 2 != 0):
+            raise ValueError("ZERO_PAD: source and target shapes "
+                             "have different parity.")
+        offx, offy = dshape // 2
+    else:
+        offx, offy = (0, 0)
+    pad_img[idx + offx, idy + offy] = image
+    return pad_img
+
+
+def upsample_np(x, sf=3, center=False):
+    st = (sf-1)//2 if center else 0
+    z = np.zeros((x.shape[0]*sf, x.shape[1]*sf, x.shape[2]))
+    z[st::sf, st::sf, ...] = x
+    return z
+
+
+def downsample_np(x, sf=3, center=False):
+    st = (sf-1)//2 if center else 0
+    return x[st::sf, st::sf, ...]
+
+
+def imfilter_np(x, k):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+    '''
+    x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+    return x
+
+
+def G_np(x, k, sf=3, center=False):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+
+    Matlab function:
+    tmp = imfilter(x,h,'circular');
+    y = downsample2(tmp,K);
+    '''
+    x = downsample_np(imfilter_np(x, k), sf=sf, center=center)
+    return x
+
+
+def Gt_np(x, k, sf=3, center=False):
+    '''
+    x: image, NxcxHxW
+    k: kernel, cx1xhxw
+
+    Matlab function:
+    tmp = upsample2(x,K);
+    y = imfilter(tmp,h,'circular');
+    '''
+    x = imfilter_np(upsample_np(x, sf=sf, center=center), k)
+    return x
+
+
+if __name__ == '__main__':
+    img = util.imread_uint('test.bmp', 3)
+
+    img = util.uint2single(img)
+    k = anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6)
+    util.imshow(k*10)
+
+
+    for sf in [2, 3, 4]:
+
+        # modcrop
+        img = modcrop_np(img, sf=sf)
+
+        # 1) bicubic degradation
+        img_b = bicubic_degradation(img, sf=sf)
+        print(img_b.shape)
+
+        # 2) srmd degradation
+        img_s = srmd_degradation(img, k, sf=sf)
+        print(img_s.shape)
+
+        # 3) dpsr degradation
+        img_d = dpsr_degradation(img, k, sf=sf)
+        print(img_d.shape)
+
+        # 4) classical degradation
+        img_d = classical_degradation(img, k, sf=sf)
+        print(img_d.shape)
+
+    k = anisotropic_Gaussian(ksize=7, theta=0.25*np.pi, l1=0.01, l2=0.01)
+    #print(k)
+#    util.imshow(k*10)
+
+    k = shifted_anisotropic_Gaussian(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.8, max_var=10.8, noise_level=0.0)
+#    util.imshow(k*10)
+
+
+    # PCA
+#    pca_matrix = cal_pca_matrix(ksize=15, l_max=10.0, dim_pca=15, num_samples=12500)
+#    print(pca_matrix.shape)
+#    show_pca(pca_matrix)
+    # run utils/utils_sisr.py
+    # run utils_sisr.py
+    
+    
+    
+    
+    
+    
+    

core/data/deg_kair_utils/utils_video.py

@@ -0,0 +1,493 @@
+import os
+import cv2
+import numpy as np
+import torch
+import random
+from os import path as osp
+from torch.nn import functional as F
+from abc import ABCMeta, abstractmethod
+
+
+def scandir(dir_path, suffix=None, recursive=False, full_path=False):
+    """Scan a directory to find the interested files.
+
+    Args:
+        dir_path (str): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        full_path (bool, optional): If set to True, include the dir_path.
+            Default: False.
+
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith('.') and entry.is_file():
+                if full_path:
+                    return_path = entry.path
+                else:
+                    return_path = osp.relpath(entry.path, root)
+
+                if suffix is None:
+                    yield return_path
+                elif return_path.endswith(suffix):
+                    yield return_path
+            else:
+                if recursive:
+                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
+                else:
+                    continue
+
+    return _scandir(dir_path, suffix=suffix, recursive=recursive)
+
+
+def read_img_seq(path, require_mod_crop=False, scale=1, return_imgname=False):
+    """Read a sequence of images from a given folder path.
+
+    Args:
+        path (list[str] | str): List of image paths or image folder path.
+        require_mod_crop (bool): Require mod crop for each image.
+            Default: False.
+        scale (int): Scale factor for mod_crop. Default: 1.
+        return_imgname(bool): Whether return image names. Default False.
+
+    Returns:
+        Tensor: size (t, c, h, w), RGB, [0, 1].
+        list[str]: Returned image name list.
+    """
+    if isinstance(path, list):
+        img_paths = path
+    else:
+        img_paths = sorted(list(scandir(path, full_path=True)))
+    imgs = [cv2.imread(v).astype(np.float32) / 255. for v in img_paths]
+
+    if require_mod_crop:
+        imgs = [mod_crop(img, scale) for img in imgs]
+    imgs = img2tensor(imgs, bgr2rgb=True, float32=True)
+    imgs = torch.stack(imgs, dim=0)
+
+    if return_imgname:
+        imgnames = [osp.splitext(osp.basename(path))[0] for path in img_paths]
+        return imgs, imgnames
+    else:
+        return imgs
+
+
+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 augment(imgs, hflip=True, rotation=True, flows=None, return_status=False):
+    """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees).
+
+    We use vertical flip and transpose for rotation implementation.
+    All the images in the list use the same augmentation.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Images to be augmented. If the input
+            is an ndarray, it will be transformed to a list.
+        hflip (bool): Horizontal flip. Default: True.
+        rotation (bool): Ratotation. Default: True.
+        flows (list[ndarray]: Flows to be augmented. If the input is an
+            ndarray, it will be transformed to a list.
+            Dimension is (h, w, 2). Default: None.
+        return_status (bool): Return the status of flip and rotation.
+            Default: False.
+
+    Returns:
+        list[ndarray] | ndarray: Augmented images and flows. If returned
+            results only have one element, just return ndarray.
+
+    """
+    hflip = hflip and random.random() < 0.5
+    vflip = rotation and random.random() < 0.5
+    rot90 = rotation and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:  # horizontal
+            cv2.flip(img, 1, img)
+        if vflip:  # vertical
+            cv2.flip(img, 0, img)
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    def _augment_flow(flow):
+        if hflip:  # horizontal
+            cv2.flip(flow, 1, flow)
+            flow[:, :, 0] *= -1
+        if vflip:  # vertical
+            cv2.flip(flow, 0, flow)
+            flow[:, :, 1] *= -1
+        if rot90:
+            flow = flow.transpose(1, 0, 2)
+            flow = flow[:, :, [1, 0]]
+        return flow
+
+    if not isinstance(imgs, list):
+        imgs = [imgs]
+    imgs = [_augment(img) for img in imgs]
+    if len(imgs) == 1:
+        imgs = imgs[0]
+
+    if flows is not None:
+        if not isinstance(flows, list):
+            flows = [flows]
+        flows = [_augment_flow(flow) for flow in flows]
+        if len(flows) == 1:
+            flows = flows[0]
+        return imgs, flows
+    else:
+        if return_status:
+            return imgs, (hflip, vflip, rot90)
+        else:
+            return imgs
+
+
+def paired_random_crop(img_gts, img_lqs, gt_patch_size, scale, gt_path=None):
+    """Paired random crop. Support Numpy array and Tensor inputs.
+
+    It crops lists of lq and gt images with corresponding locations.
+
+    Args:
+        img_gts (list[ndarray] | ndarray | list[Tensor] | Tensor): GT images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        img_lqs (list[ndarray] | ndarray): LQ images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        gt_patch_size (int): GT patch size.
+        scale (int): Scale factor.
+        gt_path (str): Path to ground-truth. Default: None.
+
+    Returns:
+        list[ndarray] | ndarray: GT images and LQ images. If returned results
+            only have one element, just return ndarray.
+    """
+
+    if not isinstance(img_gts, list):
+        img_gts = [img_gts]
+    if not isinstance(img_lqs, list):
+        img_lqs = [img_lqs]
+
+    # determine input type: Numpy array or Tensor
+    input_type = 'Tensor' if torch.is_tensor(img_gts[0]) else 'Numpy'
+
+    if input_type == 'Tensor':
+        h_lq, w_lq = img_lqs[0].size()[-2:]
+        h_gt, w_gt = img_gts[0].size()[-2:]
+    else:
+        h_lq, w_lq = img_lqs[0].shape[0:2]
+        h_gt, w_gt = img_gts[0].shape[0:2]
+    lq_patch_size = gt_patch_size // scale
+
+    if h_gt != h_lq * scale or w_gt != w_lq * scale:
+        raise ValueError(f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
+                         f'multiplication of LQ ({h_lq}, {w_lq}).')
+    if h_lq < lq_patch_size or w_lq < lq_patch_size:
+        raise ValueError(f'LQ ({h_lq}, {w_lq}) is smaller than patch size '
+                         f'({lq_patch_size}, {lq_patch_size}). '
+                         f'Please remove {gt_path}.')
+
+    # randomly choose top and left coordinates for lq patch
+    top = random.randint(0, h_lq - lq_patch_size)
+    left = random.randint(0, w_lq - lq_patch_size)
+
+    # crop lq patch
+    if input_type == 'Tensor':
+        img_lqs = [v[:, :, top:top + lq_patch_size, left:left + lq_patch_size] for v in img_lqs]
+    else:
+        img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]
+
+    # crop corresponding gt patch
+    top_gt, left_gt = int(top * scale), int(left * scale)
+    if input_type == 'Tensor':
+        img_gts = [v[:, :, top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size] for v in img_gts]
+    else:
+        img_gts = [v[top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size, ...] for v in img_gts]
+    if len(img_gts) == 1:
+        img_gts = img_gts[0]
+    if len(img_lqs) == 1:
+        img_lqs = img_lqs[0]
+    return img_gts, img_lqs
+
+
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py  # noqa: E501
+class BaseStorageBackend(metaclass=ABCMeta):
+    """Abstract class of storage backends.
+
+    All backends need to implement two apis: ``get()`` and ``get_text()``.
+    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file
+    as texts.
+    """
+
+    @abstractmethod
+    def get(self, filepath):
+        pass
+
+    @abstractmethod
+    def get_text(self, filepath):
+        pass
+
+
+class MemcachedBackend(BaseStorageBackend):
+    """Memcached storage backend.
+
+    Attributes:
+        server_list_cfg (str): Config file for memcached server list.
+        client_cfg (str): Config file for memcached client.
+        sys_path (str | None): Additional path to be appended to `sys.path`.
+            Default: None.
+    """
+
+    def __init__(self, server_list_cfg, client_cfg, sys_path=None):
+        if sys_path is not None:
+            import sys
+            sys.path.append(sys_path)
+        try:
+            import mc
+        except ImportError:
+            raise ImportError('Please install memcached to enable MemcachedBackend.')
+
+        self.server_list_cfg = server_list_cfg
+        self.client_cfg = client_cfg
+        self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)
+        # mc.pyvector servers as a point which points to a memory cache
+        self._mc_buffer = mc.pyvector()
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        import mc
+        self._client.Get(filepath, self._mc_buffer)
+        value_buf = mc.ConvertBuffer(self._mc_buffer)
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class HardDiskBackend(BaseStorageBackend):
+    """Raw hard disks storage backend."""
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'rb') as f:
+            value_buf = f.read()
+        return value_buf
+
+    def get_text(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'r') as f:
+            value_buf = f.read()
+        return value_buf
+
+
+class LmdbBackend(BaseStorageBackend):
+    """Lmdb storage backend.
+
+    Args:
+        db_paths (str | list[str]): Lmdb database paths.
+        client_keys (str | list[str]): Lmdb client keys. Default: 'default'.
+        readonly (bool, optional): Lmdb environment parameter. If True,
+            disallow any write operations. Default: True.
+        lock (bool, optional): Lmdb environment parameter. If False, when
+            concurrent access occurs, do not lock the database. Default: False.
+        readahead (bool, optional): Lmdb environment parameter. If False,
+            disable the OS filesystem readahead mechanism, which may improve
+            random read performance when a database is larger than RAM.
+            Default: False.
+
+    Attributes:
+        db_paths (list): Lmdb database path.
+        _client (list): A list of several lmdb envs.
+    """
+
+    def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):
+        try:
+            import lmdb
+        except ImportError:
+            raise ImportError('Please install lmdb to enable LmdbBackend.')
+
+        if isinstance(client_keys, str):
+            client_keys = [client_keys]
+
+        if isinstance(db_paths, list):
+            self.db_paths = [str(v) for v in db_paths]
+        elif isinstance(db_paths, str):
+            self.db_paths = [str(db_paths)]
+        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
+                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')
+
+        self._client = {}
+        for client, path in zip(client_keys, self.db_paths):
+            self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)
+
+    def get(self, filepath, client_key):
+        """Get values according to the filepath from one lmdb named client_key.
+
+        Args:
+            filepath (str | obj:`Path`): Here, filepath is the lmdb key.
+            client_key (str): Used for distinguishing different lmdb envs.
+        """
+        filepath = str(filepath)
+        assert client_key in self._client, (f'client_key {client_key} is not ' 'in lmdb clients.')
+        client = self._client[client_key]
+        with client.begin(write=False) as txn:
+            value_buf = txn.get(filepath.encode('ascii'))
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class FileClient(object):
+    """A general file client to access files in different backend.
+
+    The client loads a file or text in a specified backend from its path
+    and return it as a binary file. it can also register other backend
+    accessor with a given name and backend class.
+
+    Attributes:
+        backend (str): The storage backend type. Options are "disk",
+            "memcached" and "lmdb".
+        client (:obj:`BaseStorageBackend`): The backend object.
+    """
+
+    _backends = {
+        'disk': HardDiskBackend,
+        'memcached': MemcachedBackend,
+        'lmdb': LmdbBackend,
+    }
+
+    def __init__(self, backend='disk', **kwargs):
+        if backend not in self._backends:
+            raise ValueError(f'Backend {backend} is not supported. Currently supported ones'
+                             f' are {list(self._backends.keys())}')
+        self.backend = backend
+        self.client = self._backends[backend](**kwargs)
+
+    def get(self, filepath, client_key='default'):
+        # client_key is used only for lmdb, where different fileclients have
+        # different lmdb environments.
+        if self.backend == 'lmdb':
+            return self.client.get(filepath, client_key)
+        else:
+            return self.client.get(filepath)
+
+    def get_text(self, filepath):
+        return self.client.get_text(filepath)
+
+
+def imfrombytes(content, flag='color', float32=False):
+    """Read an image from bytes.
+
+    Args:
+        content (bytes): Image bytes got from files or other streams.
+        flag (str): Flags specifying the color type of a loaded image,
+            candidates are `color`, `grayscale` and `unchanged`.
+        float32 (bool): Whether to change to float32., If True, will also norm
+            to [0, 1]. Default: False.
+
+    Returns:
+        ndarray: Loaded image array.
+    """
+    img_np = np.frombuffer(content, np.uint8)
+    imread_flags = {'color': cv2.IMREAD_COLOR, 'grayscale': cv2.IMREAD_GRAYSCALE, 'unchanged': cv2.IMREAD_UNCHANGED}
+    img = cv2.imdecode(img_np, imread_flags[flag])
+    if float32:
+        img = img.astype(np.float32) / 255.
+    return img
+

core/data/deg_kair_utils/utils_videoio.py

@@ -0,0 +1,555 @@
+import os
+import cv2
+import numpy as np
+import torch
+import random
+from os import path as osp
+from torchvision.utils import make_grid
+import sys
+from pathlib import Path
+import six
+from collections import OrderedDict
+import math
+import glob
+import av
+import io
+from cv2 import (CAP_PROP_FOURCC, CAP_PROP_FPS, CAP_PROP_FRAME_COUNT,
+                 CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH,
+                 CAP_PROP_POS_FRAMES, VideoWriter_fourcc)
+
+if sys.version_info <= (3, 3):
+    FileNotFoundError = IOError
+else:
+    FileNotFoundError = FileNotFoundError
+
+
+def is_str(x):
+    """Whether the input is an string instance."""
+    return isinstance(x, six.string_types)
+
+
+def is_filepath(x):
+    return is_str(x) or isinstance(x, Path)
+
+
+def fopen(filepath, *args, **kwargs):
+    if is_str(filepath):
+        return open(filepath, *args, **kwargs)
+    elif isinstance(filepath, Path):
+        return filepath.open(*args, **kwargs)
+    raise ValueError('`filepath` should be a string or a Path')
+
+
+def check_file_exist(filename, msg_tmpl='file "{}" does not exist'):
+    if not osp.isfile(filename):
+        raise FileNotFoundError(msg_tmpl.format(filename))
+
+
+def mkdir_or_exist(dir_name, mode=0o777):
+    if dir_name == '':
+        return
+    dir_name = osp.expanduser(dir_name)
+    os.makedirs(dir_name, mode=mode, exist_ok=True)
+
+
+def symlink(src, dst, overwrite=True, **kwargs):
+    if os.path.lexists(dst) and overwrite:
+        os.remove(dst)
+    os.symlink(src, dst, **kwargs)
+
+
+def scandir(dir_path, suffix=None, recursive=False, case_sensitive=True):
+    """Scan a directory to find the interested files.
+    Args:
+        dir_path (str | :obj:`Path`): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        case_sensitive (bool, optional) : If set to False, ignore the case of
+            suffix. Default: True.
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+    if isinstance(dir_path, (str, Path)):
+        dir_path = str(dir_path)
+    else:
+        raise TypeError('"dir_path" must be a string or Path object')
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    if suffix is not None and not case_sensitive:
+        suffix = suffix.lower() if isinstance(suffix, str) else tuple(
+            item.lower() for item in suffix)
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive, case_sensitive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith('.') and entry.is_file():
+                rel_path = osp.relpath(entry.path, root)
+                _rel_path = rel_path if case_sensitive else rel_path.lower()
+                if suffix is None or _rel_path.endswith(suffix):
+                    yield rel_path
+            elif recursive and os.path.isdir(entry.path):
+                # scan recursively if entry.path is a directory
+                yield from _scandir(entry.path, suffix, recursive,
+                                    case_sensitive)
+
+    return _scandir(dir_path, suffix, recursive, case_sensitive)
+
+
+class Cache:
+
+    def __init__(self, capacity):
+        self._cache = OrderedDict()
+        self._capacity = int(capacity)
+        if capacity <= 0:
+            raise ValueError('capacity must be a positive integer')
+
+    @property
+    def capacity(self):
+        return self._capacity
+
+    @property
+    def size(self):
+        return len(self._cache)
+
+    def put(self, key, val):
+        if key in self._cache:
+            return
+        if len(self._cache) >= self.capacity:
+            self._cache.popitem(last=False)
+        self._cache[key] = val
+
+    def get(self, key, default=None):
+        val = self._cache[key] if key in self._cache else default
+        return val
+
+
+class VideoReader:
+    """Video class with similar usage to a list object.
+
+    This video warpper class provides convenient apis to access frames.
+    There exists an issue of OpenCV's VideoCapture class that jumping to a
+    certain frame may be inaccurate. It is fixed in this class by checking
+    the position after jumping each time.
+    Cache is used when decoding videos. So if the same frame is visited for
+    the second time, there is no need to decode again if it is stored in the
+    cache.
+
+    """
+
+    def __init__(self, filename, cache_capacity=10):
+        # Check whether the video path is a url
+        if not filename.startswith(('https://', 'http://')):
+            check_file_exist(filename, 'Video file not found: ' + filename)
+        self._vcap = cv2.VideoCapture(filename)
+        assert cache_capacity > 0
+        self._cache = Cache(cache_capacity)
+        self._position = 0
+        # get basic info
+        self._width = int(self._vcap.get(CAP_PROP_FRAME_WIDTH))
+        self._height = int(self._vcap.get(CAP_PROP_FRAME_HEIGHT))
+        self._fps = self._vcap.get(CAP_PROP_FPS)
+        self._frame_cnt = int(self._vcap.get(CAP_PROP_FRAME_COUNT))
+        self._fourcc = self._vcap.get(CAP_PROP_FOURCC)
+
+    @property
+    def vcap(self):
+        """:obj:`cv2.VideoCapture`: The raw VideoCapture object."""
+        return self._vcap
+
+    @property
+    def opened(self):
+        """bool: Indicate whether the video is opened."""
+        return self._vcap.isOpened()
+
+    @property
+    def width(self):
+        """int: Width of video frames."""
+        return self._width
+
+    @property
+    def height(self):
+        """int: Height of video frames."""
+        return self._height
+
+    @property
+    def resolution(self):
+        """tuple: Video resolution (width, height)."""
+        return (self._width, self._height)
+
+    @property
+    def fps(self):
+        """float: FPS of the video."""
+        return self._fps
+
+    @property
+    def frame_cnt(self):
+        """int: Total frames of the video."""
+        return self._frame_cnt
+
+    @property
+    def fourcc(self):
+        """str: "Four character code" of the video."""
+        return self._fourcc
+
+    @property
+    def position(self):
+        """int: Current cursor position, indicating frame decoded."""
+        return self._position
+
+    def _get_real_position(self):
+        return int(round(self._vcap.get(CAP_PROP_POS_FRAMES)))
+
+    def _set_real_position(self, frame_id):
+        self._vcap.set(CAP_PROP_POS_FRAMES, frame_id)
+        pos = self._get_real_position()
+        for _ in range(frame_id - pos):
+            self._vcap.read()
+        self._position = frame_id
+
+    def read(self):
+        """Read the next frame.
+
+        If the next frame have been decoded before and in the cache, then
+        return it directly, otherwise decode, cache and return it.
+
+        Returns:
+            ndarray or None: Return the frame if successful, otherwise None.
+        """
+        # pos = self._position
+        if self._cache:
+            img = self._cache.get(self._position)
+            if img is not None:
+                ret = True
+            else:
+                if self._position != self._get_real_position():
+                    self._set_real_position(self._position)
+                ret, img = self._vcap.read()
+                if ret:
+                    self._cache.put(self._position, img)
+        else:
+            ret, img = self._vcap.read()
+        if ret:
+            self._position += 1
+        return img
+
+    def get_frame(self, frame_id):
+        """Get frame by index.
+
+        Args:
+            frame_id (int): Index of the expected frame, 0-based.
+
+        Returns:
+            ndarray or None: Return the frame if successful, otherwise None.
+        """
+        if frame_id < 0 or frame_id >= self._frame_cnt:
+            raise IndexError(
+                f'"frame_id" must be between 0 and {self._frame_cnt - 1}')
+        if frame_id == self._position:
+            return self.read()
+        if self._cache:
+            img = self._cache.get(frame_id)
+            if img is not None:
+                self._position = frame_id + 1
+                return img
+        self._set_real_position(frame_id)
+        ret, img = self._vcap.read()
+        if ret:
+            if self._cache:
+                self._cache.put(self._position, img)
+            self._position += 1
+        return img
+
+    def current_frame(self):
+        """Get the current frame (frame that is just visited).
+
+        Returns:
+            ndarray or None: If the video is fresh, return None, otherwise
+            return the frame.
+        """
+        if self._position == 0:
+            return None
+        return self._cache.get(self._position - 1)
+
+    def cvt2frames(self,
+                   frame_dir,
+                   file_start=0,
+                   filename_tmpl='{:06d}.jpg',
+                   start=0,
+                   max_num=0,
+                   show_progress=False):
+        """Convert a video to frame images.
+
+        Args:
+            frame_dir (str): Output directory to store all the frame images.
+            file_start (int): Filenames will start from the specified number.
+            filename_tmpl (str): Filename template with the index as the
+                placeholder.
+            start (int): The starting frame index.
+            max_num (int): Maximum number of frames to be written.
+            show_progress (bool): Whether to show a progress bar.
+        """
+        mkdir_or_exist(frame_dir)
+        if max_num == 0:
+            task_num = self.frame_cnt - start
+        else:
+            task_num = min(self.frame_cnt - start, max_num)
+        if task_num <= 0:
+            raise ValueError('start must be less than total frame number')
+        if start > 0:
+            self._set_real_position(start)
+
+        def write_frame(file_idx):
+            img = self.read()
+            if img is None:
+                return
+            filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
+            cv2.imwrite(filename, img)
+
+        if show_progress:
+            pass
+            #track_progress(write_frame, range(file_start,file_start + task_num))
+        else:
+            for i in range(task_num):
+                write_frame(file_start + i)
+
+    def __len__(self):
+        return self.frame_cnt
+
+    def __getitem__(self, index):
+        if isinstance(index, slice):
+            return [
+                self.get_frame(i)
+                for i in range(*index.indices(self.frame_cnt))
+            ]
+        # support negative indexing
+        if index < 0:
+            index += self.frame_cnt
+            if index < 0:
+                raise IndexError('index out of range')
+        return self.get_frame(index)
+
+    def __iter__(self):
+        self._set_real_position(0)
+        return self
+
+    def __next__(self):
+        img = self.read()
+        if img is not None:
+            return img
+        else:
+            raise StopIteration
+
+    next = __next__
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._vcap.release()
+
+
+def frames2video(frame_dir,
+                 video_file,
+                 fps=30,
+                 fourcc='XVID',
+                 filename_tmpl='{:06d}.jpg',
+                 start=0,
+                 end=0,
+                 show_progress=False):
+    """Read the frame images from a directory and join them as a video.
+
+    Args:
+        frame_dir (str): The directory containing video frames.
+        video_file (str): Output filename.
+        fps (float): FPS of the output video.
+        fourcc (str): Fourcc of the output video, this should be compatible
+            with the output file type.
+        filename_tmpl (str): Filename template with the index as the variable.
+        start (int): Starting frame index.
+        end (int): Ending frame index.
+        show_progress (bool): Whether to show a progress bar.
+    """
+    if end == 0:
+        ext = filename_tmpl.split('.')[-1]
+        end = len([name for name in scandir(frame_dir, ext)])
+    first_file = osp.join(frame_dir, filename_tmpl.format(start))
+    check_file_exist(first_file, 'The start frame not found: ' + first_file)
+    img = cv2.imread(first_file)
+    height, width = img.shape[:2]
+    resolution = (width, height)
+    vwriter = cv2.VideoWriter(video_file, VideoWriter_fourcc(*fourcc), fps,
+                              resolution)
+
+    def write_frame(file_idx):
+        filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
+        img = cv2.imread(filename)
+        vwriter.write(img)
+
+    if show_progress:
+        pass
+        # track_progress(write_frame, range(start, end))
+    else:
+        for i in range(start, end):
+            write_frame(i)
+    vwriter.release()
+
+
+def video2images(video_path, output_dir):
+    vidcap = cv2.VideoCapture(video_path)
+    in_fps = vidcap.get(cv2.CAP_PROP_FPS)
+    print('video fps:', in_fps)
+    if not os.path.isdir(output_dir):
+        os.makedirs(output_dir)
+    loaded, frame = vidcap.read()
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    print(f'number of total frames is: {total_frames:06}')
+    for i_frame in range(total_frames):
+        if i_frame % 100 == 0:
+            print(f'{i_frame:06} / {total_frames:06}')
+        frame_name = os.path.join(output_dir, f'{i_frame:06}' + '.png')
+        cv2.imwrite(frame_name, frame)
+        loaded, frame = vidcap.read()
+
+
+def images2video(image_dir, video_path, fps=24, image_ext='png'):
+    '''
+    #codec = cv2.VideoWriter_fourcc(*'XVID')
+    #codec = cv2.VideoWriter_fourcc('A','V','C','1')
+    #codec = cv2.VideoWriter_fourcc('Y','U','V','1')
+    #codec = cv2.VideoWriter_fourcc('P','I','M','1')
+    #codec = cv2.VideoWriter_fourcc('M','J','P','G')
+    codec = cv2.VideoWriter_fourcc('M','P','4','2')
+    #codec = cv2.VideoWriter_fourcc('D','I','V','3')
+    #codec =  cv2.VideoWriter_fourcc('D','I','V','X')
+    #codec = cv2.VideoWriter_fourcc('U','2','6','3')
+    #codec = cv2.VideoWriter_fourcc('I','2','6','3')
+    #codec = cv2.VideoWriter_fourcc('F','L','V','1')
+    #codec = cv2.VideoWriter_fourcc('H','2','6','4')
+    #codec = cv2.VideoWriter_fourcc('A','Y','U','V')
+    #codec = cv2.VideoWriter_fourcc('I','U','Y','V')
+    编码器常用的几种：
+    cv2.VideoWriter_fourcc("I", "4", "2", "0") 
+        压缩的yuv颜色编码器，4:2:0色彩度子采样 兼容性好，产生很大的视频 avi
+    cv2.VideoWriter_fourcc("P", I", "M", "1")
+        采用mpeg-1编码，文件为avi
+    cv2.VideoWriter_fourcc("X", "V", "T", "D")
+        采用mpeg-4编码，得到视频大小平均 拓展名avi
+    cv2.VideoWriter_fourcc("T", "H", "E", "O")
+        Ogg Vorbis， 拓展名为ogv
+    cv2.VideoWriter_fourcc("F", "L", "V", "1")
+        FLASH视频，拓展名为.flv
+    '''
+    image_files = sorted(glob.glob(os.path.join(image_dir, '*.{}'.format(image_ext))))
+    print(len(image_files))
+    height, width, _ = cv2.imread(image_files[0]).shape
+    out_fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')  # cv2.VideoWriter_fourcc(*'MP4V')
+    out_video = cv2.VideoWriter(video_path, out_fourcc, fps, (width, height))
+
+    for image_file in image_files:
+        img = cv2.imread(image_file)
+        img = cv2.resize(img, (width, height), interpolation=3)
+        out_video.write(img)
+    out_video.release()
+
+
+def add_video_compression(imgs):
+    codec_type = ['libx264', 'h264', 'mpeg4']
+    codec_prob = [1 / 3., 1 / 3., 1 / 3.]
+    codec = random.choices(codec_type, codec_prob)[0]
+    # codec = 'mpeg4'
+    bitrate = [1e4, 1e5]
+    bitrate = np.random.randint(bitrate[0], bitrate[1] + 1)
+
+    buf = io.BytesIO()
+    with av.open(buf, 'w', 'mp4') as container:
+        stream = container.add_stream(codec, rate=1)
+        stream.height = imgs[0].shape[0]
+        stream.width = imgs[0].shape[1]
+        stream.pix_fmt = 'yuv420p'
+        stream.bit_rate = bitrate
+        
+        for img in imgs:
+            img = np.uint8((img.clip(0, 1)*255.).round())
+            frame = av.VideoFrame.from_ndarray(img, format='rgb24')
+            frame.pict_type = 'NONE'
+            # pdb.set_trace()
+            for packet in stream.encode(frame):
+                container.mux(packet)
+
+        # Flush stream
+        for packet in stream.encode():
+            container.mux(packet)
+
+    outputs = []
+    with av.open(buf, 'r', 'mp4') as container:
+        if container.streams.video:
+            for frame in container.decode(**{'video': 0}):
+                outputs.append(
+                    frame.to_rgb().to_ndarray().astype(np.float32) / 255.)
+
+    #outputs = np.stack(outputs, axis=0)
+    return outputs
+
+
+if __name__ == '__main__':
+
+    # -----------------------------------
+    # test VideoReader(filename, cache_capacity=10)
+    # -----------------------------------
+#    video_reader = VideoReader('utils/test.mp4')
+#    from utils import utils_image as util
+#    inputs = []
+#    for frame in video_reader:
+#        print(frame.dtype)
+#        util.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+#        #util.imshow(np.flip(frame, axis=2))
+
+    # -----------------------------------
+    # test video2images(video_path, output_dir)
+    # -----------------------------------
+#    video2images('utils/test.mp4', 'frames')
+
+    # -----------------------------------
+    # test images2video(image_dir, video_path, fps=24, image_ext='png')
+    # -----------------------------------
+#    images2video('frames', 'video_02.mp4', fps=30, image_ext='png')
+
+
+    # -----------------------------------
+    # test frames2video(frame_dir, video_file, fps=30, fourcc='XVID', filename_tmpl='{:06d}.png')
+    # -----------------------------------
+#    frames2video('frames', 'video_01.mp4', filename_tmpl='{:06d}.png')
+
+
+    # -----------------------------------
+    # test add_video_compression(imgs)
+    # -----------------------------------
+#    imgs = []
+#    image_ext = 'png'
+#    frames = 'frames'
+#    from utils import utils_image as util
+#    image_files = sorted(glob.glob(os.path.join(frames, '*.{}'.format(image_ext))))
+#    for i, image_file in enumerate(image_files):
+#        if i < 7:
+#            img = util.imread_uint(image_file, 3)
+#            img = util.uint2single(img)
+#            imgs.append(img)
+#
+#    results = add_video_compression(imgs)
+#    for i, img in enumerate(results):
+#        util.imshow(util.single2uint(img))
+#        util.imsave(util.single2uint(img),f'{i:05}.png')
+
+    # run utils/utils_video.py
+
+
+
+
+
+
+

core/scripts/cli.py

@@ -0,0 +1,41 @@
+import sys
+import argparse
+from .. import WarpCore
+from .. import templates
+
+
+def template_init(args):
+    return ''''
+
+
+    '''.strip()
+
+
+def init_template(args):
+    parser = argparse.ArgumentParser(description='WarpCore template init tool')
+    parser.add_argument('-t', '--template', type=str, default='WarpCore')
+    args = parser.parse_args(args)
+
+    if args.template == 'WarpCore':
+        template_cls = WarpCore
+    else:
+        try:
+            template_cls = __import__(args.template)
+        except ModuleNotFoundError:
+            template_cls = getattr(templates, args.template)
+    print(template_cls)
+
+
+def main():
+    if len(sys.argv) < 2:
+        print('Usage: core <command>')
+        sys.exit(1)
+    if sys.argv[1] == 'init':
+        init_template(sys.argv[2:])
+    else:
+        print('Unknown command')
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    main()

core/templates/__init__.py

@@ -0,0 +1 @@
+from .diffusion import DiffusionCore

core/templates/diffusion.py

@@ -0,0 +1,236 @@
+from .. import WarpCore
+from ..utils import EXPECTED, EXPECTED_TRAIN, update_weights_ema, create_folder_if_necessary
+from abc import abstractmethod
+from dataclasses import dataclass
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from gdf import GDF
+import numpy as np
+from tqdm import tqdm
+import wandb
+
+import webdataset as wds
+from webdataset.handlers import warn_and_continue
+from torch.distributed import barrier
+from enum import Enum
+
+class TargetReparametrization(Enum):
+    EPSILON = 'epsilon'
+    X0 = 'x0'
+
+class DiffusionCore(WarpCore):
+    @dataclass(frozen=True)
+    class Config(WarpCore.Config):
+        # TRAINING PARAMS
+        lr: float = EXPECTED_TRAIN
+        grad_accum_steps: int = EXPECTED_TRAIN
+        batch_size: int = EXPECTED_TRAIN
+        updates: int = EXPECTED_TRAIN
+        warmup_updates: int = EXPECTED_TRAIN
+        save_every: int = 500
+        backup_every: int = 20000
+        use_fsdp: bool = True
+
+        # EMA UPDATE
+        ema_start_iters: int = None
+        ema_iters: int = None
+        ema_beta: float = None
+
+        # GDF setting
+        gdf_target_reparametrization: TargetReparametrization = None # epsilon or x0
+    
+    @dataclass() # not frozen, means that fields are mutable. Doesn't support EXPECTED
+    class Info(WarpCore.Info):
+        ema_loss: float = None
+
+    @dataclass(frozen=True)
+    class Models(WarpCore.Models):
+        generator : nn.Module = EXPECTED
+        generator_ema : nn.Module = None # optional
+
+    @dataclass(frozen=True)
+    class Optimizers(WarpCore.Optimizers):
+        generator : any = EXPECTED
+
+    @dataclass(frozen=True)
+    class Schedulers(WarpCore.Schedulers):
+        generator: any = None
+
+    @dataclass(frozen=True)
+    class Extras(WarpCore.Extras):
+        gdf: GDF = EXPECTED
+        sampling_configs: dict = EXPECTED
+
+    # --------------------------------------------
+    info: Info
+    config: Config
+
+    @abstractmethod
+    def encode_latents(self, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
+        raise NotImplementedError("This method needs to be overriden")
+
+    @abstractmethod
+    def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
+        raise NotImplementedError("This method needs to be overriden")
+
+    @abstractmethod
+    def get_conditions(self, batch: dict, models: Models, extras: Extras, is_eval=False, is_unconditional=False):
+        raise NotImplementedError("This method needs to be overriden")
+
+    @abstractmethod
+    def webdataset_path(self, extras: Extras):
+        raise NotImplementedError("This method needs to be overriden")
+
+    @abstractmethod
+    def webdataset_filters(self, extras: Extras):
+        raise NotImplementedError("This method needs to be overriden")
+    
+    @abstractmethod
+    def webdataset_preprocessors(self, extras: Extras):
+        raise NotImplementedError("This method needs to be overriden")
+
+    @abstractmethod
+    def sample(self, models: Models, data: WarpCore.Data, extras: Extras):
+        raise NotImplementedError("This method needs to be overriden")
+    # -------------
+
+    def setup_data(self, extras: Extras) -> WarpCore.Data:
+        # SETUP DATASET
+        dataset_path = self.webdataset_path(extras)
+        preprocessors = self.webdataset_preprocessors(extras)
+        filters = self.webdataset_filters(extras)
+
+        handler = warn_and_continue # None
+        # handler = None
+        dataset = wds.WebDataset(
+            dataset_path, resampled=True, handler=handler
+        ).select(filters).shuffle(690, handler=handler).decode(
+            "pilrgb", handler=handler
+        ).to_tuple(
+            *[p[0] for p in preprocessors], handler=handler
+        ).map_tuple(
+            *[p[1] for p in preprocessors], handler=handler
+        ).map(lambda x: {p[2]:x[i] for i, p in enumerate(preprocessors)})
+
+        # SETUP DATALOADER
+        real_batch_size = self.config.batch_size//(self.world_size*self.config.grad_accum_steps)
+        dataloader = DataLoader(
+            dataset, batch_size=real_batch_size, num_workers=8, pin_memory=True
+        )
+
+        return self.Data(dataset=dataset, dataloader=dataloader, iterator=iter(dataloader))
+
+    def forward_pass(self, data: WarpCore.Data, extras: Extras, models: Models):
+        batch = next(data.iterator)
+
+        with torch.no_grad():
+            conditions = self.get_conditions(batch, models, extras)
+            latents = self.encode_latents(batch, models, extras)
+            noised, noise, target, logSNR, noise_cond, loss_weight = extras.gdf.diffuse(latents, shift=1, loss_shift=1)
+
+        # FORWARD PASS
+        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            pred = models.generator(noised, noise_cond, **conditions)
+            if self.config.gdf_target_reparametrization == TargetReparametrization.EPSILON:
+                pred = extras.gdf.undiffuse(noised, logSNR, pred)[1] # transform whatever prediction to epsilon to use in the loss
+                target = noise
+            elif self.config.gdf_target_reparametrization == TargetReparametrization.X0:
+                pred = extras.gdf.undiffuse(noised, logSNR, pred)[0] # transform whatever prediction to x0 to use in the loss
+                target = latents
+            loss = nn.functional.mse_loss(pred, target, reduction='none').mean(dim=[1, 2, 3])
+            loss_adjusted = (loss * loss_weight).mean() / self.config.grad_accum_steps
+
+        return loss, loss_adjusted
+    
+    def train(self, data: WarpCore.Data, extras: Extras, models: Models, optimizers: Optimizers, schedulers: Schedulers):
+        start_iter = self.info.iter+1
+        max_iters = self.config.updates * self.config.grad_accum_steps
+        if self.is_main_node:
+            print(f"STARTING AT STEP: {start_iter}/{max_iters}")
+
+        pbar = tqdm(range(start_iter, max_iters+1)) if self.is_main_node else range(start_iter, max_iters+1) # <--- DDP
+        models.generator.train()
+        for i in pbar:
+            # FORWARD PASS
+            loss, loss_adjusted = self.forward_pass(data, extras, models)
+
+            # BACKWARD PASS
+            if i % self.config.grad_accum_steps == 0 or i == max_iters:
+                loss_adjusted.backward()
+                grad_norm = nn.utils.clip_grad_norm_(models.generator.parameters(), 1.0)
+                optimizers_dict = optimizers.to_dict()
+                for k in optimizers_dict:
+                    optimizers_dict[k].step()
+                schedulers_dict = schedulers.to_dict()
+                for k in schedulers_dict:
+                    schedulers_dict[k].step()
+                models.generator.zero_grad(set_to_none=True)
+                self.info.total_steps += 1
+            else:
+                with models.generator.no_sync():
+                    loss_adjusted.backward()
+            self.info.iter = i
+
+            # UPDATE EMA
+            if models.generator_ema is not None and i % self.config.ema_iters == 0:
+                update_weights_ema(
+                    models.generator_ema, models.generator,
+                    beta=(self.config.ema_beta if i > self.config.ema_start_iters else 0)
+                )
+
+            # UPDATE LOSS METRICS
+            self.info.ema_loss = loss.mean().item() if self.info.ema_loss is None else self.info.ema_loss * 0.99 + loss.mean().item() * 0.01
+
+            if self.is_main_node and self.config.wandb_project is not None and np.isnan(loss.mean().item()) or np.isnan(grad_norm.item()):
+                wandb.alert(
+                    title=f"NaN value encountered in training run {self.info.wandb_run_id}", 
+                    text=f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}",
+                    wait_duration=60*30
+                )
+
+            if self.is_main_node:
+                logs = {
+                    'loss': self.info.ema_loss, 
+                    'raw_loss': loss.mean().item(),
+                    'grad_norm': grad_norm.item(),
+                    'lr': optimizers.generator.param_groups[0]['lr'],
+                    'total_steps': self.info.total_steps,
+                }
+
+                pbar.set_postfix(logs)
+                if self.config.wandb_project is not None:
+                    wandb.log(logs)
+
+            if i == 1 or i % (self.config.save_every*self.config.grad_accum_steps) == 0 or i == max_iters:
+                # SAVE AND CHECKPOINT STUFF
+                if np.isnan(loss.mean().item()):
+                    if self.is_main_node and self.config.wandb_project is not None:
+                        tqdm.write("Skipping sampling & checkpoint because the loss is NaN")
+                        wandb.alert(title=f"Skipping sampling & checkpoint for training run {self.config.run_id}", text=f"Skipping sampling & checkpoint at {self.info.total_steps} for training run {self.info.wandb_run_id} iters because loss is NaN")
+                else:
+                    self.save_checkpoints(models, optimizers)
+                    if self.is_main_node:
+                        create_folder_if_necessary(f'{self.config.output_path}/{self.config.experiment_id}/')
+                    self.sample(models, data, extras)
+
+    def models_to_save(self):
+        return ['generator', 'generator_ema']
+
+    def save_checkpoints(self, models: Models, optimizers: Optimizers, suffix=None):
+        barrier()
+        suffix = '' if suffix is None else suffix
+        self.save_info(self.info, suffix=suffix)
+        models_dict = models.to_dict()
+        optimizers_dict = optimizers.to_dict()
+        for key in self.models_to_save():
+            model = models_dict[key]
+            if model is not None:
+                self.save_model(model, f"{key}{suffix}", is_fsdp=self.config.use_fsdp)
+        for key in optimizers_dict:
+            optimizer = optimizers_dict[key]
+            if optimizer is not None:
+                self.save_optimizer(optimizer, f'{key}_optim{suffix}', fsdp_model=models.generator if self.config.use_fsdp else None)
+        if suffix == '' and self.info.total_steps > 1 and self.info.total_steps % self.config.backup_every == 0:
+            self.save_checkpoints(models, optimizers, suffix=f"_{self.info.total_steps//1000}k")
+        torch.cuda.empty_cache()

core/utils/__init__.py

@@ -0,0 +1,9 @@
+from .base_dto import Base, nested_dto, EXPECTED, EXPECTED_TRAIN
+from .save_and_load import create_folder_if_necessary, safe_save, load_or_fail
+
+# MOVE IT SOMERWHERE ELSE
+def update_weights_ema(tgt_model, src_model, beta=0.999):
+    for self_params, src_params in zip(tgt_model.parameters(), src_model.parameters()):
+        self_params.data = self_params.data * beta + src_params.data.clone().to(self_params.device) * (1-beta)
+    for self_buffers, src_buffers in zip(tgt_model.buffers(), src_model.buffers()):
+        self_buffers.data = self_buffers.data * beta + src_buffers.data.clone().to(self_buffers.device) * (1-beta)

core/utils/base_dto.py

@@ -0,0 +1,56 @@
+import dataclasses
+from dataclasses import dataclass, _MISSING_TYPE
+from munch import Munch
+
+EXPECTED = "___REQUIRED___"
+EXPECTED_TRAIN = "___REQUIRED_TRAIN___"
+
+# pylint: disable=invalid-field-call
+def nested_dto(x, raw=False):
+    return dataclasses.field(default_factory=lambda: x if raw else Munch.fromDict(x))
+
+@dataclass(frozen=True)
+class Base:
+    training: bool = None
+    def __new__(cls, **kwargs):
+        training = kwargs.get('training', True)
+        setteable_fields = cls.setteable_fields(**kwargs)
+        mandatory_fields = cls.mandatory_fields(**kwargs)
+        invalid_kwargs = [
+            {k: v} for k, v in kwargs.items() if k not in setteable_fields or v == EXPECTED or (v == EXPECTED_TRAIN and training is not False)
+        ]
+        print(mandatory_fields)
+        assert (
+            len(invalid_kwargs) == 0
+        ), f"Invalid fields detected when initializing this DTO: {invalid_kwargs}.\nDeclare this field and set it to None or EXPECTED in order to make it setteable."
+        missing_kwargs = [f for f in mandatory_fields if f not in kwargs]
+        assert (
+            len(missing_kwargs) == 0
+        ), f"Required fields missing initializing this DTO: {missing_kwargs}."
+        return object.__new__(cls)
+
+
+    @classmethod
+    def setteable_fields(cls, **kwargs):
+        return [f.name for f in dataclasses.fields(cls) if f.default is None or isinstance(f.default, _MISSING_TYPE) or f.default == EXPECTED or f.default == EXPECTED_TRAIN]
+
+    @classmethod
+    def mandatory_fields(cls, **kwargs):
+        training = kwargs.get('training', True)
+        return [f.name for f in dataclasses.fields(cls) if isinstance(f.default, _MISSING_TYPE) and isinstance(f.default_factory, _MISSING_TYPE) or f.default == EXPECTED or (f.default == EXPECTED_TRAIN and training is not False)]
+
+    @classmethod
+    def from_dict(cls, kwargs):
+        for k in kwargs:
+            if isinstance(kwargs[k], (dict, list, tuple)):
+                kwargs[k] = Munch.fromDict(kwargs[k])
+        return cls(**kwargs)
+
+    def to_dict(self):
+        # selfdict = dataclasses.asdict(self) # needs to pickle stuff, doesn't support some more complex classes
+        selfdict = {}
+        for k in dataclasses.fields(self):
+            selfdict[k.name] = getattr(self, k.name)
+            if isinstance(selfdict[k.name], Munch):
+                selfdict[k.name] = selfdict[k.name].toDict()
+        return selfdict

core/utils/save_and_load.py

@@ -0,0 +1,59 @@
+import os
+import torch
+import json
+from pathlib import Path
+import safetensors
+import wandb
+
+
+def create_folder_if_necessary(path):
+    path = "/".join(path.split("/")[:-1])
+    Path(path).mkdir(parents=True, exist_ok=True)
+
+
+def safe_save(ckpt, path):
+    try:
+        os.remove(f"{path}.bak")
+    except OSError:
+        pass
+    try:
+        os.rename(path, f"{path}.bak")
+    except OSError:
+        pass
+    if path.endswith(".pt") or path.endswith(".ckpt"):
+        torch.save(ckpt, path)
+    elif path.endswith(".json"):
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(ckpt, f, indent=4)
+    elif path.endswith(".safetensors"):
+        safetensors.torch.save_file(ckpt, path)
+    else:
+        raise ValueError(f"File extension not supported: {path}")
+
+
+def load_or_fail(path, wandb_run_id=None):
+    accepted_extensions = [".pt", ".ckpt", ".json", ".safetensors"]
+    try:
+        assert any(
+            [path.endswith(ext) for ext in accepted_extensions]
+        ), f"Automatic loading not supported for this extension: {path}"
+        if not os.path.exists(path):
+            checkpoint = None
+        elif path.endswith(".pt") or path.endswith(".ckpt"):
+            checkpoint = torch.load(path, map_location="cpu")
+        elif path.endswith(".json"):
+            with open(path, "r", encoding="utf-8") as f:
+                checkpoint = json.load(f)
+        elif path.endswith(".safetensors"):
+            checkpoint = {}
+            with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+                for key in f.keys():
+                    checkpoint[key] = f.get_tensor(key)
+        return checkpoint
+    except Exception as e:
+        if wandb_run_id is not None:
+            wandb.alert(
+                title=f"Corrupt checkpoint for run {wandb_run_id}",
+                text=f"Training {wandb_run_id} tried to load checkpoint {path} and failed",
+            )
+        raise e