├── LICENSE ├── README.md ├── image ├── image1.png ├── image2.png ├── image3.png ├── image4.png └── table.png ├── requirements.txt ├── test └── texture │ └── config │ └── inpainting │ ├── data │ ├── BokehLQGT_dataset.py │ ├── BokehLQ_dataset.py │ ├── GT_dataset.py │ ├── LQGT_dataset.py │ ├── LQ_dataset.py │ ├── StereoLQGT_dataset.py │ ├── StereoLQ_dataset.py │ ├── __init__.py │ ├── canny.py │ ├── data_sampler.py │ └── util.py │ ├── file_utils.py │ ├── models │ ├── __init__.py │ ├── base_model.py │ ├── canny.py │ ├── denoising_model.py │ ├── dense_layer.py │ ├── layers.py │ ├── lr_scheduler.py │ ├── modules │ │ ├── DenoisingUNet_arch.py │ │ ├── __init__.py │ │ ├── loss.py │ │ └── module_util.py │ ├── networks.py │ ├── op │ │ ├── LICENSE_MIT │ │ ├── __init__.py │ │ ├── fused_act.py │ │ ├── fused_bias_act.cpp │ │ ├── fused_bias_act_kernel.cu │ │ ├── upfirdn2d.cpp │ │ ├── upfirdn2d.py │ │ └── upfirdn2d_kernel.cu │ ├── optimizer.py │ └── up_or_down_sampling.py │ ├── options.py │ ├── options │ └── test │ │ └── ir-sde.yml │ ├── str_utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py │ ├── test.py │ └── utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py └── train ├── discriminator └── config │ └── inpainting │ ├── data │ ├── BokehLQGT_dataset.py │ ├── BokehLQ_dataset.py │ ├── GT_dataset.py │ ├── LQGT_dataset.py │ ├── LQ_dataset.py │ ├── StereoLQGT_dataset.py │ ├── StereoLQ_dataset.py │ ├── __init__.py │ ├── canny.py │ ├── data_sampler.py │ └── util.py │ ├── demo.sh │ ├── file_utils.py │ ├── models │ ├── __init__.py │ ├── base_model.py │ ├── canny.py │ ├── denoising_model.py │ ├── dense_layer.py │ ├── layers.py │ ├── lr_scheduler.py │ ├── modules │ │ ├── DenoisingUNet_arch.py │ │ ├── __init__.py │ │ ├── loss.py │ │ └── module_util.py │ ├── networks.py │ ├── op │ │ ├── LICENSE_MIT │ │ ├── __init__.py │ │ ├── cuda-installer.log │ │ ├── fused_act.py │ │ ├── fused_bias_act.cpp │ │ ├── fused_bias_act_kernel.cu │ │ ├── libcudart.so │ │ ├── upfirdn2d.cpp │ │ ├── upfirdn2d.py │ │ └── upfirdn2d_kernel.cu │ ├── optimizer.py │ └── up_or_down_sampling.py │ ├── options.py │ ├── options │ └── train │ │ └── ir-sde.yml │ ├── str_utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py │ ├── train.py │ └── utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py ├── structure └── config │ └── inpainting │ ├── data │ ├── BokehLQGT_dataset.py │ ├── BokehLQ_dataset.py │ ├── GT_dataset.py │ ├── LQGT_dataset.py │ ├── LQ_dataset.py │ ├── StereoLQGT_dataset.py │ ├── StereoLQ_dataset.py │ ├── __init__.py │ ├── canny.py │ ├── data_sampler.py │ └── util.py │ ├── demo.sh │ ├── file_utils.py │ ├── models │ ├── __init__.py │ ├── base_model.py │ ├── canny.py │ ├── denoising_model.py │ ├── dense_layer.py │ ├── layers.py │ ├── lr_scheduler.py │ ├── modules │ │ ├── DenoisingUNet_arch.py │ │ ├── __init__.py │ │ ├── loss.py │ │ └── module_util.py │ ├── networks.py │ ├── op │ │ ├── LICENSE_MIT │ │ ├── __init__.py │ │ ├── cuda-installer.log │ │ ├── fused_act.py │ │ ├── fused_bias_act.cpp │ │ ├── fused_bias_act_kernel.cu │ │ ├── libcudart.so │ │ ├── upfirdn2d.cpp │ │ ├── upfirdn2d.py │ │ └── upfirdn2d_kernel.cu │ ├── optimizer.py │ └── up_or_down_sampling.py │ ├── options.py │ ├── options │ └── train │ │ └── ir-sde.yml │ ├── str_utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py │ ├── train.py │ └── utils │ ├── __init__.py │ ├── deg_utils.py │ ├── file_utils.py │ ├── img_utils.py │ └── sde_utils.py └── texture └── config └── inpainting ├── data ├── BokehLQGT_dataset.py ├── BokehLQ_dataset.py ├── GT_dataset.py ├── LQGT_dataset.py ├── LQ_dataset.py ├── StereoLQGT_dataset.py ├── StereoLQ_dataset.py ├── __init__.py ├── canny.py ├── data_sampler.py └── util.py ├── demo.sh ├── file_utils.py ├── models ├── __init__.py ├── base_model.py ├── canny.py ├── denoising_model.py ├── dense_layer.py ├── layers.py ├── lr_scheduler.py ├── modules │ ├── DenoisingUNet_arch.py │ ├── __init__.py │ ├── loss.py │ └── module_util.py ├── networks.py ├── op │ ├── LICENSE_MIT │ ├── __init__.py │ ├── cuda-installer.log │ ├── fused_act.py │ ├── fused_bias_act.cpp │ ├── fused_bias_act_kernel.cu │ ├── libcudart.so │ ├── upfirdn2d.cpp │ ├── upfirdn2d.py │ └── upfirdn2d_kernel.cu ├── optimizer.py └── up_or_down_sampling.py ├── options.py ├── options └── train │ └── ir-sde.yml ├── str_utils ├── __init__.py ├── deg_utils.py ├── file_utils.py ├── img_utils.py └── sde_utils.py ├── train.py └── utils ├── __init__.py ├── deg_utils.py ├── file_utils.py ├── img_utils.py └── sde_utils.py /image/image1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/image/image1.png -------------------------------------------------------------------------------- /image/image2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/image/image2.png -------------------------------------------------------------------------------- /image/image3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/image/image3.png -------------------------------------------------------------------------------- /image/image4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/image/image4.png -------------------------------------------------------------------------------- /image/table.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/image/table.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | einops==0.6.0 2 | lmdb==1.3.0 3 | lpips==0.1.4 4 | numpy==1.23.5 5 | opencv-python==4.6.0.66 6 | Pillow==9.3.0 7 | PyYAML==6.0 8 | scipy==1.9.3 9 | tensorboardX==2.5.1 10 | timm==0.6.12 11 | torch==1.13.0 12 | torchsummaryX==1.3.0 13 | torchvision==0.14.0 14 | tqdm 15 | gradio 16 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/BokehLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class BokehLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and GT image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LR_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from image files 32 | if opt["data_type"] == "img": 33 | self.LR_paths = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) # LR list 36 | self.metas = self._read_meta_data(opt["dataroot_meta"]) 37 | else: 38 | print("Error: data_type is not matched in Dataset") 39 | 40 | def _read_meta_data(self, meta_file_path: str): 41 | """Read the meta file containing source / target lens and disparity for each image. 42 | Args: 43 | meta_file_path (str): File path 44 | Raises: 45 | ValueError: File not found. 46 | Returns: 47 | dict: Meta dict of tuples like {id: (id, src_lens, tgt_lens, disparity)}. 48 | """ 49 | if not os.path.isfile(meta_file_path): 50 | raise ValueError(f"Meta file missing under {meta_file_path}.") 51 | 52 | meta = {} 53 | with open(meta_file_path, "r") as f: 54 | lines = f.readlines() 55 | 56 | for line in lines: 57 | id, src_lens, tgt_lens, disparity = [part.strip() for part in line.split(",")] 58 | meta[id] = (src_lens, tgt_lens, disparity) 59 | return meta 60 | 61 | def lenstr2tensor(self, lenstr, scale=1.): 62 | # Canon50mm -> -1, Sony50mm -> 1 63 | lenstr = lenstr.replace('Canon50mmf', '-') 64 | lenstr = lenstr.replace('Sony50mmf', '') 65 | lenstr = lenstr.replace('BS', '') 66 | return torch.tensor(float(lenstr)) * scale 67 | 68 | def __getitem__(self, index): 69 | 70 | # get LR image 71 | LR_path = self.LR_paths[index] 72 | img_LR = util.read_img(self.LR_env, LR_path, None) 73 | 74 | id = os.path.basename(LR_path).split(".")[0] 75 | src_lens, tgt_lens, disparity = self.metas[id] 76 | 77 | src_lens = self.lenstr2tensor(src_lens, 10) 78 | tgt_lens = self.lenstr2tensor(tgt_lens, 10) 79 | disparity = self.lenstr2tensor(disparity) 80 | 81 | # change color space if necessary 82 | if self.opt["color"]: 83 | H, W, C = img_LR.shape 84 | img_LR = util.channel_convert(C, self.opt["color"], [img_LR])[ 85 | 0 86 | ] # TODO during val no definition 87 | 88 | 89 | # BGR to RGB, HWC to CHW, numpy to tensor 90 | if img_LR.shape[2] == 3: 91 | img_LR = img_LR[:, :, [2, 1, 0]] 92 | 93 | img_LR = torch.from_numpy( 94 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 95 | ).float() 96 | 97 | return { 98 | "LQ": img_LR, 99 | "src_lens": src_lens, 100 | "tgt_lens": tgt_lens, 101 | "disparity": disparity, 102 | "LQ_path": LR_path, 103 | } 104 | 105 | def __len__(self): 106 | return len(self.LR_paths) 107 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/GT_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | 12 | 13 | import numpy as np 14 | from PIL import Image 15 | 16 | from torchvision import transforms 17 | 18 | from skimage.feature import canny 19 | from skimage.color import gray2rgb, rgb2gray 20 | 21 | 22 | def tensor_to_image(): 23 | 24 | return transforms.ToPILImage() 25 | 26 | 27 | def image_to_tensor(): 28 | 29 | return transforms.ToTensor() 30 | 31 | 32 | def image_to_edge(image, sigma): 33 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 34 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 35 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 36 | 37 | return edge, gray_image 38 | 39 | try: 40 | sys.path.append("..") 41 | import data.util as util 42 | except ImportError: 43 | pass 44 | 45 | 46 | class GTDataset(data.Dataset): 47 | """ 48 | Read LR (Low Quality, here is LR) and GT image pairs. 49 | The pair is ensured by 'sorted' function, so please check the name convention. 50 | """ 51 | 52 | def __init__(self, opt): 53 | super().__init__() 54 | self.opt = opt 55 | self.GT_paths = None 56 | self.GT_env = None # environment for lmdb 57 | self.GT_size = opt["GT_size"] 58 | 59 | # read image list from lmdb or image files 60 | if opt["data_type"] == "lmdb": 61 | self.GT_paths, self.GT_sizes = util.get_image_paths( 62 | opt["data_type"], opt["dataroot_GT"] 63 | ) 64 | elif opt["data_type"] == "img": 65 | self.GT_paths = util.get_image_paths( 66 | opt["data_type"], opt["dataroot_GT"] 67 | ) # GT list 68 | else: 69 | print("Error: data_type is not matched in Dataset") 70 | assert self.GT_paths, "Error: GT paths are empty." 71 | print("dataset length: {}".format(len(self.GT_paths))) 72 | self.random_scale_list = [1] 73 | 74 | def _init_lmdb(self): 75 | # https://github.com/chainer/chainermn/issues/129 76 | self.GT_env = lmdb.open( 77 | self.opt["dataroot_GT"], 78 | readonly=True, 79 | lock=False, 80 | readahead=False, 81 | meminit=False, 82 | ) 83 | 84 | def __getitem__(self, index): 85 | if self.opt["data_type"] == "lmdb": 86 | if self.GT_env is None: 87 | self._init_lmdb() 88 | 89 | GT_path = None 90 | GT_size = self.opt["GT_size"] 91 | 92 | # get GT image 93 | GT_path = self.GT_paths[index] 94 | if self.opt["data_type"] == "lmdb": 95 | resolution = [int(s) for s in self.GT_sizes[index].split("_")] 96 | else: 97 | resolution = None 98 | img_GT = util.read_img( 99 | self.GT_env, GT_path, resolution 100 | ) # return: Numpy float32, HWC, BGR, [0,1] 101 | 102 | if self.opt["phase"] == "train": 103 | H, W, C = img_GT.shape 104 | 105 | rnd_h = random.randint(0, max(0, H - GT_size)) 106 | rnd_w = random.randint(0, max(0, W - GT_size)) 107 | img_GT = img_GT[rnd_h : rnd_h + GT_size, rnd_w : rnd_w + GT_size, :] 108 | 109 | # augmentation - flip, rotate 110 | img_GT = util.augment( 111 | img_GT, 112 | self.opt["use_flip"], 113 | self.opt["use_rot"], 114 | self.opt["mode"], 115 | ) 116 | else: 117 | img_GT = cv2.resize( 118 | np.copy(img_GT), (256, 256), interpolation=cv2.INTER_LINEAR 119 | ) 120 | 121 | # change color space if necessary 122 | if self.opt["color"]: 123 | img_GT = util.channel_convert(img_GT.shape[2], self.opt["color"], [img_GT])[ 124 | 0 125 | ] 126 | print(GT_path) 127 | # BGR to RGB, HWC to CHW, numpy to tensor 128 | if img_GT.shape[2] == 3: 129 | img_GT = img_GT[:, :, [2, 1, 0]] 130 | img_GT = torch.from_numpy( 131 | np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1))) 132 | ).float() 133 | 134 | GT_edge,GT_gray = image_to_edge(img_GT, sigma=2.) 135 | return {"GT": img_GT, "GT_path": GT_path, "GT_edge": GT_edge, "GT_gray": GT_gray} 136 | 137 | def __len__(self): 138 | return len(self.GT_paths) 139 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/LQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class LQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path = None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path = self.LQ_paths[index] 67 | if self.opt["data_type"] == "lmdb": 68 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 69 | else: 70 | resolution = None 71 | img_LR = util.read_img( 72 | self.LR_env, LR_path, resolution 73 | ) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # modcrop in the validation / test phase 76 | if self.opt["phase"] != "train": 77 | img_LR = util.modcrop(img_LR, scale) 78 | 79 | if self.opt["phase"] == "train": 80 | H, W, C = img_LR.shape 81 | 82 | rnd_h = random.randint(0, max(0, H - LR_size)) 83 | rnd_w = random.randint(0, max(0, W - LR_size)) 84 | img_LR = img_LR[rnd_h : rnd_h + LR_size, rnd_w : rnd_w + LR_size, :] 85 | 86 | # augmentation - flip, rotate 87 | img_LR = util.augment( 88 | img_LR, 89 | self.opt["use_flip"], 90 | self.opt["use_rot"], 91 | self.opt["mode"], 92 | ) 93 | 94 | # change color space if necessary 95 | if self.opt["color"]: 96 | img_LR = util.channel_convert(img_LR.shape[2], self.opt["color"], [img_LR])[ 97 | 0 98 | ] 99 | 100 | # BGR to RGB, HWC to CHW, numpy to tensor 101 | if img_LR.shape[2] == 3: 102 | img_LR = img_LR[:, :, [2, 1, 0]] 103 | img_LR = torch.from_numpy( 104 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 105 | ).float() 106 | 107 | return {"LQ": img_LR, "LQ_path": LR_path} 108 | 109 | def __len__(self): 110 | return len(self.LQ_paths) 111 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/StereoLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class StereoLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path_L, LR_path_R = None, None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path_L = self.LQ_paths[index*2] 67 | LR_path_R = self.LQ_paths[index*2+1] 68 | if self.opt["data_type"] == "lmdb": 69 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 70 | else: 71 | resolution = None 72 | imgL_LR = util.read_img(self.LR_env, LR_path_L, resolution) # return: Numpy float32, HWC, BGR, [0,1] 73 | imgR_LR = util.read_img(self.LR_env, LR_path_R, resolution) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # change color space if necessary 76 | if self.opt["color"]: 77 | imgL_LR, imgR_LR = util.channel_convert( 78 | imgL_LR.shape[2], self.opt["color"], [imgL_LR, imgR_LR]) 79 | 80 | # BGR to RGB, HWC to CHW, numpy to tensor 81 | if imgL_LR.shape[2] == 3: 82 | imgL_LR = imgL_LR[:, :, [2, 1, 0]] 83 | imgR_LR = imgR_LR[:, :, [2, 1, 0]] 84 | imgL_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgL_LR, (2, 0, 1)))).float() 85 | imgR_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgR_LR, (2, 0, 1)))).float() 86 | 87 | img_LR = torch.cat([imgL_LR, imgR_LR], dim=0) 88 | 89 | return {"LQ": img_LR, "LQ_path": LR_path_L} 90 | 91 | def __len__(self): 92 | return len(self.LQ_paths) // 2 93 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/__init__.py: -------------------------------------------------------------------------------- 1 | """create dataset and dataloader""" 2 | import logging 3 | 4 | import torch 5 | import torch.utils.data 6 | 7 | 8 | def create_dataloader(dataset, dataset_opt, opt=None, sampler=None): 9 | phase = dataset_opt["phase"] 10 | if phase == "train": 11 | if opt["dist"]: 12 | world_size = torch.distributed.get_world_size() 13 | num_workers = dataset_opt["n_workers"] 14 | assert dataset_opt["batch_size"] % world_size == 0 15 | batch_size = dataset_opt["batch_size"] // world_size 16 | shuffle = True 17 | else: 18 | num_workers = dataset_opt["n_workers"] * len(opt["gpu_ids"]) 19 | batch_size = dataset_opt["batch_size"] 20 | shuffle = True 21 | return torch.utils.data.DataLoader( 22 | dataset, 23 | batch_size=batch_size, 24 | shuffle=shuffle, 25 | num_workers=num_workers, 26 | sampler=sampler, 27 | drop_last=True, 28 | pin_memory=False, 29 | ) 30 | else: 31 | return torch.utils.data.DataLoader( 32 | dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=(phase=="val") 33 | ) 34 | 35 | 36 | def create_dataset(dataset_opt): 37 | mode = dataset_opt["mode"] 38 | if mode == "LQ": # Predictor 39 | from data.LQ_dataset import LQDataset as D 40 | dataset = D(dataset_opt) 41 | elif mode == "LQGT": # SFTMD 42 | from data.LQGT_dataset import LQGTDataset as D 43 | dataset = D(dataset_opt) 44 | elif mode == "GT": # Corrector 45 | from data.GT_dataset import GTDataset as D 46 | dataset = D(dataset_opt) 47 | elif mode == 'SteLQGT': 48 | from data.StereoLQGT_dataset import StereoLQGTDataset as D 49 | dataset = D(dataset_opt) 50 | elif mode == 'SteLQ': 51 | from data.StereoLQ_dataset import StereoLQDataset as D 52 | dataset = D(dataset_opt) 53 | elif mode == 'BokehLQGT': 54 | from data.BokehLQGT_dataset import BokehLQGTDataset as D 55 | dataset = D(dataset_opt) 56 | elif mode == 'BokehLQ': 57 | from data.BokehLQ_dataset import BokehLQDataset as D 58 | dataset = D(dataset_opt) 59 | else: 60 | raise NotImplementedError("Dataset [{:s}] is not recognized.".format(mode)) 61 | 62 | logger = logging.getLogger("base") 63 | logger.info( 64 | "Dataset [{:s} - {:s}] is created.".format( 65 | dataset.__class__.__name__, dataset_opt["name"] 66 | ) 67 | ) 68 | return dataset 69 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def image_to_edge(image, sigma): 21 | 22 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 25 | 26 | return edge, gray_image 27 | 28 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/data/data_sampler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from torch.utils.data.distributed.DistributedSampler 3 | Support enlarging the dataset for *iter-oriented* training, for saving time when restart the 4 | dataloader after each epoch 5 | """ 6 | import math 7 | 8 | import torch 9 | import torch.distributed as dist 10 | from torch.utils.data.sampler import Sampler 11 | 12 | 13 | class DistIterSampler(Sampler): 14 | """Sampler that restricts data loading to a subset of the dataset. 15 | 16 | It is especially useful in conjunction with 17 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 18 | process can pass a DistributedSampler instance as a DataLoader sampler, 19 | and load a subset of the original dataset that is exclusive to it. 20 | 21 | .. note:: 22 | Dataset is assumed to be of constant size. 23 | 24 | Arguments: 25 | dataset: Dataset used for sampling. 26 | num_replicas (optional): Number of processes participating in 27 | distributed training. 28 | rank (optional): Rank of the current process within num_replicas. 29 | """ 30 | 31 | def __init__(self, dataset, num_replicas=None, rank=None, ratio=100): 32 | if num_replicas is None: 33 | if not dist.is_available(): 34 | raise RuntimeError("Requires distributed package to be available") 35 | num_replicas = dist.get_world_size() 36 | if rank is None: 37 | if not dist.is_available(): 38 | raise RuntimeError("Requires distributed package to be available") 39 | rank = dist.get_rank() 40 | self.dataset = dataset 41 | self.num_replicas = num_replicas 42 | self.rank = rank 43 | self.epoch = 0 44 | self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas)) 45 | self.total_size = self.num_samples * self.num_replicas 46 | 47 | def __iter__(self): 48 | # deterministically shuffle based on epoch 49 | g = torch.Generator() 50 | g.manual_seed(self.epoch) 51 | indices = torch.randperm( 52 | self.total_size, generator=g 53 | ).tolist() # Returns a random permutation of integers from 0 to n - 1 54 | 55 | dsize = len(self.dataset) 56 | indices = [v % dsize for v in indices] 57 | 58 | # subsample 59 | indices = indices[self.rank : self.total_size : self.num_replicas] 60 | assert len(indices) == self.num_samples 61 | 62 | return iter(indices) 63 | 64 | def __len__(self): 65 | return self.num_samples 66 | 67 | def set_epoch(self, epoch): 68 | self.epoch = epoch 69 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logger = logging.getLogger("base") 4 | 5 | 6 | def create_model(opt): 7 | model = opt["model"] 8 | 9 | if model == "denoising": 10 | from .denoising_model import DenoisingModel as M 11 | else: 12 | raise NotImplementedError("Model [{:s}] not recognized.".format(model)) 13 | m = M(opt) 14 | logger.info("Model [{:s}] is created.".format(m.__class__.__name__)) 15 | return m 16 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import OrderedDict 3 | 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.parallel import DistributedDataParallel 7 | 8 | 9 | class BaseModel: 10 | def __init__(self, opt): 11 | self.opt = opt 12 | self.device = torch.device("cuda" if opt["gpu_ids"] is not None else "cpu") 13 | self.is_train = opt["is_train"] 14 | self.schedulers = [] 15 | self.optimizers = [] 16 | 17 | def feed_data(self, data): 18 | pass 19 | 20 | def optimize_parameters(self): 21 | pass 22 | 23 | def get_current_visuals(self): 24 | pass 25 | 26 | def get_current_losses(self): 27 | pass 28 | 29 | def print_network(self): 30 | pass 31 | 32 | def save(self, label): 33 | pass 34 | 35 | def load(self): 36 | pass 37 | 38 | def _set_lr(self, lr_groups_l): 39 | """set learning rate for warmup, 40 | lr_groups_l: list for lr_groups. each for a optimizer""" 41 | for optimizer, lr_groups in zip(self.optimizers, lr_groups_l): 42 | for param_group, lr in zip(optimizer.param_groups, lr_groups): 43 | param_group["lr"] = lr 44 | 45 | def _get_init_lr(self): 46 | # get the initial lr, which is set by the scheduler 47 | init_lr_groups_l = [] 48 | for optimizer in self.optimizers: 49 | init_lr_groups_l.append([v["initial_lr"] for v in optimizer.param_groups]) 50 | return init_lr_groups_l 51 | 52 | def update_learning_rate(self, cur_iter, warmup_iter=-1): 53 | for scheduler in self.schedulers: 54 | scheduler.step() 55 | #### set up warm up learning rate 56 | if cur_iter < warmup_iter: 57 | # get initial lr for each group 58 | init_lr_g_l = self._get_init_lr() 59 | # modify warming-up learning rates 60 | warm_up_lr_l = [] 61 | for init_lr_g in init_lr_g_l: 62 | warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g]) 63 | # set learning rate 64 | self._set_lr(warm_up_lr_l) 65 | 66 | def get_current_learning_rate(self): 67 | # return self.schedulers[0].get_lr()[0] 68 | return self.optimizers[0].param_groups[0]["lr"] 69 | 70 | def get_network_description(self, network): 71 | """Get the string and total parameters of the network""" 72 | if isinstance(network, nn.DataParallel) or isinstance( 73 | network, DistributedDataParallel 74 | ): 75 | network = network.module 76 | s = str(network) 77 | n = sum(map(lambda x: x.numel(), network.parameters())) 78 | return s, n 79 | 80 | def save_network(self, network, network_label, iter_label): 81 | save_filename = "{}_{}.pth".format(iter_label, network_label) 82 | save_path = os.path.join(self.opt["path"]["models"], save_filename) 83 | if isinstance(network, nn.DataParallel) or isinstance( 84 | network, DistributedDataParallel 85 | ): 86 | network = network.module 87 | state_dict = network.state_dict() 88 | for key, param in state_dict.items(): 89 | state_dict[key] = param.cpu() 90 | torch.save(state_dict, save_path) 91 | 92 | def load_network(self, load_path, network, strict=True): 93 | if isinstance(network, nn.DataParallel) or isinstance( 94 | network, DistributedDataParallel 95 | ): 96 | network = network.module 97 | print(network) 98 | load_net = torch.load(load_path) 99 | load_net_clean = OrderedDict() # remove unnecessary 'module.' 100 | for k, v in load_net.items(): 101 | if k.startswith("module."): 102 | load_net_clean[k[7:]] = v 103 | else: 104 | load_net_clean[k] = v 105 | 106 | network.load_state_dict(load_net_clean, strict=strict) 107 | 108 | def save_training_state(self, epoch, iter_step): 109 | """Saves training state during training, which will be used for resuming""" 110 | state = {"epoch": epoch, "iter": iter_step, "schedulers": [], "optimizers": []} 111 | for s in self.schedulers: 112 | state["schedulers"].append(s.state_dict()) 113 | for o in self.optimizers: 114 | state["optimizers"].append(o.state_dict()) 115 | save_filename = "{}.state".format(iter_step) 116 | save_path = os.path.join(self.opt["path"]["training_state"], save_filename) 117 | torch.save(state, save_path) 118 | 119 | def resume_training(self, resume_state): 120 | """Resume the optimizers and schedulers for training""" 121 | resume_optimizers = resume_state["optimizers"] 122 | resume_schedulers = resume_state["schedulers"] 123 | assert len(resume_optimizers) == len( 124 | self.optimizers 125 | ), "Wrong lengths of optimizers" 126 | assert len(resume_schedulers) == len( 127 | self.schedulers 128 | ), "Wrong lengths of schedulers" 129 | for i, o in enumerate(resume_optimizers): 130 | self.optimizers[i].load_state_dict(o) 131 | for i, s in enumerate(resume_schedulers): 132 | self.schedulers[i].load_state_dict(s) 133 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def gray_to_edge(image, sigma): 21 | 22 | gray_image = np.array(tensor_to_image()(image)) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | 25 | return edge 26 | 27 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/dense_layer.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from a file released under the MIT License. 5 | # 6 | # Source: 7 | # https://github.com/CW-Huang/sdeflow-light/blob/524650bc5ad69522b3e0905672deef0650374512/lib/models/unet.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_MIT). The modifications 11 | # to this file are subject to the same MIT License. 12 | # --------------------------------------------------------------- 13 | 14 | 15 | import math 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | from torch.nn.init import _calculate_fan_in_and_fan_out 20 | import numpy as np 21 | 22 | 23 | def _calculate_correct_fan(tensor, mode): 24 | """ 25 | copied and modified from https://github.com/pytorch/pytorch/blob/master/torch/nn/init.py#L337 26 | """ 27 | mode = mode.lower() 28 | valid_modes = ['fan_in', 'fan_out', 'fan_avg'] 29 | if mode not in valid_modes: 30 | raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes)) 31 | 32 | fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) 33 | return fan_in if mode == 'fan_in' else fan_out 34 | 35 | 36 | def kaiming_uniform_(tensor, gain=1., mode='fan_in'): 37 | r"""Fills the input `Tensor` with values according to the method 38 | described in `Delving deep into rectifiers: Surpassing human-level 39 | performance on ImageNet classification` - He, K. et al. (2015), using a 40 | uniform distribution. The resulting tensor will have values sampled from 41 | :math:`\mathcal{U}(-\text{bound}, \text{bound})` where 42 | .. math:: 43 | \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} 44 | Also known as He initialization. 45 | Args: 46 | tensor: an n-dimensional `torch.Tensor` 47 | gain: multiplier to the dispersion 48 | mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'`` 49 | preserves the magnitude of the variance of the weights in the 50 | forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the 51 | backwards pass. 52 | Examples: 53 | >>> w = torch.empty(3, 5) 54 | >>> nn.init.kaiming_uniform_(w, mode='fan_in') 55 | """ 56 | fan = _calculate_correct_fan(tensor, mode) 57 | var = gain / max(1., fan) 58 | bound = math.sqrt(3.0 * var) # Calculate uniform bounds from standard deviation 59 | with torch.no_grad(): 60 | return tensor.uniform_(-bound, bound) 61 | 62 | 63 | def variance_scaling_init_(tensor, scale): 64 | return kaiming_uniform_(tensor, gain=1e-10 if scale == 0 else scale, mode='fan_avg') 65 | 66 | 67 | def dense(in_channels, out_channels, init_scale=1.): 68 | lin = nn.Linear(in_channels, out_channels) 69 | variance_scaling_init_(lin.weight, scale=init_scale) 70 | nn.init.zeros_(lin.bias) 71 | return lin 72 | 73 | def conv2d(in_planes, out_planes, kernel_size=(3, 3), stride=1, dilation=1, padding=1, bias=True, padding_mode='zeros', 74 | init_scale=1.): 75 | conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 76 | bias=bias, padding_mode=padding_mode) 77 | variance_scaling_init_(conv.weight, scale=init_scale) 78 | if bias: 79 | nn.init.zeros_(conv.bias) 80 | return conv 81 | 82 | 83 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .DenoisingUNet_arch import ConditionalUNet,ConditionalUNets -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/modules/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import einops 5 | import numpy as np 6 | import sys 7 | 8 | class MatchingLoss(nn.Module): 9 | def __init__(self, loss_type='l1', is_weighted=False): 10 | super().__init__() 11 | self.is_weighted = is_weighted 12 | 13 | if loss_type == 'l1': 14 | self.loss_fn = F.l1_loss 15 | elif loss_type == 'l2': 16 | self.loss_fn = F.mse_loss 17 | else: 18 | raise ValueError(f'invalid loss type {loss_type}') 19 | 20 | def forward(self, predict, target, mask, weights=None): 21 | 22 | lossm = self.loss_fn(predict * (1 - mask), target * (1 - mask), reduction='none') 23 | lossm = einops.reduce(lossm, 'b ... -> b (...)', 'mean') 24 | 25 | lossu = self.loss_fn(predict * mask, target * mask, reduction='none') 26 | lossu = einops.reduce(lossu, 'b ... -> b (...)', 'mean') 27 | 28 | loss = lossu + 10 * lossm 29 | if self.is_weighted and weights is not None: 30 | loss = weights * loss 31 | 32 | return loss.mean() -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/LICENSE_MIT: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Kim Seonghyeon 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/__init__.py: -------------------------------------------------------------------------------- 1 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 2 | from .upfirdn2d import upfirdn2d 3 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/fused_act.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # --------------------------------------------------------------- 4 | 5 | """ Originated from https://github.com/rosinality/stylegan2-pytorch 6 | The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | """ 8 | 9 | import os 10 | 11 | import torch 12 | from torch import nn 13 | from torch.nn import functional as F 14 | from torch.autograd import Function 15 | from torch.utils.cpp_extension import load 16 | 17 | 18 | module_path = os.path.dirname(__file__) 19 | fused = load( 20 | "fused", 21 | sources=[ 22 | os.path.join(module_path, "fused_bias_act.cpp"), 23 | os.path.join(module_path, "fused_bias_act_kernel.cu"), 24 | ], 25 | ) 26 | 27 | 28 | class FusedLeakyReLUFunctionBackward(Function): 29 | @staticmethod 30 | def forward(ctx, grad_output, out, negative_slope, scale): 31 | ctx.save_for_backward(out) 32 | ctx.negative_slope = negative_slope 33 | ctx.scale = scale 34 | 35 | empty = grad_output.new_empty(0) 36 | 37 | grad_input = fused.fused_bias_act( 38 | grad_output, empty, out, 3, 1, negative_slope, scale 39 | ) 40 | 41 | dim = [0] 42 | 43 | if grad_input.ndim > 2: 44 | dim += list(range(2, grad_input.ndim)) 45 | 46 | grad_bias = grad_input.sum(dim).detach() 47 | 48 | return grad_input, grad_bias 49 | 50 | @staticmethod 51 | def backward(ctx, gradgrad_input, gradgrad_bias): 52 | out, = ctx.saved_tensors 53 | gradgrad_out = fused.fused_bias_act( 54 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 55 | ) 56 | 57 | return gradgrad_out, None, None, None 58 | 59 | 60 | class FusedLeakyReLUFunction(Function): 61 | @staticmethod 62 | def forward(ctx, input, bias, negative_slope, scale): 63 | empty = input.new_empty(0) 64 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 65 | ctx.save_for_backward(out) 66 | ctx.negative_slope = negative_slope 67 | ctx.scale = scale 68 | 69 | return out 70 | 71 | @staticmethod 72 | def backward(ctx, grad_output): 73 | out, = ctx.saved_tensors 74 | 75 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 76 | grad_output, out, ctx.negative_slope, ctx.scale 77 | ) 78 | 79 | return grad_input, grad_bias, None, None 80 | 81 | 82 | class FusedLeakyReLU(nn.Module): 83 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): 84 | super().__init__() 85 | 86 | self.bias = nn.Parameter(torch.zeros(channel)) 87 | self.negative_slope = negative_slope 88 | self.scale = scale 89 | 90 | def forward(self, input): 91 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 92 | 93 | 94 | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): 95 | if input.device.type == "cpu": 96 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 97 | return ( 98 | F.leaky_relu( 99 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 100 | ) 101 | * scale 102 | ) 103 | 104 | else: 105 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 106 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | #include 9 | 10 | 11 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale); 13 | 14 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 15 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 17 | 18 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 19 | int act, int grad, float alpha, float scale) { 20 | CHECK_CUDA(input); 21 | CHECK_CUDA(bias); 22 | 23 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 24 | } 25 | 26 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 27 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 28 | } -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved. 3 | // --------------------------------------------------------------- 4 | // 5 | // This work is made available under the Nvidia Source Code License-NC. 6 | // To view a copy of this license, visit 7 | // https://nvlabs.github.io/stylegan2/license.html 8 | 9 | #include 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | #include 17 | #include 18 | 19 | 20 | template 21 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 22 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 23 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 24 | 25 | scalar_t zero = 0.0; 26 | 27 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 28 | scalar_t x = p_x[xi]; 29 | 30 | if (use_bias) { 31 | x += p_b[(xi / step_b) % size_b]; 32 | } 33 | 34 | scalar_t ref = use_ref ? p_ref[xi] : zero; 35 | 36 | scalar_t y; 37 | 38 | switch (act * 10 + grad) { 39 | default: 40 | case 10: y = x; break; 41 | case 11: y = x; break; 42 | case 12: y = 0.0; break; 43 | 44 | case 30: y = (x > 0.0) ? x : x * alpha; break; 45 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 46 | case 32: y = 0.0; break; 47 | } 48 | 49 | out[xi] = y * scale; 50 | } 51 | } 52 | 53 | 54 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 55 | int act, int grad, float alpha, float scale) { 56 | int curDevice = -1; 57 | cudaGetDevice(&curDevice); 58 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 59 | 60 | auto x = input.contiguous(); 61 | auto b = bias.contiguous(); 62 | auto ref = refer.contiguous(); 63 | 64 | int use_bias = b.numel() ? 1 : 0; 65 | int use_ref = ref.numel() ? 1 : 0; 66 | 67 | int size_x = x.numel(); 68 | int size_b = b.numel(); 69 | int step_b = 1; 70 | 71 | for (int i = 1 + 1; i < x.dim(); i++) { 72 | step_b *= x.size(i); 73 | } 74 | 75 | int loop_x = 4; 76 | int block_size = 4 * 32; 77 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 78 | 79 | auto y = torch::empty_like(x); 80 | 81 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 82 | fused_bias_act_kernel<<>>( 83 | y.data_ptr(), 84 | x.data_ptr(), 85 | b.data_ptr(), 86 | ref.data_ptr(), 87 | act, 88 | grad, 89 | alpha, 90 | scale, 91 | loop_x, 92 | size_x, 93 | step_b, 94 | size_b, 95 | use_bias, 96 | use_ref 97 | ); 98 | }); 99 | 100 | return y; 101 | } -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | 9 | #include 10 | 11 | 12 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 15 | 16 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 17 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 18 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 19 | 20 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 21 | int up_x, int up_y, int down_x, int down_y, 22 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 23 | CHECK_CUDA(input); 24 | CHECK_CUDA(kernel); 25 | 26 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 27 | } 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 31 | } -------------------------------------------------------------------------------- /test/texture/config/inpainting/models/optimizer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google Research. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """PyTorch implementation of the Lion optimizer.""" 16 | 17 | import torch 18 | from torch.optim.optimizer import Optimizer 19 | 20 | 21 | class Lion(Optimizer): 22 | r"""Implements Lion algorithm.""" 23 | 24 | def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0): 25 | """Initialize the hyperparameters. 26 | 27 | Args: 28 | params (iterable): iterable of parameters to optimize or dicts defining 29 | parameter groups 30 | lr (float, optional): learning rate (default: 1e-4) 31 | betas (Tuple[float, float], optional): coefficients used for computing 32 | running averages of gradient and its square (default: (0.9, 0.99)) 33 | weight_decay (float, optional): weight decay coefficient (default: 0) 34 | """ 35 | 36 | if not 0.0 <= lr: 37 | raise ValueError('Invalid learning rate: {}'.format(lr)) 38 | if not 0.0 <= betas[0] < 1.0: 39 | raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) 40 | if not 0.0 <= betas[1] < 1.0: 41 | raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) 42 | defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) 43 | super().__init__(params, defaults) 44 | 45 | @torch.no_grad() 46 | def step(self, closure=None): 47 | """Performs a single optimization step. 48 | 49 | Args: 50 | closure (callable, optional): A closure that reevaluates the model 51 | and returns the loss. 52 | 53 | Returns: 54 | the loss. 55 | """ 56 | loss = None 57 | if closure is not None: 58 | with torch.enable_grad(): 59 | loss = closure() 60 | 61 | for group in self.param_groups: 62 | for p in group['params']: 63 | if p.grad is None: 64 | continue 65 | 66 | # Perform stepweight decay 67 | p.data.mul_(1 - group['lr'] * group['weight_decay']) 68 | 69 | grad = p.grad 70 | state = self.state[p] 71 | # State initialization 72 | if len(state) == 0: 73 | # Exponential moving average of gradient values 74 | state['exp_avg'] = torch.zeros_like(p) 75 | 76 | exp_avg = state['exp_avg'] 77 | beta1, beta2 = group['betas'] 78 | 79 | # Weight update 80 | update = exp_avg * beta1 + grad * (1 - beta1) 81 | p.add_(torch.sign(update), alpha=-group['lr']) 82 | # Decay the momentum running average coefficient 83 | exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) 84 | 85 | return loss 86 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/options.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import os.path as osp 4 | import sys 5 | import math 6 | 7 | import yaml 8 | 9 | from file_utils import OrderedYaml 10 | 11 | Loader, Dumper = OrderedYaml() 12 | 13 | def parse(opt_path, is_train=True): 14 | with open(opt_path, mode="r") as f: 15 | opt = yaml.load(f, Loader=Loader) 16 | # export CUDA_VISIBLE_DEVICES 17 | gpu_list = ",".join(str(x) for x in opt["gpu_ids"]) 18 | os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list 19 | print("export CUDA_VISIBLE_DEVICES=" + gpu_list) 20 | 21 | opt["is_train"] = is_train 22 | 23 | scale = 1 24 | if opt['distortion'] == 'sr': 25 | scale = opt['degradation']['scale'] 26 | 27 | ##### sr network #### 28 | opt["network_G"]["setting"]["upscale"] = scale 29 | # opt["network_G"]["setting"]["in_nc"] *= scale**2 30 | 31 | # datasets 32 | for phase, dataset in opt["datasets"].items(): 33 | phase = phase.split("_")[0] 34 | print(dataset) 35 | dataset["phase"] = phase 36 | dataset["scale"] = scale 37 | 38 | is_lmdb = False 39 | if dataset.get("dataroot_GT", None) is not None: 40 | dataset["dataroot_GT"] = osp.expanduser(dataset["dataroot_GT"]) 41 | if dataset["dataroot_GT"].endswith("lmdb"): 42 | is_lmdb = True 43 | # if dataset.get('dataroot_GT_bg', None) is not None: 44 | # dataset['dataroot_GT_bg'] = osp.expanduser(dataset['dataroot_GT_bg']) 45 | if dataset.get("dataroot_LQ", None) is not None: 46 | dataset["dataroot_LQ"] = osp.expanduser(dataset["dataroot_LQ"]) 47 | if dataset["dataroot_LQ"].endswith("lmdb"): 48 | is_lmdb = True 49 | dataset["data_type"] = "lmdb" if is_lmdb else "img" 50 | if dataset["mode"].endswith("mc"): # for memcached 51 | dataset["data_type"] = "mc" 52 | dataset["mode"] = dataset["mode"].replace("_mc", "") 53 | 54 | # path 55 | for key, path in opt["path"].items(): 56 | if path and key in opt["path"] and key != "strict_load": 57 | opt["path"][key] = osp.expanduser(path) 58 | opt["path"]["root"] = osp.abspath( 59 | osp.join(__file__, osp.pardir, osp.pardir, osp.pardir, osp.pardir) 60 | ) 61 | path = osp.abspath(__file__) 62 | config_dir = path.split("/")[-2] 63 | if is_train: 64 | experiments_root = osp.join( 65 | opt["path"]["root"], "experiments", config_dir, opt["name"] 66 | ) 67 | opt["path"]["experiments_root"] = experiments_root 68 | opt["path"]["models"] = osp.join(experiments_root, "models") 69 | opt["path"]["training_state"] = osp.join(experiments_root, "training_state") 70 | opt["path"]["log"] = experiments_root 71 | opt["path"]["val_images"] = osp.join(experiments_root, "val_images") 72 | 73 | # change some options for debug mode 74 | if "debug" in opt["name"]: 75 | opt["train"]["val_freq"] = 8 76 | opt["logger"]["print_freq"] = 1 77 | opt["logger"]["save_checkpoint_freq"] = 8 78 | else: # test 79 | results_root = osp.join(opt["path"]["root"], "results", config_dir) 80 | opt["path"]["results_root"] = osp.join(results_root, opt["name"]) 81 | opt["path"]["log"] = osp.join(results_root, opt["name"]) 82 | 83 | return opt 84 | 85 | 86 | def dict2str(opt, indent_l=1): 87 | """dict to string for logger""" 88 | msg = "" 89 | for k, v in opt.items(): 90 | if isinstance(v, dict): 91 | msg += " " * (indent_l * 2) + k + ":[\n" 92 | msg += dict2str(v, indent_l + 1) 93 | msg += " " * (indent_l * 2) + "]\n" 94 | else: 95 | msg += " " * (indent_l * 2) + k + ": " + str(v) + "\n" 96 | return msg 97 | 98 | 99 | class NoneDict(dict): 100 | def __missing__(self, key): 101 | return None 102 | 103 | 104 | # convert to NoneDict, which return None for missing key. 105 | def dict_to_nonedict(opt): 106 | if isinstance(opt, dict): 107 | new_opt = dict() 108 | for key, sub_opt in opt.items(): 109 | new_opt[key] = dict_to_nonedict(sub_opt) 110 | return NoneDict(**new_opt) 111 | elif isinstance(opt, list): 112 | return [dict_to_nonedict(sub_opt) for sub_opt in opt] 113 | else: 114 | return opt 115 | 116 | 117 | def check_resume(opt, resume_iter): 118 | """Check resume states and pretrain_model paths""" 119 | logger = logging.getLogger("base") 120 | if opt["path"]["resume_state"]: 121 | if ( 122 | opt["path"].get("pretrain_model_G", None) is not None 123 | or opt["path"].get("pretrain_model_D", None) is not None 124 | ): 125 | logger.warning( 126 | "pretrain_model path will be ignored when resuming training." 127 | ) 128 | 129 | opt["path"]["pretrain_model_G"] = osp.join( 130 | opt["path"]["models"], "{}_G.pth".format(resume_iter) 131 | ) 132 | logger.info("Set [pretrain_model_G] to " + opt["path"]["pretrain_model_G"]) 133 | if "gan" in opt["model"]: 134 | opt["path"]["pretrain_model_D"] = osp.join( 135 | opt["path"]["models"], "{}_D.pth".format(resume_iter) 136 | ) 137 | logger.info("Set [pretrain_model_D] to " + opt["path"]["pretrain_model_D"]) 138 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/options/test/ir-sde.yml: -------------------------------------------------------------------------------- 1 | #### general settings 2 | name: ir-sde 3 | suffix: ~ # add suffix to saved images 4 | model: denoising 5 | distortion: derain 6 | gpu_ids: [3] 7 | 8 | sde: 9 | max_sigma: 30 10 | T: 100 11 | schedule: cosine 12 | eps: 0.005 13 | 14 | degradation: 15 | mask_root: your mask path 16 | 17 | 18 | #### datasets 19 | datasets: 20 | test1: 21 | name: Val_Dataset 22 | mode: GT 23 | dataroot_GT: your image path 24 | 25 | #### network structures 26 | network_G: 27 | which_model_G: ConditionalUNet 28 | setting: 29 | in_nc: 3 30 | out_nc: 3 31 | nf: 64 32 | depth: 4 33 | 34 | network_Gs: 35 | which_model_G: ConditionalUNets 36 | setting: 37 | in_nc: 1 38 | out_nc: 1 39 | nf: 64 40 | depth: 4 41 | 42 | 43 | path: 44 | pretrain_model_G: your path of texture.pth 45 | pretrain_model_Gs: your path of structure.pth 46 | pretrain_model_D: your path of discriminator.pth 47 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/str_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /test/texture/config/inpainting/str_utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import abc 4 | import math 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torchvision.utils as tvutils 10 | 11 | 12 | ########### denoising ############### 13 | def add_noise(tensor, sigma): 14 | sigma = sigma / 255 if sigma > 1 else sigma 15 | return tensor + torch.randn_like(tensor) * sigma 16 | 17 | 18 | ######## inpainting ########### 19 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 20 | batch = tensor.shape[0] 21 | if mask_id < 0: 22 | mask_id = np.random.randint(0, n, batch) 23 | masks = [] 24 | for i in range(batch): 25 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 26 | mask = np.concatenate(masks, axis=0) 27 | else: 28 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 29 | 30 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 31 | # for images are clipped or scaled 32 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 33 | masked_tensor = mask * tensor 34 | return masked_tensor + (1. - mask) 35 | 36 | ######## super-resolution ########### 37 | 38 | def upscale(tensor, scale=4, mode='bicubic'): 39 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 40 | return tensor 41 | 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /test/texture/config/inpainting/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /test/texture/config/inpainting/utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import abc 3 | import math 4 | import numpy as np 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torchvision.utils as tvutils 9 | import cv2 10 | 11 | ########### denoising ############### 12 | def add_noise(tensor, sigma): 13 | sigma = sigma / 255 if sigma > 1 else sigma 14 | return tensor + torch.randn_like(tensor) * sigma 15 | 16 | 17 | ######## inpainting ########### 18 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 19 | batch = tensor.shape[0] 20 | if mask_id < 0: 21 | mask_id = np.random.randint(0, n, batch) 22 | masks = [] 23 | for i in range(batch): 24 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 25 | mask = np.concatenate(masks, axis=0) 26 | else: 27 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 28 | 29 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 30 | # for images are clipped or scaled 31 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 32 | masked_tensor = mask * tensor 33 | return masked_tensor + (1. - mask) 34 | 35 | ######## super-resolution ########### 36 | 37 | def upscale(tensor, scale=4, mode='bicubic'): 38 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 39 | return tensor 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/BokehLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class BokehLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and GT image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LR_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from image files 32 | if opt["data_type"] == "img": 33 | self.LR_paths = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) # LR list 36 | self.metas = self._read_meta_data(opt["dataroot_meta"]) 37 | else: 38 | print("Error: data_type is not matched in Dataset") 39 | 40 | def _read_meta_data(self, meta_file_path: str): 41 | """Read the meta file containing source / target lens and disparity for each image. 42 | Args: 43 | meta_file_path (str): File path 44 | Raises: 45 | ValueError: File not found. 46 | Returns: 47 | dict: Meta dict of tuples like {id: (id, src_lens, tgt_lens, disparity)}. 48 | """ 49 | if not os.path.isfile(meta_file_path): 50 | raise ValueError(f"Meta file missing under {meta_file_path}.") 51 | 52 | meta = {} 53 | with open(meta_file_path, "r") as f: 54 | lines = f.readlines() 55 | 56 | for line in lines: 57 | id, src_lens, tgt_lens, disparity = [part.strip() for part in line.split(",")] 58 | meta[id] = (src_lens, tgt_lens, disparity) 59 | return meta 60 | 61 | def lenstr2tensor(self, lenstr, scale=1.): 62 | # Canon50mm -> -1, Sony50mm -> 1 63 | lenstr = lenstr.replace('Canon50mmf', '-') 64 | lenstr = lenstr.replace('Sony50mmf', '') 65 | lenstr = lenstr.replace('BS', '') 66 | return torch.tensor(float(lenstr)) * scale 67 | 68 | def __getitem__(self, index): 69 | 70 | # get LR image 71 | LR_path = self.LR_paths[index] 72 | img_LR = util.read_img(self.LR_env, LR_path, None) 73 | 74 | id = os.path.basename(LR_path).split(".")[0] 75 | src_lens, tgt_lens, disparity = self.metas[id] 76 | 77 | src_lens = self.lenstr2tensor(src_lens, 10) 78 | tgt_lens = self.lenstr2tensor(tgt_lens, 10) 79 | disparity = self.lenstr2tensor(disparity) 80 | 81 | # change color space if necessary 82 | if self.opt["color"]: 83 | H, W, C = img_LR.shape 84 | img_LR = util.channel_convert(C, self.opt["color"], [img_LR])[ 85 | 0 86 | ] # TODO during val no definition 87 | 88 | 89 | # BGR to RGB, HWC to CHW, numpy to tensor 90 | if img_LR.shape[2] == 3: 91 | img_LR = img_LR[:, :, [2, 1, 0]] 92 | 93 | img_LR = torch.from_numpy( 94 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 95 | ).float() 96 | 97 | return { 98 | "LQ": img_LR, 99 | "src_lens": src_lens, 100 | "tgt_lens": tgt_lens, 101 | "disparity": disparity, 102 | "LQ_path": LR_path, 103 | } 104 | 105 | def __len__(self): 106 | return len(self.LR_paths) 107 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/GT_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | 12 | 13 | import numpy as np 14 | from PIL import Image 15 | 16 | from torchvision import transforms 17 | 18 | from skimage.feature import canny 19 | from skimage.color import gray2rgb, rgb2gray 20 | 21 | 22 | def tensor_to_image(): 23 | 24 | return transforms.ToPILImage() 25 | 26 | 27 | def image_to_tensor(): 28 | 29 | return transforms.ToTensor() 30 | 31 | 32 | def image_to_edge(image, sigma): 33 | 34 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 35 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 36 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 37 | 38 | return edge, gray_image 39 | 40 | 41 | 42 | 43 | 44 | try: 45 | sys.path.append("..") 46 | import data.util as util 47 | except ImportError: 48 | pass 49 | 50 | 51 | class GTDataset(data.Dataset): 52 | """ 53 | Read LR (Low Quality, here is LR) and GT image pairs. 54 | The pair is ensured by 'sorted' function, so please check the name convention. 55 | """ 56 | 57 | def __init__(self, opt): 58 | super().__init__() 59 | self.opt = opt 60 | self.GT_paths = None 61 | self.GT_env = None # environment for lmdb 62 | self.GT_size = opt["GT_size"] 63 | 64 | # read image list from lmdb or image files 65 | if opt["data_type"] == "lmdb": 66 | self.GT_paths, self.GT_sizes = util.get_image_paths( 67 | opt["data_type"], opt["dataroot_GT"] 68 | ) 69 | elif opt["data_type"] == "img": 70 | self.GT_paths = util.get_image_paths( 71 | opt["data_type"], opt["dataroot_GT"] 72 | ) # GT list 73 | else: 74 | print("Error: data_type is not matched in Dataset") 75 | assert self.GT_paths, "Error: GT paths are empty." 76 | print("dataset length: {}".format(len(self.GT_paths))) 77 | self.random_scale_list = [1] 78 | 79 | def _init_lmdb(self): 80 | # https://github.com/chainer/chainermn/issues/129 81 | self.GT_env = lmdb.open( 82 | self.opt["dataroot_GT"], 83 | readonly=True, 84 | lock=False, 85 | readahead=False, 86 | meminit=False, 87 | ) 88 | 89 | def __getitem__(self, index): 90 | if self.opt["data_type"] == "lmdb": 91 | if self.GT_env is None: 92 | self._init_lmdb() 93 | 94 | GT_path = None 95 | GT_size = self.opt["GT_size"] 96 | 97 | # get GT image 98 | GT_path = self.GT_paths[index] 99 | if self.opt["data_type"] == "lmdb": 100 | resolution = [int(s) for s in self.GT_sizes[index].split("_")] 101 | else: 102 | resolution = None 103 | img_GT = util.read_img( 104 | self.GT_env, GT_path, resolution 105 | ) # return: Numpy float32, HWC, BGR, [0,1] 106 | 107 | if self.opt["phase"] == "train": 108 | H, W, C = img_GT.shape 109 | 110 | rnd_h = random.randint(0, max(0, H - GT_size)) 111 | rnd_w = random.randint(0, max(0, W - GT_size)) 112 | img_GT = img_GT[rnd_h : rnd_h + GT_size, rnd_w : rnd_w + GT_size, :] 113 | 114 | # augmentation - flip, rotate 115 | img_GT = util.augment( 116 | img_GT, 117 | self.opt["use_flip"], 118 | self.opt["use_rot"], 119 | self.opt["mode"], 120 | ) 121 | 122 | # change color space if necessary 123 | if self.opt["color"]: 124 | img_GT = util.channel_convert(img_GT.shape[2], self.opt["color"], [img_GT])[ 125 | 0 126 | ] 127 | 128 | # BGR to RGB, HWC to CHW, numpy to tensor 129 | if img_GT.shape[2] == 3: 130 | img_GT = img_GT[:, :, [2, 1, 0]] 131 | img_GT = torch.from_numpy( 132 | np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1))) 133 | ).float() 134 | 135 | GT_edge,GT_gray = image_to_edge(img_GT, sigma=3.) 136 | return {"GT": img_GT, "GT_path": GT_path, "GT_edge": GT_edge, "GT_gray": GT_gray} 137 | 138 | def __len__(self): 139 | return len(self.GT_paths) 140 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/LQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class LQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path = None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path = self.LQ_paths[index] 67 | if self.opt["data_type"] == "lmdb": 68 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 69 | else: 70 | resolution = None 71 | img_LR = util.read_img( 72 | self.LR_env, LR_path, resolution 73 | ) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # modcrop in the validation / test phase 76 | if self.opt["phase"] != "train": 77 | img_LR = util.modcrop(img_LR, scale) 78 | 79 | if self.opt["phase"] == "train": 80 | H, W, C = img_LR.shape 81 | 82 | rnd_h = random.randint(0, max(0, H - LR_size)) 83 | rnd_w = random.randint(0, max(0, W - LR_size)) 84 | img_LR = img_LR[rnd_h : rnd_h + LR_size, rnd_w : rnd_w + LR_size, :] 85 | 86 | # augmentation - flip, rotate 87 | img_LR = util.augment( 88 | img_LR, 89 | self.opt["use_flip"], 90 | self.opt["use_rot"], 91 | self.opt["mode"], 92 | ) 93 | 94 | # change color space if necessary 95 | if self.opt["color"]: 96 | img_LR = util.channel_convert(img_LR.shape[2], self.opt["color"], [img_LR])[ 97 | 0 98 | ] 99 | 100 | # BGR to RGB, HWC to CHW, numpy to tensor 101 | if img_LR.shape[2] == 3: 102 | img_LR = img_LR[:, :, [2, 1, 0]] 103 | img_LR = torch.from_numpy( 104 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 105 | ).float() 106 | 107 | return {"LQ": img_LR, "LQ_path": LR_path} 108 | 109 | def __len__(self): 110 | return len(self.LQ_paths) 111 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/StereoLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class StereoLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path_L, LR_path_R = None, None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path_L = self.LQ_paths[index*2] 67 | LR_path_R = self.LQ_paths[index*2+1] 68 | if self.opt["data_type"] == "lmdb": 69 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 70 | else: 71 | resolution = None 72 | imgL_LR = util.read_img(self.LR_env, LR_path_L, resolution) # return: Numpy float32, HWC, BGR, [0,1] 73 | imgR_LR = util.read_img(self.LR_env, LR_path_R, resolution) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # change color space if necessary 76 | if self.opt["color"]: 77 | imgL_LR, imgR_LR = util.channel_convert( 78 | imgL_LR.shape[2], self.opt["color"], [imgL_LR, imgR_LR]) 79 | 80 | # BGR to RGB, HWC to CHW, numpy to tensor 81 | if imgL_LR.shape[2] == 3: 82 | imgL_LR = imgL_LR[:, :, [2, 1, 0]] 83 | imgR_LR = imgR_LR[:, :, [2, 1, 0]] 84 | imgL_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgL_LR, (2, 0, 1)))).float() 85 | imgR_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgR_LR, (2, 0, 1)))).float() 86 | 87 | img_LR = torch.cat([imgL_LR, imgR_LR], dim=0) 88 | 89 | return {"LQ": img_LR, "LQ_path": LR_path_L} 90 | 91 | def __len__(self): 92 | return len(self.LQ_paths) // 2 93 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/__init__.py: -------------------------------------------------------------------------------- 1 | """create dataset and dataloader""" 2 | import logging 3 | 4 | import torch 5 | import torch.utils.data 6 | 7 | 8 | def create_dataloader(dataset, dataset_opt, opt=None, sampler=None): 9 | phase = dataset_opt["phase"] 10 | if phase == "train": 11 | if opt["dist"]: 12 | world_size = torch.distributed.get_world_size() 13 | num_workers = dataset_opt["n_workers"] 14 | assert dataset_opt["batch_size"] % world_size == 0 15 | batch_size = dataset_opt["batch_size"] // world_size 16 | shuffle = False 17 | else: 18 | num_workers = dataset_opt["n_workers"] * len(opt["gpu_ids"]) 19 | batch_size = dataset_opt["batch_size"] 20 | shuffle = True 21 | return torch.utils.data.DataLoader( 22 | dataset, 23 | batch_size=batch_size, 24 | shuffle=shuffle, 25 | num_workers=num_workers, 26 | sampler=sampler, 27 | drop_last=True, 28 | pin_memory=False, 29 | ) 30 | else: 31 | return torch.utils.data.DataLoader( 32 | dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=(phase=="val") 33 | ) 34 | 35 | 36 | def create_dataset(dataset_opt): 37 | mode = dataset_opt["mode"] 38 | if mode == "LQ": # Predictor 39 | from data.LQ_dataset import LQDataset as D 40 | dataset = D(dataset_opt) 41 | elif mode == "LQGT": # SFTMD 42 | from data.LQGT_dataset import LQGTDataset as D 43 | dataset = D(dataset_opt) 44 | elif mode == "GT": # Corrector 45 | from data.GT_dataset import GTDataset as D 46 | dataset = D(dataset_opt) 47 | elif mode == 'SteLQGT': 48 | from data.StereoLQGT_dataset import StereoLQGTDataset as D 49 | dataset = D(dataset_opt) 50 | elif mode == 'SteLQ': 51 | from data.StereoLQ_dataset import StereoLQDataset as D 52 | dataset = D(dataset_opt) 53 | elif mode == 'BokehLQGT': 54 | from data.BokehLQGT_dataset import BokehLQGTDataset as D 55 | dataset = D(dataset_opt) 56 | elif mode == 'BokehLQ': 57 | from data.BokehLQ_dataset import BokehLQDataset as D 58 | dataset = D(dataset_opt) 59 | else: 60 | raise NotImplementedError("Dataset [{:s}] is not recognized.".format(mode)) 61 | 62 | logger = logging.getLogger("base") 63 | logger.info( 64 | "Dataset [{:s} - {:s}] is created.".format( 65 | dataset.__class__.__name__, dataset_opt["name"] 66 | ) 67 | ) 68 | return dataset 69 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def image_to_edge(image, sigma): 21 | 22 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 25 | 26 | return edge, gray_image 27 | 28 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/data/data_sampler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from torch.utils.data.distributed.DistributedSampler 3 | Support enlarging the dataset for *iter-oriented* training, for saving time when restart the 4 | dataloader after each epoch 5 | """ 6 | import math 7 | 8 | import torch 9 | import torch.distributed as dist 10 | from torch.utils.data.sampler import Sampler 11 | 12 | 13 | class DistIterSampler(Sampler): 14 | """Sampler that restricts data loading to a subset of the dataset. 15 | 16 | It is especially useful in conjunction with 17 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 18 | process can pass a DistributedSampler instance as a DataLoader sampler, 19 | and load a subset of the original dataset that is exclusive to it. 20 | 21 | .. note:: 22 | Dataset is assumed to be of constant size. 23 | 24 | Arguments: 25 | dataset: Dataset used for sampling. 26 | num_replicas (optional): Number of processes participating in 27 | distributed training. 28 | rank (optional): Rank of the current process within num_replicas. 29 | """ 30 | 31 | def __init__(self, dataset, num_replicas=None, rank=None, ratio=100): 32 | if num_replicas is None: 33 | if not dist.is_available(): 34 | raise RuntimeError("Requires distributed package to be available") 35 | num_replicas = dist.get_world_size() 36 | if rank is None: 37 | if not dist.is_available(): 38 | raise RuntimeError("Requires distributed package to be available") 39 | rank = dist.get_rank() 40 | self.dataset = dataset 41 | self.num_replicas = num_replicas 42 | self.rank = rank 43 | self.epoch = 0 44 | self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas)) 45 | self.total_size = self.num_samples * self.num_replicas 46 | 47 | def __iter__(self): 48 | # deterministically shuffle based on epoch 49 | g = torch.Generator() 50 | g.manual_seed(self.epoch) 51 | indices = torch.randperm( 52 | self.total_size, generator=g 53 | ).tolist() # Returns a random permutation of integers from 0 to n - 1 54 | 55 | dsize = len(self.dataset) 56 | indices = [v % dsize for v in indices] 57 | 58 | # subsample 59 | indices = indices[self.rank : self.total_size : self.num_replicas] 60 | assert len(indices) == self.num_samples 61 | 62 | return iter(indices) 63 | 64 | def __len__(self): 65 | return self.num_samples 66 | 67 | def set_epoch(self, epoch): 68 | self.epoch = epoch 69 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/demo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ############################################################# 4 | ### training ### 5 | 6 | # for single GPU 7 | python train.py -opt=options/train/ir-sde.yml 8 | # python train.py -opt=options/train/refusion.yml 9 | 10 | # for multiple GPUs 11 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/ir-sde.yml --launcher pytorch 12 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/refusion.yml --launcher pytorch 13 | 14 | ############################################################# 15 | 16 | ### testing ### 17 | # python test.py -opt=options/test/ir-sde.yml 18 | # python test.py -opt=options/test/refusion.yml 19 | 20 | ############################################################# -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logger = logging.getLogger("base") 4 | 5 | 6 | def create_model(opt): 7 | model = opt["model"] 8 | 9 | if model == "denoising": 10 | from .denoising_model import DenoisingModel as M 11 | else: 12 | raise NotImplementedError("Model [{:s}] not recognized.".format(model)) 13 | m = M(opt) 14 | logger.info("Model [{:s}] is created.".format(m.__class__.__name__)) 15 | return m 16 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import OrderedDict 3 | 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.parallel import DistributedDataParallel 7 | 8 | 9 | class BaseModel: 10 | def __init__(self, opt): 11 | self.opt = opt 12 | self.device = torch.device("cuda" if opt["gpu_ids"] is not None else "cpu") 13 | self.is_train = opt["is_train"] 14 | self.schedulers = [] 15 | self.optimizers = [] 16 | 17 | def feed_data(self, data): 18 | pass 19 | 20 | def optimize_parameters(self): 21 | pass 22 | 23 | def get_current_visuals(self): 24 | pass 25 | 26 | def get_current_losses(self): 27 | pass 28 | 29 | def print_network(self): 30 | pass 31 | 32 | def save(self, label): 33 | pass 34 | 35 | def load(self): 36 | pass 37 | 38 | def _set_lr(self, lr_groups_l): 39 | """set learning rate for warmup, 40 | lr_groups_l: list for lr_groups. each for a optimizer""" 41 | for optimizer, lr_groups in zip(self.optimizers, lr_groups_l): 42 | for param_group, lr in zip(optimizer.param_groups, lr_groups): 43 | param_group["lr"] = lr 44 | 45 | def _get_init_lr(self): 46 | # get the initial lr, which is set by the scheduler 47 | init_lr_groups_l = [] 48 | for optimizer in self.optimizers: 49 | init_lr_groups_l.append([v["initial_lr"] for v in optimizer.param_groups]) 50 | return init_lr_groups_l 51 | 52 | def update_learning_rate(self, cur_iter, warmup_iter=-1): 53 | for scheduler in self.schedulers: 54 | scheduler.step() 55 | #### set up warm up learning rate 56 | if cur_iter < warmup_iter: 57 | # get initial lr for each group 58 | init_lr_g_l = self._get_init_lr() 59 | # modify warming-up learning rates 60 | warm_up_lr_l = [] 61 | for init_lr_g in init_lr_g_l: 62 | warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g]) 63 | # set learning rate 64 | self._set_lr(warm_up_lr_l) 65 | 66 | def get_current_learning_rate(self): 67 | # return self.schedulers[0].get_lr()[0] 68 | return self.optimizers[0].param_groups[0]["lr"] 69 | 70 | def get_network_description(self, network): 71 | """Get the string and total parameters of the network""" 72 | if isinstance(network, nn.DataParallel) or isinstance( 73 | network, DistributedDataParallel 74 | ): 75 | network = network.module 76 | s = str(network) 77 | n = sum(map(lambda x: x.numel(), network.parameters())) 78 | return s, n 79 | 80 | def save_network(self, network, network_label, iter_label): 81 | save_filename = "{}_{}.pth".format(iter_label, network_label) 82 | save_path = os.path.join(self.opt["path"]["models"], save_filename) 83 | if isinstance(network, nn.DataParallel) or isinstance( 84 | network, DistributedDataParallel 85 | ): 86 | network = network.module 87 | state_dict = network.state_dict() 88 | for key, param in state_dict.items(): 89 | state_dict[key] = param.cpu() 90 | torch.save(state_dict, save_path) 91 | 92 | def load_network(self, load_path, network, strict=True): 93 | if isinstance(network, nn.DataParallel) or isinstance( 94 | network, DistributedDataParallel 95 | ): 96 | network = network.module 97 | print(network) 98 | load_net = torch.load(load_path) 99 | load_net_clean = OrderedDict() # remove unnecessary 'module.' 100 | for k, v in load_net.items(): 101 | if k.startswith("module."): 102 | load_net_clean[k[7:]] = v 103 | else: 104 | load_net_clean[k] = v 105 | 106 | network.load_state_dict(load_net_clean, strict=strict) 107 | 108 | def save_training_state(self, epoch, iter_step): 109 | """Saves training state during training, which will be used for resuming""" 110 | state = {"epoch": epoch, "iter": iter_step, "schedulers": [], "optimizers": []} 111 | for s in self.schedulers: 112 | state["schedulers"].append(s.state_dict()) 113 | for o in self.optimizers: 114 | state["optimizers"].append(o.state_dict()) 115 | save_filename = "{}.state".format(iter_step) 116 | save_path = os.path.join(self.opt["path"]["training_state"], save_filename) 117 | torch.save(state, save_path) 118 | 119 | def resume_training(self, resume_state): 120 | """Resume the optimizers and schedulers for training""" 121 | resume_optimizers = resume_state["optimizers"] 122 | resume_schedulers = resume_state["schedulers"] 123 | assert len(resume_optimizers) == len( 124 | self.optimizers 125 | ), "Wrong lengths of optimizers" 126 | assert len(resume_schedulers) == len( 127 | self.schedulers 128 | ), "Wrong lengths of schedulers" 129 | for i, o in enumerate(resume_optimizers): 130 | self.optimizers[i].load_state_dict(o) 131 | for i, s in enumerate(resume_schedulers): 132 | self.schedulers[i].load_state_dict(s) 133 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def gray_to_edge(image, sigma): 21 | 22 | gray_image = np.array(tensor_to_image()(image)) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | 25 | return edge 26 | 27 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/dense_layer.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from a file released under the MIT License. 5 | # 6 | # Source: 7 | # https://github.com/CW-Huang/sdeflow-light/blob/524650bc5ad69522b3e0905672deef0650374512/lib/models/unet.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_MIT). The modifications 11 | # to this file are subject to the same MIT License. 12 | # --------------------------------------------------------------- 13 | 14 | 15 | import math 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | from torch.nn.init import _calculate_fan_in_and_fan_out 20 | import numpy as np 21 | 22 | 23 | def _calculate_correct_fan(tensor, mode): 24 | """ 25 | copied and modified from https://github.com/pytorch/pytorch/blob/master/torch/nn/init.py#L337 26 | """ 27 | mode = mode.lower() 28 | valid_modes = ['fan_in', 'fan_out', 'fan_avg'] 29 | if mode not in valid_modes: 30 | raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes)) 31 | 32 | fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) 33 | return fan_in if mode == 'fan_in' else fan_out 34 | 35 | 36 | def kaiming_uniform_(tensor, gain=1., mode='fan_in'): 37 | r"""Fills the input `Tensor` with values according to the method 38 | described in `Delving deep into rectifiers: Surpassing human-level 39 | performance on ImageNet classification` - He, K. et al. (2015), using a 40 | uniform distribution. The resulting tensor will have values sampled from 41 | :math:`\mathcal{U}(-\text{bound}, \text{bound})` where 42 | .. math:: 43 | \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} 44 | Also known as He initialization. 45 | Args: 46 | tensor: an n-dimensional `torch.Tensor` 47 | gain: multiplier to the dispersion 48 | mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'`` 49 | preserves the magnitude of the variance of the weights in the 50 | forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the 51 | backwards pass. 52 | Examples: 53 | >>> w = torch.empty(3, 5) 54 | >>> nn.init.kaiming_uniform_(w, mode='fan_in') 55 | """ 56 | fan = _calculate_correct_fan(tensor, mode) 57 | var = gain / max(1., fan) 58 | bound = math.sqrt(3.0 * var) # Calculate uniform bounds from standard deviation 59 | with torch.no_grad(): 60 | return tensor.uniform_(-bound, bound) 61 | 62 | 63 | def variance_scaling_init_(tensor, scale): 64 | return kaiming_uniform_(tensor, gain=1e-10 if scale == 0 else scale, mode='fan_avg') 65 | 66 | 67 | def dense(in_channels, out_channels, init_scale=1.): 68 | lin = nn.Linear(in_channels, out_channels) 69 | variance_scaling_init_(lin.weight, scale=init_scale) 70 | nn.init.zeros_(lin.bias) 71 | return lin 72 | 73 | def conv2d(in_planes, out_planes, kernel_size=(3, 3), stride=1, dilation=1, padding=1, bias=True, padding_mode='zeros', 74 | init_scale=1.): 75 | conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 76 | bias=bias, padding_mode=padding_mode) 77 | variance_scaling_init_(conv.weight, scale=init_scale) 78 | if bias: 79 | nn.init.zeros_(conv.bias) 80 | return conv 81 | 82 | 83 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .DenoisingUNet_arch import ConditionalUNet -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/modules/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import einops 5 | import numpy as np 6 | import sys 7 | 8 | class MatchingLoss(nn.Module): 9 | def __init__(self, loss_type='l1', is_weighted=False): 10 | super().__init__() 11 | self.is_weighted = is_weighted 12 | 13 | if loss_type == 'l1': 14 | self.loss_fn = F.l1_loss 15 | elif loss_type == 'l2': 16 | self.loss_fn = F.mse_loss 17 | else: 18 | raise ValueError(f'invalid loss type {loss_type}') 19 | 20 | def forward(self, predict, target, mask, weights=None): 21 | 22 | lossm = self.loss_fn(predict * (1 - mask), target * (1 - mask), reduction='none') 23 | lossm = einops.reduce(lossm, 'b ... -> b (...)', 'mean') 24 | 25 | lossu = self.loss_fn(predict * mask, target * mask, reduction='none') 26 | lossu = einops.reduce(lossu, 'b ... -> b (...)', 'mean') 27 | 28 | loss = lossu + 10 * lossm 29 | if self.is_weighted and weights is not None: 30 | loss = weights * loss 31 | 32 | return loss.mean() 33 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/LICENSE_MIT: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Kim Seonghyeon 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/__init__.py: -------------------------------------------------------------------------------- 1 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 2 | from .upfirdn2d import upfirdn2d 3 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/cuda-installer.log: -------------------------------------------------------------------------------- 1 | [INFO]: Driver not installed. 2 | [INFO]: Checking compiler version... 3 | [INFO]: gcc location: /usr/bin/gcc 4 | 5 | [INFO]: gcc version: gcc version 9.4.0 (Ubuntu 9.4.0-1ubuntu1~20.04.2) 6 | 7 | [INFO]: Initializing menu 8 | [INFO]: nvidia-fs.setKOVersion(2.13.5) 9 | [INFO]: Setup complete 10 | [INFO]: Components to install: 11 | [INFO]: Driver 12 | [INFO]: 520.61.05 13 | [INFO]: Executing NVIDIA-Linux-x86_64-520.61.05.run --ui=none --no-questions --accept-license --disable-nouveau --no-cc-version-check --install-libglvnd 2>&1 14 | [INFO]: Finished with code: 256 15 | [ERROR]: Install of driver component failed. Consult the driver log at /var/log/nvidia-installer.log for more details. 16 | [ERROR]: Install of 520.61.05 failed, quitting 17 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/fused_act.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # --------------------------------------------------------------- 4 | 5 | """ Originated from https://github.com/rosinality/stylegan2-pytorch 6 | The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | """ 8 | 9 | import os 10 | 11 | import torch 12 | from torch import nn 13 | from torch.nn import functional as F 14 | from torch.autograd import Function 15 | from torch.utils.cpp_extension import load 16 | 17 | module_path = os.path.dirname(__file__) 18 | fused = load( 19 | "fused", 20 | sources=[ 21 | os.path.join(module_path, "fused_bias_act.cpp"), 22 | os.path.join(module_path, "fused_bias_act_kernel.cu"), 23 | ], 24 | ) 25 | 26 | 27 | class FusedLeakyReLUFunctionBackward(Function): 28 | @staticmethod 29 | def forward(ctx, grad_output, out, negative_slope, scale): 30 | ctx.save_for_backward(out) 31 | ctx.negative_slope = negative_slope 32 | ctx.scale = scale 33 | 34 | empty = grad_output.new_empty(0) 35 | 36 | grad_input = fused.fused_bias_act( 37 | grad_output, empty, out, 3, 1, negative_slope, scale 38 | ) 39 | 40 | dim = [0] 41 | 42 | if grad_input.ndim > 2: 43 | dim += list(range(2, grad_input.ndim)) 44 | 45 | grad_bias = grad_input.sum(dim).detach() 46 | 47 | return grad_input, grad_bias 48 | 49 | @staticmethod 50 | def backward(ctx, gradgrad_input, gradgrad_bias): 51 | out, = ctx.saved_tensors 52 | gradgrad_out = fused.fused_bias_act( 53 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 54 | ) 55 | 56 | return gradgrad_out, None, None, None 57 | 58 | 59 | class FusedLeakyReLUFunction(Function): 60 | @staticmethod 61 | def forward(ctx, input, bias, negative_slope, scale): 62 | empty = input.new_empty(0) 63 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 64 | ctx.save_for_backward(out) 65 | ctx.negative_slope = negative_slope 66 | ctx.scale = scale 67 | 68 | return out 69 | 70 | @staticmethod 71 | def backward(ctx, grad_output): 72 | out, = ctx.saved_tensors 73 | 74 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 75 | grad_output, out, ctx.negative_slope, ctx.scale 76 | ) 77 | 78 | return grad_input, grad_bias, None, None 79 | 80 | 81 | class FusedLeakyReLU(nn.Module): 82 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): 83 | super().__init__() 84 | 85 | self.bias = nn.Parameter(torch.zeros(channel)) 86 | self.negative_slope = negative_slope 87 | self.scale = scale 88 | 89 | def forward(self, input): 90 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 91 | 92 | 93 | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): 94 | if input.device.type == "cpu": 95 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 96 | return ( 97 | F.leaky_relu( 98 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 99 | ) 100 | * scale 101 | ) 102 | 103 | else: 104 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 105 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | #include 9 | 10 | 11 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale); 13 | 14 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 15 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 17 | 18 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 19 | int act, int grad, float alpha, float scale) { 20 | CHECK_CUDA(input); 21 | CHECK_CUDA(bias); 22 | 23 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 24 | } 25 | 26 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 27 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 28 | } -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved. 3 | // --------------------------------------------------------------- 4 | // 5 | // This work is made available under the Nvidia Source Code License-NC. 6 | // To view a copy of this license, visit 7 | // https://nvlabs.github.io/stylegan2/license.html 8 | 9 | #include 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | #include 17 | #include 18 | 19 | 20 | template 21 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 22 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 23 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 24 | 25 | scalar_t zero = 0.0; 26 | 27 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 28 | scalar_t x = p_x[xi]; 29 | 30 | if (use_bias) { 31 | x += p_b[(xi / step_b) % size_b]; 32 | } 33 | 34 | scalar_t ref = use_ref ? p_ref[xi] : zero; 35 | 36 | scalar_t y; 37 | 38 | switch (act * 10 + grad) { 39 | default: 40 | case 10: y = x; break; 41 | case 11: y = x; break; 42 | case 12: y = 0.0; break; 43 | 44 | case 30: y = (x > 0.0) ? x : x * alpha; break; 45 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 46 | case 32: y = 0.0; break; 47 | } 48 | 49 | out[xi] = y * scale; 50 | } 51 | } 52 | 53 | 54 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 55 | int act, int grad, float alpha, float scale) { 56 | int curDevice = -1; 57 | cudaGetDevice(&curDevice); 58 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 59 | 60 | auto x = input.contiguous(); 61 | auto b = bias.contiguous(); 62 | auto ref = refer.contiguous(); 63 | 64 | int use_bias = b.numel() ? 1 : 0; 65 | int use_ref = ref.numel() ? 1 : 0; 66 | 67 | int size_x = x.numel(); 68 | int size_b = b.numel(); 69 | int step_b = 1; 70 | 71 | for (int i = 1 + 1; i < x.dim(); i++) { 72 | step_b *= x.size(i); 73 | } 74 | 75 | int loop_x = 4; 76 | int block_size = 4 * 32; 77 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 78 | 79 | auto y = torch::empty_like(x); 80 | 81 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 82 | fused_bias_act_kernel<<>>( 83 | y.data_ptr(), 84 | x.data_ptr(), 85 | b.data_ptr(), 86 | ref.data_ptr(), 87 | act, 88 | grad, 89 | alpha, 90 | scale, 91 | loop_x, 92 | size_x, 93 | step_b, 94 | size_b, 95 | use_bias, 96 | use_ref 97 | ); 98 | }); 99 | 100 | return y; 101 | } -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/libcudart.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/train/discriminator/config/inpainting/models/op/libcudart.so -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | 9 | #include 10 | 11 | 12 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 15 | 16 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 17 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 18 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 19 | 20 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 21 | int up_x, int up_y, int down_x, int down_y, 22 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 23 | CHECK_CUDA(input); 24 | CHECK_CUDA(kernel); 25 | 26 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 27 | } 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 31 | } -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/models/optimizer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google Research. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """PyTorch implementation of the Lion optimizer.""" 16 | 17 | import torch 18 | from torch.optim.optimizer import Optimizer 19 | 20 | 21 | class Lion(Optimizer): 22 | r"""Implements Lion algorithm.""" 23 | 24 | def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0): 25 | """Initialize the hyperparameters. 26 | 27 | Args: 28 | params (iterable): iterable of parameters to optimize or dicts defining 29 | parameter groups 30 | lr (float, optional): learning rate (default: 1e-4) 31 | betas (Tuple[float, float], optional): coefficients used for computing 32 | running averages of gradient and its square (default: (0.9, 0.99)) 33 | weight_decay (float, optional): weight decay coefficient (default: 0) 34 | """ 35 | 36 | if not 0.0 <= lr: 37 | raise ValueError('Invalid learning rate: {}'.format(lr)) 38 | if not 0.0 <= betas[0] < 1.0: 39 | raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) 40 | if not 0.0 <= betas[1] < 1.0: 41 | raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) 42 | defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) 43 | super().__init__(params, defaults) 44 | 45 | @torch.no_grad() 46 | def step(self, closure=None): 47 | """Performs a single optimization step. 48 | 49 | Args: 50 | closure (callable, optional): A closure that reevaluates the model 51 | and returns the loss. 52 | 53 | Returns: 54 | the loss. 55 | """ 56 | loss = None 57 | if closure is not None: 58 | with torch.enable_grad(): 59 | loss = closure() 60 | 61 | for group in self.param_groups: 62 | for p in group['params']: 63 | if p.grad is None: 64 | continue 65 | 66 | # Perform stepweight decay 67 | p.data.mul_(1 - group['lr'] * group['weight_decay']) 68 | 69 | grad = p.grad 70 | state = self.state[p] 71 | # State initialization 72 | if len(state) == 0: 73 | # Exponential moving average of gradient values 74 | state['exp_avg'] = torch.zeros_like(p) 75 | 76 | exp_avg = state['exp_avg'] 77 | beta1, beta2 = group['betas'] 78 | 79 | # Weight update 80 | update = exp_avg * beta1 + grad * (1 - beta1) 81 | p.add_(torch.sign(update), alpha=-group['lr']) 82 | # Decay the momentum running average coefficient 83 | exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) 84 | 85 | return loss 86 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/options/train/ir-sde.yml: -------------------------------------------------------------------------------- 1 | #### general settings 2 | name: ir-sde 3 | use_tb_logger: true 4 | model: denoising 5 | distortion: derain 6 | gpu_ids: [2] 7 | 8 | sde: 9 | max_sigma: 30 10 | T: 400 11 | schedule: cosine 12 | eps: 0.005 13 | 14 | degradation: 15 | mask_root: your mask path 16 | 17 | #### datasets 18 | datasets: 19 | train: 20 | name: Train_Dataset 21 | mode: GT 22 | dataroot_GT: your image path 23 | 24 | use_shuffle: true 25 | n_workers: 4 26 | batch_size: 16 27 | GT_size: 256 28 | use_flip: true 29 | use_rot: true 30 | color: RGB 31 | 32 | 33 | #### network structures 34 | network_G: 35 | which_model_G: ConditionalUNet 36 | setting: 37 | in_nc: 1 38 | out_nc: 1 39 | nf: 64 40 | depth: 4 41 | 42 | #### path 43 | path: 44 | pretrain_model_G: ~ 45 | strict_load: true 46 | resume_state: ~ 47 | 48 | #### training settings: learning rate scheme, loss 49 | train: 50 | optimizer: Adam 51 | lr_G: !!float 1e-4 52 | lr_scheme: MultiStepLR 53 | beta1: 0.9 54 | beta2: 0.99 55 | niter: 8000000 56 | warmup_iter: -1 # no warm up 57 | lr_steps: [2000000, 4000000, 6000000] 58 | lr_gamma: 0.5 59 | eta_min: !!float 1e-7 60 | 61 | # criterion 62 | is_weighted: False 63 | loss_type: l1 64 | weight: 1.0 65 | 66 | manual_seed: 0 67 | 68 | #### logger 69 | logger: 70 | print_freq: 100 71 | save_checkpoint_freq: !!float 5e3 72 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/str_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/str_utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import abc 4 | import math 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torchvision.utils as tvutils 10 | 11 | 12 | ########### denoising ############### 13 | def add_noise(tensor, sigma): 14 | sigma = sigma / 255 if sigma > 1 else sigma 15 | return tensor + torch.randn_like(tensor) * sigma 16 | 17 | 18 | ######## inpainting ########### 19 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 20 | batch = tensor.shape[0] 21 | if mask_id < 0: 22 | mask_id = np.random.randint(0, n, batch) 23 | masks = [] 24 | for i in range(batch): 25 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 26 | mask = np.concatenate(masks, axis=0) 27 | else: 28 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 29 | 30 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 31 | # for images are clipped or scaled 32 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 33 | masked_tensor = mask * tensor 34 | return masked_tensor + (1. - mask) 35 | 36 | ######## super-resolution ########### 37 | 38 | def upscale(tensor, scale=4, mode='bicubic'): 39 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 40 | return tensor 41 | 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/discriminator/config/inpainting/utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import abc 3 | import math 4 | import numpy as np 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torchvision.utils as tvutils 9 | import cv2 10 | 11 | ########### denoising ############### 12 | def add_noise(tensor, sigma): 13 | sigma = sigma / 255 if sigma > 1 else sigma 14 | return tensor + torch.randn_like(tensor) * sigma 15 | 16 | 17 | ######## inpainting ########### 18 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 19 | batch = tensor.shape[0] 20 | if mask_id < 0: 21 | mask_id = np.random.randint(0, n, batch) 22 | masks = [] 23 | for i in range(batch): 24 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 25 | mask = np.concatenate(masks, axis=0) 26 | else: 27 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 28 | 29 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 30 | # for images are clipped or scaled 31 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 32 | masked_tensor = mask * tensor 33 | return masked_tensor + (1. - mask) 34 | 35 | ######## super-resolution ########### 36 | 37 | def upscale(tensor, scale=4, mode='bicubic'): 38 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 39 | return tensor 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/BokehLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class BokehLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and GT image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LR_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from image files 32 | if opt["data_type"] == "img": 33 | self.LR_paths = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) # LR list 36 | self.metas = self._read_meta_data(opt["dataroot_meta"]) 37 | else: 38 | print("Error: data_type is not matched in Dataset") 39 | 40 | def _read_meta_data(self, meta_file_path: str): 41 | """Read the meta file containing source / target lens and disparity for each image. 42 | Args: 43 | meta_file_path (str): File path 44 | Raises: 45 | ValueError: File not found. 46 | Returns: 47 | dict: Meta dict of tuples like {id: (id, src_lens, tgt_lens, disparity)}. 48 | """ 49 | if not os.path.isfile(meta_file_path): 50 | raise ValueError(f"Meta file missing under {meta_file_path}.") 51 | 52 | meta = {} 53 | with open(meta_file_path, "r") as f: 54 | lines = f.readlines() 55 | 56 | for line in lines: 57 | id, src_lens, tgt_lens, disparity = [part.strip() for part in line.split(",")] 58 | meta[id] = (src_lens, tgt_lens, disparity) 59 | return meta 60 | 61 | def lenstr2tensor(self, lenstr, scale=1.): 62 | # Canon50mm -> -1, Sony50mm -> 1 63 | lenstr = lenstr.replace('Canon50mmf', '-') 64 | lenstr = lenstr.replace('Sony50mmf', '') 65 | lenstr = lenstr.replace('BS', '') 66 | return torch.tensor(float(lenstr)) * scale 67 | 68 | def __getitem__(self, index): 69 | 70 | # get LR image 71 | LR_path = self.LR_paths[index] 72 | img_LR = util.read_img(self.LR_env, LR_path, None) 73 | 74 | id = os.path.basename(LR_path).split(".")[0] 75 | src_lens, tgt_lens, disparity = self.metas[id] 76 | 77 | src_lens = self.lenstr2tensor(src_lens, 10) 78 | tgt_lens = self.lenstr2tensor(tgt_lens, 10) 79 | disparity = self.lenstr2tensor(disparity) 80 | 81 | # change color space if necessary 82 | if self.opt["color"]: 83 | H, W, C = img_LR.shape 84 | img_LR = util.channel_convert(C, self.opt["color"], [img_LR])[ 85 | 0 86 | ] # TODO during val no definition 87 | 88 | 89 | # BGR to RGB, HWC to CHW, numpy to tensor 90 | if img_LR.shape[2] == 3: 91 | img_LR = img_LR[:, :, [2, 1, 0]] 92 | 93 | img_LR = torch.from_numpy( 94 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 95 | ).float() 96 | 97 | return { 98 | "LQ": img_LR, 99 | "src_lens": src_lens, 100 | "tgt_lens": tgt_lens, 101 | "disparity": disparity, 102 | "LQ_path": LR_path, 103 | } 104 | 105 | def __len__(self): 106 | return len(self.LR_paths) 107 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/GT_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | 12 | 13 | import numpy as np 14 | from PIL import Image 15 | 16 | from torchvision import transforms 17 | 18 | from skimage.feature import canny 19 | from skimage.color import gray2rgb, rgb2gray 20 | 21 | 22 | def tensor_to_image(): 23 | 24 | return transforms.ToPILImage() 25 | 26 | 27 | def image_to_tensor(): 28 | 29 | return transforms.ToTensor() 30 | 31 | 32 | def image_to_edge(image, sigma): 33 | 34 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 35 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 36 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 37 | 38 | return edge, gray_image 39 | 40 | 41 | 42 | 43 | 44 | try: 45 | sys.path.append("..") 46 | import data.util as util 47 | except ImportError: 48 | pass 49 | 50 | 51 | class GTDataset(data.Dataset): 52 | """ 53 | Read LR (Low Quality, here is LR) and GT image pairs. 54 | The pair is ensured by 'sorted' function, so please check the name convention. 55 | """ 56 | 57 | def __init__(self, opt): 58 | super().__init__() 59 | self.opt = opt 60 | self.GT_paths = None 61 | self.GT_env = None # environment for lmdb 62 | self.GT_size = opt["GT_size"] 63 | 64 | # read image list from lmdb or image files 65 | if opt["data_type"] == "lmdb": 66 | self.GT_paths, self.GT_sizes = util.get_image_paths( 67 | opt["data_type"], opt["dataroot_GT"] 68 | ) 69 | elif opt["data_type"] == "img": 70 | self.GT_paths = util.get_image_paths( 71 | opt["data_type"], opt["dataroot_GT"] 72 | ) # GT list 73 | else: 74 | print("Error: data_type is not matched in Dataset") 75 | assert self.GT_paths, "Error: GT paths are empty." 76 | print("dataset length: {}".format(len(self.GT_paths))) 77 | self.random_scale_list = [1] 78 | 79 | def _init_lmdb(self): 80 | # https://github.com/chainer/chainermn/issues/129 81 | self.GT_env = lmdb.open( 82 | self.opt["dataroot_GT"], 83 | readonly=True, 84 | lock=False, 85 | readahead=False, 86 | meminit=False, 87 | ) 88 | 89 | def __getitem__(self, index): 90 | if self.opt["data_type"] == "lmdb": 91 | if self.GT_env is None: 92 | self._init_lmdb() 93 | 94 | GT_path = None 95 | GT_size = self.opt["GT_size"] 96 | 97 | # get GT image 98 | GT_path = self.GT_paths[index] 99 | if self.opt["data_type"] == "lmdb": 100 | resolution = [int(s) for s in self.GT_sizes[index].split("_")] 101 | else: 102 | resolution = None 103 | img_GT = util.read_img( 104 | self.GT_env, GT_path, resolution 105 | ) # return: Numpy float32, HWC, BGR, [0,1] 106 | 107 | if self.opt["phase"] == "train": 108 | H, W, C = img_GT.shape 109 | 110 | rnd_h = random.randint(0, max(0, H - GT_size)) 111 | rnd_w = random.randint(0, max(0, W - GT_size)) 112 | img_GT = img_GT[rnd_h : rnd_h + GT_size, rnd_w : rnd_w + GT_size, :] 113 | 114 | # augmentation - flip, rotate 115 | img_GT = util.augment( 116 | img_GT, 117 | self.opt["use_flip"], 118 | self.opt["use_rot"], 119 | self.opt["mode"], 120 | ) 121 | 122 | # change color space if necessary 123 | if self.opt["color"]: 124 | img_GT = util.channel_convert(img_GT.shape[2], self.opt["color"], [img_GT])[ 125 | 0 126 | ] 127 | 128 | # BGR to RGB, HWC to CHW, numpy to tensor 129 | if img_GT.shape[2] == 3: 130 | img_GT = img_GT[:, :, [2, 1, 0]] 131 | img_GT = torch.from_numpy( 132 | np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1))) 133 | ).float() 134 | 135 | GT_edge,GT_gray = image_to_edge(img_GT, sigma=3.) 136 | return {"GT": img_GT, "GT_path": GT_path, "GT_edge": GT_edge, "GT_gray": GT_gray} 137 | 138 | def __len__(self): 139 | return len(self.GT_paths) 140 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/LQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class LQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path = None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path = self.LQ_paths[index] 67 | if self.opt["data_type"] == "lmdb": 68 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 69 | else: 70 | resolution = None 71 | img_LR = util.read_img( 72 | self.LR_env, LR_path, resolution 73 | ) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # modcrop in the validation / test phase 76 | if self.opt["phase"] != "train": 77 | img_LR = util.modcrop(img_LR, scale) 78 | 79 | if self.opt["phase"] == "train": 80 | H, W, C = img_LR.shape 81 | 82 | rnd_h = random.randint(0, max(0, H - LR_size)) 83 | rnd_w = random.randint(0, max(0, W - LR_size)) 84 | img_LR = img_LR[rnd_h : rnd_h + LR_size, rnd_w : rnd_w + LR_size, :] 85 | 86 | # augmentation - flip, rotate 87 | img_LR = util.augment( 88 | img_LR, 89 | self.opt["use_flip"], 90 | self.opt["use_rot"], 91 | self.opt["mode"], 92 | ) 93 | 94 | # change color space if necessary 95 | if self.opt["color"]: 96 | img_LR = util.channel_convert(img_LR.shape[2], self.opt["color"], [img_LR])[ 97 | 0 98 | ] 99 | 100 | # BGR to RGB, HWC to CHW, numpy to tensor 101 | if img_LR.shape[2] == 3: 102 | img_LR = img_LR[:, :, [2, 1, 0]] 103 | img_LR = torch.from_numpy( 104 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 105 | ).float() 106 | 107 | return {"LQ": img_LR, "LQ_path": LR_path} 108 | 109 | def __len__(self): 110 | return len(self.LQ_paths) 111 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/StereoLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class StereoLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path_L, LR_path_R = None, None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path_L = self.LQ_paths[index*2] 67 | LR_path_R = self.LQ_paths[index*2+1] 68 | if self.opt["data_type"] == "lmdb": 69 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 70 | else: 71 | resolution = None 72 | imgL_LR = util.read_img(self.LR_env, LR_path_L, resolution) # return: Numpy float32, HWC, BGR, [0,1] 73 | imgR_LR = util.read_img(self.LR_env, LR_path_R, resolution) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # change color space if necessary 76 | if self.opt["color"]: 77 | imgL_LR, imgR_LR = util.channel_convert( 78 | imgL_LR.shape[2], self.opt["color"], [imgL_LR, imgR_LR]) 79 | 80 | # BGR to RGB, HWC to CHW, numpy to tensor 81 | if imgL_LR.shape[2] == 3: 82 | imgL_LR = imgL_LR[:, :, [2, 1, 0]] 83 | imgR_LR = imgR_LR[:, :, [2, 1, 0]] 84 | imgL_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgL_LR, (2, 0, 1)))).float() 85 | imgR_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgR_LR, (2, 0, 1)))).float() 86 | 87 | img_LR = torch.cat([imgL_LR, imgR_LR], dim=0) 88 | 89 | return {"LQ": img_LR, "LQ_path": LR_path_L} 90 | 91 | def __len__(self): 92 | return len(self.LQ_paths) // 2 93 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/__init__.py: -------------------------------------------------------------------------------- 1 | """create dataset and dataloader""" 2 | import logging 3 | 4 | import torch 5 | import torch.utils.data 6 | 7 | 8 | def create_dataloader(dataset, dataset_opt, opt=None, sampler=None): 9 | phase = dataset_opt["phase"] 10 | if phase == "train": 11 | if opt["dist"]: 12 | world_size = torch.distributed.get_world_size() 13 | num_workers = dataset_opt["n_workers"] 14 | assert dataset_opt["batch_size"] % world_size == 0 15 | batch_size = dataset_opt["batch_size"] // world_size 16 | shuffle = False 17 | else: 18 | num_workers = dataset_opt["n_workers"] * len(opt["gpu_ids"]) 19 | batch_size = dataset_opt["batch_size"] 20 | shuffle = True 21 | return torch.utils.data.DataLoader( 22 | dataset, 23 | batch_size=batch_size, 24 | shuffle=shuffle, 25 | num_workers=num_workers, 26 | sampler=sampler, 27 | drop_last=True, 28 | pin_memory=False, 29 | ) 30 | else: 31 | return torch.utils.data.DataLoader( 32 | dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=(phase=="val") 33 | ) 34 | 35 | 36 | def create_dataset(dataset_opt): 37 | mode = dataset_opt["mode"] 38 | if mode == "LQ": # Predictor 39 | from data.LQ_dataset import LQDataset as D 40 | dataset = D(dataset_opt) 41 | elif mode == "LQGT": # SFTMD 42 | from data.LQGT_dataset import LQGTDataset as D 43 | dataset = D(dataset_opt) 44 | elif mode == "GT": # Corrector 45 | from data.GT_dataset import GTDataset as D 46 | dataset = D(dataset_opt) 47 | elif mode == 'SteLQGT': 48 | from data.StereoLQGT_dataset import StereoLQGTDataset as D 49 | dataset = D(dataset_opt) 50 | elif mode == 'SteLQ': 51 | from data.StereoLQ_dataset import StereoLQDataset as D 52 | dataset = D(dataset_opt) 53 | elif mode == 'BokehLQGT': 54 | from data.BokehLQGT_dataset import BokehLQGTDataset as D 55 | dataset = D(dataset_opt) 56 | elif mode == 'BokehLQ': 57 | from data.BokehLQ_dataset import BokehLQDataset as D 58 | dataset = D(dataset_opt) 59 | else: 60 | raise NotImplementedError("Dataset [{:s}] is not recognized.".format(mode)) 61 | 62 | logger = logging.getLogger("base") 63 | logger.info( 64 | "Dataset [{:s} - {:s}] is created.".format( 65 | dataset.__class__.__name__, dataset_opt["name"] 66 | ) 67 | ) 68 | return dataset 69 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def image_to_edge(image, sigma): 21 | 22 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 25 | 26 | return edge, gray_image 27 | 28 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/data/data_sampler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from torch.utils.data.distributed.DistributedSampler 3 | Support enlarging the dataset for *iter-oriented* training, for saving time when restart the 4 | dataloader after each epoch 5 | """ 6 | import math 7 | 8 | import torch 9 | import torch.distributed as dist 10 | from torch.utils.data.sampler import Sampler 11 | 12 | 13 | class DistIterSampler(Sampler): 14 | """Sampler that restricts data loading to a subset of the dataset. 15 | 16 | It is especially useful in conjunction with 17 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 18 | process can pass a DistributedSampler instance as a DataLoader sampler, 19 | and load a subset of the original dataset that is exclusive to it. 20 | 21 | .. note:: 22 | Dataset is assumed to be of constant size. 23 | 24 | Arguments: 25 | dataset: Dataset used for sampling. 26 | num_replicas (optional): Number of processes participating in 27 | distributed training. 28 | rank (optional): Rank of the current process within num_replicas. 29 | """ 30 | 31 | def __init__(self, dataset, num_replicas=None, rank=None, ratio=100): 32 | if num_replicas is None: 33 | if not dist.is_available(): 34 | raise RuntimeError("Requires distributed package to be available") 35 | num_replicas = dist.get_world_size() 36 | if rank is None: 37 | if not dist.is_available(): 38 | raise RuntimeError("Requires distributed package to be available") 39 | rank = dist.get_rank() 40 | self.dataset = dataset 41 | self.num_replicas = num_replicas 42 | self.rank = rank 43 | self.epoch = 0 44 | self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas)) 45 | self.total_size = self.num_samples * self.num_replicas 46 | 47 | def __iter__(self): 48 | # deterministically shuffle based on epoch 49 | g = torch.Generator() 50 | g.manual_seed(self.epoch) 51 | indices = torch.randperm( 52 | self.total_size, generator=g 53 | ).tolist() # Returns a random permutation of integers from 0 to n - 1 54 | 55 | dsize = len(self.dataset) 56 | indices = [v % dsize for v in indices] 57 | 58 | # subsample 59 | indices = indices[self.rank : self.total_size : self.num_replicas] 60 | assert len(indices) == self.num_samples 61 | 62 | return iter(indices) 63 | 64 | def __len__(self): 65 | return self.num_samples 66 | 67 | def set_epoch(self, epoch): 68 | self.epoch = epoch 69 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/demo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ############################################################# 4 | ### training ### 5 | 6 | # for single GPU 7 | python train.py -opt=options/train/ir-sde.yml 8 | # python train.py -opt=options/train/refusion.yml 9 | 10 | # for multiple GPUs 11 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/ir-sde.yml --launcher pytorch 12 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/refusion.yml --launcher pytorch 13 | 14 | ############################################################# 15 | 16 | ### testing ### 17 | # python test.py -opt=options/test/ir-sde.yml 18 | # python test.py -opt=options/test/refusion.yml 19 | 20 | ############################################################# -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logger = logging.getLogger("base") 4 | 5 | 6 | def create_model(opt): 7 | model = opt["model"] 8 | 9 | if model == "denoising": 10 | from .denoising_model import DenoisingModel as M 11 | else: 12 | raise NotImplementedError("Model [{:s}] not recognized.".format(model)) 13 | m = M(opt) 14 | logger.info("Model [{:s}] is created.".format(m.__class__.__name__)) 15 | return m 16 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import OrderedDict 3 | 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.parallel import DistributedDataParallel 7 | 8 | 9 | class BaseModel: 10 | def __init__(self, opt): 11 | self.opt = opt 12 | self.device = torch.device("cuda" if opt["gpu_ids"] is not None else "cpu") 13 | self.is_train = opt["is_train"] 14 | self.schedulers = [] 15 | self.optimizers = [] 16 | 17 | def feed_data(self, data): 18 | pass 19 | 20 | def optimize_parameters(self): 21 | pass 22 | 23 | def get_current_visuals(self): 24 | pass 25 | 26 | def get_current_losses(self): 27 | pass 28 | 29 | def print_network(self): 30 | pass 31 | 32 | def save(self, label): 33 | pass 34 | 35 | def load(self): 36 | pass 37 | 38 | def _set_lr(self, lr_groups_l): 39 | """set learning rate for warmup, 40 | lr_groups_l: list for lr_groups. each for a optimizer""" 41 | for optimizer, lr_groups in zip(self.optimizers, lr_groups_l): 42 | for param_group, lr in zip(optimizer.param_groups, lr_groups): 43 | param_group["lr"] = lr 44 | 45 | def _get_init_lr(self): 46 | # get the initial lr, which is set by the scheduler 47 | init_lr_groups_l = [] 48 | for optimizer in self.optimizers: 49 | init_lr_groups_l.append([v["initial_lr"] for v in optimizer.param_groups]) 50 | return init_lr_groups_l 51 | 52 | def update_learning_rate(self, cur_iter, warmup_iter=-1): 53 | for scheduler in self.schedulers: 54 | scheduler.step() 55 | #### set up warm up learning rate 56 | if cur_iter < warmup_iter: 57 | # get initial lr for each group 58 | init_lr_g_l = self._get_init_lr() 59 | # modify warming-up learning rates 60 | warm_up_lr_l = [] 61 | for init_lr_g in init_lr_g_l: 62 | warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g]) 63 | # set learning rate 64 | self._set_lr(warm_up_lr_l) 65 | 66 | def get_current_learning_rate(self): 67 | # return self.schedulers[0].get_lr()[0] 68 | return self.optimizers[0].param_groups[0]["lr"] 69 | 70 | def get_network_description(self, network): 71 | """Get the string and total parameters of the network""" 72 | if isinstance(network, nn.DataParallel) or isinstance( 73 | network, DistributedDataParallel 74 | ): 75 | network = network.module 76 | s = str(network) 77 | n = sum(map(lambda x: x.numel(), network.parameters())) 78 | return s, n 79 | 80 | def save_network(self, network, network_label, iter_label): 81 | save_filename = "{}_{}.pth".format(iter_label, network_label) 82 | save_path = os.path.join(self.opt["path"]["models"], save_filename) 83 | if isinstance(network, nn.DataParallel) or isinstance( 84 | network, DistributedDataParallel 85 | ): 86 | network = network.module 87 | state_dict = network.state_dict() 88 | for key, param in state_dict.items(): 89 | state_dict[key] = param.cpu() 90 | torch.save(state_dict, save_path) 91 | 92 | def load_network(self, load_path, network, strict=True): 93 | if isinstance(network, nn.DataParallel) or isinstance( 94 | network, DistributedDataParallel 95 | ): 96 | network = network.module 97 | print(network) 98 | load_net = torch.load(load_path) 99 | load_net_clean = OrderedDict() # remove unnecessary 'module.' 100 | for k, v in load_net.items(): 101 | if k.startswith("module."): 102 | load_net_clean[k[7:]] = v 103 | else: 104 | load_net_clean[k] = v 105 | 106 | network.load_state_dict(load_net_clean, strict=strict) 107 | 108 | def save_training_state(self, epoch, iter_step): 109 | """Saves training state during training, which will be used for resuming""" 110 | state = {"epoch": epoch, "iter": iter_step, "schedulers": [], "optimizers": []} 111 | for s in self.schedulers: 112 | state["schedulers"].append(s.state_dict()) 113 | for o in self.optimizers: 114 | state["optimizers"].append(o.state_dict()) 115 | save_filename = "{}.state".format(iter_step) 116 | save_path = os.path.join(self.opt["path"]["training_state"], save_filename) 117 | torch.save(state, save_path) 118 | 119 | def resume_training(self, resume_state): 120 | """Resume the optimizers and schedulers for training""" 121 | resume_optimizers = resume_state["optimizers"] 122 | resume_schedulers = resume_state["schedulers"] 123 | assert len(resume_optimizers) == len( 124 | self.optimizers 125 | ), "Wrong lengths of optimizers" 126 | assert len(resume_schedulers) == len( 127 | self.schedulers 128 | ), "Wrong lengths of schedulers" 129 | for i, o in enumerate(resume_optimizers): 130 | self.optimizers[i].load_state_dict(o) 131 | for i, s in enumerate(resume_schedulers): 132 | self.schedulers[i].load_state_dict(s) 133 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def gray_to_edge(image, sigma): 21 | 22 | gray_image = np.array(tensor_to_image()(image)) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | 25 | return edge 26 | 27 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/dense_layer.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from a file released under the MIT License. 5 | # 6 | # Source: 7 | # https://github.com/CW-Huang/sdeflow-light/blob/524650bc5ad69522b3e0905672deef0650374512/lib/models/unet.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_MIT). The modifications 11 | # to this file are subject to the same MIT License. 12 | # --------------------------------------------------------------- 13 | 14 | 15 | import math 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | from torch.nn.init import _calculate_fan_in_and_fan_out 20 | import numpy as np 21 | 22 | 23 | def _calculate_correct_fan(tensor, mode): 24 | """ 25 | copied and modified from https://github.com/pytorch/pytorch/blob/master/torch/nn/init.py#L337 26 | """ 27 | mode = mode.lower() 28 | valid_modes = ['fan_in', 'fan_out', 'fan_avg'] 29 | if mode not in valid_modes: 30 | raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes)) 31 | 32 | fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) 33 | return fan_in if mode == 'fan_in' else fan_out 34 | 35 | 36 | def kaiming_uniform_(tensor, gain=1., mode='fan_in'): 37 | r"""Fills the input `Tensor` with values according to the method 38 | described in `Delving deep into rectifiers: Surpassing human-level 39 | performance on ImageNet classification` - He, K. et al. (2015), using a 40 | uniform distribution. The resulting tensor will have values sampled from 41 | :math:`\mathcal{U}(-\text{bound}, \text{bound})` where 42 | .. math:: 43 | \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} 44 | Also known as He initialization. 45 | Args: 46 | tensor: an n-dimensional `torch.Tensor` 47 | gain: multiplier to the dispersion 48 | mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'`` 49 | preserves the magnitude of the variance of the weights in the 50 | forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the 51 | backwards pass. 52 | Examples: 53 | >>> w = torch.empty(3, 5) 54 | >>> nn.init.kaiming_uniform_(w, mode='fan_in') 55 | """ 56 | fan = _calculate_correct_fan(tensor, mode) 57 | var = gain / max(1., fan) 58 | bound = math.sqrt(3.0 * var) # Calculate uniform bounds from standard deviation 59 | with torch.no_grad(): 60 | return tensor.uniform_(-bound, bound) 61 | 62 | 63 | def variance_scaling_init_(tensor, scale): 64 | return kaiming_uniform_(tensor, gain=1e-10 if scale == 0 else scale, mode='fan_avg') 65 | 66 | 67 | def dense(in_channels, out_channels, init_scale=1.): 68 | lin = nn.Linear(in_channels, out_channels) 69 | variance_scaling_init_(lin.weight, scale=init_scale) 70 | nn.init.zeros_(lin.bias) 71 | return lin 72 | 73 | def conv2d(in_planes, out_planes, kernel_size=(3, 3), stride=1, dilation=1, padding=1, bias=True, padding_mode='zeros', 74 | init_scale=1.): 75 | conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 76 | bias=bias, padding_mode=padding_mode) 77 | variance_scaling_init_(conv.weight, scale=init_scale) 78 | if bias: 79 | nn.init.zeros_(conv.bias) 80 | return conv 81 | 82 | 83 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .DenoisingUNet_arch import ConditionalUNet -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/modules/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import einops 5 | import numpy as np 6 | import sys 7 | 8 | class MatchingLoss(nn.Module): 9 | def __init__(self, loss_type='l1', is_weighted=False): 10 | super().__init__() 11 | self.is_weighted = is_weighted 12 | 13 | if loss_type == 'l1': 14 | self.loss_fn = F.l1_loss 15 | elif loss_type == 'l2': 16 | self.loss_fn = F.mse_loss 17 | else: 18 | raise ValueError(f'invalid loss type {loss_type}') 19 | 20 | def forward(self, predict, target, mask, weights=None): 21 | 22 | lossm = self.loss_fn(predict * (1 - mask), target * (1 - mask), reduction='none') 23 | lossm = einops.reduce(lossm, 'b ... -> b (...)', 'mean') 24 | 25 | lossu = self.loss_fn(predict * mask, target * mask, reduction='none') 26 | lossu = einops.reduce(lossu, 'b ... -> b (...)', 'mean') 27 | 28 | loss = lossu + 10 * lossm 29 | if self.is_weighted and weights is not None: 30 | loss = weights * loss 31 | 32 | return loss.mean() 33 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/LICENSE_MIT: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Kim Seonghyeon 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/__init__.py: -------------------------------------------------------------------------------- 1 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 2 | from .upfirdn2d import upfirdn2d 3 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/cuda-installer.log: -------------------------------------------------------------------------------- 1 | [INFO]: Driver not installed. 2 | [INFO]: Checking compiler version... 3 | [INFO]: gcc location: /usr/bin/gcc 4 | 5 | [INFO]: gcc version: gcc version 9.4.0 (Ubuntu 9.4.0-1ubuntu1~20.04.2) 6 | 7 | [INFO]: Initializing menu 8 | [INFO]: nvidia-fs.setKOVersion(2.13.5) 9 | [INFO]: Setup complete 10 | [INFO]: Components to install: 11 | [INFO]: Driver 12 | [INFO]: 520.61.05 13 | [INFO]: Executing NVIDIA-Linux-x86_64-520.61.05.run --ui=none --no-questions --accept-license --disable-nouveau --no-cc-version-check --install-libglvnd 2>&1 14 | [INFO]: Finished with code: 256 15 | [ERROR]: Install of driver component failed. Consult the driver log at /var/log/nvidia-installer.log for more details. 16 | [ERROR]: Install of 520.61.05 failed, quitting 17 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/fused_act.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # --------------------------------------------------------------- 4 | 5 | """ Originated from https://github.com/rosinality/stylegan2-pytorch 6 | The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | """ 8 | 9 | import os 10 | 11 | import torch 12 | from torch import nn 13 | from torch.nn import functional as F 14 | from torch.autograd import Function 15 | from torch.utils.cpp_extension import load 16 | 17 | module_path = os.path.dirname(__file__) 18 | fused = load( 19 | "fused", 20 | sources=[ 21 | os.path.join(module_path, "fused_bias_act.cpp"), 22 | os.path.join(module_path, "fused_bias_act_kernel.cu"), 23 | ], 24 | ) 25 | 26 | 27 | class FusedLeakyReLUFunctionBackward(Function): 28 | @staticmethod 29 | def forward(ctx, grad_output, out, negative_slope, scale): 30 | ctx.save_for_backward(out) 31 | ctx.negative_slope = negative_slope 32 | ctx.scale = scale 33 | 34 | empty = grad_output.new_empty(0) 35 | 36 | grad_input = fused.fused_bias_act( 37 | grad_output, empty, out, 3, 1, negative_slope, scale 38 | ) 39 | 40 | dim = [0] 41 | 42 | if grad_input.ndim > 2: 43 | dim += list(range(2, grad_input.ndim)) 44 | 45 | grad_bias = grad_input.sum(dim).detach() 46 | 47 | return grad_input, grad_bias 48 | 49 | @staticmethod 50 | def backward(ctx, gradgrad_input, gradgrad_bias): 51 | out, = ctx.saved_tensors 52 | gradgrad_out = fused.fused_bias_act( 53 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 54 | ) 55 | 56 | return gradgrad_out, None, None, None 57 | 58 | 59 | class FusedLeakyReLUFunction(Function): 60 | @staticmethod 61 | def forward(ctx, input, bias, negative_slope, scale): 62 | empty = input.new_empty(0) 63 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 64 | ctx.save_for_backward(out) 65 | ctx.negative_slope = negative_slope 66 | ctx.scale = scale 67 | 68 | return out 69 | 70 | @staticmethod 71 | def backward(ctx, grad_output): 72 | out, = ctx.saved_tensors 73 | 74 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 75 | grad_output, out, ctx.negative_slope, ctx.scale 76 | ) 77 | 78 | return grad_input, grad_bias, None, None 79 | 80 | 81 | class FusedLeakyReLU(nn.Module): 82 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): 83 | super().__init__() 84 | 85 | self.bias = nn.Parameter(torch.zeros(channel)) 86 | self.negative_slope = negative_slope 87 | self.scale = scale 88 | 89 | def forward(self, input): 90 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 91 | 92 | 93 | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): 94 | if input.device.type == "cpu": 95 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 96 | return ( 97 | F.leaky_relu( 98 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 99 | ) 100 | * scale 101 | ) 102 | 103 | else: 104 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 105 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | #include 9 | 10 | 11 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale); 13 | 14 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 15 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 17 | 18 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 19 | int act, int grad, float alpha, float scale) { 20 | CHECK_CUDA(input); 21 | CHECK_CUDA(bias); 22 | 23 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 24 | } 25 | 26 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 27 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 28 | } -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved. 3 | // --------------------------------------------------------------- 4 | // 5 | // This work is made available under the Nvidia Source Code License-NC. 6 | // To view a copy of this license, visit 7 | // https://nvlabs.github.io/stylegan2/license.html 8 | 9 | #include 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | #include 17 | #include 18 | 19 | 20 | template 21 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 22 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 23 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 24 | 25 | scalar_t zero = 0.0; 26 | 27 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 28 | scalar_t x = p_x[xi]; 29 | 30 | if (use_bias) { 31 | x += p_b[(xi / step_b) % size_b]; 32 | } 33 | 34 | scalar_t ref = use_ref ? p_ref[xi] : zero; 35 | 36 | scalar_t y; 37 | 38 | switch (act * 10 + grad) { 39 | default: 40 | case 10: y = x; break; 41 | case 11: y = x; break; 42 | case 12: y = 0.0; break; 43 | 44 | case 30: y = (x > 0.0) ? x : x * alpha; break; 45 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 46 | case 32: y = 0.0; break; 47 | } 48 | 49 | out[xi] = y * scale; 50 | } 51 | } 52 | 53 | 54 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 55 | int act, int grad, float alpha, float scale) { 56 | int curDevice = -1; 57 | cudaGetDevice(&curDevice); 58 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 59 | 60 | auto x = input.contiguous(); 61 | auto b = bias.contiguous(); 62 | auto ref = refer.contiguous(); 63 | 64 | int use_bias = b.numel() ? 1 : 0; 65 | int use_ref = ref.numel() ? 1 : 0; 66 | 67 | int size_x = x.numel(); 68 | int size_b = b.numel(); 69 | int step_b = 1; 70 | 71 | for (int i = 1 + 1; i < x.dim(); i++) { 72 | step_b *= x.size(i); 73 | } 74 | 75 | int loop_x = 4; 76 | int block_size = 4 * 32; 77 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 78 | 79 | auto y = torch::empty_like(x); 80 | 81 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 82 | fused_bias_act_kernel<<>>( 83 | y.data_ptr(), 84 | x.data_ptr(), 85 | b.data_ptr(), 86 | ref.data_ptr(), 87 | act, 88 | grad, 89 | alpha, 90 | scale, 91 | loop_x, 92 | size_x, 93 | step_b, 94 | size_b, 95 | use_bias, 96 | use_ref 97 | ); 98 | }); 99 | 100 | return y; 101 | } -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/libcudart.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/train/structure/config/inpainting/models/op/libcudart.so -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | 9 | #include 10 | 11 | 12 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 15 | 16 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 17 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 18 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 19 | 20 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 21 | int up_x, int up_y, int down_x, int down_y, 22 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 23 | CHECK_CUDA(input); 24 | CHECK_CUDA(kernel); 25 | 26 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 27 | } 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 31 | } -------------------------------------------------------------------------------- /train/structure/config/inpainting/models/optimizer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google Research. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """PyTorch implementation of the Lion optimizer.""" 16 | 17 | import torch 18 | from torch.optim.optimizer import Optimizer 19 | 20 | 21 | class Lion(Optimizer): 22 | r"""Implements Lion algorithm.""" 23 | 24 | def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0): 25 | """Initialize the hyperparameters. 26 | 27 | Args: 28 | params (iterable): iterable of parameters to optimize or dicts defining 29 | parameter groups 30 | lr (float, optional): learning rate (default: 1e-4) 31 | betas (Tuple[float, float], optional): coefficients used for computing 32 | running averages of gradient and its square (default: (0.9, 0.99)) 33 | weight_decay (float, optional): weight decay coefficient (default: 0) 34 | """ 35 | 36 | if not 0.0 <= lr: 37 | raise ValueError('Invalid learning rate: {}'.format(lr)) 38 | if not 0.0 <= betas[0] < 1.0: 39 | raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) 40 | if not 0.0 <= betas[1] < 1.0: 41 | raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) 42 | defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) 43 | super().__init__(params, defaults) 44 | 45 | @torch.no_grad() 46 | def step(self, closure=None): 47 | """Performs a single optimization step. 48 | 49 | Args: 50 | closure (callable, optional): A closure that reevaluates the model 51 | and returns the loss. 52 | 53 | Returns: 54 | the loss. 55 | """ 56 | loss = None 57 | if closure is not None: 58 | with torch.enable_grad(): 59 | loss = closure() 60 | 61 | for group in self.param_groups: 62 | for p in group['params']: 63 | if p.grad is None: 64 | continue 65 | 66 | # Perform stepweight decay 67 | p.data.mul_(1 - group['lr'] * group['weight_decay']) 68 | 69 | grad = p.grad 70 | state = self.state[p] 71 | # State initialization 72 | if len(state) == 0: 73 | # Exponential moving average of gradient values 74 | state['exp_avg'] = torch.zeros_like(p) 75 | 76 | exp_avg = state['exp_avg'] 77 | beta1, beta2 = group['betas'] 78 | 79 | # Weight update 80 | update = exp_avg * beta1 + grad * (1 - beta1) 81 | p.add_(torch.sign(update), alpha=-group['lr']) 82 | # Decay the momentum running average coefficient 83 | exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) 84 | 85 | return loss 86 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/options/train/ir-sde.yml: -------------------------------------------------------------------------------- 1 | #### general settings 2 | name: ir-sde 3 | use_tb_logger: true 4 | model: denoising 5 | distortion: derain 6 | gpu_ids: [2] 7 | 8 | sde: 9 | max_sigma: 30 10 | T: 400 11 | schedule: cosine 12 | eps: 0.005 13 | 14 | degradation: 15 | mask_root: your mask path 16 | 17 | #### datasets 18 | datasets: 19 | train: 20 | name: Train_Dataset 21 | mode: GT 22 | dataroot_GT: your image path 23 | 24 | use_shuffle: true 25 | n_workers: 4 26 | batch_size: 16 27 | GT_size: 256 28 | use_flip: true 29 | use_rot: true 30 | color: RGB 31 | 32 | 33 | #### network structures 34 | network_G: 35 | which_model_G: ConditionalUNet 36 | setting: 37 | in_nc: 1 38 | out_nc: 1 39 | nf: 64 40 | depth: 4 41 | 42 | #### path 43 | path: 44 | pretrain_model_G: ~ 45 | strict_load: true 46 | resume_state: ~ 47 | 48 | #### training settings: learning rate scheme, loss 49 | train: 50 | optimizer: Adam 51 | lr_G: !!float 1e-4 52 | lr_scheme: MultiStepLR 53 | beta1: 0.9 54 | beta2: 0.99 55 | niter: 8000000 56 | warmup_iter: -1 # no warm up 57 | lr_steps: [2000000, 4000000, 6000000] 58 | lr_gamma: 0.5 59 | eta_min: !!float 1e-7 60 | 61 | # criterion 62 | is_weighted: False 63 | loss_type: l1 64 | weight: 1.0 65 | 66 | manual_seed: 0 67 | 68 | #### logger 69 | logger: 70 | print_freq: 100 71 | save_checkpoint_freq: !!float 5e3 72 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/str_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/structure/config/inpainting/str_utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import abc 4 | import math 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torchvision.utils as tvutils 10 | 11 | 12 | ########### denoising ############### 13 | def add_noise(tensor, sigma): 14 | sigma = sigma / 255 if sigma > 1 else sigma 15 | return tensor + torch.randn_like(tensor) * sigma 16 | 17 | 18 | ######## inpainting ########### 19 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 20 | batch = tensor.shape[0] 21 | if mask_id < 0: 22 | mask_id = np.random.randint(0, n, batch) 23 | masks = [] 24 | for i in range(batch): 25 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 26 | mask = np.concatenate(masks, axis=0) 27 | else: 28 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 29 | 30 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 31 | # for images are clipped or scaled 32 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 33 | masked_tensor = mask * tensor 34 | return masked_tensor + (1. - mask) 35 | 36 | ######## super-resolution ########### 37 | 38 | def upscale(tensor, scale=4, mode='bicubic'): 39 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 40 | return tensor 41 | 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /train/structure/config/inpainting/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/structure/config/inpainting/utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import abc 3 | import math 4 | import numpy as np 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torchvision.utils as tvutils 9 | import cv2 10 | 11 | ########### denoising ############### 12 | def add_noise(tensor, sigma): 13 | sigma = sigma / 255 if sigma > 1 else sigma 14 | return tensor + torch.randn_like(tensor) * sigma 15 | 16 | 17 | ######## inpainting ########### 18 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 19 | batch = tensor.shape[0] 20 | if mask_id < 0: 21 | mask_id = np.random.randint(0, n, batch) 22 | masks = [] 23 | for i in range(batch): 24 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 25 | mask = np.concatenate(masks, axis=0) 26 | else: 27 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 28 | 29 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 30 | # for images are clipped or scaled 31 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 32 | masked_tensor = mask * tensor 33 | return masked_tensor + (1. - mask) 34 | 35 | ######## super-resolution ########### 36 | 37 | def upscale(tensor, scale=4, mode='bicubic'): 38 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 39 | return tensor 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/BokehLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class BokehLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and GT image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LR_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from image files 32 | if opt["data_type"] == "img": 33 | self.LR_paths = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) # LR list 36 | self.metas = self._read_meta_data(opt["dataroot_meta"]) 37 | else: 38 | print("Error: data_type is not matched in Dataset") 39 | 40 | def _read_meta_data(self, meta_file_path: str): 41 | """Read the meta file containing source / target lens and disparity for each image. 42 | Args: 43 | meta_file_path (str): File path 44 | Raises: 45 | ValueError: File not found. 46 | Returns: 47 | dict: Meta dict of tuples like {id: (id, src_lens, tgt_lens, disparity)}. 48 | """ 49 | if not os.path.isfile(meta_file_path): 50 | raise ValueError(f"Meta file missing under {meta_file_path}.") 51 | 52 | meta = {} 53 | with open(meta_file_path, "r") as f: 54 | lines = f.readlines() 55 | 56 | for line in lines: 57 | id, src_lens, tgt_lens, disparity = [part.strip() for part in line.split(",")] 58 | meta[id] = (src_lens, tgt_lens, disparity) 59 | return meta 60 | 61 | def lenstr2tensor(self, lenstr, scale=1.): 62 | # Canon50mm -> -1, Sony50mm -> 1 63 | lenstr = lenstr.replace('Canon50mmf', '-') 64 | lenstr = lenstr.replace('Sony50mmf', '') 65 | lenstr = lenstr.replace('BS', '') 66 | return torch.tensor(float(lenstr)) * scale 67 | 68 | def __getitem__(self, index): 69 | 70 | # get LR image 71 | LR_path = self.LR_paths[index] 72 | img_LR = util.read_img(self.LR_env, LR_path, None) 73 | 74 | id = os.path.basename(LR_path).split(".")[0] 75 | src_lens, tgt_lens, disparity = self.metas[id] 76 | 77 | src_lens = self.lenstr2tensor(src_lens, 10) 78 | tgt_lens = self.lenstr2tensor(tgt_lens, 10) 79 | disparity = self.lenstr2tensor(disparity) 80 | 81 | # change color space if necessary 82 | if self.opt["color"]: 83 | H, W, C = img_LR.shape 84 | img_LR = util.channel_convert(C, self.opt["color"], [img_LR])[ 85 | 0 86 | ] # TODO during val no definition 87 | 88 | 89 | # BGR to RGB, HWC to CHW, numpy to tensor 90 | if img_LR.shape[2] == 3: 91 | img_LR = img_LR[:, :, [2, 1, 0]] 92 | 93 | img_LR = torch.from_numpy( 94 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 95 | ).float() 96 | 97 | return { 98 | "LQ": img_LR, 99 | "src_lens": src_lens, 100 | "tgt_lens": tgt_lens, 101 | "disparity": disparity, 102 | "LQ_path": LR_path, 103 | } 104 | 105 | def __len__(self): 106 | return len(self.LR_paths) 107 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/GT_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | 12 | 13 | import numpy as np 14 | from PIL import Image 15 | 16 | from torchvision import transforms 17 | 18 | from skimage.feature import canny 19 | from skimage.color import gray2rgb, rgb2gray 20 | 21 | 22 | def tensor_to_image(): 23 | 24 | return transforms.ToPILImage() 25 | 26 | 27 | def image_to_tensor(): 28 | 29 | return transforms.ToTensor() 30 | 31 | 32 | def image_to_edge(image, sigma): 33 | 34 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 35 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 36 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 37 | 38 | return edge, gray_image 39 | 40 | 41 | 42 | 43 | 44 | try: 45 | sys.path.append("..") 46 | import data.util as util 47 | except ImportError: 48 | pass 49 | 50 | 51 | class GTDataset(data.Dataset): 52 | """ 53 | Read LR (Low Quality, here is LR) and GT image pairs. 54 | The pair is ensured by 'sorted' function, so please check the name convention. 55 | """ 56 | 57 | def __init__(self, opt): 58 | super().__init__() 59 | self.opt = opt 60 | self.GT_paths = None 61 | self.GT_env = None # environment for lmdb 62 | self.GT_size = opt["GT_size"] 63 | 64 | # read image list from lmdb or image files 65 | if opt["data_type"] == "lmdb": 66 | self.GT_paths, self.GT_sizes = util.get_image_paths( 67 | opt["data_type"], opt["dataroot_GT"] 68 | ) 69 | elif opt["data_type"] == "img": 70 | self.GT_paths = util.get_image_paths( 71 | opt["data_type"], opt["dataroot_GT"] 72 | ) # GT list 73 | else: 74 | print("Error: data_type is not matched in Dataset") 75 | assert self.GT_paths, "Error: GT paths are empty." 76 | print("dataset length: {}".format(len(self.GT_paths))) 77 | self.random_scale_list = [1] 78 | 79 | def _init_lmdb(self): 80 | # https://github.com/chainer/chainermn/issues/129 81 | self.GT_env = lmdb.open( 82 | self.opt["dataroot_GT"], 83 | readonly=True, 84 | lock=False, 85 | readahead=False, 86 | meminit=False, 87 | ) 88 | 89 | def __getitem__(self, index): 90 | if self.opt["data_type"] == "lmdb": 91 | if self.GT_env is None: 92 | self._init_lmdb() 93 | 94 | GT_path = None 95 | GT_size = self.opt["GT_size"] 96 | 97 | # get GT image 98 | GT_path = self.GT_paths[index] 99 | if self.opt["data_type"] == "lmdb": 100 | resolution = [int(s) for s in self.GT_sizes[index].split("_")] 101 | else: 102 | resolution = None 103 | img_GT = util.read_img( 104 | self.GT_env, GT_path, resolution 105 | ) # return: Numpy float32, HWC, BGR, [0,1] 106 | 107 | if self.opt["phase"] == "train": 108 | H, W, C = img_GT.shape 109 | 110 | rnd_h = random.randint(0, max(0, H - GT_size)) 111 | rnd_w = random.randint(0, max(0, W - GT_size)) 112 | img_GT = img_GT[rnd_h : rnd_h + GT_size, rnd_w : rnd_w + GT_size, :] 113 | 114 | # augmentation - flip, rotate 115 | img_GT = util.augment( 116 | img_GT, 117 | self.opt["use_flip"], 118 | self.opt["use_rot"], 119 | self.opt["mode"], 120 | ) 121 | 122 | # change color space if necessary 123 | if self.opt["color"]: 124 | img_GT = util.channel_convert(img_GT.shape[2], self.opt["color"], [img_GT])[ 125 | 0 126 | ] 127 | 128 | # BGR to RGB, HWC to CHW, numpy to tensor 129 | if img_GT.shape[2] == 3: 130 | img_GT = img_GT[:, :, [2, 1, 0]] 131 | img_GT = torch.from_numpy( 132 | np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1))) 133 | ).float() 134 | 135 | GT_edge,GT_gray = image_to_edge(img_GT, sigma=3.) 136 | return {"GT": img_GT, "GT_path": GT_path, "GT_edge": GT_edge, "GT_gray": GT_gray} 137 | 138 | def __len__(self): 139 | return len(self.GT_paths) 140 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/LQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class LQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path = None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path = self.LQ_paths[index] 67 | if self.opt["data_type"] == "lmdb": 68 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 69 | else: 70 | resolution = None 71 | img_LR = util.read_img( 72 | self.LR_env, LR_path, resolution 73 | ) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # modcrop in the validation / test phase 76 | if self.opt["phase"] != "train": 77 | img_LR = util.modcrop(img_LR, scale) 78 | 79 | if self.opt["phase"] == "train": 80 | H, W, C = img_LR.shape 81 | 82 | rnd_h = random.randint(0, max(0, H - LR_size)) 83 | rnd_w = random.randint(0, max(0, W - LR_size)) 84 | img_LR = img_LR[rnd_h : rnd_h + LR_size, rnd_w : rnd_w + LR_size, :] 85 | 86 | # augmentation - flip, rotate 87 | img_LR = util.augment( 88 | img_LR, 89 | self.opt["use_flip"], 90 | self.opt["use_rot"], 91 | self.opt["mode"], 92 | ) 93 | 94 | # change color space if necessary 95 | if self.opt["color"]: 96 | img_LR = util.channel_convert(img_LR.shape[2], self.opt["color"], [img_LR])[ 97 | 0 98 | ] 99 | 100 | # BGR to RGB, HWC to CHW, numpy to tensor 101 | if img_LR.shape[2] == 3: 102 | img_LR = img_LR[:, :, [2, 1, 0]] 103 | img_LR = torch.from_numpy( 104 | np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1))) 105 | ).float() 106 | 107 | return {"LQ": img_LR, "LQ_path": LR_path} 108 | 109 | def __len__(self): 110 | return len(self.LQ_paths) 111 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/StereoLQ_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | 5 | import cv2 6 | import lmdb 7 | import numpy as np 8 | import torch 9 | import torch.utils.data as data 10 | 11 | try: 12 | sys.path.append("..") 13 | import data.util as util 14 | except ImportError: 15 | pass 16 | 17 | 18 | class StereoLQDataset(data.Dataset): 19 | """ 20 | Read LR (Low Quality, here is LR) and LR image pairs. 21 | The pair is ensured by 'sorted' function, so please check the name convention. 22 | """ 23 | 24 | def __init__(self, opt): 25 | super().__init__() 26 | self.opt = opt 27 | self.LQ_paths = None 28 | self.LR_env = None # environment for lmdb 29 | self.LR_size = opt["LR_size"] 30 | 31 | # read image list from lmdb or image files 32 | if opt["data_type"] == "lmdb": 33 | self.LQ_paths, self.LR_sizes = util.get_image_paths( 34 | opt["data_type"], opt["dataroot_LQ"] 35 | ) 36 | elif opt["data_type"] == "img": 37 | self.LQ_paths = util.get_image_paths( 38 | opt["data_type"], opt["dataroot_LQ"] 39 | ) # LR list 40 | else: 41 | print("Error: data_type is not matched in Dataset") 42 | assert self.LQ_paths, "Error: LQ paths are empty." 43 | 44 | self.random_scale_list = [1] 45 | 46 | def _init_lmdb(self): 47 | # https://github.com/chainer/chainermn/issues/129 48 | self.LR_env = lmdb.open( 49 | self.opt["dataroot_LR"], 50 | readonly=True, 51 | lock=False, 52 | readahead=False, 53 | meminit=False, 54 | ) 55 | 56 | def __getitem__(self, index): 57 | if self.opt["data_type"] == "lmdb": 58 | if self.LR_env is None: 59 | self._init_lmdb() 60 | 61 | LR_path_L, LR_path_R = None, None 62 | scale = self.opt["scale"] 63 | LR_size = self.opt["LR_size"] 64 | 65 | # get LR image 66 | LR_path_L = self.LQ_paths[index*2] 67 | LR_path_R = self.LQ_paths[index*2+1] 68 | if self.opt["data_type"] == "lmdb": 69 | resolution = [int(s) for s in self.LR_sizes[index].split("_")] 70 | else: 71 | resolution = None 72 | imgL_LR = util.read_img(self.LR_env, LR_path_L, resolution) # return: Numpy float32, HWC, BGR, [0,1] 73 | imgR_LR = util.read_img(self.LR_env, LR_path_R, resolution) # return: Numpy float32, HWC, BGR, [0,1] 74 | 75 | # change color space if necessary 76 | if self.opt["color"]: 77 | imgL_LR, imgR_LR = util.channel_convert( 78 | imgL_LR.shape[2], self.opt["color"], [imgL_LR, imgR_LR]) 79 | 80 | # BGR to RGB, HWC to CHW, numpy to tensor 81 | if imgL_LR.shape[2] == 3: 82 | imgL_LR = imgL_LR[:, :, [2, 1, 0]] 83 | imgR_LR = imgR_LR[:, :, [2, 1, 0]] 84 | imgL_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgL_LR, (2, 0, 1)))).float() 85 | imgR_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(imgR_LR, (2, 0, 1)))).float() 86 | 87 | img_LR = torch.cat([imgL_LR, imgR_LR], dim=0) 88 | 89 | return {"LQ": img_LR, "LQ_path": LR_path_L} 90 | 91 | def __len__(self): 92 | return len(self.LQ_paths) // 2 93 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/__init__.py: -------------------------------------------------------------------------------- 1 | """create dataset and dataloader""" 2 | import logging 3 | 4 | import torch 5 | import torch.utils.data 6 | 7 | 8 | def create_dataloader(dataset, dataset_opt, opt=None, sampler=None): 9 | phase = dataset_opt["phase"] 10 | if phase == "train": 11 | if opt["dist"]: 12 | world_size = torch.distributed.get_world_size() 13 | num_workers = dataset_opt["n_workers"] 14 | assert dataset_opt["batch_size"] % world_size == 0 15 | batch_size = dataset_opt["batch_size"] // world_size 16 | shuffle = False 17 | else: 18 | num_workers = dataset_opt["n_workers"] * len(opt["gpu_ids"]) 19 | batch_size = dataset_opt["batch_size"] 20 | shuffle = True 21 | return torch.utils.data.DataLoader( 22 | dataset, 23 | batch_size=batch_size, 24 | shuffle=shuffle, 25 | num_workers=num_workers, 26 | sampler=sampler, 27 | drop_last=True, 28 | pin_memory=False, 29 | ) 30 | else: 31 | return torch.utils.data.DataLoader( 32 | dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=(phase=="val") 33 | ) 34 | 35 | 36 | def create_dataset(dataset_opt): 37 | mode = dataset_opt["mode"] 38 | if mode == "LQ": # Predictor 39 | from data.LQ_dataset import LQDataset as D 40 | dataset = D(dataset_opt) 41 | elif mode == "LQGT": # SFTMD 42 | from data.LQGT_dataset import LQGTDataset as D 43 | dataset = D(dataset_opt) 44 | elif mode == "GT": # Corrector 45 | from data.GT_dataset import GTDataset as D 46 | dataset = D(dataset_opt) 47 | elif mode == 'SteLQGT': 48 | from data.StereoLQGT_dataset import StereoLQGTDataset as D 49 | dataset = D(dataset_opt) 50 | elif mode == 'SteLQ': 51 | from data.StereoLQ_dataset import StereoLQDataset as D 52 | dataset = D(dataset_opt) 53 | elif mode == 'BokehLQGT': 54 | from data.BokehLQGT_dataset import BokehLQGTDataset as D 55 | dataset = D(dataset_opt) 56 | elif mode == 'BokehLQ': 57 | from data.BokehLQ_dataset import BokehLQDataset as D 58 | dataset = D(dataset_opt) 59 | else: 60 | raise NotImplementedError("Dataset [{:s}] is not recognized.".format(mode)) 61 | 62 | logger = logging.getLogger("base") 63 | logger.info( 64 | "Dataset [{:s} - {:s}] is created.".format( 65 | dataset.__class__.__name__, dataset_opt["name"] 66 | ) 67 | ) 68 | return dataset 69 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def image_to_edge(image, sigma): 21 | 22 | gray_image = rgb2gray(np.array(tensor_to_image()(image))) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | gray_image = image_to_tensor()(Image.fromarray(gray_image)) 25 | 26 | return edge, gray_image 27 | 28 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/data/data_sampler.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modified from torch.utils.data.distributed.DistributedSampler 3 | Support enlarging the dataset for *iter-oriented* training, for saving time when restart the 4 | dataloader after each epoch 5 | """ 6 | import math 7 | 8 | import torch 9 | import torch.distributed as dist 10 | from torch.utils.data.sampler import Sampler 11 | 12 | 13 | class DistIterSampler(Sampler): 14 | """Sampler that restricts data loading to a subset of the dataset. 15 | 16 | It is especially useful in conjunction with 17 | :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each 18 | process can pass a DistributedSampler instance as a DataLoader sampler, 19 | and load a subset of the original dataset that is exclusive to it. 20 | 21 | .. note:: 22 | Dataset is assumed to be of constant size. 23 | 24 | Arguments: 25 | dataset: Dataset used for sampling. 26 | num_replicas (optional): Number of processes participating in 27 | distributed training. 28 | rank (optional): Rank of the current process within num_replicas. 29 | """ 30 | 31 | def __init__(self, dataset, num_replicas=None, rank=None, ratio=100): 32 | if num_replicas is None: 33 | if not dist.is_available(): 34 | raise RuntimeError("Requires distributed package to be available") 35 | num_replicas = dist.get_world_size() 36 | if rank is None: 37 | if not dist.is_available(): 38 | raise RuntimeError("Requires distributed package to be available") 39 | rank = dist.get_rank() 40 | self.dataset = dataset 41 | self.num_replicas = num_replicas 42 | self.rank = rank 43 | self.epoch = 0 44 | self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas)) 45 | self.total_size = self.num_samples * self.num_replicas 46 | 47 | def __iter__(self): 48 | # deterministically shuffle based on epoch 49 | g = torch.Generator() 50 | g.manual_seed(self.epoch) 51 | indices = torch.randperm( 52 | self.total_size, generator=g 53 | ).tolist() # Returns a random permutation of integers from 0 to n - 1 54 | 55 | dsize = len(self.dataset) 56 | indices = [v % dsize for v in indices] 57 | 58 | # subsample 59 | indices = indices[self.rank : self.total_size : self.num_replicas] 60 | assert len(indices) == self.num_samples 61 | 62 | return iter(indices) 63 | 64 | def __len__(self): 65 | return self.num_samples 66 | 67 | def set_epoch(self, epoch): 68 | self.epoch = epoch 69 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/demo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ############################################################# 4 | ### training ### 5 | 6 | # for single GPU 7 | python train.py -opt=options/train/ir-sde.yml 8 | # python train.py -opt=options/train/refusion.yml 9 | 10 | # for multiple GPUs 11 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/ir-sde.yml --launcher pytorch 12 | # python -m torch.distributed.launch --nproc_per_node=2 --master_port=6512 train.py -opt=options/train/refusion.yml --launcher pytorch 13 | 14 | ############################################################# 15 | 16 | ### testing ### 17 | # python test.py -opt=options/test/ir-sde.yml 18 | # python test.py -opt=options/test/refusion.yml 19 | 20 | ############################################################# -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logger = logging.getLogger("base") 4 | 5 | 6 | def create_model(opt): 7 | model = opt["model"] 8 | 9 | if model == "denoising": 10 | from .denoising_model import DenoisingModel as M 11 | else: 12 | raise NotImplementedError("Model [{:s}] not recognized.".format(model)) 13 | m = M(opt) 14 | logger.info("Model [{:s}] is created.".format(m.__class__.__name__)) 15 | return m 16 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/base_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | from collections import OrderedDict 3 | 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.parallel import DistributedDataParallel 7 | 8 | 9 | class BaseModel: 10 | def __init__(self, opt): 11 | self.opt = opt 12 | self.device = torch.device("cuda" if opt["gpu_ids"] is not None else "cpu") 13 | self.is_train = opt["is_train"] 14 | self.schedulers = [] 15 | self.optimizers = [] 16 | 17 | def feed_data(self, data): 18 | pass 19 | 20 | def optimize_parameters(self): 21 | pass 22 | 23 | def get_current_visuals(self): 24 | pass 25 | 26 | def get_current_losses(self): 27 | pass 28 | 29 | def print_network(self): 30 | pass 31 | 32 | def save(self, label): 33 | pass 34 | 35 | def load(self): 36 | pass 37 | 38 | def _set_lr(self, lr_groups_l): 39 | """set learning rate for warmup, 40 | lr_groups_l: list for lr_groups. each for a optimizer""" 41 | for optimizer, lr_groups in zip(self.optimizers, lr_groups_l): 42 | for param_group, lr in zip(optimizer.param_groups, lr_groups): 43 | param_group["lr"] = lr 44 | 45 | def _get_init_lr(self): 46 | # get the initial lr, which is set by the scheduler 47 | init_lr_groups_l = [] 48 | for optimizer in self.optimizers: 49 | init_lr_groups_l.append([v["initial_lr"] for v in optimizer.param_groups]) 50 | return init_lr_groups_l 51 | 52 | def update_learning_rate(self, cur_iter, warmup_iter=-1): 53 | for scheduler in self.schedulers: 54 | scheduler.step() 55 | #### set up warm up learning rate 56 | if cur_iter < warmup_iter: 57 | # get initial lr for each group 58 | init_lr_g_l = self._get_init_lr() 59 | # modify warming-up learning rates 60 | warm_up_lr_l = [] 61 | for init_lr_g in init_lr_g_l: 62 | warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g]) 63 | # set learning rate 64 | self._set_lr(warm_up_lr_l) 65 | 66 | def get_current_learning_rate(self): 67 | # return self.schedulers[0].get_lr()[0] 68 | return self.optimizers[0].param_groups[0]["lr"] 69 | 70 | def get_network_description(self, network): 71 | """Get the string and total parameters of the network""" 72 | if isinstance(network, nn.DataParallel) or isinstance( 73 | network, DistributedDataParallel 74 | ): 75 | network = network.module 76 | s = str(network) 77 | n = sum(map(lambda x: x.numel(), network.parameters())) 78 | return s, n 79 | 80 | def save_network(self, network, network_label, iter_label): 81 | save_filename = "{}_{}.pth".format(iter_label, network_label) 82 | save_path = os.path.join(self.opt["path"]["models"], save_filename) 83 | if isinstance(network, nn.DataParallel) or isinstance( 84 | network, DistributedDataParallel 85 | ): 86 | network = network.module 87 | state_dict = network.state_dict() 88 | for key, param in state_dict.items(): 89 | state_dict[key] = param.cpu() 90 | torch.save(state_dict, save_path) 91 | 92 | def load_network(self, load_path, network, strict=True): 93 | if isinstance(network, nn.DataParallel) or isinstance( 94 | network, DistributedDataParallel 95 | ): 96 | network = network.module 97 | print(network) 98 | load_net = torch.load(load_path) 99 | load_net_clean = OrderedDict() # remove unnecessary 'module.' 100 | for k, v in load_net.items(): 101 | if k.startswith("module."): 102 | load_net_clean[k[7:]] = v 103 | else: 104 | load_net_clean[k] = v 105 | 106 | network.load_state_dict(load_net_clean, strict=strict) 107 | 108 | def save_training_state(self, epoch, iter_step): 109 | """Saves training state during training, which will be used for resuming""" 110 | state = {"epoch": epoch, "iter": iter_step, "schedulers": [], "optimizers": []} 111 | for s in self.schedulers: 112 | state["schedulers"].append(s.state_dict()) 113 | for o in self.optimizers: 114 | state["optimizers"].append(o.state_dict()) 115 | save_filename = "{}.state".format(iter_step) 116 | save_path = os.path.join(self.opt["path"]["training_state"], save_filename) 117 | torch.save(state, save_path) 118 | 119 | def resume_training(self, resume_state): 120 | """Resume the optimizers and schedulers for training""" 121 | resume_optimizers = resume_state["optimizers"] 122 | resume_schedulers = resume_state["schedulers"] 123 | assert len(resume_optimizers) == len( 124 | self.optimizers 125 | ), "Wrong lengths of optimizers" 126 | assert len(resume_schedulers) == len( 127 | self.schedulers 128 | ), "Wrong lengths of schedulers" 129 | for i, o in enumerate(resume_optimizers): 130 | self.optimizers[i].load_state_dict(o) 131 | for i, s in enumerate(resume_schedulers): 132 | self.schedulers[i].load_state_dict(s) 133 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/canny.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from PIL import Image 3 | 4 | from torchvision import transforms 5 | 6 | from skimage.feature import canny 7 | from skimage.color import gray2rgb, rgb2gray 8 | 9 | 10 | def tensor_to_image(): 11 | 12 | return transforms.ToPILImage() 13 | 14 | 15 | def image_to_tensor(): 16 | 17 | return transforms.ToTensor() 18 | 19 | 20 | def gray_to_edge(image, sigma): 21 | 22 | gray_image = np.array(tensor_to_image()(image)) 23 | edge = image_to_tensor()(Image.fromarray(canny(gray_image, sigma=sigma))) 24 | 25 | return edge 26 | 27 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/dense_layer.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # This file has been modified from a file released under the MIT License. 5 | # 6 | # Source: 7 | # https://github.com/CW-Huang/sdeflow-light/blob/524650bc5ad69522b3e0905672deef0650374512/lib/models/unet.py 8 | # 9 | # The license for the original version of this file can be 10 | # found in this directory (LICENSE_MIT). The modifications 11 | # to this file are subject to the same MIT License. 12 | # --------------------------------------------------------------- 13 | 14 | 15 | import math 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | from torch.nn.init import _calculate_fan_in_and_fan_out 20 | import numpy as np 21 | 22 | 23 | def _calculate_correct_fan(tensor, mode): 24 | """ 25 | copied and modified from https://github.com/pytorch/pytorch/blob/master/torch/nn/init.py#L337 26 | """ 27 | mode = mode.lower() 28 | valid_modes = ['fan_in', 'fan_out', 'fan_avg'] 29 | if mode not in valid_modes: 30 | raise ValueError("Mode {} not supported, please use one of {}".format(mode, valid_modes)) 31 | 32 | fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor) 33 | return fan_in if mode == 'fan_in' else fan_out 34 | 35 | 36 | def kaiming_uniform_(tensor, gain=1., mode='fan_in'): 37 | r"""Fills the input `Tensor` with values according to the method 38 | described in `Delving deep into rectifiers: Surpassing human-level 39 | performance on ImageNet classification` - He, K. et al. (2015), using a 40 | uniform distribution. The resulting tensor will have values sampled from 41 | :math:`\mathcal{U}(-\text{bound}, \text{bound})` where 42 | .. math:: 43 | \text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}} 44 | Also known as He initialization. 45 | Args: 46 | tensor: an n-dimensional `torch.Tensor` 47 | gain: multiplier to the dispersion 48 | mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'`` 49 | preserves the magnitude of the variance of the weights in the 50 | forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the 51 | backwards pass. 52 | Examples: 53 | >>> w = torch.empty(3, 5) 54 | >>> nn.init.kaiming_uniform_(w, mode='fan_in') 55 | """ 56 | fan = _calculate_correct_fan(tensor, mode) 57 | var = gain / max(1., fan) 58 | bound = math.sqrt(3.0 * var) # Calculate uniform bounds from standard deviation 59 | with torch.no_grad(): 60 | return tensor.uniform_(-bound, bound) 61 | 62 | 63 | def variance_scaling_init_(tensor, scale): 64 | return kaiming_uniform_(tensor, gain=1e-10 if scale == 0 else scale, mode='fan_avg') 65 | 66 | 67 | def dense(in_channels, out_channels, init_scale=1.): 68 | lin = nn.Linear(in_channels, out_channels) 69 | variance_scaling_init_(lin.weight, scale=init_scale) 70 | nn.init.zeros_(lin.bias) 71 | return lin 72 | 73 | def conv2d(in_planes, out_planes, kernel_size=(3, 3), stride=1, dilation=1, padding=1, bias=True, padding_mode='zeros', 74 | init_scale=1.): 75 | conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 76 | bias=bias, padding_mode=padding_mode) 77 | variance_scaling_init_(conv.weight, scale=init_scale) 78 | if bias: 79 | nn.init.zeros_(conv.bias) 80 | return conv 81 | 82 | 83 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .DenoisingUNet_arch import ConditionalUNet -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/modules/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import einops 5 | import numpy as np 6 | import sys 7 | 8 | class MatchingLoss(nn.Module): 9 | def __init__(self, loss_type='l1', is_weighted=False): 10 | super().__init__() 11 | self.is_weighted = is_weighted 12 | 13 | if loss_type == 'l1': 14 | self.loss_fn = F.l1_loss 15 | elif loss_type == 'l2': 16 | self.loss_fn = F.mse_loss 17 | else: 18 | raise ValueError(f'invalid loss type {loss_type}') 19 | 20 | def forward(self, predict, target, mask, weights=None): 21 | 22 | lossm = self.loss_fn(predict * (1 - mask), target * (1 - mask), reduction='none') 23 | lossm = einops.reduce(lossm, 'b ... -> b (...)', 'mean') 24 | 25 | lossu = self.loss_fn(predict * mask, target * mask, reduction='none') 26 | lossu = einops.reduce(lossu, 'b ... -> b (...)', 'mean') 27 | 28 | loss = lossu + 10 * lossm 29 | if self.is_weighted and weights is not None: 30 | loss = weights * loss 31 | 32 | return loss.mean() 33 | 34 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/LICENSE_MIT: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Kim Seonghyeon 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/__init__.py: -------------------------------------------------------------------------------- 1 | from .fused_act import FusedLeakyReLU, fused_leaky_relu 2 | from .upfirdn2d import upfirdn2d 3 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/cuda-installer.log: -------------------------------------------------------------------------------- 1 | [INFO]: Driver not installed. 2 | [INFO]: Checking compiler version... 3 | [INFO]: gcc location: /usr/bin/gcc 4 | 5 | [INFO]: gcc version: gcc version 9.4.0 (Ubuntu 9.4.0-1ubuntu1~20.04.2) 6 | 7 | [INFO]: Initializing menu 8 | [INFO]: nvidia-fs.setKOVersion(2.13.5) 9 | [INFO]: Setup complete 10 | [INFO]: Components to install: 11 | [INFO]: Driver 12 | [INFO]: 520.61.05 13 | [INFO]: Executing NVIDIA-Linux-x86_64-520.61.05.run --ui=none --no-questions --accept-license --disable-nouveau --no-cc-version-check --install-libglvnd 2>&1 14 | [INFO]: Finished with code: 256 15 | [ERROR]: Install of driver component failed. Consult the driver log at /var/log/nvidia-installer.log for more details. 16 | [ERROR]: Install of 520.61.05 failed, quitting 17 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/fused_act.py: -------------------------------------------------------------------------------- 1 | # --------------------------------------------------------------- 2 | # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | # --------------------------------------------------------------- 4 | 5 | """ Originated from https://github.com/rosinality/stylegan2-pytorch 6 | The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | """ 8 | 9 | import os 10 | 11 | import torch 12 | from torch import nn 13 | from torch.nn import functional as F 14 | from torch.autograd import Function 15 | from torch.utils.cpp_extension import load 16 | 17 | module_path = os.path.dirname(__file__) 18 | fused = load( 19 | "fused", 20 | sources=[ 21 | os.path.join(module_path, "fused_bias_act.cpp"), 22 | os.path.join(module_path, "fused_bias_act_kernel.cu"), 23 | ], 24 | ) 25 | 26 | 27 | class FusedLeakyReLUFunctionBackward(Function): 28 | @staticmethod 29 | def forward(ctx, grad_output, out, negative_slope, scale): 30 | ctx.save_for_backward(out) 31 | ctx.negative_slope = negative_slope 32 | ctx.scale = scale 33 | 34 | empty = grad_output.new_empty(0) 35 | 36 | grad_input = fused.fused_bias_act( 37 | grad_output, empty, out, 3, 1, negative_slope, scale 38 | ) 39 | 40 | dim = [0] 41 | 42 | if grad_input.ndim > 2: 43 | dim += list(range(2, grad_input.ndim)) 44 | 45 | grad_bias = grad_input.sum(dim).detach() 46 | 47 | return grad_input, grad_bias 48 | 49 | @staticmethod 50 | def backward(ctx, gradgrad_input, gradgrad_bias): 51 | out, = ctx.saved_tensors 52 | gradgrad_out = fused.fused_bias_act( 53 | gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale 54 | ) 55 | 56 | return gradgrad_out, None, None, None 57 | 58 | 59 | class FusedLeakyReLUFunction(Function): 60 | @staticmethod 61 | def forward(ctx, input, bias, negative_slope, scale): 62 | empty = input.new_empty(0) 63 | out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale) 64 | ctx.save_for_backward(out) 65 | ctx.negative_slope = negative_slope 66 | ctx.scale = scale 67 | 68 | return out 69 | 70 | @staticmethod 71 | def backward(ctx, grad_output): 72 | out, = ctx.saved_tensors 73 | 74 | grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply( 75 | grad_output, out, ctx.negative_slope, ctx.scale 76 | ) 77 | 78 | return grad_input, grad_bias, None, None 79 | 80 | 81 | class FusedLeakyReLU(nn.Module): 82 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): 83 | super().__init__() 84 | 85 | self.bias = nn.Parameter(torch.zeros(channel)) 86 | self.negative_slope = negative_slope 87 | self.scale = scale 88 | 89 | def forward(self, input): 90 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) 91 | 92 | 93 | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): 94 | if input.device.type == "cpu": 95 | rest_dim = [1] * (input.ndim - bias.ndim - 1) 96 | return ( 97 | F.leaky_relu( 98 | input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2 99 | ) 100 | * scale 101 | ) 102 | 103 | else: 104 | return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale) 105 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/fused_bias_act.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | #include 9 | 10 | 11 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 12 | int act, int grad, float alpha, float scale); 13 | 14 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 15 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 16 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 17 | 18 | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 19 | int act, int grad, float alpha, float scale) { 20 | CHECK_CUDA(input); 21 | CHECK_CUDA(bias); 22 | 23 | return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale); 24 | } 25 | 26 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 27 | m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)"); 28 | } -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/fused_bias_act_kernel.cu: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2019-2022, NVIDIA Corporation. All rights reserved. 3 | // --------------------------------------------------------------- 4 | // 5 | // This work is made available under the Nvidia Source Code License-NC. 6 | // To view a copy of this license, visit 7 | // https://nvlabs.github.io/stylegan2/license.html 8 | 9 | #include 10 | 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | #include 17 | #include 18 | 19 | 20 | template 21 | static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref, 22 | int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) { 23 | int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x; 24 | 25 | scalar_t zero = 0.0; 26 | 27 | for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) { 28 | scalar_t x = p_x[xi]; 29 | 30 | if (use_bias) { 31 | x += p_b[(xi / step_b) % size_b]; 32 | } 33 | 34 | scalar_t ref = use_ref ? p_ref[xi] : zero; 35 | 36 | scalar_t y; 37 | 38 | switch (act * 10 + grad) { 39 | default: 40 | case 10: y = x; break; 41 | case 11: y = x; break; 42 | case 12: y = 0.0; break; 43 | 44 | case 30: y = (x > 0.0) ? x : x * alpha; break; 45 | case 31: y = (ref > 0.0) ? x : x * alpha; break; 46 | case 32: y = 0.0; break; 47 | } 48 | 49 | out[xi] = y * scale; 50 | } 51 | } 52 | 53 | 54 | torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer, 55 | int act, int grad, float alpha, float scale) { 56 | int curDevice = -1; 57 | cudaGetDevice(&curDevice); 58 | cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice); 59 | 60 | auto x = input.contiguous(); 61 | auto b = bias.contiguous(); 62 | auto ref = refer.contiguous(); 63 | 64 | int use_bias = b.numel() ? 1 : 0; 65 | int use_ref = ref.numel() ? 1 : 0; 66 | 67 | int size_x = x.numel(); 68 | int size_b = b.numel(); 69 | int step_b = 1; 70 | 71 | for (int i = 1 + 1; i < x.dim(); i++) { 72 | step_b *= x.size(i); 73 | } 74 | 75 | int loop_x = 4; 76 | int block_size = 4 * 32; 77 | int grid_size = (size_x - 1) / (loop_x * block_size) + 1; 78 | 79 | auto y = torch::empty_like(x); 80 | 81 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] { 82 | fused_bias_act_kernel<<>>( 83 | y.data_ptr(), 84 | x.data_ptr(), 85 | b.data_ptr(), 86 | ref.data_ptr(), 87 | act, 88 | grad, 89 | alpha, 90 | scale, 91 | loop_x, 92 | size_x, 93 | step_b, 94 | size_b, 95 | use_bias, 96 | use_ref 97 | ); 98 | }); 99 | 100 | return y; 101 | } -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/libcudart.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/htyjers/StrDiffusion/7f3ba4ad55add146741adea609759e7b12567f46/train/texture/config/inpainting/models/op/libcudart.so -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/op/upfirdn2d.cpp: -------------------------------------------------------------------------------- 1 | // --------------------------------------------------------------- 2 | // Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. 3 | // --------------------------------------------------------------- 4 | 5 | // Originated from https://github.com/rosinality/stylegan2-pytorch 6 | // The license for the original version of this file can be found in this directory (LICENSE_MIT). 7 | 8 | 9 | #include 10 | 11 | 12 | torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel, 13 | int up_x, int up_y, int down_x, int down_y, 14 | int pad_x0, int pad_x1, int pad_y0, int pad_y1); 15 | 16 | #define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 17 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") 18 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 19 | 20 | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel, 21 | int up_x, int up_y, int down_x, int down_y, 22 | int pad_x0, int pad_x1, int pad_y0, int pad_y1) { 23 | CHECK_CUDA(input); 24 | CHECK_CUDA(kernel); 25 | 26 | return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1); 27 | } 28 | 29 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 30 | m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)"); 31 | } -------------------------------------------------------------------------------- /train/texture/config/inpainting/models/optimizer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2023 Google Research. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """PyTorch implementation of the Lion optimizer.""" 16 | 17 | import torch 18 | from torch.optim.optimizer import Optimizer 19 | 20 | 21 | class Lion(Optimizer): 22 | r"""Implements Lion algorithm.""" 23 | 24 | def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0): 25 | """Initialize the hyperparameters. 26 | 27 | Args: 28 | params (iterable): iterable of parameters to optimize or dicts defining 29 | parameter groups 30 | lr (float, optional): learning rate (default: 1e-4) 31 | betas (Tuple[float, float], optional): coefficients used for computing 32 | running averages of gradient and its square (default: (0.9, 0.99)) 33 | weight_decay (float, optional): weight decay coefficient (default: 0) 34 | """ 35 | 36 | if not 0.0 <= lr: 37 | raise ValueError('Invalid learning rate: {}'.format(lr)) 38 | if not 0.0 <= betas[0] < 1.0: 39 | raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0])) 40 | if not 0.0 <= betas[1] < 1.0: 41 | raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1])) 42 | defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay) 43 | super().__init__(params, defaults) 44 | 45 | @torch.no_grad() 46 | def step(self, closure=None): 47 | """Performs a single optimization step. 48 | 49 | Args: 50 | closure (callable, optional): A closure that reevaluates the model 51 | and returns the loss. 52 | 53 | Returns: 54 | the loss. 55 | """ 56 | loss = None 57 | if closure is not None: 58 | with torch.enable_grad(): 59 | loss = closure() 60 | 61 | for group in self.param_groups: 62 | for p in group['params']: 63 | if p.grad is None: 64 | continue 65 | 66 | # Perform stepweight decay 67 | p.data.mul_(1 - group['lr'] * group['weight_decay']) 68 | 69 | grad = p.grad 70 | state = self.state[p] 71 | # State initialization 72 | if len(state) == 0: 73 | # Exponential moving average of gradient values 74 | state['exp_avg'] = torch.zeros_like(p) 75 | 76 | exp_avg = state['exp_avg'] 77 | beta1, beta2 = group['betas'] 78 | 79 | # Weight update 80 | update = exp_avg * beta1 + grad * (1 - beta1) 81 | p.add_(torch.sign(update), alpha=-group['lr']) 82 | # Decay the momentum running average coefficient 83 | exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2) 84 | 85 | return loss 86 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/options/train/ir-sde.yml: -------------------------------------------------------------------------------- 1 | #### general settings 2 | name: ir-sde 3 | use_tb_logger: true 4 | model: denoising 5 | distortion: derain 6 | gpu_ids: [2] 7 | 8 | sde: 9 | max_sigma: 30 10 | T: 400 11 | schedule: cosine 12 | eps: 0.005 13 | 14 | degradation: 15 | mask_root: your mask path 16 | 17 | #### datasets 18 | datasets: 19 | train: 20 | name: Train_Dataset 21 | mode: GT 22 | dataroot_GT: your image path 23 | 24 | use_shuffle: true 25 | n_workers: 4 26 | batch_size: 16 27 | GT_size: 256 28 | use_flip: true 29 | use_rot: true 30 | color: RGB 31 | 32 | 33 | #### network structures 34 | network_G: 35 | which_model_G: ConditionalUNet 36 | setting: 37 | in_nc: 3 38 | out_nc: 3 39 | nf: 64 40 | depth: 4 41 | 42 | #### path 43 | path: 44 | pretrain_model_G: ~ 45 | strict_load: true 46 | resume_state: ~ 47 | 48 | #### training settings: learning rate scheme, loss 49 | train: 50 | optimizer: Adam 51 | lr_G: !!float 1e-4 52 | lr_scheme: MultiStepLR 53 | beta1: 0.9 54 | beta2: 0.99 55 | niter: 8000000 56 | warmup_iter: -1 # no warm up 57 | lr_steps: [2000000, 4000000, 6000000] 58 | lr_gamma: 0.5 59 | eta_min: !!float 1e-7 60 | 61 | # criterion 62 | is_weighted: False 63 | loss_type: l1 64 | weight: 1.0 65 | 66 | manual_seed: 0 67 | 68 | #### logger 69 | logger: 70 | print_freq: 100 71 | save_checkpoint_freq: !!float 5e3 72 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/str_utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/texture/config/inpainting/str_utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import abc 4 | import math 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | import torchvision.utils as tvutils 10 | 11 | 12 | ########### denoising ############### 13 | def add_noise(tensor, sigma): 14 | sigma = sigma / 255 if sigma > 1 else sigma 15 | return tensor + torch.randn_like(tensor) * sigma 16 | 17 | 18 | ######## inpainting ########### 19 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 20 | batch = tensor.shape[0] 21 | if mask_id < 0: 22 | mask_id = np.random.randint(0, n, batch) 23 | masks = [] 24 | for i in range(batch): 25 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 26 | mask = np.concatenate(masks, axis=0) 27 | else: 28 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 29 | 30 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 31 | # for images are clipped or scaled 32 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 33 | masked_tensor = mask * tensor 34 | return masked_tensor + (1. - mask) 35 | 36 | ######## super-resolution ########### 37 | 38 | def upscale(tensor, scale=4, mode='bicubic'): 39 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 40 | return tensor 41 | 42 | 43 | 44 | 45 | -------------------------------------------------------------------------------- /train/texture/config/inpainting/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .deg_utils import * 2 | from .file_utils import * 3 | from .img_utils import * 4 | from .sde_utils import * -------------------------------------------------------------------------------- /train/texture/config/inpainting/utils/deg_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import abc 3 | import math 4 | import numpy as np 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torchvision.utils as tvutils 9 | import cv2 10 | 11 | ########### denoising ############### 12 | def add_noise(tensor, sigma): 13 | sigma = sigma / 255 if sigma > 1 else sigma 14 | return tensor + torch.randn_like(tensor) * sigma 15 | 16 | 17 | ######## inpainting ########### 18 | def mask_to(tensor, mask_root='data/datasets/gt_keep_masks/genhalf', mask_id=-1, n=100): 19 | batch = tensor.shape[0] 20 | if mask_id < 0: 21 | mask_id = np.random.randint(0, n, batch) 22 | masks = [] 23 | for i in range(batch): 24 | masks.append(cv2.imread(os.path.join(mask_root, f'{mask_id[i]:06d}.png'))[None, ...] / 255.) 25 | mask = np.concatenate(masks, axis=0) 26 | else: 27 | mask = cv2.imread(os.path.join(mask_root, f'{mask_id:06d}.png'))[None, ...] / 255. 28 | 29 | mask = torch.tensor(mask).permute(0, 3, 1, 2).float() 30 | # for images are clipped or scaled 31 | mask = F.interpolate(mask, size=tensor.shape[2:], mode='nearest') 32 | masked_tensor = mask * tensor 33 | return masked_tensor + (1. - mask) 34 | 35 | ######## super-resolution ########### 36 | 37 | def upscale(tensor, scale=4, mode='bicubic'): 38 | tensor = F.interpolate(tensor, scale_factor=scale, mode=mode) 39 | return tensor 40 | 41 | 42 | 43 | 44 | --------------------------------------------------------------------------------