Upload 100 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+core/data/deg_kair_utils/test.png filter=lfs diff=lfs merge=lfs -text
+figures/California_000490.jpg filter=lfs diff=lfs merge=lfs -text
+figures/example_dataset/000008.jpg filter=lfs diff=lfs merge=lfs -text
+figures/example_dataset/000012.jpg filter=lfs diff=lfs merge=lfs -text
+figures/teaser.jpg filter=lfs diff=lfs merge=lfs -text

core/__init__.py

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

core/data/__init__.py

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

core/data/bucketeer.py

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

core/data/bucketeer_deg.py

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

core/data/deg_kair_utils/utils_alignfaces.py

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

core/data/deg_kair_utils/utils_blindsr.py

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

core/data/deg_kair_utils/utils_bnorm.py

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

core/data/deg_kair_utils/utils_deblur.py

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

core/data/deg_kair_utils/utils_dist.py

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

core/data/deg_kair_utils/utils_googledownload.py

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

core/data/deg_kair_utils/utils_image.py

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

core/data/deg_kair_utils/utils_lmdb.py

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

core/data/deg_kair_utils/utils_logger.py

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

core/data/deg_kair_utils/utils_mat.py

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

core/data/deg_kair_utils/utils_matconvnet.py

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

core/data/deg_kair_utils/utils_model.py

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

core/data/deg_kair_utils/utils_modelsummary.py

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

core/data/deg_kair_utils/utils_option.py

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

core/data/deg_kair_utils/utils_params.py

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

core/data/deg_kair_utils/utils_receptivefield.py

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

core/data/deg_kair_utils/utils_regularizers.py

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

core/data/deg_kair_utils/utils_sisr.py

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

core/data/deg_kair_utils/utils_video.py

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

core/data/deg_kair_utils/utils_videoio.py

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

core/scripts/cli.py

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

core/templates/__init__.py

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

core/templates/diffusion.py

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

core/utils/__init__.py

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

core/utils/base_dto.py

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

core/utils/save_and_load.py

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

figures/example_dataset/000008.json

@@ -0,0 +1,2 @@
+{ "caption": "The image captures the iconic Shard, a modern skyscraper that stands as the tallest building in the United Kingdom. The Shard, with its glass and steel structure, pierces the sky, its pointed top reaching towards the heavens. The photograph is taken from a low angle, which emphasizes the height and grandeur of the building. The sky forms a beautiful backdrop, painted in hues of pinkish-orange, suggesting that the photo was taken at sunset. The Shard is nestled between two other buildings, their presence subtly hinted at in the foreground. The image does not contain any discernible text or countable objects, and there are no visible actions taking place. The relative positions of the objects confirm that the Shard is the central focus of the image, with the other buildings and the sky providing context to its location. The image is devoid of any aesthetic descriptions, focusing solely on the factual representation of the scene."    
+}

figures/example_dataset/000012.json

@@ -0,0 +1 @@
+{"caption": "cars in a road during daytime"}

gdf/__init__.py

