Florence2Flux

app.py

@@ -1,161 +1,74 @@
-import spaces
-import os
-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
-
-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/", "stage_a.safetensors"),
-        "STABLEWURST_PREVIEWER": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/previewer.safetensors?download=true", "models/", "previewer.safetensors"),
-        "STABLEWURST_EFFNET": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/effnet_encoder.safetensors?download=true", "models/", "effnet_encoder.safetensors"),
-        "STABLEWURST_B_LITE": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_b_lite_bf16.safetensors?download=true", "models/", "stage_b_lite_bf16.safetensors"),
-        "STABLEWURST_C": ("https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_c_bf16.safetensors?download=true", "models/", "stage_c_bf16.safetensors"),
-        "ULTRAPIXEL_T2I": ("https://huggingface.co/roubaofeipi/UltraPixel/resolve/main/ultrapixel_t2i.safetensors?download=true", "models/", "ultrapixel_t2i.safetensors"),
-        "ULTRAPIXEL_LORA_CAT": ("https://huggingface.co/roubaofeipi/UltraPixel/resolve/main/lora_cat.safetensors?download=true", "models/", "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}
-)
-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(duration=180)
-def generate_images(prompt, height, width, seed, num_images):
-    torch.manual_seed(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-
-    batch_size = num_images
-    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)    
-    
-    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)
-    
-    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
-
-iface = gr.Interface(
-    fn=generate_images,
-    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),
-        gr.Slider(minimum=1, maximum=10, step=1, label="Number of Images", value=1)
-    ],
-    outputs=gr.Gallery(label="Generated Images", columns=5, rows=2),
-    title="UltraPixel Image Generation",
-    description="Generate high-resolution images using UltraPixel model.",
-    theme='bethecloud/storj_theme',
-    examples=[
-        ["The image features a snow-covered mountain range with a large, snow-covered mountain in the background. The mountain is surrounded by a forest of trees, and the sky is filled with clouds. The scene is set during the winter season, with snow covering the ground and the trees.", 1024, 1024, 42, 1]
-    ],
-    cache_examples=True
-)
-
-iface.launch()
+from transformers import AutoProcessor, AutoModelForCausalLM
+import spaces
+from PIL import Image 
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+models = {
+    'gokaygokay/Florence-2-Flux-Large': AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-Flux-Large', trust_remote_code=True).eval(),
+    'gokaygokay/Florence-2-Flux': AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-Flux', trust_remote_code=True).eval(),
+}
+
+processors = {
+    'gokaygokay/Florence-2-Flux-Large': AutoProcessor.from_pretrained('gokaygokay/Florence-2-Flux-Large', trust_remote_code=True),
+    'gokaygokay/Florence-2-Flux': AutoProcessor.from_pretrained('gokaygokay/Florence-2-Flux', trust_remote_code=True),
+}
+
+title = """<h1 align="center">Florence-2 Captioner for Flux Prompts</h1>
+<p><center>
+<a href="https://huggingface.co/gokaygokay/Florence-2-Flux-Large" target="_blank">[Florence-2 Flux Large]</a>
+<a href="https://huggingface.co/gokaygokay/Florence-2-Flux" target="_blank">[Florence-2 Flux Base]</a>
+</center></p>
+"""
+
+@spaces.GPU
+def run_example(image, model_name='gokaygokay/Florence-2-Flux-Large'):
+    image = Image.fromarray(image)
+    task_prompt = "<DESCRIPTION>"
+    prompt = task_prompt + "Describe this image in great detail."
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    model = models[model_name]
+    processor = processors[model_name]
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3,
+        repetition_penalty=1.10,
+    )
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+    return parsed_answer["<DESCRIPTION>"]
+
+with gr.Blocks(theme='bethecloud/storj_theme') as demo:
+    gr.HTML(title)
+
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="Input Picture")
+            model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='gokaygokay/Florence-2-Flux-Large')
+            submit_btn = gr.Button(value="Submit")
+        with gr.Column():
+            output_text = gr.Textbox(label="Output Text")
+
+    gr.Examples(
+        [["image1.jpg"], 
+         ["image2.jpg"], 
+         ["image3.png"], 
+         ["image5.jpg"]],
+        inputs=[input_img, model_selector],
+        outputs=[output_text],
+        fn=run_example,
+        label='Try captioning on below examples'
+    )
+
+    submit_btn.click(run_example, [input_img, model_selector], [output_text])
+
+demo.launch(debug=True)

configs/inference/controlnet_c_3b_canny.yaml

@@ -1,14 +0,0 @@
-# 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

@@ -1,17 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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

@@ -1,15 +0,0 @@
-# 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

@@ -1,13 +0,0 @@
-# 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

@@ -1,13 +0,0 @@
-# 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

@@ -1,7 +0,0 @@
-# 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

@@ -1,7 +0,0 @@
-# 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

@@ -1,47 +0,0 @@
-# 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

@@ -1,37 +0,0 @@
-# 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

@@ -1,29 +0,0 @@
-# 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

@@ -1,372 +0,0 @@
-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

@@ -1,69 +0,0 @@
-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

@@ -1,88 +0,0 @@
-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

@@ -1,91 +0,0 @@
-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

@@ -1,263 +0,0 @@
-# -*- 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

@@ -1,631 +0,0 @@
-# -*- 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 ([email protected])
-# 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

@@ -1,91 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-"""
-# --------------------------------------------
-# Batch Normalization
-# --------------------------------------------
-
-# Kai Zhang ([email protected])
-# 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

@@ -1,655 +0,0 @@
-# -*- 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 ([email protected])
-    # 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

@@ -1,201 +0,0 @@
-# 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

@@ -1,93 +0,0 @@
-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

@@ -1,1016 +0,0 @@
-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

@@ -1,205 +0,0 @@
-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

@@ -1,66 +0,0 @@
-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

@@ -1,88 +0,0 @@
-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

@@ -1,197 +0,0 @@
-# -*- 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 ([email protected])
-# 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

@@ -1,330 +0,0 @@
-# -*- 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

@@ -1,485 +0,0 @@
-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

@@ -1,255 +0,0 @@
-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

@@ -1,135 +0,0 @@
-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

@@ -1,62 +0,0 @@
-# -*- 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

@@ -1,104 +0,0 @@
-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

@@ -1,848 +0,0 @@
-# -*- 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 ([email protected])
-# 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

@@ -1,493 +0,0 @@
-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

@@ -1,555 +0,0 @@
-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

@@ -1,41 +0,0 @@
-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

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

core/templates/diffusion.py

@@ -1,236 +0,0 @@
-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

@@ -1,9 +0,0 @@
-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

@@ -1,56 +0,0 @@
-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

@@ -1,59 +0,0 @@
-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