@@ -0,0 +1,205 @@
+import torch
+from .scalers import *
+from .targets import *
+from .schedulers import *
+from .noise_conditions import *
+from .loss_weights import *
+from .samplers import *
+import torch.nn.functional as F
+import math
+class GDF():
+    def __init__(self, schedule, input_scaler, target, noise_cond, loss_weight, offset_noise=0):
+        self.schedule = schedule
+        self.input_scaler = input_scaler
+        self.target = target
+        self.noise_cond = noise_cond
+        self.loss_weight = loss_weight
+        self.offset_noise = offset_noise
+
+    def setup_limits(self, stretch_max=True, stretch_min=True, shift=1):
+        stretched_limits = self.input_scaler.setup_limits(self.schedule, self.input_scaler, stretch_max, stretch_min, shift)
+        return stretched_limits
+
+    def diffuse(self, x0, epsilon=None, t=None, shift=1, loss_shift=1, offset=None):
+        if epsilon is None:
+            epsilon = torch.randn_like(x0)
+        if self.offset_noise > 0:
+            if offset is None:
+                offset = torch.randn([x0.size(0), x0.size(1)] + [1]*(len(x0.shape)-2)).to(x0.device)
+            epsilon = epsilon + offset * self.offset_noise
+        logSNR = self.schedule(x0.size(0) if t is None else t, shift=shift).to(x0.device)
+        a, b = self.input_scaler(logSNR) # B
+        if len(a.shape) == 1:
+            a, b = a.view(-1, *[1]*(len(x0.shape)-1)), b.view(-1, *[1]*(len(x0.shape)-1)) # BxCxHxW
+        #print('in line 33 a b', a.shape, b.shape, x0.shape, logSNR.shape, logSNR, self.noise_cond(logSNR))
+        target = self.target(x0, epsilon, logSNR, a, b)
+
+        # noised, noise, logSNR, t_cond
+        #noised, noise, target, logSNR, noise_cond, loss_weight
+        return x0 * a + epsilon * b, epsilon, target, logSNR, self.noise_cond(logSNR), self.loss_weight(logSNR, shift=loss_shift)
+
+    def undiffuse(self, x, logSNR, pred):
+        a, b = self.input_scaler(logSNR)
+        if len(a.shape) == 1:
+            a, b = a.view(-1, *[1]*(len(x.shape)-1)), b.view(-1, *[1]*(len(x.shape)-1))
+        return self.target.x0(x, pred, logSNR, a, b), self.target.epsilon(x, pred, logSNR, a, b)
+
+    def sample(self, model, model_inputs, shape, unconditional_inputs=None, sampler=None, schedule=None, t_start=1.0, t_end=0.0, timesteps=20, x_init=None, cfg=3.0, cfg_t_stop=None, cfg_t_start=None, cfg_rho=0.7, sampler_params=None, shift=1, device="cpu"):
+        sampler_params = {} if sampler_params is None else sampler_params
+        if sampler is None:
+            sampler = DDPMSampler(self)
+        r_range = torch.linspace(t_start, t_end, timesteps+1)
+        schedule = self.schedule if schedule is None else schedule
+        logSNR_range = schedule(r_range, shift=shift)[:, None].expand(
+            -1, shape[0] if x_init is None else x_init.size(0)
+        ).to(device)
+
+        x = sampler.init_x(shape).to(device) if x_init is None else x_init.clone()
+       
+        if cfg is not None:
+            if unconditional_inputs is None:
+                unconditional_inputs = {k: torch.zeros_like(v) for k, v in model_inputs.items()}
+            model_inputs = {
+                k: torch.cat([v, v_u], dim=0) if isinstance(v, torch.Tensor) 
+                else [torch.cat([vi, vi_u], dim=0) if isinstance(vi, torch.Tensor) and isinstance(vi_u, torch.Tensor) else None for vi, vi_u in zip(v, v_u)] if isinstance(v, list)
+                else {vk: torch.cat([v[vk], v_u.get(vk, torch.zeros_like(v[vk]))], dim=0) for vk in v} if isinstance(v, dict)
+                else None for (k, v), (k_u, v_u) in zip(model_inputs.items(), unconditional_inputs.items())
+            }
+      
+        for i in range(0, timesteps):
+            noise_cond = self.noise_cond(logSNR_range[i])
+            if cfg is not None and (cfg_t_stop is None or r_range[i].item() >= cfg_t_stop) and (cfg_t_start is None or r_range[i].item() <= cfg_t_start):
+                cfg_val = cfg
+                if isinstance(cfg_val, (list, tuple)):
+                    assert len(cfg_val) == 2, "cfg must be a float or a list/tuple of length 2"
+                    cfg_val = cfg_val[0] * r_range[i].item() + cfg_val[1] * (1-r_range[i].item())
+                
+                pred, pred_unconditional = model(torch.cat([x, x], dim=0), noise_cond.repeat(2), **model_inputs).chunk(2)
+                
+                pred_cfg = torch.lerp(pred_unconditional, pred, cfg_val)
+                if cfg_rho > 0:
+                    std_pos, std_cfg = pred.std(),  pred_cfg.std()
+                    pred = cfg_rho * (pred_cfg * std_pos/(std_cfg+1e-9)) + pred_cfg * (1-cfg_rho)
+                else:
+                    pred = pred_cfg
+            else:
+                pred = model(x, noise_cond, **model_inputs)
+            x0, epsilon = self.undiffuse(x, logSNR_range[i], pred)
+            x = sampler(x, x0, epsilon, logSNR_range[i], logSNR_range[i+1], **sampler_params)
+            #print('in line 86', x0.shape, x.shape, i, )
+            altered_vars = yield (x0, x, pred)
+
+            # Update some running variables if the user wants
+            if altered_vars is not None:
+                cfg = altered_vars.get('cfg', cfg)
+                cfg_rho = altered_vars.get('cfg_rho', cfg_rho)
+                sampler = altered_vars.get('sampler', sampler)
+                model_inputs = altered_vars.get('model_inputs', model_inputs)
+                x = altered_vars.get('x', x)
+                x_init = altered_vars.get('x_init', x_init)
+                
+class GDF_dual_fixlrt(GDF):
+    def ref_noise(self, noised, x0, logSNR):
+        a, b = self.input_scaler(logSNR)
+        if len(a.shape) == 1:
+            a, b = a.view(-1, *[1]*(len(x0.shape)-1)), b.view(-1, *[1]*(len(x0.shape)-1)) 
+        #print('in line 210', a.shape, b.shape, x0.shape, noised.shape)
+        return self.target.noise_givenx0_noised(x0, noised, logSNR, a, b)
+   
+    def sample(self, model,  model_inputs, shape, shape_lr, unconditional_inputs=None, sampler=None, 
+    schedule=None, t_start=1.0, t_end=0.0, timesteps=20, x_init=None, cfg=3.0, cfg_t_stop=None, 
+    cfg_t_start=None, cfg_rho=0.7, sampler_params=None, shift=1, device="cpu"):
+        sampler_params = {} if sampler_params is None else sampler_params
+        if sampler is None:
+            sampler = DDPMSampler(self)
+        r_range = torch.linspace(t_start, t_end, timesteps+1)
+        schedule = self.schedule if schedule is None else schedule
+        logSNR_range = schedule(r_range, shift=shift)[:, None].expand(
+            -1, shape[0] if x_init is None else x_init.size(0)
+        ).to(device)
+
+        x = sampler.init_x(shape).to(device) if x_init is None else x_init.clone()
+        x_lr = sampler.init_x(shape_lr).to(device) if x_init is None else x_init.clone()
+        if cfg is not None:
+            if unconditional_inputs is None:
+                unconditional_inputs = {k: torch.zeros_like(v) for k, v in model_inputs.items()}
+            model_inputs = {
+                k: torch.cat([v, v_u], dim=0) if isinstance(v, torch.Tensor) 
+                else [torch.cat([vi, vi_u], dim=0) if isinstance(vi, torch.Tensor) and isinstance(vi_u, torch.Tensor) else None for vi, vi_u in zip(v, v_u)] if isinstance(v, list)
+                else {vk: torch.cat([v[vk], v_u.get(vk, torch.zeros_like(v[vk]))], dim=0) for vk in v} if isinstance(v, dict)
+                else None for (k, v), (k_u, v_u) in zip(model_inputs.items(), unconditional_inputs.items())
+            }
+            
+        ###############################################lr sampling   
+        
+        guide_feas = [None] * timesteps
+       
+        for i in range(0, timesteps):
+            noise_cond = self.noise_cond(logSNR_range[i])
+            if cfg is not None and (cfg_t_stop is None or r_range[i].item() >= cfg_t_stop) and (cfg_t_start is None or r_range[i].item() <= cfg_t_start):
+                cfg_val = cfg
+                if isinstance(cfg_val, (list, tuple)):
+                    assert len(cfg_val) == 2, "cfg must be a float or a list/tuple of length 2"
+                    cfg_val = cfg_val[0] * r_range[i].item() + cfg_val[1] * (1-r_range[i].item())
+                
+              
+                
+                if i == timesteps -1 :
+                    output, guide_lr_enc, guide_lr_dec = model(torch.cat([x_lr, x_lr], dim=0), noise_cond.repeat(2), reuire_f=True, **model_inputs)
+                    guide_feas[i] = ([f.chunk(2)[0].repeat(2, 1, 1, 1) for f in guide_lr_enc], [f.chunk(2)[0].repeat(2, 1, 1, 1) for f in guide_lr_dec])
+                else:
+                    output, _, _ = model(torch.cat([x_lr, x_lr], dim=0), noise_cond.repeat(2), reuire_f=True, **model_inputs)
+                
+                pred, pred_unconditional = output.chunk(2)
+
+                
+                pred_cfg = torch.lerp(pred_unconditional, pred, cfg_val)
+                if cfg_rho > 0:
+                    std_pos, std_cfg = pred.std(),  pred_cfg.std()
+                    pred = cfg_rho * (pred_cfg * std_pos/(std_cfg+1e-9)) + pred_cfg * (1-cfg_rho)
+                else:
+                    pred = pred_cfg
+            else:
+                pred = model(x_lr, noise_cond, **model_inputs)
+            x0_lr, epsilon_lr = self.undiffuse(x_lr, logSNR_range[i], pred)
+            x_lr = sampler(x_lr, x0_lr, epsilon_lr, logSNR_range[i], logSNR_range[i+1], **sampler_params)
+
+       ###############################################hr HR sampling   
+        for i in range(0, timesteps):
+            noise_cond = self.noise_cond(logSNR_range[i])
+            if cfg is not None and (cfg_t_stop is None or r_range[i].item() >= cfg_t_stop) and (cfg_t_start is None or r_range[i].item() <= cfg_t_start):
+                cfg_val = cfg
+                if isinstance(cfg_val, (list, tuple)):
+                    assert len(cfg_val) == 2, "cfg must be a float or a list/tuple of length 2"
+                    cfg_val = cfg_val[0] * r_range[i].item() + cfg_val[1] * (1-r_range[i].item())
+
+                out_pred, t_emb = model(torch.cat([x, x], dim=0), noise_cond.repeat(2), \
+                lr_guide=guide_feas[timesteps -1] if i <=19 else None  , **model_inputs, require_t=True, guide_weight=1 - i/timesteps)
+                pred, pred_unconditional = out_pred.chunk(2)
+                pred_cfg = torch.lerp(pred_unconditional, pred, cfg_val)
+                if cfg_rho > 0:
+                    std_pos, std_cfg = pred.std(),  pred_cfg.std()
+                    pred = cfg_rho * (pred_cfg * std_pos/(std_cfg+1e-9)) + pred_cfg * (1-cfg_rho)
+                else:
+                    pred = pred_cfg
+            else:
+                pred = model(x, noise_cond, guide_lr=(guide_lr_enc, guide_lr_dec), **model_inputs)
+            x0, epsilon = self.undiffuse(x, logSNR_range[i], pred)
+
+            x = sampler(x, x0, epsilon, logSNR_range[i], logSNR_range[i+1], **sampler_params)
+            altered_vars = yield (x0, x, pred, x_lr)
+            
+         
+
+            # Update some running variables if the user wants
+            if altered_vars is not None:
+                cfg = altered_vars.get('cfg', cfg)
+                cfg_rho = altered_vars.get('cfg_rho', cfg_rho)
+                sampler = altered_vars.get('sampler', sampler)
+                model_inputs = altered_vars.get('model_inputs', model_inputs)
+                x = altered_vars.get('x', x)
+                x_init = altered_vars.get('x_init', x_init)
+
+
+
+

gdf/loss_weights.py

@@ -0,0 +1,101 @@
+import torch
+import numpy as np
+
+# --- Loss Weighting
+class BaseLossWeight():
+    def weight(self, logSNR):
+        raise NotImplementedError("this method needs to be overridden")
+
+    def __call__(self, logSNR, *args, shift=1, clamp_range=None, **kwargs):
+        clamp_range = [-1e9, 1e9] if clamp_range is None else clamp_range
+        if shift != 1:
+            logSNR = logSNR.clone() + 2 * np.log(shift)
+        return self.weight(logSNR, *args, **kwargs).clamp(*clamp_range)
+
+class ComposedLossWeight(BaseLossWeight):
+    def __init__(self, div, mul):
+        self.mul = [mul] if isinstance(mul, BaseLossWeight) else mul
+        self.div = [div] if isinstance(div, BaseLossWeight) else div
+
+    def weight(self, logSNR):
+        prod, div = 1, 1
+        for m in self.mul:
+            prod *= m.weight(logSNR)
+        for d in self.div:
+            div *= d.weight(logSNR)
+        return prod/div
+
+class ConstantLossWeight(BaseLossWeight):
+    def __init__(self, v=1):
+        self.v = v
+
+    def weight(self, logSNR):
+        return torch.ones_like(logSNR) * self.v
+
+class SNRLossWeight(BaseLossWeight):
+    def weight(self, logSNR):
+        return logSNR.exp()
+
+class P2LossWeight(BaseLossWeight):
+    def __init__(self, k=1.0, gamma=1.0, s=1.0):
+        self.k, self.gamma, self.s = k, gamma, s
+
+    def weight(self, logSNR):
+        return (self.k + (logSNR * self.s).exp()) ** -self.gamma
+
+class SNRPlusOneLossWeight(BaseLossWeight):
+    def weight(self, logSNR):
+        return logSNR.exp() + 1
+
+class MinSNRLossWeight(BaseLossWeight):
+    def __init__(self, max_snr=5):
+        self.max_snr = max_snr
+
+    def weight(self, logSNR):
+        return logSNR.exp().clamp(max=self.max_snr)
+
+class MinSNRPlusOneLossWeight(BaseLossWeight):
+    def __init__(self, max_snr=5):
+        self.max_snr = max_snr
+
+    def weight(self, logSNR):
+        return (logSNR.exp() + 1).clamp(max=self.max_snr)
+
+class TruncatedSNRLossWeight(BaseLossWeight):
+    def __init__(self, min_snr=1):
+        self.min_snr = min_snr
+
+    def weight(self, logSNR):
+        return logSNR.exp().clamp(min=self.min_snr)
+
+class SechLossWeight(BaseLossWeight):
+    def __init__(self, div=2):
+        self.div = div
+
+    def weight(self, logSNR):
+        return 1/(logSNR/self.div).cosh()
+
+class DebiasedLossWeight(BaseLossWeight):
+    def weight(self, logSNR):
+        return 1/logSNR.exp().sqrt()
+
+class SigmoidLossWeight(BaseLossWeight):
+    def __init__(self, s=1):
+        self.s = s
+
+    def weight(self, logSNR):
+        return (logSNR * self.s).sigmoid()
+
+class AdaptiveLossWeight(BaseLossWeight):
+    def __init__(self, logsnr_range=[-10, 10], buckets=300, weight_range=[1e-7, 1e7]):
+        self.bucket_ranges = torch.linspace(logsnr_range[0], logsnr_range[1], buckets-1)
+        self.bucket_losses = torch.ones(buckets)
+        self.weight_range = weight_range
+
+    def weight(self, logSNR):
+        indices = torch.searchsorted(self.bucket_ranges.to(logSNR.device), logSNR)
+        return (1/self.bucket_losses.to(logSNR.device)[indices]).clamp(*self.weight_range)
+
+    def update_buckets(self, logSNR, loss, beta=0.99):
+        indices = torch.searchsorted(self.bucket_ranges.to(logSNR.device), logSNR).cpu()
+        self.bucket_losses[indices] = self.bucket_losses[indices]*beta + loss.detach().cpu() * (1-beta)

gdf/noise_conditions.py

@@ -0,0 +1,102 @@
+import torch
+import numpy as np
+
+class BaseNoiseCond():
+    def __init__(self, *args, shift=1, clamp_range=None, **kwargs):
+        clamp_range = [-1e9, 1e9] if clamp_range is None else clamp_range
+        self.shift = shift
+        self.clamp_range = clamp_range
+        self.setup(*args, **kwargs)
+
+    def setup(self, *args, **kwargs):
+        pass # this method is optional, override it if required
+
+    def cond(self, logSNR):
+        raise NotImplementedError("this method needs to be overriden")
+
+    def __call__(self, logSNR):
+        if self.shift != 1:
+            logSNR = logSNR.clone() + 2 * np.log(self.shift)
+        return self.cond(logSNR).clamp(*self.clamp_range)
+
+class CosineTNoiseCond(BaseNoiseCond):
+    def setup(self, s=0.008, clamp_range=[0, 1]): # [0.0001, 0.9999]
+        self.s = torch.tensor([s])
+        self.clamp_range = clamp_range
+        self.min_var = torch.cos(self.s / (1 + self.s) * torch.pi * 0.5) ** 2
+
+    def cond(self, logSNR):
+        var = logSNR.sigmoid()
+        var = var.clamp(*self.clamp_range)
+        s, min_var = self.s.to(var.device), self.min_var.to(var.device)
+        t = (((var * min_var) ** 0.5).acos() / (torch.pi * 0.5)) * (1 + s) - s
+        return t
+
+class EDMNoiseCond(BaseNoiseCond):
+    def cond(self, logSNR):
+        return -logSNR/8
+
+class SigmoidNoiseCond(BaseNoiseCond):
+    def cond(self, logSNR):
+        return (-logSNR).sigmoid()
+
+class LogSNRNoiseCond(BaseNoiseCond):
+    def cond(self, logSNR):
+        return logSNR
+
+class EDMSigmaNoiseCond(BaseNoiseCond):
+    def setup(self, sigma_data=1):
+        self.sigma_data = sigma_data
+
+    def cond(self, logSNR):
+        return torch.exp(-logSNR / 2) * self.sigma_data
+
+class RectifiedFlowsNoiseCond(BaseNoiseCond):
+    def cond(self, logSNR):
+        _a = logSNR.exp() - 1
+        _a[_a == 0] = 1e-3 # Avoid division by zero
+        a = 1 + (2-(2**2 + 4*_a)**0.5) / (2*_a)
+        return a
+
+# Any NoiseCond that cannot be described easily as a continuous function of t
+# It needs to define self.x and self.y in the setup() method
+class PiecewiseLinearNoiseCond(BaseNoiseCond):
+    def setup(self):
+        self.x = None
+        self.y = None
+
+    def piecewise_linear(self, y, xs, ys):
+        indices = (len(xs)-2) - torch.searchsorted(ys.flip(dims=(-1,))[:-2], y)  
+        x_min, x_max = xs[indices], xs[indices+1]
+        y_min, y_max = ys[indices], ys[indices+1]
+        x = x_min + (x_max - x_min) * (y - y_min) / (y_max - y_min)
+        return x
+
+    def cond(self, logSNR):
+        var = logSNR.sigmoid()
+        t = self.piecewise_linear(var, self.x.to(var.device), self.y.to(var.device)) # .mul(1000).round().clamp(min=0)
+        return t
+
+class StableDiffusionNoiseCond(PiecewiseLinearNoiseCond):
+    def setup(self, linear_range=[0.00085, 0.012], total_steps=1000):
+        self.total_steps = total_steps
+        linear_range_sqrt = [r**0.5 for r in linear_range]
+        self.x = torch.linspace(0, 1, total_steps+1)
+
+        alphas = 1-(linear_range_sqrt[0]*(1-self.x) + linear_range_sqrt[1]*self.x)**2
+        self.y = alphas.cumprod(dim=-1)
+
+    def cond(self, logSNR):
+        return super().cond(logSNR).clamp(0, 1)
+
+class DiscreteNoiseCond(BaseNoiseCond):
+    def setup(self, noise_cond, steps=1000, continuous_range=[0, 1]):
+        self.noise_cond = noise_cond
+        self.steps = steps
+        self.continuous_range = continuous_range
+
+    def cond(self, logSNR):
+        cond = self.noise_cond(logSNR)
+        cond = (cond-self.continuous_range[0]) / (self.continuous_range[1]-self.continuous_range[0])
+        return cond.mul(self.steps).long()
+    

gdf/readme.md

@@ -0,0 +1,86 @@
+# Generic Diffusion Framework (GDF)
+
+# Basic usage
+GDF is a simple framework for working with diffusion models. It implements most common diffusion frameworks (DDPM / DDIM
+, EDM, Rectified Flows, etc.) and makes it very easy to switch between them or combine different parts of different
+frameworks
+
+Using GDF is very straighforward, first of all just define an instance of the GDF class:
+
+```python
+from gdf import GDF
+from gdf import CosineSchedule
+from gdf import VPScaler, EpsilonTarget, CosineTNoiseCond, P2LossWeight
+
+gdf = GDF(
+    schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]),
+    input_scaler=VPScaler(), target=EpsilonTarget(),
+    noise_cond=CosineTNoiseCond(),
+    loss_weight=P2LossWeight(),
+)
+```
+
+You need to define the following components: 
+* **Train Schedule**: This will return the logSNR schedule that will be used during training, some of the schedulers can be configured. A train schedule will then be called with a batch size and will randomly sample some values from the defined distribution.
+* **Sample Schedule**: This is the schedule that will be used later on when sampling. It might be different from the training schedule. 
+* **Input Scaler**: If you want to use Variance Preserving or LERP (rectified flows)
+* **Target**: What the target is during training, usually: epsilon, x0 or v
+* **Noise Conditioning**: You could directly pass the logSNR to your model but usually a normalized value is used instead, for example the EDM framework proposes to use `-logSNR/8`
+* **Loss Weight**: There are many proposed loss weighting strategies, here you define which one you'll use
+
+All of those classes are actually very simple logSNR centric definitions, for example the VPScaler is defined as just:
+```python 
+class VPScaler():
+    def __call__(self, logSNR): 
+        a_squared = logSNR.sigmoid()
+        a = a_squared.sqrt()
+        b = (1-a_squared).sqrt()
+        return a, b
+
+```
+
+So it's very easy to extend this framework with custom schedulers, scalers, targets, loss weights, etc...
+
+### Training
+
+When you define your training loop you can get all you need by just doing:
+```python
+shift, loss_shift = 1, 1 # this can be set to higher values as per what the Simple Diffusion paper sugested for high resolution
+for inputs, extra_conditions in dataloader_iterator:
+	noised, noise, target, logSNR, noise_cond, loss_weight = gdf.diffuse(inputs, shift=shift, loss_shift=loss_shift) 
+	pred = diffusion_model(noised, noise_cond, extra_conditions)
+
+	loss = nn.functional.mse_loss(pred, target, reduction='none').mean(dim=[1, 2, 3])
+	loss_adjusted = (loss * loss_weight).mean()
+
+	loss_adjusted.backward()
+	optimizer.step()
+	optimizer.zero_grad(set_to_none=True)
+```
+
+And that's all, you have a diffusion model training, where it's very easy to customize the different elements of the 
+training from the GDF class.
+
+### Sampling
+
+The other important part is sampling, when you want to use this framework to sample you can just do the following:
+
+```python
+from gdf import DDPMSampler
+
+shift = 1
+sampling_configs = {
+	"timesteps": 30, "cfg": 7,  "sampler": DDPMSampler(gdf), "shift": shift,
+	"schedule": CosineSchedule(clamp_range=[0.0001, 0.9999])
+}
+
+*_, (sampled, _, _) = gdf.sample(
+	diffusion_model, {"cond": extra_conditions}, latents.shape, 
+	unconditional_inputs= {"cond": torch.zeros_like(extra_conditions)}, 
+	device=device, **sampling_configs
+)
+```
+
+# Available modules
+
+TODO