update

=2.0

@@ -0,0 +1,26 @@
+Requirement already satisfied: torch in /usr/local/lib/python3.10/site-packages (2.6.0)
+Requirement already satisfied: torchvision in /usr/local/lib/python3.10/site-packages (0.21.0)
+Requirement already satisfied: filelock in /usr/local/lib/python3.10/site-packages (from torch) (3.17.0)
+Requirement already satisfied: typing-extensions>=4.10.0 in /usr/local/lib/python3.10/site-packages (from torch) (4.12.2)
+Requirement already satisfied: networkx in /usr/local/lib/python3.10/site-packages (from torch) (3.4.2)
+Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/site-packages (from torch) (3.1.5)
+Requirement already satisfied: fsspec in /usr/local/lib/python3.10/site-packages (from torch) (2024.12.0)
+Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.127)
+Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.127)
+Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.127)
+Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.10/site-packages (from torch) (9.1.0.70)
+Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.5.8)
+Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /usr/local/lib/python3.10/site-packages (from torch) (11.2.1.3)
+Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /usr/local/lib/python3.10/site-packages (from torch) (10.3.5.147)
+Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /usr/local/lib/python3.10/site-packages (from torch) (11.6.1.9)
+Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /usr/local/lib/python3.10/site-packages (from torch) (12.3.1.170)
+Requirement already satisfied: nvidia-cusparselt-cu12==0.6.2 in /usr/local/lib/python3.10/site-packages (from torch) (0.6.2)
+Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.10/site-packages (from torch) (2.21.5)
+Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.127)
+Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /usr/local/lib/python3.10/site-packages (from torch) (12.4.127)
+Requirement already satisfied: triton==3.2.0 in /usr/local/lib/python3.10/site-packages (from torch) (3.2.0)
+Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/site-packages (from torch) (1.13.1)
+Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/site-packages (from sympy==1.13.1->torch) (1.3.0)
+Requirement already satisfied: numpy in /usr/local/lib/python3.10/site-packages (from torchvision) (2.2.3)
+Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/site-packages (from torchvision) (10.4.0)
+Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/site-packages (from jinja2->torch) (2.1.5)

README.md

@@ -1,5 +1,5 @@
 ---
-title: Lime Evangelion
+title: RDNet
 emoji: 💻
 colorFrom: indigo
 colorTo: blue

data/dataset_sir.py

@@ -1,332 +0,0 @@
-import math
-import os.path
-import os.path
-import random
-from os.path import join
-
-import cv2
-import numpy as np
-import torch.utils.data
-import torchvision.transforms.functional as TF
-from PIL import Image
-from scipy.signal import convolve2d
-
-from data.image_folder import make_dataset
-from data.torchdata import Dataset as BaseDataset
-from data.transforms import to_tensor
-
-
-def __scale_width(img, target_width):
-    ow, oh = img.size
-    if (ow == target_width):
-        return img
-    w = target_width
-    h = int(target_width * oh / ow)
-    h = math.ceil(h / 2.) * 2  # round up to even
-    return img.resize((w, h), Image.BICUBIC)
-
-
-def __scale_height(img, target_height):
-    ow, oh = img.size
-    if (oh == target_height):
-        return img
-    h = target_height
-    w = int(target_height * ow / oh)
-    w = math.ceil(w / 2.) * 2
-    return img.resize((w, h), Image.BICUBIC)
-
-
-def paired_data_transforms(img_1, img_2, unaligned_transforms=False):
-    def get_params(img, output_size):
-        w, h = img.size
-        th, tw = output_size
-        if w == tw and h == th:
-            return 0, 0, h, w
-
-        i = random.randint(0, h - th)
-        j = random.randint(0, w - tw)
-        return i, j, th, tw
-
-    target_size = int(random.randint(320, 640) / 2.) * 2
-    ow, oh = img_1.size
-    if ow >= oh:
-        img_1 = __scale_height(img_1, target_size)
-        img_2 = __scale_height(img_2, target_size)
-    else:
-        img_1 = __scale_width(img_1, target_size)
-        img_2 = __scale_width(img_2, target_size)
-
-    if random.random() < 0.5:
-        img_1 = TF.hflip(img_1)
-        img_2 = TF.hflip(img_2)
-
-    if random.random() < 0.5:
-        angle = random.choice([90, 180, 270])
-        img_1 = TF.rotate(img_1, angle)
-        img_2 = TF.rotate(img_2, angle)
-
-    i, j, h, w = get_params(img_1, (320, 320))
-    img_1 = TF.crop(img_1, i, j, h, w)
-
-    if unaligned_transforms:
-        # print('random shift')
-        i_shift = random.randint(-10, 10)
-        j_shift = random.randint(-10, 10)
-        i += i_shift
-        j += j_shift
-
-    img_2 = TF.crop(img_2, i, j, h, w)
-
-    return img_1, img_2
-
-
-class ReflectionSynthesis(object):
-    def __init__(self):
-        # Kernel Size of the Gaussian Blurry
-        self.kernel_sizes = [5, 7, 9, 11]
-        self.kernel_probs = [0.1, 0.2, 0.3, 0.4]
-
-        # Sigma of the Gaussian Blurry
-        self.sigma_range = [2, 5]
-        self.alpha_range = [0.8, 1.0]
-        self.beta_range = [0.4, 1.0]
-
-    def __call__(self, T_, R_):
-        T_ = np.asarray(T_, np.float32) / 255.
-        R_ = np.asarray(R_, np.float32) / 255.
-
-        kernel_size = np.random.choice(self.kernel_sizes, p=self.kernel_probs)
-        sigma = np.random.uniform(self.sigma_range[0], self.sigma_range[1])
-        kernel = cv2.getGaussianKernel(kernel_size, sigma)
-        kernel2d = np.dot(kernel, kernel.T)
-        for i in range(3):
-            R_[..., i] = convolve2d(R_[..., i], kernel2d, mode='same')
-
-        a = np.random.uniform(self.alpha_range[0], self.alpha_range[1])
-        b = np.random.uniform(self.beta_range[0], self.beta_range[1])
-        T, R = a * T_, b * R_
-
-        if random.random() < 0.7:
-            I = T + R - T * R
-            
-        else:
-            I = T + R
-            if np.max(I) > 1:
-                m = I[I > 1]
-                m = (np.mean(m) - 1) * 1.3
-                I = np.clip(T + np.clip(R - m, 0, 1), 0, 1)
-
-        return T_, R_, I
-
-
-class DataLoader(torch.utils.data.DataLoader):
-    def __init__(self, dataset, batch_size, shuffle, *args, **kwargs):
-        super(DataLoader, self).__init__(dataset, batch_size, shuffle, *args, **kwargs)
-        self.shuffle = shuffle
-
-    def reset(self):
-        if self.shuffle:
-            print('Reset Dataset...')
-            self.dataset.reset()
-
-
-class DSRDataset(BaseDataset):
-    def __init__(self, datadir, fns=None, size=None, enable_transforms=True):
-        super(DSRDataset, self).__init__()
-        self.size = size
-        self.datadir = datadir
-        self.enable_transforms = enable_transforms
-        sortkey = lambda key: os.path.split(key)[-1]
-        self.paths = sorted(make_dataset(datadir, fns), key=sortkey)
-        if size is not None:
-            self.paths = np.random.choice(self.paths, size)
-
-        self.syn_model = ReflectionSynthesis()
-        self.reset(shuffle=False)
-
-    def reset(self, shuffle=True):
-        if shuffle:
-            random.shuffle(self.paths)
-        num_paths = len(self.paths) // 2
-        self.B_paths = self.paths[0:num_paths]
-        self.R_paths = self.paths[num_paths:2 * num_paths]
-
-    def data_synthesis(self, t_img, r_img):
-        if self.enable_transforms:
-            t_img, r_img = paired_data_transforms(t_img, r_img)
-
-        t_img, r_img, m_img = self.syn_model(t_img, r_img)
-
-        B = TF.to_tensor(t_img)
-        R = TF.to_tensor(r_img)
-        M = TF.to_tensor(m_img)
-
-        return B, R, M
-
-    def __getitem__(self, index):
-        index_B = index % len(self.B_paths)
-        index_R = index % len(self.R_paths)
-
-        B_path = self.B_paths[index_B]
-        R_path = self.R_paths[index_R]
-
-        t_img = Image.open(B_path).convert('RGB')
-        r_img = Image.open(R_path).convert('RGB')
-
-        B, R, M = self.data_synthesis(t_img, r_img)
-        fn = os.path.basename(B_path)
-        return {'input': M, 'target_t': B, 'target_r': M-B, 'fn': fn, 'real': False}
-
-    def __len__(self):
-        if self.size is not None:
-            return min(max(len(self.B_paths), len(self.R_paths)), self.size)
-        else:
-            return max(len(self.B_paths), len(self.R_paths))
-
-
-class DSRTestDataset(BaseDataset):
-    def __init__(self, datadir, fns=None, size=None, enable_transforms=False, unaligned_transforms=False,
-                 round_factor=1, flag=None, if_align=True):
-        super(DSRTestDataset, self).__init__()
-        self.size = size
-        self.datadir = datadir
-        self.fns = fns or os.listdir(join(datadir, 'blended'))
-        self.enable_transforms = enable_transforms
-        self.unaligned_transforms = unaligned_transforms
-        self.round_factor = round_factor
-        self.flag = flag
-        self.if_align = True # if_align
-
-        if size is not None:
-            self.fns = self.fns[:size]
-
-    def align(self, x1, x2):
-        h, w = x1.height, x1.width
-        h, w = h // 32 * 32, w // 32 * 32
-        x1 = x1.resize((w, h))
-        x2 = x2.resize((w, h))
-        return x1, x2
-
-    def __getitem__(self, index):
-        fn = self.fns[index]
-
-        t_img = Image.open(join(self.datadir, 'transmission_layer', fn)).convert('RGB')
-        m_img = Image.open(join(self.datadir, 'blended', fn)).convert('RGB')
-
-        if self.if_align:
-            t_img, m_img = self.align(t_img, m_img)
-
-        if self.enable_transforms:
-            t_img, m_img = paired_data_transforms(t_img, m_img, self.unaligned_transforms)
-
-        B = TF.to_tensor(t_img)
-        M = TF.to_tensor(m_img)
-
-        dic = {'input': M, 'target_t': B, 'fn': fn, 'real': True, 'target_r': M - B}
-        if self.flag is not None:
-            dic.update(self.flag)
-        return dic
-
-    def __len__(self):
-        if self.size is not None:
-            return min(len(self.fns), self.size)
-        else:
-            return len(self.fns)
-
-
-class SIRTestDataset(BaseDataset):
-    def __init__(self, datadir, fns=None, size=None, if_align=True):
-        super(SIRTestDataset, self).__init__()
-        self.size = size
-        self.datadir = datadir
-        self.fns = fns or os.listdir(join(datadir, 'blended'))
-        self.if_align = if_align
-
-        if size is not None:
-            self.fns = self.fns[:size]
-
-    def align(self, x1, x2, x3):
-        h, w = x1.height, x1.width
-        h, w = h // 32 * 32, w // 32 * 32
-        x1 = x1.resize((w, h))
-        x2 = x2.resize((w, h))
-        x3 = x3.resize((w, h))
-        return x1, x2, x3
-
-    def __getitem__(self, index):
-        fn = self.fns[index]
-
-        t_img = Image.open(join(self.datadir, 'transmission_layer', fn)).convert('RGB')
-        r_img = Image.open(join(self.datadir, 'reflection_layer', fn)).convert('RGB')
-        m_img = Image.open(join(self.datadir, 'blended', fn)).convert('RGB')
-
-        if self.if_align:
-            t_img, r_img, m_img = self.align(t_img, r_img, m_img)
-
-        B = TF.to_tensor(t_img)
-        R = TF.to_tensor(r_img)
-        M = TF.to_tensor(m_img)
-
-        dic = {'input': M, 'target_t': B, 'fn': fn, 'real': True, 'target_r': R, 'target_r_hat': M - B}
-        return dic
-
-    def __len__(self):
-        if self.size is not None:
-            return min(len(self.fns), self.size)
-        else:
-            return len(self.fns)
-
-
-class RealDataset(BaseDataset):
-    def __init__(self, datadir, fns=None, size=None):
-        super(RealDataset, self).__init__()
-        self.size = size
-        self.datadir = datadir
-        self.fns = fns or os.listdir(join(datadir))
-
-        if size is not None:
-            self.fns = self.fns[:size]
-
-    def align(self, x):
-        h, w = x.height, x.width
-        h, w = h // 32 * 32, w // 32 * 32
-        x = x.resize((w, h))
-        return x
-
-    def __getitem__(self, index):
-        fn = self.fns[index]
-        B = -1
-        m_img = Image.open(join(self.datadir, fn)).convert('RGB')
-        M = to_tensor(self.align(m_img))
-        data = {'input': M, 'target_t': B, 'fn': fn}
-        return data
-
-    def __len__(self):
-        if self.size is not None:
-            return min(len(self.fns), self.size)
-        else:
-            return len(self.fns)
-
-
-class FusionDataset(BaseDataset):
-    def __init__(self, datasets, fusion_ratios=None):
-        self.datasets = datasets
-        self.size = sum([len(dataset) for dataset in datasets])
-        self.fusion_ratios = fusion_ratios or [1. / len(datasets)] * len(datasets)
-        print('[i] using a fusion dataset: %d %s imgs fused with ratio %s' % (
-            self.size, [len(dataset) for dataset in datasets], self.fusion_ratios))
-
-    def reset(self):
-        for dataset in self.datasets:
-            dataset.reset()
-
-    def __getitem__(self, index):
-        residual = 1
-        for i, ratio in enumerate(self.fusion_ratios):
-            if random.random() < ratio / residual or i == len(self.fusion_ratios) - 1:
-                dataset = self.datasets[i]
-                return dataset[index % len(dataset)]
-            residual -= ratio
-
-    def __len__(self):
-        return self.size

data/image_folder.py

@@ -1,51 +0,0 @@
-###############################################################################
-# Code from
-# https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
-# Modified the original code so that it also loads images from the current
-# directory as well as the subdirectories
-###############################################################################
-
-import torch.utils.data as data
-
-from PIL import Image
-import os
-import os.path
-
-IMG_EXTENSIONS = [
-    '.jpg', '.JPG', '.jpeg', '.JPEG',
-    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
-]
-
-
-def read_fns(filename):
-    with open(filename) as f:
-        fns = f.readlines()
-        fns = [fn.strip() for fn in fns]
-    return fns
-
-
-def is_image_file(filename):
-    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
-
-
-def make_dataset(dir, fns=None):
-    images = []
-    assert os.path.isdir(dir), '%s is not a valid directory' % dir
-
-    if fns is None:
-        for root, _, fnames in sorted(os.walk(dir)):
-            for fname in fnames:
-                if is_image_file(fname):                
-                    path = os.path.join(root, fname)
-                    images.append(path)
-    else:
-        for fname in fns:
-            if is_image_file(fname):
-                path = os.path.join(dir, fname)
-                images.append(path)
-
-    return images
-
-
-def default_loader(path):
-    return Image.open(path).convert('RGB')

data/real_test.txt

@@ -1,20 +0,0 @@
-3.jpg
-4.jpg
-9.jpg
-12.jpg
-15.jpg
-22.jpg
-23.jpg
-25.jpg
-29.jpg
-39.jpg
-46.jpg
-47.jpg
-58.jpg
-86.jpg
-87.jpg
-89.jpg
-93.jpg
-103.jpg
-107.jpg
-110.jpg

data/torchdata.py

@@ -1,67 +0,0 @@
-import bisect
-import warnings
-
-
-class Dataset(object):
-    """An abstract class representing a Dataset.
-
-    All other datasets should subclass it. All subclasses should override
-    ``__len__``, that provides the size of the dataset, and ``__getitem__``,
-    supporting integer indexing in range from 0 to len(self) exclusive.
-    """
-
-    def __getitem__(self, index):
-        raise NotImplementedError
-
-    def __len__(self):
-        raise NotImplementedError
-
-    def __add__(self, other):
-        return ConcatDataset([self, other])
-
-    def reset(self):
-        return 
-
-
-class ConcatDataset(Dataset):
-    """
-    Dataset to concatenate multiple datasets.
-    Purpose: useful to assemble different existing datasets, possibly
-    large-scale datasets as the concatenation operation is done in an
-    on-the-fly manner.
-
-    Arguments:
-        datasets (sequence): List of datasets to be concatenated
-    """
-
-    @staticmethod
-    def cumsum(sequence):
-        r, s = [], 0
-        for e in sequence:
-            l = len(e)
-            r.append(l + s)
-            s += l
-        return r
-
-    def __init__(self, datasets):
-        super(ConcatDataset, self).__init__()
-        assert len(datasets) > 0, 'datasets should not be an empty iterable'
-        self.datasets = list(datasets)
-        self.cumulative_sizes = self.cumsum(self.datasets)
-
-    def __len__(self):
-        return self.cumulative_sizes[-1]
-
-    def __getitem__(self, idx):
-        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
-        if dataset_idx == 0:
-            sample_idx = idx
-        else:
-            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
-        return self.datasets[dataset_idx][sample_idx]
-
-    @property
-    def cummulative_sizes(self):
-        warnings.warn("cummulative_sizes attribute is renamed to "
-                      "cumulative_sizes", DeprecationWarning, stacklevel=2)
-        return self.cumulative_sizes

data/transforms.py

@@ -1,301 +0,0 @@
-from __future__ import division
-
-import math
-import random
-
-import torch
-from PIL import Image
-
-try:
-    import accimage
-except ImportError:
-    accimage = None
-import numpy as np
-import scipy.stats as st
-import cv2
-import collections
-import torchvision.transforms as transforms
-import util.util as util
-from scipy.signal import convolve2d
-
-
-# utility
-def _is_pil_image(img):
-    if accimage is not None:
-        return isinstance(img, (Image.Image, accimage.Image))
-    else:
-        return isinstance(img, Image.Image)
-
-
-def _is_tensor_image(img):
-    return torch.is_tensor(img) and img.ndimension() == 3
-
-
-def _is_numpy_image(img):
-    return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
-
-
-def arrshow(arr):
-    Image.fromarray(arr.astype(np.uint8)).show()
-
-
-def get_transform(opt):
-    transform_list = []
-    osizes = util.parse_args(opt.loadSize)
-    fineSize = util.parse_args(opt.fineSize)
-    if opt.resize_or_crop == 'resize_and_crop':
-        transform_list.append(
-            transforms.RandomChoice([
-                transforms.Resize([osize, osize], Image.BICUBIC) for osize in osizes
-            ]))
-        transform_list.append(transforms.RandomCrop(fineSize))
-    elif opt.resize_or_crop == 'crop':
-        transform_list.append(transforms.RandomCrop(fineSize))
-    elif opt.resize_or_crop == 'scale_width':
-        transform_list.append(transforms.Lambda(
-            lambda img: __scale_width(img, fineSize)))
-    elif opt.resize_or_crop == 'scale_width_and_crop':
-        transform_list.append(transforms.Lambda(
-            lambda img: __scale_width(img, opt.loadSize)))
-        transform_list.append(transforms.RandomCrop(opt.fineSize))
-
-    if opt.isTrain and not opt.no_flip:
-        transform_list.append(transforms.RandomHorizontalFlip())
-
-    return transforms.Compose(transform_list)
-
-
-to_norm_tensor = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Normalize(
-        (0.5, 0.5, 0.5),
-        (0.5, 0.5, 0.5)
-    )
-])
-
-to_tensor = transforms.ToTensor()
-
-
-def __scale_width(img, target_width):
-    ow, oh = img.size
-    if (ow == target_width):
-        return img
-    w = target_width
-    h = int(target_width * oh / ow)
-    h = math.ceil(h / 2.) * 2  # round up to even
-    return img.resize((w, h), Image.BICUBIC)
-
-
-# functional 
-def gaussian_blur(img, kernel_size, sigma):
-    if not _is_pil_image(img):
-        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
-
-    img = np.asarray(img)
-    # the 3rd dimension (i.e. inter-band) would be filtered which is unwanted for our purpose
-    # new = gaussian_filter(img, sigma=sigma, truncate=truncate)
-    if isinstance(kernel_size, int):
-        kernel_size = (kernel_size, kernel_size)
-    elif isinstance(kernel_size, collections.Sequence):
-        assert len(kernel_size) == 2
-    new = cv2.GaussianBlur(img, kernel_size, sigma)  # apply gaussian filter band by band    
-    return Image.fromarray(new)
-
-
-# transforms
-class GaussianBlur(object):
-    def __init__(self, kernel_size=11, sigma=3):
-        self.kernel_size = kernel_size
-        self.sigma = sigma
-
-    def __call__(self, img):
-        return gaussian_blur(img, self.kernel_size, self.sigma)
-
-
-class ReflectionSythesis_0(object):
-    """Reflection image data synthesis for weakly-supervised learning
-    of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*
-    """
-
-    def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3,
-                 high_gamma=1.3, low_delta=0.4, high_delta=1.8):
-        self.kernel_sizes = kernel_sizes or [11]
-        self.low_sigma = low_sigma
-        self.high_sigma = high_sigma
-        self.low_gamma = low_gamma
-        self.high_gamma = high_gamma
-        self.low_delta = low_delta
-        self.high_delta = high_delta
-        print('[i] reflection sythesis model: {}'.format({
-            'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
-            'low_gamma': low_gamma, 'high_gamma': high_gamma}))
-
-    def __call__(self, B, R):
-        if not _is_pil_image(B):
-            raise TypeError('B should be PIL Image. Got {}'.format(type(B)))
-        if not _is_pil_image(R):
-            raise TypeError('R should be PIL Image. Got {}'.format(type(R)))
-        B_ = np.asarray(B, np.float32)
-        if random.random() < 0.4:
-            B_ = np.tile(np.random.uniform(0, 30, (1, 1, 1)), B_.shape) / 255.
-        else:
-            B_ = np.tile(np.random.normal(50, 50, (1, 1, 3)), (B_.shape[0], B_.shape[1], 1)).clip(0, 255) / 255.
-        R_ = np.asarray(R, np.float32) / 255.
-
-        kernel_size = np.random.choice(self.kernel_sizes)
-        sigma = np.random.uniform(self.low_sigma, self.high_sigma)
-        gamma = np.random.uniform(self.low_gamma, self.high_gamma)
-        delta = np.random.uniform(self.low_delta, self.high_delta)
-        R_blur = R_
-        kernel = cv2.getGaussianKernel(11, sigma)
-        kernel2d = np.dot(kernel, kernel.T)
-
-        for i in range(3):
-            R_blur[..., i] = convolve2d(R_blur[..., i], kernel2d, mode='same')
-
-        R_blur = np.clip(R_blur - np.mean(R_blur) * gamma, 0, 1)
-        R_blur = np.clip(R_blur * delta, 0, 1)
-        M_ = np.clip(R_blur + B_, 0, 1)
-
-        return B_, R_blur, M_
-
-
-class ReflectionSythesis_1(object):
-    """Reflection image data synthesis for weakly-supervised learning 
-    of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*    
-    """
-
-    def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3, high_gamma=1.3):
-        self.kernel_sizes = kernel_sizes or [11]
-        self.low_sigma = low_sigma
-        self.high_sigma = high_sigma
-        self.low_gamma = low_gamma
-        self.high_gamma = high_gamma
-        print('[i] reflection sythesis model: {}'.format({
-            'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
-            'low_gamma': low_gamma, 'high_gamma': high_gamma}))
-
-    def __call__(self, B, R):
-        if not _is_pil_image(B):
-            raise TypeError('B should be PIL Image. Got {}'.format(type(B)))
-        if not _is_pil_image(R):
-            raise TypeError('R should be PIL Image. Got {}'.format(type(R)))
-
-        B_ = np.asarray(B, np.float32) / 255.
-        R_ = np.asarray(R, np.float32) / 255.
-
-        kernel_size = np.random.choice(self.kernel_sizes)
-        sigma = np.random.uniform(self.low_sigma, self.high_sigma)
-        gamma = np.random.uniform(self.low_gamma, self.high_gamma)
-        R_blur = R_
-        kernel = cv2.getGaussianKernel(11, sigma)
-        kernel2d = np.dot(kernel, kernel.T)
-
-        for i in range(3):
-            R_blur[..., i] = convolve2d(R_blur[..., i], kernel2d, mode='same')
-
-        M_ = B_ + R_blur
-
-        if np.max(M_) > 1:
-            m = M_[M_ > 1]
-            m = (np.mean(m) - 1) * gamma
-            R_blur = np.clip(R_blur - m, 0, 1)
-            M_ = np.clip(R_blur + B_, 0, 1)
-
-        return B_, R_blur, M_
-
-
-class Sobel(object):
-    def __call__(self, img):
-        if not _is_pil_image(img):
-            raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
-
-        gray_img = np.array(img.convert('L'))
-        x = cv2.Sobel(gray_img, cv2.CV_16S, 1, 0)
-        y = cv2.Sobel(gray_img, cv2.CV_16S, 0, 1)
-
-        absX = cv2.convertScaleAbs(x)
-        absY = cv2.convertScaleAbs(y)
-
-        dst = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
-        return Image.fromarray(dst)
-
-
-class ReflectionSythesis_2(object):
-    """Reflection image data synthesis for weakly-supervised learning 
-    of CVPR 2018 paper *"Single Image Reflection Separation with Perceptual Losses"*
-    """
-
-    def __init__(self, kernel_sizes=None):
-        self.kernel_sizes = kernel_sizes or np.linspace(1, 5, 80)
-
-    @staticmethod
-    def gkern(kernlen=100, nsig=1):
-        """Returns a 2D Gaussian kernel array."""
-        interval = (2 * nsig + 1.) / (kernlen)
-        x = np.linspace(-nsig - interval / 2., nsig + interval / 2., kernlen + 1)
-        kern1d = np.diff(st.norm.cdf(x))
-        kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
-        kernel = kernel_raw / kernel_raw.sum()
-        kernel = kernel / kernel.max()
-        return kernel
-
-    def __call__(self, t, r):
-        t = np.float32(t) / 255.
-        r = np.float32(r) / 255.
-        ori_t = t
-        # create a vignetting mask
-        g_mask = self.gkern(560, 3)
-        g_mask = np.dstack((g_mask, g_mask, g_mask))
-        sigma = self.kernel_sizes[np.random.randint(0, len(self.kernel_sizes))]
-
-        t = np.power(t, 2.2)
-        r = np.power(r, 2.2)
-
-        sz = int(2 * np.ceil(2 * sigma) + 1)
-
-        r_blur = cv2.GaussianBlur(r, (sz, sz), sigma, sigma, 0)
-        blend = r_blur + t
-
-        att = 1.08 + np.random.random() / 10.0
-
-        for i in range(3):
-            maski = blend[:, :, i] > 1
-            mean_i = max(1., np.sum(blend[:, :, i] * maski) / (maski.sum() + 1e-6))
-            r_blur[:, :, i] = r_blur[:, :, i] - (mean_i - 1) * att
-        r_blur[r_blur >= 1] = 1
-        r_blur[r_blur <= 0] = 0
-
-        h, w = r_blur.shape[0:2]
-        neww = np.random.randint(0, 560 - w - 10)
-        newh = np.random.randint(0, 560 - h - 10)
-        alpha1 = g_mask[newh:newh + h, neww:neww + w, :]
-        alpha2 = 1 - np.random.random() / 5.0
-        r_blur_mask = np.multiply(r_blur, alpha1)
-        blend = r_blur_mask + t * alpha2
-
-        t = np.power(t, 1 / 2.2)
-        r_blur_mask = np.power(r_blur_mask, 1 / 2.2)
-        blend = np.power(blend, 1 / 2.2)
-        blend[blend >= 1] = 1
-        blend[blend <= 0] = 0
-
-        return np.float32(ori_t), np.float32(r_blur_mask), np.float32(blend)
-
-
-# Examples
-if __name__ == '__main__':
-    """cv2 imread"""
-    # img = cv2.imread('testdata_reflection_real/19-input.png')
-    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    # img2 = cv2.GaussianBlur(img, (11,11), 3)
-
-    """Sobel Operator"""
-    # img = np.array(Image.open('datasets/VOC224/train/B/2007_000250.png').convert('L'))
-
-    """Reflection Sythesis"""
-    b = Image.open('')
-    r = Image.open('')
-    G = ReflectionSythesis_0()
-    m, r = G(b, r)
-    r.show()

engine.py

@@ -1,178 +0,0 @@
-import torch
-import util.util as util
-from models import make_model
-import time
-import os
-import sys
-from os.path import join
-from util.visualizer import Visualizer
-import tqdm
-import visdom
-import numpy as np
-from tools import mutils
-
-class Engine(object):
-    def __init__(self, opt,eval_dataset_real,eval_dataset_solidobject,eval_dataset_postcard,eval_dataloader_wild):
-        self.opt = opt
-        self.writer = None
-        self.visualizer = None
-        self.model = None
-        self.best_val_loss = 1e6
-        self.eval_dataset_real = eval_dataset_real
-        self.eval_dataset_solidobject = eval_dataset_solidobject
-        self.eval_dataset_postcard = eval_dataset_postcard
-        self.eval_dataloader_wild = eval_dataloader_wild
-        self.result_dir = os.path.join(f'./experiment/{self.opt.name}/results',
-                          mutils.get_formatted_time())
-        self.biggest_psnr=0
-        self.__setup()
-
-    def __setup(self):
-        self.basedir = join('experiment', self.opt.name)
-        os.makedirs(self.basedir, exist_ok=True)
-
-        opt = self.opt
-
-        """Model"""
-        self.model = make_model(self.opt.model)  # models.__dict__[self.opt.model]()
-        self.model.initialize(opt)
-        if True:
-            print("IN")
-            self.writer = util.get_summary_writer(os.path.join(self.basedir, 'logs'))
-            self.visualizer = Visualizer(opt)
-
-    def train(self, train_loader, **kwargs):
-        print('\nEpoch: %d' % self.epoch)
-        avg_meters = util.AverageMeters()
-        opt = self.opt
-        model = self.model
-        epoch = self.epoch
-
-        epoch_start_time = time.time()
-        for i, data in tqdm.tqdm(enumerate(train_loader)):
-
-            iter_start_time = time.time()
-            iterations = self.iterations
-
-            model.set_input(data, mode='train')
-            model.optimize_parameters(**kwargs)
-
-            errors = model.get_current_errors()
-            avg_meters.update(errors)
-            util.progress_bar(i, len(train_loader), str(avg_meters))
-            util.write_loss(self.writer, 'train', avg_meters, iterations)
-            if iterations%100==0:
-                imgs=[]
-                output_clean,output_reflection,input=model.return_output()
-                # output_clean,input=model.return_output()
-                
-                output_clean=np.transpose(output_clean,(2,0,1))/255
-                #output_reflection = np.transpose(output_reflection, (2, 0, 1))/255
-                input = np.transpose(input, (2, 0, 1))/255
-                imgs.append(output_clean)
-                #imgs.append(output_reflection)
-                imgs.append(input)
-                util.get_visual(self.writer,iterations,imgs)
-                if iterations % opt.print_freq == 0 and opt.display_id != 0:
-                    t = (time.time() - iter_start_time)
-
-            self.iterations += 1
-
-        self.epoch += 1
-
-        if True:#not self.opt.no_log:
-            if self.epoch % opt.save_epoch_freq == 0:
-                save_dir = os.path.join(self.result_dir, '%03d' % self.epoch)
-                os.makedirs(save_dir, exist_ok=True)
-                matrix_real=self.eval(self.eval_dataset_real, dataset_name='testdata_real20', savedir=save_dir, suffix='real20')
-                matrix_solid=self.eval(self.eval_dataset_solidobject, dataset_name='testdata_solidobject', savedir=save_dir,
-                    suffix='solidobject')
-                matrix_post=self.eval(self.eval_dataset_postcard, dataset_name='testdata_postcard', savedir=save_dir, suffix='postcard')
-                matrix_wild=self.eval(self.eval_dataloader_wild, dataset_name='testdata_wild', savedir=save_dir, suffix='wild')
-                sum_PSNR_real=matrix_real['PSNR']*20
-                sum_PSNR_solid=matrix_solid['PSNR']*200
-                sum_PSNR_post=matrix_post['PSNR']*199
-                sum_PSNR_wild=matrix_wild['PSNR']*55
-                print("sum_PSNR_real: ",matrix_real['PSNR'],"sum_PSNR_solid: ",matrix_solid['PSNR'],"sum_PSNR_post: ",matrix_post['PSNR'],"sum_PSNR_wild: ",matrix_wild['PSNR'])
-                sum_PSNR = float(sum_PSNR_real + sum_PSNR_solid + sum_PSNR_post + sum_PSNR_wild)/474.0
-                print('总PSNR:', sum_PSNR)
-                if sum_PSNR>self.biggest_psnr:
-                    self.biggest_psnr=sum_PSNR
-                    print('saving the model at epoch %d, iters %d' %(self.epoch, self.iterations))
-                    model.save()
-                print('highest： ',self.biggest_psnr,' name: ',opt.name)
-
-            print('saving the latest model at the end of epoch %d, iters %d' %
-                  (self.epoch, self.iterations))
-            model.save(label='latest')
-
-            print('Time Taken: %d sec' %
-                  (time.time() - epoch_start_time))
-
-        # model.update_learning_rate()
-        try:
-            train_loader.reset()
-        except:
-            pass
-
-    def eval(self, val_loader, dataset_name, savedir='./tmp', loss_key=None, **kwargs):
-        # print(dataset_name)
-        if savedir is not None:
-            os.makedirs(savedir, exist_ok=True)
-            self.f = open(os.path.join(savedir, 'metrics.txt'), 'w+')
-            self.f.write(dataset_name + '\n')
-        avg_meters = util.AverageMeters()
-        model = self.model
-        opt = self.opt
-        with torch.no_grad():
-            for i, data in enumerate(val_loader):
-                if self.opt.select is not None and data['fn'][0] not in [f'{self.opt.select}.jpg']:
-                    continue
-                #print(data.shape())
-                index = model.eval(data, savedir=savedir, **kwargs)
-
-                # print(data['fn'][0], index)
-                if savedir is not None:
-                    self.f.write(f"{data['fn'][0]} {index['PSNR']} {index['SSIM']}\n")
-                avg_meters.update(index)
-                util.progress_bar(i, len(val_loader), str(avg_meters))
-
-        if not opt.no_log:
-            util.write_loss(self.writer, join('eval', dataset_name), avg_meters, self.epoch)
-
-        if loss_key is not None:
-            val_loss = avg_meters[loss_key]
-            if val_loss < self.best_val_loss:
-                self.best_val_loss = val_loss
-                print('saving the best model at the end of epoch %d, iters %d' %
-                      (self.epoch, self.iterations))
-                model.save(label='best_{}_{}'.format(loss_key, dataset_name))
-
-        return avg_meters
-
-    def test(self, test_loader, savedir=None, **kwargs):
-        model = self.model
-        opt = self.opt
-        with torch.no_grad():
-            for i, data in enumerate(test_loader):
-                model.test(data, savedir=savedir, **kwargs)
-                util.progress_bar(i, len(test_loader))
-
-    def save_eval(self, label):
-        self.model.save_eval(label)
-
-    @property
-    def iterations(self):
-        return self.model.iterations
-
-    @iterations.setter
-    def iterations(self, i):
-        self.model.iterations = i
-
-    @property
-    def epoch(self):
-        return self.model.epoch
-
-    @epoch.setter
-    def epoch(self, e):
-        self.model.epoch = e

models/__init__.py

@@ -1,11 +0,0 @@
-import importlib
-
-from models.arch import *
-
-from models.cls_model_eval_nocls_reg import ClsModel as ClsReg
-
-
-def make_model(name: str):
-
-    model = ClsReg()
-    return model

models/base_model.py

@@ -1,71 +0,0 @@
-import os
-import torch
-import util.util as util
-
-
-class BaseModel:
-    def name(self):
-        return self.__class__.__name__.lower()
-
-    def initialize(self, opt):
-        self.opt = opt
-        self.gpu_ids = opt.gpu_ids
-        self.isTrain = opt.isTrain
-        self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
-        last_split = opt.checkpoints_dir.split('/')[-1]
-        if opt.resume and last_split != 'checkpoints' and (last_split != opt.name or opt.supp_eval):
-
-            self.save_dir = opt.checkpoints_dir
-            self.model_save_dir = os.path.join(opt.checkpoints_dir.replace(opt.checkpoints_dir.split('/')[-1], ''),
-                                               opt.name)
-        else:
-            self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
-            self.model_save_dir = os.path.join(opt.checkpoints_dir, opt.name)
-        self._count = 0
-
-    def set_input(self, input):
-        self.input = input
-
-    def forward(self, mode='train'):
-        pass
-
-    # used in test time, no backprop
-    def test(self):
-        pass
-
-    def get_image_paths(self):
-        pass
-
-    def optimize_parameters(self):
-        pass
-
-    def get_current_visuals(self):
-        return self.input
-
-    def get_current_errors(self):
-        return {}
-
-    def print_optimizer_param(self):
-        print(self.optimizers[-1])
-
-    def save(self, label=None):
-        epoch = self.epoch
-        iterations = self.iterations
-
-        if label is None:
-            model_name = os.path.join(self.model_save_dir, self.opt.name + '_%03d_%08d.pt' % ((epoch), (iterations)))
-        else:
-            model_name = os.path.join(self.model_save_dir, self.opt.name + '_' + label + '.pt')
-
-        torch.save(self.state_dict(), model_name)
-
-    def save_eval(self, label=None):
-        model_name = os.path.join(self.model_save_dir, label + '.pt')
-
-        torch.save(self.state_dict_eval(), model_name)
-
-    def _init_optimizer(self, optimizers):
-        self.optimizers = optimizers
-        for optimizer in self.optimizers:
-            util.set_opt_param(optimizer, 'initial_lr', self.opt.lr)
-            util.set_opt_param(optimizer, 'weight_decay', self.opt.wd)

models/cls_model_eval_nocls_reg.py

@@ -1,517 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from models.losses import DINOLoss
-import os
-import numpy as np
-from collections import OrderedDict
-from ema_pytorch import EMA
-from models.arch.classifier import PretrainedConvNext
-import util.util as util
-import util.index as index
-import models.networks as networks
-import models.losses as losses
-from models import arch
-#from models.arch.dncnn import effnetv2_s
-from .base_model import BaseModel
-from PIL import Image
-from os.path import join
-#from torchviz import make_dot
-from models.arch.RDnet_ import FullNet_NLP
-import timm
-
-def tensor2im(image_tensor, imtype=np.uint8):
-    image_tensor = image_tensor.detach()
-    image_numpy = image_tensor[0].cpu().float().numpy()
-    image_numpy = np.clip(image_numpy, 0, 1)
-    if image_numpy.shape[0] == 1:
-        image_numpy = np.tile(image_numpy, (3, 1, 1))
-    image_numpy = (np.transpose(image_numpy, (1, 2, 0))) * 255.0
-    # image_numpy = image_numpy.astype(imtype)
-    return image_numpy
-
-
-class EdgeMap(nn.Module):
-    def __init__(self, scale=1):
-        super(EdgeMap, self).__init__()
-        self.scale = scale
-        self.requires_grad = False
-
-    def forward(self, img):
-        img = img / self.scale
-
-        N, C, H, W = img.shape
-        gradX = torch.zeros(N, 1, H, W, dtype=img.dtype, device=img.device)
-        gradY = torch.zeros(N, 1, H, W, dtype=img.dtype, device=img.device)
-
-        gradx = (img[..., 1:, :] - img[..., :-1, :]).abs().sum(dim=1, keepdim=True)
-        grady = (img[..., 1:] - img[..., :-1]).abs().sum(dim=1, keepdim=True)
-
-        gradX[..., :-1, :] += gradx
-        gradX[..., 1:, :] += gradx
-        gradX[..., 1:-1, :] /= 2
-
-        gradY[..., :-1] += grady
-        gradY[..., 1:] += grady
-        gradY[..., 1:-1] /= 2
-
-        # edge = (gradX + gradY) / 2
-        edge = (gradX + gradY)
-
-        return edge
-
-
-class YTMTNetBase(BaseModel):
-    def _init_optimizer(self, optimizers):
-        self.optimizers = optimizers
-        for optimizer in self.optimizers:
-            util.set_opt_param(optimizer, 'initial_lr', self.opt.lr)
-            util.set_opt_param(optimizer, 'weight_decay', self.opt.wd)
-
-    def set_input(self, data, mode='train'):
-        target_t = None
-        target_r = None
-        data_name = None
-        identity = False
-        mode = mode.lower()
-        if mode == 'train':
-            input, target_t, target_r = data['input'], data['target_t'], data['target_r']
-        elif mode == 'eval':
-            input, target_t, target_r, data_name = data['input'], data['target_t'], data['target_r'], data['fn']
-        elif mode == 'test':
-            input, data_name = data['input'], data['fn']
-        else:
-            raise NotImplementedError('Mode [%s] is not implemented' % mode)
-
-        if len(self.gpu_ids) > 0:  # transfer data into gpu
-            input = input.to(device=self.gpu_ids[0])
-            if target_t is not None:
-                target_t = target_t.to(device=self.gpu_ids[0])
-            if target_r is not None:
-                target_r = target_r.to(device=self.gpu_ids[0])
-
-        self.input = input
-        self.identity = identity
-        self.input_edge = self.edge_map(self.input)
-        self.target_t = target_t
-        self.target_r = target_r
-        self.data_name = data_name
-
-        self.issyn = False if 'real' in data else True
-        self.aligned = False if 'unaligned' in data else True
-
-        if target_t is not None:
-            self.target_edge = self.edge_map(self.target_t)
-
-    def eval(self, data, savedir=None, suffix=None, pieapp=None):
-        self._eval()
-        self.set_input(data, 'eval')
-        with torch.no_grad():
-            self.forward_eval()
-
-            output_i = tensor2im(self.output_j[6])
-            output_j = tensor2im(self.output_j[7])
-            target = tensor2im(self.target_t)
-            target_r = tensor2im(self.target_r)
-
-            if self.aligned:
-                res = index.quality_assess(output_i, target)
-            else:
-                res = {}
-
-            if savedir is not None:
-                if self.data_name is not None:
-                    name = os.path.splitext(os.path.basename(self.data_name[0]))[0]
-                    savedir = join(savedir, suffix, name)
-                    os.makedirs(savedir, exist_ok=True)
-                    Image.fromarray(output_i.astype(np.uint8)).save(
-                        join(savedir, '{}_t.png'.format(self.opt.name)))
-                    Image.fromarray(output_j.astype(np.uint8)).save(
-                        join(savedir, '{}_r.png'.format(self.opt.name)))
-                    Image.fromarray(target.astype(np.uint8)).save(join(savedir, 't_label.png'))
-                    Image.fromarray(tensor2im(self.input).astype(np.uint8)).save(join(savedir, 'm_input.png'))
-                else:
-                    if not os.path.exists(join(savedir, 'transmission_layer')):
-                        os.makedirs(join(savedir, 'transmission_layer'))
-                        os.makedirs(join(savedir, 'blended'))
-                    Image.fromarray(target.astype(np.uint8)).save(
-                        join(savedir, 'transmission_layer', str(self._count) + '.png'))
-                    Image.fromarray(tensor2im(self.input).astype(np.uint8)).save(
-                        join(savedir, 'blended', str(self._count) + '.png'))
-                    self._count += 1
-
-            return res
-
-    def test(self, data, savedir=None):
-        # only the 1st input of the whole minibatch would be evaluated
-        self._eval()
-        self.set_input(data, 'test')
-
-        if self.data_name is not None and savedir is not None:
-            name = os.path.splitext(os.path.basename(self.data_name[0]))[0]
-            if not os.path.exists(join(savedir, name)):
-                os.makedirs(join(savedir, name))
-
-            if os.path.exists(join(savedir, name, '{}.png'.format(self.opt.name))):
-                return
-
-        with torch.no_grad():
-            output_i, output_j = self.forward()
-            output_i = tensor2im(output_i)
-            output_j = tensor2im(output_j)
-            if self.data_name is not None and savedir is not None:
-                Image.fromarray(output_i.astype(np.uint8)).save(join(savedir, name, '{}_l.png'.format(self.opt.name)))
-                Image.fromarray(output_j.astype(np.uint8)).save(join(savedir, name, '{}_r.png'.format(self.opt.name)))
-                Image.fromarray(tensor2im(self.input).astype(np.uint8)).save(join(savedir, name, 'm_input.png'))
-
-
-class ClsModel(YTMTNetBase):
-    def name(self):
-        return 'ytmtnet'
-
-    def __init__(self):
-        self.epoch = 0
-        self.iterations = 0
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.net_c = None
-
-    def print_network(self):
-        print('--------------------- Model ---------------------')
-        print('##################### NetG #####################')
-        networks.print_network(self.net_i)
-        if self.isTrain and self.opt.lambda_gan > 0:
-            print('##################### NetD #####################')
-            networks.print_network(self.netD)
-
-    def _eval(self):
-        self.net_i.eval()
-        self.net_c.eval()
-
-    def _train(self):
-        self.net_i.train()
-        self.net_c.eval()
-    def initialize(self, opt):
-        self.opt=opt
-        BaseModel.initialize(self, opt)
-
-        in_channels = 3
-        self.vgg = None
-
-        if opt.hyper:
-            self.vgg = losses.Vgg19(requires_grad=False).to(self.device)
-            in_channels += 1472
-        channels = [64, 128, 256, 512]
-        layers = [2, 2, 4, 2]
-        num_subnet = opt.num_subnet
-        self.net_c = PretrainedConvNext("convnext_small_in22k").cuda()
-        
-        self.net_c.load_state_dict(torch.load('pretrained/cls_model.pth')['icnn'])
-
-        self.net_i = FullNet_NLP(channels, layers, num_subnet, opt.loss_col,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None, kernel_size=3).to(self.device)
-    
-        self.edge_map = EdgeMap(scale=1).to(self.device)
-    
-        if self.isTrain:
-            self.loss_dic = losses.init_loss(opt, self.Tensor)
-            vggloss = losses.ContentLoss()
-            vggloss.initialize(losses.VGGLoss(self.vgg))
-            self.loss_dic['t_vgg'] = vggloss
-
-            cxloss = losses.ContentLoss()
-            if opt.unaligned_loss == 'vgg':
-                cxloss.initialize(losses.VGGLoss(self.vgg, weights=[0.1], indices=[opt.vgg_layer]))
-            elif opt.unaligned_loss == 'ctx':
-                cxloss.initialize(losses.CXLoss(self.vgg, weights=[0.1, 0.1, 0.1], indices=[8, 13, 22]))
-            elif opt.unaligned_loss == 'mse':
-                cxloss.initialize(nn.MSELoss())
-            elif opt.unaligned_loss == 'ctx_vgg':
-                cxloss.initialize(losses.CXLoss(self.vgg, weights=[0.1, 0.1, 0.1, 0.1], indices=[8, 13, 22, 31],
-                                                criterions=[losses.CX_loss] * 3 + [nn.L1Loss()]))
-            else:
-                raise NotImplementedError
-            self.scaler=torch.cuda.amp.GradScaler()
-            with torch.autocast(device_type='cuda',dtype=torch.float16):
-                self.dinoloss=DINOLoss()
-            self.loss_dic['t_cx'] = cxloss
-        
-            self.optimizer_G = torch.optim.Adam(self.net_i.parameters(),
-                                                lr=opt.lr, betas=(0.9, 0.999), weight_decay=opt.wd)
-
-
-            self._init_optimizer([self.optimizer_G])
-
-        if opt.resume:
-            self.load(self, opt.resume_epoch)
-
-
-    def backward_D(self):
-        loss_D=[]
-        weight=self.opt.weight_loss
-        for p in self.netD.parameters():
-            p.requires_grad = True
-        for i in range(4):
-            loss_D_1, pred_fake_1, pred_real_1 = self.loss_dic['gan'].get_loss(
-                self.netD, self.input, self.output_j[2*i], self.target_t)
-            loss_D.append(loss_D_1*weight)
-            weight+=self.opt.weight_loss
-        loss_sum=sum(loss_D)
-
-        self.loss_D, self.pred_fake, self.pred_real = (loss_sum, pred_fake_1, pred_real_1)
-
-        (self.loss_D * self.opt.lambda_gan).backward(retain_graph=True)
-
-    def get_loss(self, out_l, out_r):
-        loss_G_GAN_sum=[]
-        loss_icnn_pixel_sum=[]
-        loss_rcnn_pixel_sum=[]
-        loss_icnn_vgg_sum=[]
-        weight=self.opt.weight_loss
-        for i in range(self.opt.loss_col):
-            out_r_clean=out_r[2*i]
-            out_r_reflection=out_r[2*i+1]
-            if i != self.opt.loss_col -1:
-                loss_G_GAN = 0
-                loss_icnn_pixel = self.loss_dic['t_pixel'].get_loss(out_r_clean, self.target_t)
-                loss_rcnn_pixel = self.loss_dic['r_pixel'].get_loss(out_r_reflection, self.target_r) * 1.5 * self.opt.r_pixel_weight
-                loss_icnn_vgg = self.loss_dic['t_vgg'].get_loss(out_r_clean, self.target_t) * self.opt.lambda_vgg
-            else:
-                if self.opt.lambda_gan>0:
-
-                    loss_G_GAN=0
-                else:
-                    loss_G_GAN=0
-                loss_icnn_pixel = self.loss_dic['t_pixel'].get_loss(out_r_clean, self.target_t)
-                loss_rcnn_pixel = self.loss_dic['r_pixel'].get_loss(out_r_reflection, self.target_r) * 1.5 * self.opt.r_pixel_weight
-                loss_icnn_vgg = self.loss_dic['t_vgg'].get_loss(out_r_clean, self.target_t) * self.opt.lambda_vgg
-
-            loss_G_GAN_sum.append(loss_G_GAN*weight)
-            loss_icnn_pixel_sum.append(loss_icnn_pixel*weight)
-            loss_rcnn_pixel_sum.append(loss_rcnn_pixel*weight)
-            loss_icnn_vgg_sum.append(loss_icnn_vgg*weight)
-            weight=weight+self.opt.weight_loss
-        return sum(loss_G_GAN_sum), sum(loss_icnn_pixel_sum), sum(loss_rcnn_pixel_sum), sum(loss_icnn_vgg_sum)
-
-    def backward_G(self):
-
-        self.loss_G_GAN,self.loss_icnn_pixel, self.loss_rcnn_pixel, \
-        self.loss_icnn_vgg = self.get_loss(self.output_i, self.output_j)
-
-        self.loss_exclu = self.exclusion_loss(self.output_i, self.output_j, 3)
-
-        self.loss_recons = self.loss_dic['recons'](self.output_i, self.output_j, self.input) * 0.2
-
-        self.loss_G =  self.loss_G_GAN +self.loss_icnn_pixel + self.loss_rcnn_pixel + \
-                      self.loss_icnn_vgg
-        self.scaler.scale(self.loss_G).backward()
-
-
-
-    def hyper_column(self, input_img):
-        hypercolumn = self.vgg(input_img)
-        _, C, H, W = input_img.shape
-        hypercolumn = [F.interpolate(feature.detach(), size=(H, W), mode='bilinear', align_corners=False) for
-                       feature in hypercolumn]
-        input_i = [input_img]
-        input_i.extend(hypercolumn)
-        input_i = torch.cat(input_i, dim=1)
-        return input_i
-
-    def forward(self):
-        # without edge
-        
-        self.output_j=[]
-        input_i = self.input
-        if self.vgg is not None:
-            input_i = self.hyper_column(input_i)
-        with torch.no_grad():
-            ipt = self.net_c(input_i)
-        output_i, output_j = self.net_i(input_i,ipt,prompt=True)
-        self.output_i = output_i
-        for i in range(self.opt.loss_col):
-            out_reflection, out_clean = output_j[i][:, :3, ...], output_j[i][:, 3:, ...]
-            self.output_j.append(out_clean) 
-            self.output_j.append(out_reflection) 
-        return self.output_i, self.output_j
-
-
-    @torch.no_grad() 
-    def forward_eval(self):
-       
-        self.output_j=[]
-        input_i = self.input
-        if self.vgg is not None:
-            input_i = self.hyper_column(input_i)
-        ipt = self.net_c(input_i)
-        
-        output_i, output_j = self.net_i(input_i,ipt,prompt=True)
-        self.output_i = output_i #alpha * output_i + beta
-        for i in range(self.opt.loss_col):
-            out_reflection, out_clean = output_j[i][:, :3, ...], output_j[i][:, 3:, ...]
-            self.output_j.append(out_clean) 
-            self.output_j.append(out_reflection)
-        return self.output_i, self.output_j
-
-    def optimize_parameters(self):
-        self._train()
-        self.forward()
-        self.optimizer_G.zero_grad()
-        self.backward_G()
-        self.optimizer_G.step()
-
-    def return_output(self):
-        output_clean = self.output_j[1]
-        output_reflection = self.output_j[0]
-        output_clean = tensor2im(output_clean).astype(np.uint8)
-        output_reflection = tensor2im(output_reflection).astype(np.uint8)
-        input=tensor2im(self.input)
-        return output_clean,output_reflection,input
-    def exclusion_loss(self, img_T, img_R, level=3, eps=1e-6):
-        loss_gra=[]
-        weight=0.25
-        for i in range(4):
-            grad_x_loss = []
-            grad_y_loss = []
-            img_T=self.output_j[2*i]
-            img_R=self.output_j[2*i+1]
-            for l in range(level):
-                grad_x_T, grad_y_T = self.compute_grad(img_T)
-                grad_x_R, grad_y_R = self.compute_grad(img_R)
-
-                alphax = (2.0 * torch.mean(torch.abs(grad_x_T))) / (torch.mean(torch.abs(grad_x_R)) + eps)
-                alphay = (2.0 * torch.mean(torch.abs(grad_y_T))) / (torch.mean(torch.abs(grad_y_R)) + eps)
-
-                gradx1_s = (torch.sigmoid(grad_x_T) * 2) - 1  # mul 2 minus 1 is to change sigmoid into tanh
-                grady1_s = (torch.sigmoid(grad_y_T) * 2) - 1
-                gradx2_s = (torch.sigmoid(grad_x_R * alphax) * 2) - 1
-                grady2_s = (torch.sigmoid(grad_y_R * alphay) * 2) - 1
-
-                grad_x_loss.append(((torch.mean(torch.mul(gradx1_s.pow(2), gradx2_s.pow(2)))) + eps) ** 0.25)
-                grad_y_loss.append(((torch.mean(torch.mul(grady1_s.pow(2), grady2_s.pow(2)))) + eps) ** 0.25)
-
-                img_T = F.interpolate(img_T, scale_factor=0.5, mode='bilinear')
-                img_R = F.interpolate(img_R, scale_factor=0.5, mode='bilinear')
-            loss_gradxy = torch.sum(sum(grad_x_loss) / 3) + torch.sum(sum(grad_y_loss) / 3)
-            loss_gra.append(loss_gradxy*weight)
-            weight+=0.25
-
-
-        return sum(loss_gra) / 2
-
-    def contain_loss(self, img_T, img_R, img_I, eps=1e-6):
-        pix_num = np.prod(img_I.shape)
-        predict_tx, predict_ty = self.compute_grad(img_T)
-        predict_tx, predict_ty = self.compute_grad(img_T)
-        predict_rx, predict_ry = self.compute_grad(img_R)
-        input_x, input_y = self.compute_grad(img_I)
-
-        out = torch.norm(predict_tx / (input_x + eps), 2) ** 2 + \
-              torch.norm(predict_ty / (input_y + eps), 2) ** 2 + \
-              torch.norm(predict_rx / (input_x + eps), 2) ** 2 + \
-              torch.norm(predict_ry / (input_y + eps), 2) ** 2
-
-        return out / pix_num
-
-    def compute_grad(self, img):
-        gradx = img[:, :, 1:, :] - img[:, :, :-1, :]
-        grady = img[:, :, :, 1:] - img[:, :, :, :-1]
-        return gradx, grady
-
-    def load(self, model, resume_epoch=None):
-        icnn_path = model.opt.icnn_path
-        state_dict = torch.load(icnn_path)
-        model.net_i.load_state_dict(state_dict['icnn'])
-        return state_dict
-
-    def state_dict(self):
-        state_dict = {
-            'icnn': self.net_i.state_dict(),
-            'opt_g': self.optimizer_G.state_dict(),
-            #'ema' : self.ema.state_dict(),
-            'epoch': self.epoch, 'iterations': self.iterations
-        }
-
-        if self.opt.lambda_gan > 0:
-            state_dict.update({
-                'opt_d': self.optimizer_D.state_dict(),
-                'netD': self.netD.state_dict(),
-            })
-
-        return state_dict
-class AvgPool2d(nn.Module):
-    def __init__(self, kernel_size=None, base_size=None, auto_pad=True, fast_imp=False, train_size=None):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.base_size = base_size
-        self.auto_pad = auto_pad
-
-        # only used for fast implementation
-        self.fast_imp = fast_imp
-        self.rs = [5, 4, 3, 2, 1]
-        self.max_r1 = self.rs[0]
-        self.max_r2 = self.rs[0]
-        self.train_size = train_size
-
-    def extra_repr(self) -> str:
-        return 'kernel_size={}, base_size={}, stride={}, fast_imp={}'.format(
-            self.kernel_size, self.base_size, self.kernel_size, self.fast_imp
-        )
-
-    def forward(self, x):
-        if self.kernel_size is None and self.base_size:
-            train_size = self.train_size
-            if isinstance(self.base_size, int):
-                self.base_size = (self.base_size, self.base_size)
-            self.kernel_size = list(self.base_size)
-            self.kernel_size[0] = x.shape[2] * self.base_size[0] // train_size[-2]
-            self.kernel_size[1] = x.shape[3] * self.base_size[1] // train_size[-1]
-
-            # only used for fast implementation
-            self.max_r1 = max(1, self.rs[0] * x.shape[2] // train_size[-2])
-            self.max_r2 = max(1, self.rs[0] * x.shape[3] // train_size[-1])
-
-        if self.kernel_size[0] >= x.size(-2) and self.kernel_size[1] >= x.size(-1):
-            return F.adaptive_avg_pool2d(x, 1)
-
-        if self.fast_imp:  # Non-equivalent implementation but faster
-            h, w = x.shape[2:]
-            if self.kernel_size[0] >= h and self.kernel_size[1] >= w:
-                out = F.adaptive_avg_pool2d(x, 1)
-            else:
-                r1 = [r for r in self.rs if h % r == 0][0]
-                r2 = [r for r in self.rs if w % r == 0][0]
-                # reduction_constraint
-                r1 = min(self.max_r1, r1)
-                r2 = min(self.max_r2, r2)
-                s = x[:, :, ::r1, ::r2].cumsum(dim=-1).cumsum(dim=-2)
-                n, c, h, w = s.shape
-                k1, k2 = min(h - 1, self.kernel_size[0] // r1), min(w - 1, self.kernel_size[1] // r2)
-                out = (s[:, :, :-k1, :-k2] - s[:, :, :-k1, k2:] - s[:, :, k1:, :-k2] + s[:, :, k1:, k2:]) / (k1 * k2)
-                out = torch.nn.functional.interpolate(out, scale_factor=(r1, r2))
-        else:
-            n, c, h, w = x.shape
-            s = x.cumsum(dim=-1).cumsum_(dim=-2)
-            s = torch.nn.functional.pad(s, (1, 0, 1, 0))  # pad 0 for convenience
-            k1, k2 = min(h, self.kernel_size[0]), min(w, self.kernel_size[1])
-            s1, s2, s3, s4 = s[:, :, :-k1, :-k2], s[:, :, :-k1, k2:], s[:, :, k1:, :-k2], s[:, :, k1:, k2:]
-            out = s4 + s1 - s2 - s3
-            out = out / (k1 * k2)
-
-        if self.auto_pad:
-            n, c, h, w = x.shape
-            _h, _w = out.shape[2:]
-            # print(x.shape, self.kernel_size)
-            pad2d = ((w - _w) // 2, (w - _w + 1) // 2, (h - _h) // 2, (h - _h + 1) // 2)
-            out = torch.nn.functional.pad(out, pad2d, mode='replicate')
-
-        return out
-
-def replace_layers(model, base_size, train_size, fast_imp, **kwargs):
-    for n, m in model.named_children():
-        if len(list(m.children())) > 0:
-            ## compound module, go inside it
-            replace_layers(m, base_size, train_size, fast_imp, **kwargs)
-
-        if isinstance(m, nn.AdaptiveAvgPool2d):
-            pool = AvgPool2d(base_size=base_size, fast_imp=fast_imp, train_size=train_size)
-            assert m.output_size == 1
-            setattr(model, n, pool)

models/losses.py

@@ -1,468 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from pytorch_msssim import SSIM
-from models.vit_feature_extractor import VitExtractor
-from models.vgg import Vgg19
-
-
-###############################################################################
-# Functions
-###############################################################################
-def compute_gradient(img):
-    gradx = img[..., 1:, :] - img[..., :-1, :]
-    grady = img[..., 1:] - img[..., :-1]
-    return gradx, grady
-
-
-class GradientLoss(nn.Module):
-    def __init__(self):
-        super(GradientLoss, self).__init__()
-        self.loss = nn.L1Loss()
-
-    def forward(self, predict, target):
-        predict_gradx, predict_grady = compute_gradient(predict)
-        target_gradx, target_grady = compute_gradient(target)
-
-        return self.loss(predict_gradx, target_gradx) + self.loss(predict_grady, target_grady)
-
-
-class ContainLoss(nn.Module):
-    def __init__(self, eps=1e-12):
-        super(ContainLoss, self).__init__()
-        self.eps = eps
-
-    def forward(self, predict_t, predict_r, input_image):
-        pix_num = np.prod(input_image.shape)
-        predict_tx, predict_ty = compute_gradient(predict_t)
-        predict_rx, predict_ry = compute_gradient(predict_r)
-        input_x, input_y = compute_gradient(input_image)
-
-        out = torch.norm(predict_tx / (input_x + self.eps), 2) ** 2 + \
-              torch.norm(predict_ty / (input_y + self.eps), 2) ** 2 + \
-              torch.norm(predict_rx / (input_x + self.eps), 2) ** 2 + \
-              torch.norm(predict_ry / (input_y + self.eps), 2) ** 2
-
-        return out / pix_num
-
-
-class MultipleLoss(nn.Module):
-    def __init__(self, losses, weight=None):
-        super(MultipleLoss, self).__init__()
-        self.losses = nn.ModuleList(losses)
-        self.weight = weight or [1 / len(self.losses)] * len(self.losses)
-
-    def forward(self, predict, target):
-        total_loss = 0
-        for weight, loss in zip(self.weight, self.losses):
-            total_loss += loss(predict, target) * weight
-        return total_loss
-
-
-class MeanShift(nn.Conv2d):
-    def __init__(self, data_mean, data_std, data_range=1, norm=True):
-        """norm (bool): normalize/denormalize the stats"""
-        c = len(data_mean)
-        super(MeanShift, self).__init__(c, c, kernel_size=1)
-        std = torch.Tensor(data_std)
-        self.weight.data = torch.eye(c).view(c, c, 1, 1)
-        if norm:
-            self.weight.data.div_(std.view(c, 1, 1, 1))
-            self.bias.data = -1 * data_range * torch.Tensor(data_mean)
-            self.bias.data.div_(std)
-        else:
-            self.weight.data.mul_(std.view(c, 1, 1, 1))
-            self.bias.data = data_range * torch.Tensor(data_mean)
-        self.requires_grad = False
-
-
-class VGGLoss(nn.Module):
-    def __init__(self, vgg=None, weights=None, indices=None, normalize=True):
-        super(VGGLoss, self).__init__()
-        if vgg is None:
-            self.vgg = torch.compile(Vgg19().cuda())
-        else:
-            self.vgg = vgg
-        self.criterion = nn.L1Loss()
-        self.weights = weights or [1.0 / 2.6, 1.0 / 4.8, 1.0 / 3.7, 1.0 / 5.6, 10 / 1.5]
-        self.indices = indices or [2, 7, 12, 21, 30]
-        if normalize:
-            self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
-        else:
-            self.normalize = None
-
-    def forward(self, x, y):
-        if self.normalize is not None:
-            x = self.normalize(x)
-            y = self.normalize(y)
-        with torch.no_grad():
-            y_vgg = self.vgg(y, self.indices)
-        x_vgg = self.vgg(x, self.indices) #, self.vgg(y, self.indices)
-        loss = 0
-        for i in range(len(x_vgg)):
-            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i]) #.detach())
-
-        return loss
-
-
-def l1_norm_dim(x, dim):
-    return torch.mean(torch.abs(x), dim=dim)
-
-
-def l1_norm(x):
-    return torch.mean(torch.abs(x))
-
-
-def l2_norm(x):
-    return torch.mean(torch.square(x))
-
-
-def gradient_norm_kernel(x, kernel_size=10):
-    out_h, out_v = compute_gradient(x)
-    shape = out_h.shape
-    out_h = F.unfold(out_h, kernel_size=(kernel_size, kernel_size), stride=(1, 1))
-    out_h = out_h.reshape(shape[0], shape[1], kernel_size * kernel_size, -1)
-    out_h = l1_norm_dim(out_h, 2)
-    out_v = F.unfold(out_v, kernel_size=(kernel_size, kernel_size), stride=(1, 1))
-    out_v = out_v.reshape(shape[0], shape[1], kernel_size * kernel_size, -1)
-    out_v = l1_norm_dim(out_v, 2)
-    return out_h, out_v
-
-
-class KTVLoss(nn.Module):
-    def __init__(self, kernel_size=10):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.criterion = nn.L1Loss()
-        self.eps = 1e-6
-
-    def forward(self, out_l, out_r, input_i):
-        out_l_normx, out_l_normy = gradient_norm_kernel(out_l, self.kernel_size)
-        out_r_normx, out_r_normy = gradient_norm_kernel(out_r, self.kernel_size)
-        input_normx, input_normy = gradient_norm_kernel(input_i, self.kernel_size)
-        norm_l = out_l_normx + out_l_normy
-        norm_r = out_r_normx + out_r_normy
-        norm_target = input_normx + input_normy + self.eps
-        norm_loss = (norm_l / norm_target + norm_r / norm_target).mean()
-
-        out_lx, out_ly = compute_gradient(out_l)
-        out_rx, out_ry = compute_gradient(out_r)
-        input_x, input_y = compute_gradient(input_i)
-        gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-        gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-        grad_loss = gradient_diffx + gradient_diffy
-
-        loss = norm_loss * 1e-4 + grad_loss
-        return loss
-
-
-class MTVLoss(nn.Module):
-    def __init__(self, kernel_size=10):
-        super().__init__()
-        self.criterion = nn.L1Loss()
-        self.norm = l1_norm
-
-    def forward(self, out_l, out_r, input_i):
-        out_lx, out_ly = compute_gradient(out_l)
-        out_rx, out_ry = compute_gradient(out_r)
-        input_x, input_y = compute_gradient(input_i)
-
-        norm_l = self.norm(out_lx) + self.norm(out_ly)
-        norm_r = self.norm(out_rx) + self.norm(out_ry)
-        norm_target = self.norm(input_x) + self.norm(input_y)
-
-        gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-        gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-
-        loss = (norm_l / norm_target + norm_r / norm_target) * 1e-5 + gradient_diffx + gradient_diffy
-
-        return loss
-
-
-class ReconsLoss(nn.Module):
-    def __init__(self, edge_recons=True):
-        super().__init__()
-        self.criterion = nn.L1Loss()
-        self.norm = l1_norm
-        self.edge_recons = edge_recons
-        self.mse_loss=nn.MSELoss()
-
-    def forward(self, out_l, out_r, input_i):
-        loss_sum=[]
-        weight=0.25
-        for i in range(4):
-            #out_res = out_l[i]
-            out_clean=out_r[2*i]
-            out_reflection=out_r[2*i+1]
-            #content_diff = self.criterion(out_clean + out_reflection, input_i)
-            # if self.edge_recons:
-            #     out_lx, out_ly = compute_gradient(out_clean)
-            #     out_rx, out_ry = compute_gradient(out_reflection)
-            #     #out_resx, out_resy = compute_gradient(out_res)
-            #     input_x, input_y = compute_gradient(input_i)
-
-            #     gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-            #     gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-
-            #     loss = content_diff + (gradient_diffx + gradient_diffy) * 5.0
-            # else:
-            #     loss = content_diff
-            loss=self.mse_loss(out_clean+out_reflection,input_i)
-            loss_sum.append(loss*weight)
-            weight=weight+0.25
-
-        return sum(loss_sum)
-
-
-class ReconsLossX(nn.Module):
-    def __init__(self, edge_recons=True):
-        super().__init__()
-        self.criterion = nn.MSELoss()
-        self.norm = l1_norm
-        self.edge_recons = edge_recons
-
-    def forward(self, out, input_i):
-        content_diff = self.criterion(out, input_i)
-        if self.edge_recons:
-            out_x, out_y = compute_gradient(out)
-            input_x, input_y = compute_gradient(input_i)
-
-            gradient_diffx = self.criterion(out_x, input_x)
-            gradient_diffy = self.criterion(out_y, input_y)
-
-            loss = content_diff + (gradient_diffx + gradient_diffy) * 1.0
-        else:
-            loss = content_diff
-        return loss
-
-
-class ContentLoss():
-    def initialize(self, loss):
-        self.criterion = loss
-
-    def get_loss(self, fakeIm, realIm):
-        return self.criterion(fakeIm, realIm)
-
-
-class GANLoss(nn.Module):
-    def __init__(self, use_l1=True, target_real_label=1.0, target_fake_label=0.0,
-                 tensor=torch.FloatTensor):
-        super(GANLoss, self).__init__()
-        self.real_label = target_real_label
-        self.fake_label = target_fake_label
-        self.real_label_var = None
-        self.fake_label_var = None
-        self.Tensor = tensor
-        if use_l1:
-            self.loss = nn.L1Loss()
-        else:
-            self.loss = nn.BCEWithLogitsLoss()  # absorb sigmoid into BCELoss
-
-    def get_target_tensor(self, input, target_is_real):
-        target_tensor = None
-        if target_is_real:
-            create_label = ((self.real_label_var is None) or
-                            (self.real_label_var.numel() != input.numel()))
-            if create_label:
-                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
-                self.real_label_var = real_tensor
-            target_tensor = self.real_label_var
-        else:
-            create_label = ((self.fake_label_var is None) or
-                            (self.fake_label_var.numel() != input.numel()))
-            if create_label:
-                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
-                self.fake_label_var = fake_tensor
-            target_tensor = self.fake_label_var
-        return target_tensor
-
-    def __call__(self, input, target_is_real):
-        if isinstance(input, list):
-            loss = 0
-            for input_i in input:
-                target_tensor = self.get_target_tensor(input_i, target_is_real)
-                loss += self.loss(input_i, target_tensor)
-            return loss
-        else:
-            target_tensor = self.get_target_tensor(input, target_is_real)
-            return self.loss(input, target_tensor)
-
-
-class DiscLoss():
-    def name(self):
-        return 'SGAN'
-
-    def initialize(self, opt, tensor):
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB):
-        # First, G(A) should fake the discriminator
-        pred_fake = net.forward(fakeB)
-        return self.criterionGAN(pred_fake, 1)
-
-    def get_loss(self, net, realA=None, fakeB=None, realB=None):
-        pred_fake = None
-        pred_real = None
-        loss_D_fake = 0
-        loss_D_real = 0
-        # Fake
-        # stop backprop to the generator by detaching fake_B
-        # Generated Image Disc Output should be close to zero
-
-        if fakeB is not None:
-            pred_fake = net.forward(fakeB.detach())
-            loss_D_fake = self.criterionGAN(pred_fake, 0)
-
-        # Real
-        if realB is not None:
-            pred_real = net.forward(realB)
-            loss_D_real = self.criterionGAN(pred_real, 1)
-
-        # Combined loss
-        loss_D = (loss_D_fake + loss_D_real) * 0.5
-        return loss_D, pred_fake, pred_real
-
-
-class DiscLossR(DiscLoss):
-    # RSGAN from 
-    # https://arxiv.org/abs/1807.00734        
-    def name(self):
-        return 'RSGAN'
-
-    def initialize(self, opt, tensor):
-        DiscLoss.initialize(self, opt, tensor)
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB, pred_real=None):
-        if pred_real is None:
-            pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB)
-        return self.criterionGAN(pred_fake - pred_real, 1)
-
-    def get_loss(self, net, realA, fakeB, realB):
-        pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB.detach())
-
-        loss_D = self.criterionGAN(pred_real - pred_fake, 1)  # BCE_stable loss
-        return loss_D, pred_fake, pred_real
-
-
-class DiscLossRa(DiscLoss):
-    # RaSGAN from 
-    # https://arxiv.org/abs/1807.00734    
-    def name(self):
-        return 'RaSGAN'
-
-    def initialize(self, opt, tensor):
-        DiscLoss.initialize(self, opt, tensor)
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB, pred_real=None):
-        if pred_real is None:
-            pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB)
-
-        loss_G = self.criterionGAN(pred_real - torch.mean(pred_fake, dim=0, keepdim=True), 0)
-        loss_G += self.criterionGAN(pred_fake - torch.mean(pred_real, dim=0, keepdim=True), 1)
-        return loss_G * 0.5
-
-    def get_loss(self, net, realA, fakeB, realB):
-        pred_real = net.forward(realB)
-
-        pred_fake = net.forward(fakeB.detach())
-
-        loss_D = self.criterionGAN(pred_real - torch.mean(pred_fake, dim=0, keepdim=True), 1)
-        loss_D += self.criterionGAN(pred_fake - torch.mean(pred_real, dim=0, keepdim=True), 0)
-        return loss_D * 0.5, pred_fake, pred_real
-
-
-class SSIM_Loss(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.ssim = SSIM(data_range=1, size_average=True, channel=3)
-
-    def forward(self, output, target):
-        return 1 - self.ssim(output, target)
-
-
-def init_loss(opt, tensor):
-    disc_loss = None
-    content_loss = None
-
-    loss_dic = {}
-
-    pixel_loss = ContentLoss()
-    pixel_loss.initialize(MultipleLoss([nn.MSELoss(), GradientLoss()], [0.3, 0.6]))
-
-    loss_dic['t_pixel'] = pixel_loss
-
-    r_loss = ContentLoss()
-    r_loss.initialize(MultipleLoss([nn.MSELoss()], [0.9]))
-    loss_dic['r_pixel'] = pixel_loss
-
-    loss_dic['t_ssim'] = SSIM_Loss()
-    loss_dic['r_ssim'] = SSIM_Loss()
-
-    loss_dic['mtv'] = MTVLoss()
-    loss_dic['ktv'] = KTVLoss()
-    loss_dic['recons'] = ReconsLoss(edge_recons=False)
-    loss_dic['reconsx'] = ReconsLossX(edge_recons=False)
-
-    if opt.lambda_gan > 0:
-        if opt.gan_type == 'sgan' or opt.gan_type == 'gan':
-            disc_loss = DiscLoss()
-        elif opt.gan_type == 'rsgan':
-            disc_loss = DiscLossR()
-        elif opt.gan_type == 'rasgan':
-            disc_loss = DiscLossRa()
-        else:
-            raise ValueError("GAN [%s] not recognized." % opt.gan_type)
-
-        disc_loss.initialize(opt, tensor)
-        loss_dic['gan'] = disc_loss
-
-    return loss_dic
-
-class DINOLoss(nn.Module):
-    '''
-    DINO-ViT as perceptual loss
-    '''
-
-    def resize_to_dino(self, feature, size = (224, 224)): 
-        return F.interpolate(feature, size = size, mode='bilinear', align_corners=False)
-
-    def calculate_crop_cls_loss(self, outputs, inputs):
-        loss = 0.0
-        for a, b in zip(outputs, inputs):  # avoid memory limitations
-            a = self.global_transform(a).unsqueeze(0)
-            b = self.global_transform(b).unsqueeze(0)
-            cls_token = self.extractor.get_feature_from_input(a)[-1][0, 0, :]
-            with torch.no_grad():
-                target_cls_token = self.extractor.get_feature_from_input(b)[-1][0, 0, :]
-            loss += F.mse_loss(cls_token, target_cls_token)
-        return loss
-
-    def __init__(self) :
-        super(DINOLoss, self).__init__()
-        self.extractor = VitExtractor(model_name = 'dino_vits8', device = 'cuda')
-        self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
-
-    def forward(self, output, target):
-        output = self.normalize(self.resize_to_dino(output))
-        output_cls_token = self.extractor.get_feature_from_input(output)[-1][0, 0, :]
-        with torch.no_grad():
-            target = self.normalize(self.resize_to_dino(target))
-            target_cls_token = self.extractor.get_feature_from_input(target)[-1][0, 0, :]
-
-        return F.mse_loss(output_cls_token, target_cls_token)
-    
-if __name__ == '__main__':
-    x = torch.randn(3, 32, 224, 224).cuda()
-    import time
-
-    s = time.time()
-    out1, out2 = gradient_norm_kernel(x)
-    t = time.time()
-    print(t - s)
-    print(out1.shape, out2.shape)

models/losses_opt.py

@@ -1,404 +0,0 @@
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from pytorch_msssim import MS_SSIM, SSIM
-
-from models.vgg import Vgg19
-
-
-###############################################################################
-# Functions
-###############################################################################
-def compute_gradient(img):
-    gradx = img[..., 1:, :] - img[..., :-1, :]
-    grady = img[..., 1:] - img[..., :-1]
-    return gradx, grady
-
-
-class GradientLoss(nn.Module):
-    def __init__(self):
-        super(GradientLoss, self).__init__()
-        self.loss = nn.L1Loss()
-
-    def forward(self, predict, target):
-        predict_gradx, predict_grady = compute_gradient(predict)
-        target_gradx, target_grady = compute_gradient(target)
-
-        return self.loss(predict_gradx, target_gradx) + self.loss(predict_grady, target_grady)
-
-
-class ContainLoss(nn.Module):
-    def __init__(self, eps=1e-12):
-        super(ContainLoss, self).__init__()
-        self.eps = eps
-
-    def forward(self, predict_t, predict_r, input_image):
-        pix_num = np.prod(input_image.shape)
-        predict_tx, predict_ty = compute_gradient(predict_t)
-        predict_rx, predict_ry = compute_gradient(predict_r)
-        input_x, input_y = compute_gradient(input_image)
-
-        out = torch.norm(predict_tx / (input_x + self.eps), 2) ** 2 + \
-              torch.norm(predict_ty / (input_y + self.eps), 2) ** 2 + \
-              torch.norm(predict_rx / (input_x + self.eps), 2) ** 2 + \
-              torch.norm(predict_ry / (input_y + self.eps), 2) ** 2
-
-        return out / pix_num
-
-
-class MultipleLoss(nn.Module):
-    def __init__(self, losses, weight=None):
-        super(MultipleLoss, self).__init__()
-        self.losses = nn.ModuleList(losses)
-        self.weight = weight or [1 / len(self.losses)] * len(self.losses)
-
-    def forward(self, predict, target):
-        total_loss = 0
-        for weight, loss in zip(self.weight, self.losses):
-            total_loss += loss(predict, target) * weight
-        return total_loss
-
-
-class MeanShift(nn.Conv2d):
-    def __init__(self, data_mean, data_std, data_range=1, norm=True):
-        """norm (bool): normalize/denormalize the stats"""
-        c = len(data_mean)
-        super(MeanShift, self).__init__(c, c, kernel_size=1)
-        std = torch.Tensor(data_std)
-        self.weight.data = torch.eye(c).view(c, c, 1, 1)
-        if norm:
-            self.weight.data.div_(std.view(c, 1, 1, 1))
-            self.bias.data = -1 * data_range * torch.Tensor(data_mean)
-            self.bias.data.div_(std)
-        else:
-            self.weight.data.mul_(std.view(c, 1, 1, 1))
-            self.bias.data = data_range * torch.Tensor(data_mean)
-        self.requires_grad = False
-
-
-class VGGLoss(nn.Module):
-    def __init__(self, vgg=None, weights=None, indices=None, normalize=True):
-        super(VGGLoss, self).__init__()
-        if vgg is None:
-            self.vgg = Vgg19().cuda()
-        else:
-            self.vgg = vgg
-        self.criterion = nn.L1Loss()
-        self.weights = weights or [1.0 / 2.6, 1.0 / 4.8, 1.0 / 3.7, 1.0 / 5.6, 10 / 1.5]
-        self.indices = indices or [2, 7, 12, 21, 30]
-        if normalize:
-            self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
-        else:
-            self.normalize = None
-
-    def forward(self, x, y):
-        if self.normalize is not None:
-            x = self.normalize(x)
-            y = self.normalize(y)
-        x_vgg, y_vgg = self.vgg(x, self.indices), self.vgg(y, self.indices)
-        loss = 0
-        for i in range(len(x_vgg)):
-            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
-
-        return loss
-
-
-def l1_norm_dim(x, dim):
-    return torch.mean(torch.abs(x), dim=dim)
-
-
-def l1_norm(x):
-    return torch.mean(torch.abs(x))
-
-
-def l2_norm(x):
-    return torch.mean(torch.square(x))
-
-
-def gradient_norm_kernel(x, kernel_size=10):
-    out_h, out_v = compute_gradient(x)
-    shape = out_h.shape
-    out_h = F.unfold(out_h, kernel_size=(kernel_size, kernel_size), stride=(1, 1))
-    out_h = out_h.reshape(shape[0], shape[1], kernel_size * kernel_size, -1)
-    out_h = l1_norm_dim(out_h, 2)
-    out_v = F.unfold(out_v, kernel_size=(kernel_size, kernel_size), stride=(1, 1))
-    out_v = out_v.reshape(shape[0], shape[1], kernel_size * kernel_size, -1)
-    out_v = l1_norm_dim(out_v, 2)
-    return out_h, out_v
-
-
-class KTVLoss(nn.Module):
-    def __init__(self, kernel_size=10):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.criterion = nn.L1Loss()
-        self.eps = 1e-6
-
-    def forward(self, out_l, out_r, input_i):
-        out_l_normx, out_l_normy = gradient_norm_kernel(out_l, self.kernel_size)
-        out_r_normx, out_r_normy = gradient_norm_kernel(out_r, self.kernel_size)
-        input_normx, input_normy = gradient_norm_kernel(input_i, self.kernel_size)
-        norm_l = out_l_normx + out_l_normy
-        norm_r = out_r_normx + out_r_normy
-        norm_target = input_normx + input_normy + self.eps
-        norm_loss = (norm_l / norm_target + norm_r / norm_target).mean()
-
-        out_lx, out_ly = compute_gradient(out_l)
-        out_rx, out_ry = compute_gradient(out_r)
-        input_x, input_y = compute_gradient(input_i)
-        gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-        gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-        grad_loss = gradient_diffx + gradient_diffy
-
-        loss = norm_loss * 1e-4 + grad_loss
-        return loss
-
-
-class MTVLoss(nn.Module):
-    def __init__(self, kernel_size=10):
-        super().__init__()
-        self.criterion = nn.L1Loss()
-        self.norm = l1_norm
-
-    def forward(self, out_l, out_r, input_i):
-        out_lx, out_ly = compute_gradient(out_l)
-        out_rx, out_ry = compute_gradient(out_r)
-        input_x, input_y = compute_gradient(input_i)
-
-        norm_l = self.norm(out_lx) + self.norm(out_ly)
-        norm_r = self.norm(out_rx) + self.norm(out_ry)
-        norm_target = self.norm(input_x) + self.norm(input_y)
-
-        gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-        gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-
-        loss = (norm_l / norm_target + norm_r / norm_target) * 1e-5 + gradient_diffx + gradient_diffy
-
-        return loss
-
-
-class ReconsLoss(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.criterion = nn.L1Loss()
-        self.norm = l1_norm
-
-    def forward(self, out_l, out_r, input_i):
-        content_diff = self.criterion(out_l + out_r, input_i)
-        out_lx, out_ly = compute_gradient(out_l)
-        out_rx, out_ry = compute_gradient(out_r)
-        input_x, input_y = compute_gradient(input_i)
-
-        gradient_diffx = self.criterion(out_lx + out_rx, input_x)
-        gradient_diffy = self.criterion(out_ly + out_ry, input_y)
-
-        loss = content_diff + (gradient_diffx + gradient_diffy) * 0.5
-
-        return loss
-
-
-class ContentLoss():
-    def initialize(self, loss):
-        self.criterion = loss
-
-    def get_loss(self, fakeIm, realIm):
-        return self.criterion(fakeIm, realIm)
-
-
-class GANLoss(nn.Module):
-    def __init__(self, use_l1=True, target_real_label=1.0, target_fake_label=0.0,
-                 tensor=torch.FloatTensor):
-        super(GANLoss, self).__init__()
-        self.real_label = target_real_label
-        self.fake_label = target_fake_label
-        self.real_label_var = None
-        self.fake_label_var = None
-        self.Tensor = tensor
-        if use_l1:
-            self.loss = nn.L1Loss()
-        else:
-            self.loss = nn.BCEWithLogitsLoss()  # absorb sigmoid into BCELoss
-
-    def get_target_tensor(self, input, target_is_real):
-        target_tensor = None
-        if target_is_real:
-            create_label = ((self.real_label_var is None) or
-                            (self.real_label_var.numel() != input.numel()))
-            if create_label:
-                real_tensor = self.Tensor(input.size()).fill_(self.real_label)
-                self.real_label_var = real_tensor
-            target_tensor = self.real_label_var
-        else:
-            create_label = ((self.fake_label_var is None) or
-                            (self.fake_label_var.numel() != input.numel()))
-            if create_label:
-                fake_tensor = self.Tensor(input.size()).fill_(self.fake_label)
-                self.fake_label_var = fake_tensor
-            target_tensor = self.fake_label_var
-        return target_tensor
-
-    def __call__(self, input, target_is_real):
-        if isinstance(input, list):
-            loss = 0
-            for input_i in input:
-                target_tensor = self.get_target_tensor(input_i, target_is_real)
-                loss += self.loss(input_i, target_tensor)
-            return loss
-        else:
-            target_tensor = self.get_target_tensor(input, target_is_real)
-            return self.loss(input, target_tensor)
-
-
-class DiscLoss():
-    def name(self):
-        return 'SGAN'
-
-    def initialize(self, opt, tensor):
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB):
-        # First, G(A) should fake the discriminator
-        pred_fake = net.forward(fakeB)
-        return self.criterionGAN(pred_fake, 1)
-
-    def get_loss(self, net, realA=None, fakeB=None, realB=None):
-        pred_fake = None
-        pred_real = None
-        loss_D_fake = 0
-        loss_D_real = 0
-        # Fake
-        # stop backprop to the generator by detaching fake_B
-        # Generated Image Disc Output should be close to zero
-
-        if fakeB is not None:
-            pred_fake = net.forward(fakeB.detach())
-            loss_D_fake = self.criterionGAN(pred_fake, 0)
-
-        # Real
-        if realB is not None:
-            pred_real = net.forward(realB)
-            loss_D_real = self.criterionGAN(pred_real, 1)
-
-        # Combined loss
-        loss_D = (loss_D_fake + loss_D_real) * 0.5
-        return loss_D, pred_fake, pred_real
-
-
-class DiscLossR(DiscLoss):
-    # RSGAN from
-    # https://arxiv.org/abs/1807.00734
-    def name(self):
-        return 'RSGAN'
-
-    def initialize(self, opt, tensor):
-        DiscLoss.initialize(self, opt, tensor)
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB, pred_real=None):
-        if pred_real is None:
-            pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB)
-        return self.criterionGAN(pred_fake - pred_real, 1)
-
-    def get_loss(self, net, realA, fakeB, realB):
-        pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB.detach())
-
-        loss_D = self.criterionGAN(pred_real - pred_fake, 1)  # BCE_stable loss
-        return loss_D, pred_fake, pred_real
-
-
-class DiscLossRa(DiscLoss):
-    # RaSGAN from
-    # https://arxiv.org/abs/1807.00734
-    def name(self):
-        return 'RaSGAN'
-
-    def initialize(self, opt, tensor):
-        DiscLoss.initialize(self, opt, tensor)
-        self.criterionGAN = GANLoss(use_l1=False, tensor=tensor)
-
-    def get_g_loss(self, net, realA, fakeB, realB, pred_real=None):
-        if pred_real is None:
-            pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB)
-
-        loss_G = self.criterionGAN(pred_real - torch.mean(pred_fake, dim=0, keepdim=True), 0)
-        loss_G += self.criterionGAN(pred_fake - torch.mean(pred_real, dim=0, keepdim=True), 1)
-        return loss_G * 0.5
-
-    def get_loss(self, net, realA, fakeB, realB):
-        pred_real = net.forward(realB)
-        pred_fake = net.forward(fakeB.detach())
-
-        loss_D = self.criterionGAN(pred_real - torch.mean(pred_fake, dim=0, keepdim=True), 1)
-        loss_D += self.criterionGAN(pred_fake - torch.mean(pred_real, dim=0, keepdim=True), 0)
-        return loss_D * 0.5, pred_fake, pred_real
-
-
-class MS_SSIM_Loss(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.ms_ssim = MS_SSIM(data_range=1, size_average=True, channel=3)
-
-    def forward(self, output, target):
-        return 1 - self.ms_ssim(output, target)
-
-
-class SSIM_Loss(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.ssim = SSIM(data_range=1, size_average=True, channel=3)
-
-    def forward(self, output, target):
-        return 1 - self.ssim(output, target)
-
-
-def init_loss(opt, tensor):
-    disc_loss = None
-    content_loss = None
-
-    loss_dic = {}
-
-    t_pixel_loss = ContentLoss()
-    t_pixel_loss.initialize(
-        MultipleLoss([nn.MSELoss(), MS_SSIM_Loss(), GradientLoss()], [1.0, 0.4, 0.6]))
-
-    loss_dic['t_pixel'] = t_pixel_loss
-
-    r_pixel_loss = ContentLoss()
-    r_pixel_loss.initialize(
-        MultipleLoss([nn.MSELoss()], [4.0]))
-
-    loss_dic['r_pixel'] = r_pixel_loss
-    loss_dic['recons'] = ReconsLoss()
-
-    loss_dic['mtv'] = MTVLoss()
-    loss_dic['ktv'] = KTVLoss()
-
-    if opt.lambda_gan > 0:
-        if opt.gan_type == 'sgan' or opt.gan_type == 'gan':
-            disc_loss = DiscLoss()
-        elif opt.gan_type == 'rsgan':
-            disc_loss = DiscLossR()
-        elif opt.gan_type == 'rasgan':
-            disc_loss = DiscLossRa()
-        else:
-            raise ValueError("GAN [%s] not recognized." % opt.gan_type)
-
-        disc_loss.initialize(opt, tensor)
-        loss_dic['gan'] = disc_loss
-
-    return loss_dic
-
-
-if __name__ == '__main__':
-    x = torch.randn(3, 32, 224, 224).cuda()
-    import time
-
-    s = time.time()
-    out1, out2 = gradient_norm_kernel(x)
-    t = time.time()
-    print(t - s)
-    print(out1.shape, out2.shape)

models/networks.py

@@ -1,335 +0,0 @@
-import functools
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import init
-from torch.nn.utils import spectral_norm
-from torch.nn import functional as F
-###############################################################################
-# Functions
-###############################################################################
-
-
-def weights_init_normal(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if isinstance(m, nn.Sequential):
-        return
-    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
-        init.normal_(m.weight.data, 0.0, 0.02)
-    elif isinstance(m, nn.Linear):
-        init.normal_(m.weight.data, 0.0, 0.02)
-    elif isinstance(m, nn.BatchNorm2d):
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_xavier(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
-        init.xavier_normal_(m.weight.data, gain=0.02)
-    elif isinstance(m, nn.Linear):
-        init.xavier_normal_(m.weight.data, gain=0.02)
-    elif isinstance(m, nn.BatchNorm2d):
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_kaiming(m):
-    classname = m.__class__.__name__
-    # print(classname)
-    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
-        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
-    elif isinstance(m, nn.Linear):
-        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
-    elif isinstance(m, nn.BatchNorm2d):
-        init.normal_(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def weights_init_orthogonal(m):
-    classname = m.__class__.__name__
-    print(classname)
-    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
-        init.orthogonal(m.weight.data, gain=1)
-    elif isinstance(m, nn.Linear):
-        init.orthogonal(m.weight.data, gain=1)
-    elif isinstance(m, nn.BatchNorm2d):
-        init.normal(m.weight.data, 1.0, 0.02)
-        init.constant_(m.bias.data, 0.0)
-
-
-def init_weights(net, init_type='normal'):
-    print('[i] initialization method [%s]' % init_type)
-    if init_type == 'normal':
-        net.apply(weights_init_normal)
-    elif init_type == 'xavier':
-        net.apply(weights_init_xavier)
-    elif init_type == 'kaiming':
-        net.apply(weights_init_kaiming)
-    elif init_type == 'orthogonal':
-        net.apply(weights_init_orthogonal)
-    elif init_type == 'edsr':
-        pass
-    else:
-        raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
-
-
-def get_norm_layer(norm_type='instance'):
-    if norm_type == 'batch':
-        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
-    elif norm_type == 'instance':
-        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
-    elif norm_type == 'none':
-        norm_layer = None
-    else:
-        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
-    return norm_layer
-
-
-def define_D(opt, in_channels=3):
-    # use_sigmoid = opt.gan_type == 'gan'
-    use_sigmoid = False # incorporate sigmoid into BCE_stable loss
-
-    if opt.which_model_D == 'disc_vgg':
-        netD = Discriminator_VGG(in_channels, use_sigmoid=use_sigmoid)
-        init_weights(netD, init_type='kaiming')
-    elif opt.which_model_D == 'disc_patch':
-        netD = NLayerDiscriminator(in_channels, 64, 3, nn.InstanceNorm2d, use_sigmoid, getIntermFeat=False)
-        init_weights(netD, init_type='normal')
-    elif opt.which_model_D == 'disc_unet':
-        netD = UNetDiscriminatorSN(in_channels) 
-    else:
-        raise NotImplementedError('%s is not implemented' %opt.which_model_D)
-
-    if len(opt.gpu_ids) > 0:
-        assert(torch.cuda.is_available())
-        netD.cuda(opt.gpu_ids[0])
-    
-    return netD
-
-
-def print_network(net):
-    num_params = 0
-    for param in net.parameters():
-        num_params += param.numel()
-    print(net)
-    print('Total number of parameters: %d' % num_params)
-    print('The size of receptive field: %d' % receptive_field(net))
-
-
-def receptive_field(net):
-    def _f(output_size, ksize, stride, dilation):
-        return (output_size - 1) * stride + ksize * dilation - dilation + 1
-
-    stats = []
-    for m in net.modules():
-        if isinstance(m, nn.Conv2d):
-            stats.append((m.kernel_size, m.stride, m.dilation))
-    
-    rsize = 1
-    for (ksize, stride, dilation) in reversed(stats):
-        if type(ksize) == tuple: ksize = ksize[0]
-        if type(stride) == tuple: stride = stride[0]
-        if type(dilation) == tuple: dilation = dilation[0]
-        rsize = _f(rsize, ksize, stride, dilation)
-    return rsize
-
-
-def debug_network(net):
-    def _hook(m, i, o):
-        print(o.size())
-    for m in net.modules():
-        m.register_forward_hook(_hook)
-
-
-##############################################################################
-# Classes
-##############################################################################
-
-# Defines the PatchGAN discriminator with the specified arguments.
-class NLayerDiscriminator(nn.Module):
-    def __init__(self, input_nc, ndf=64, n_layers=3, 
-    norm_layer=nn.BatchNorm2d, use_sigmoid=False, 
-    branch=1, bias=True, getIntermFeat=False):
-        super(NLayerDiscriminator, self).__init__()
-        self.getIntermFeat = getIntermFeat
-        self.n_layers = n_layers
-        kw = 4
-        padw = int(np.ceil((kw-1.0)/2))
-        sequence = [[nn.Conv2d(input_nc*branch, ndf*branch, kernel_size=kw, stride=2, padding=padw, groups=branch, bias=True), nn.LeakyReLU(0.2, True)]]
-
-        nf = ndf
-        for n in range(1, n_layers):
-            nf_prev = nf
-            nf = min(nf * 2, 512)
-            sequence += [[
-                nn.Conv2d(nf_prev*branch, nf*branch, groups=branch, kernel_size=kw, stride=2, padding=padw, bias=bias),
-                norm_layer(nf*branch), nn.LeakyReLU(0.2, True)
-            ]]
-
-        nf_prev = nf
-        nf = min(nf * 2, 512)
-        sequence += [[
-            nn.Conv2d(nf_prev*branch, nf*branch, groups=branch, kernel_size=kw, stride=1, padding=padw, bias=bias),
-            norm_layer(nf*branch),
-            nn.LeakyReLU(0.2, True)
-        ]]
-
-        sequence += [[nn.Conv2d(nf*branch, 1*branch, groups=branch, kernel_size=kw, stride=1, padding=padw, bias=True)]]
-
-        if use_sigmoid:
-            sequence += [[nn.Sigmoid()]]
-
-        if getIntermFeat:
-            for n in range(len(sequence)):
-                setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
-        else:
-            sequence_stream = []
-            for n in range(len(sequence)):
-                sequence_stream += sequence[n]
-            self.model = nn.Sequential(*sequence_stream)
-
-    def forward(self, input):
-        if self.getIntermFeat:
-            res = [input]
-            for n in range(self.n_layers+2):
-                model = getattr(self, 'model'+str(n))
-                res.append(model(res[-1]))
-            return res[1:]
-        else:
-            return self.model(input)
-
-
-class Discriminator_VGG(nn.Module):
-    def __init__(self, in_channels=3, use_sigmoid=True):
-        super(Discriminator_VGG, self).__init__()
-        def conv(*args, **kwargs):
-            return nn.Conv2d(*args, **kwargs)
-
-        num_groups = 32
-
-        body = [
-            conv(in_channels, 64, kernel_size=3, padding=1), # 224
-            nn.LeakyReLU(0.2),
-
-            conv(64, 64, kernel_size=3, stride=2, padding=1), # 112
-            nn.GroupNorm(num_groups, 64),
-            nn.LeakyReLU(0.2),
-
-            conv(64, 128, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups, 128),
-            nn.LeakyReLU(0.2),
-
-            conv(128, 128, kernel_size=3, stride=2, padding=1), # 56
-            nn.GroupNorm(num_groups, 128),
-            nn.LeakyReLU(0.2),
-
-            conv(128, 256, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups, 256),
-            nn.LeakyReLU(0.2),
-
-            conv(256, 256, kernel_size=3, stride=2, padding=1), # 28
-            nn.GroupNorm(num_groups, 256),
-            nn.LeakyReLU(0.2),
-
-            conv(256, 512, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups, 512),
-            nn.LeakyReLU(0.2),
-
-            conv(512, 512, kernel_size=3, stride=2, padding=1), # 14
-            nn.GroupNorm(num_groups, 512),
-            nn.LeakyReLU(0.2),
-
-            conv(512, 512, kernel_size=3, stride=1, padding=1),
-            nn.GroupNorm(num_groups, 512),
-            nn.LeakyReLU(0.2),
-
-            conv(512, 512, kernel_size=3, stride=2, padding=1), # 7
-            nn.GroupNorm(num_groups, 512),
-            nn.LeakyReLU(0.2),
-        ]
-
-        tail = [
-            nn.AdaptiveAvgPool2d(1),
-            nn.Conv2d(512, 1024, kernel_size=1),
-            nn.LeakyReLU(0.2),
-            nn.Conv2d(1024, 1, kernel_size=1)
-        ]
-
-        if use_sigmoid:
-            tail.append(nn.Sigmoid())
-        
-        self.body = nn.Sequential(*body)
-        self.tail = nn.Sequential(*tail)
-
-    def forward(self, x):
-        x = self.body(x)
-        out = self.tail(x)
-        return out
-
-class UNetDiscriminatorSN(nn.Module):
-    """Defines a U-Net discriminator with spectral normalization (SN)
-
-    It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
-
-    Arg:
-        num_in_ch (int): Channel number of inputs. Default: 3.
-        num_feat (int): Channel number of base intermediate features. Default: 64.
-        skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
-    """
-
-    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
-        super(UNetDiscriminatorSN, self).__init__()
-        self.skip_connection = skip_connection
-        norm = spectral_norm
-        # the first convolution
-        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
-        # downsample
-        self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
-        self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
-        self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
-        # upsample
-        self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
-        self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
-        self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
-        # extra convolutions
-        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
-        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
-        self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)
-
-    def forward(self, x, illu = None):
-        # downsample
-        ingress = self.conv0(x) 
-        if illu is not None : ingress = ingress * (1 - illu * 2)
-        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
-        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
-        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
-        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
-
-        # upsample
-        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
-        x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
-
-        if self.skip_connection:
-            x4 = x4 + x2
-        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
-        x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
-
-        if self.skip_connection:
-            x5 = x5 + x1
-        x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)
-        x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
-
-        if self.skip_connection:
-            x6 = x6 + x0
-
-        # extra convolutions
-        out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
-        out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
-        out = self.conv9(out)
-        
-        # print(out.shape, 'real_esrgan out shape')
-        return out #if illu is None else out * illu

models/vgg.py

@@ -1,66 +0,0 @@
-from collections import namedtuple
-
-import torch
-from torchvision import models
-
-
-class Vgg16(torch.nn.Module):
-    def __init__(self, requires_grad=False):
-        super(Vgg16, self).__init__()
-        vgg_pretrained_features = models.vgg16(pretrained=True).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        for x in range(4):
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(4, 9):
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(9, 16):
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(16, 23):
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1_2 = h
-        h = self.slice2(h)
-        h_relu2_2 = h
-        h = self.slice3(h)
-        h_relu3_3 = h
-        h = self.slice4(h)
-        h_relu4_3 = h
-        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])
-        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)
-        return out
-
-
-class Vgg19(torch.nn.Module):
-    def __init__(self, requires_grad=False):
-        super(Vgg19, self).__init__()
-        self.vgg_pretrained_features = models.vgg19(pretrained=True).features
-      
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X, indices=None):
-        if indices is None:
-            indices = [2, 7, 12, 21, 30]
-        out = []
-        for i in range(indices[-1]):
-            X = self.vgg_pretrained_features[i](X)
-            if (i + 1) in indices:
-                out.append(X)
-
-        return out
-
-
-if __name__ == '__main__':
-    vgg = Vgg19()
-    import ipdb
-
-    ipdb.set_trace()

models/vit_feature_extractor.py

@@ -1,164 +0,0 @@
-import torch
-
-
-def attn_cosine_sim(x, eps=1e-08):
-    assert x.shape[0] == 1, 'x.shape[0] must eqs 1'
-    x = x[0]  # TEMP: getting rid of redundant dimension, TBF
-    norm1 = x.norm(dim=2, keepdim=True)
-    factor = torch.clamp(norm1 @ norm1.permute(0, 2, 1), min=eps)
-    sim_matrix = (x @ x.permute(0, 2, 1)) / factor
-    return sim_matrix
-
-
-class VitExtractor:
-    BLOCK_KEY = 'block'
-    ATTN_KEY = 'attn'
-    PATCH_IMD_KEY = 'patch_imd'
-    QKV_KEY = 'qkv'
-    KEY_LIST = [BLOCK_KEY, ATTN_KEY, PATCH_IMD_KEY, QKV_KEY]
-
-    def __init__(self, model_name, device):
-        self.model = torch.hub.load('facebookresearch/dino:main', model_name).to(device)
-        self.model.eval()
-        self.model_name = model_name
-        self.hook_handlers = []
-        self.layers_dict = {}
-        self.outputs_dict = {}
-        for key in VitExtractor.KEY_LIST:
-            self.layers_dict[key] = []
-            self.outputs_dict[key] = []
-        self._init_hooks_data()
-
-    def _init_hooks_data(self):
-        self.layers_dict[VitExtractor.BLOCK_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
-        self.layers_dict[VitExtractor.ATTN_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
-        self.layers_dict[VitExtractor.QKV_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
-        self.layers_dict[VitExtractor.PATCH_IMD_KEY] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
-        for key in VitExtractor.KEY_LIST:
-            # self.layers_dict[key] = kwargs[key] if key in kwargs.keys() else []
-            self.outputs_dict[key] = []
-
-    def _register_hooks(self, **kwargs):
-        for block_idx, block in enumerate(self.model.blocks):
-            if block_idx in self.layers_dict[VitExtractor.BLOCK_KEY]:
-                self.hook_handlers.append(block.register_forward_hook(self._get_block_hook()))
-            if block_idx in self.layers_dict[VitExtractor.ATTN_KEY]:
-                self.hook_handlers.append(block.attn.attn_drop.register_forward_hook(self._get_attn_hook()))
-            if block_idx in self.layers_dict[VitExtractor.QKV_KEY]:
-                self.hook_handlers.append(block.attn.qkv.register_forward_hook(self._get_qkv_hook()))
-            if block_idx in self.layers_dict[VitExtractor.PATCH_IMD_KEY]:
-                self.hook_handlers.append(block.attn.register_forward_hook(self._get_patch_imd_hook()))
-
-    def _clear_hooks(self):
-        for handler in self.hook_handlers:
-            handler.remove()
-        self.hook_handlers = []
-
-    def _get_block_hook(self):
-        def _get_block_output(model, input, output):
-            self.outputs_dict[VitExtractor.BLOCK_KEY].append(output)
-
-        return _get_block_output
-
-    def _get_attn_hook(self):
-        def _get_attn_output(model, inp, output):
-            self.outputs_dict[VitExtractor.ATTN_KEY].append(output)
-
-        return _get_attn_output
-
-    def _get_qkv_hook(self):
-        def _get_qkv_output(model, inp, output):
-            self.outputs_dict[VitExtractor.QKV_KEY].append(output)
-
-        return _get_qkv_output
-
-    # TODO: CHECK ATTN OUTPUT TUPLE
-    def _get_patch_imd_hook(self):
-        def _get_attn_output(model, inp, output):
-            self.outputs_dict[VitExtractor.PATCH_IMD_KEY].append(output[0])
-
-        return _get_attn_output
-
-    def get_feature_from_input(self, input_img):  # List([B, N, D])
-        self._register_hooks()
-        self.model(input_img)
-        feature = self.outputs_dict[VitExtractor.BLOCK_KEY]
-        self._clear_hooks()
-        self._init_hooks_data()
-        return feature
-
-    def get_qkv_feature_from_input(self, input_img):
-        self._register_hooks()
-        self.model(input_img)
-        feature = self.outputs_dict[VitExtractor.QKV_KEY]
-        self._clear_hooks()
-        self._init_hooks_data()
-        return feature
-
-    def get_attn_feature_from_input(self, input_img):
-        self._register_hooks()
-        self.model(input_img)
-        feature = self.outputs_dict[VitExtractor.ATTN_KEY]
-        self._clear_hooks()
-        self._init_hooks_data()
-        return feature
-
-    def get_patch_size(self):
-        return 8 if "8" in self.model_name else 16
-
-    def get_width_patch_num(self, input_img_shape):
-        b, c, h, w = input_img_shape
-        patch_size = self.get_patch_size()
-        return w // patch_size
-
-    def get_height_patch_num(self, input_img_shape):
-        b, c, h, w = input_img_shape
-        patch_size = self.get_patch_size()
-        return h // patch_size
-
-    def get_patch_num(self, input_img_shape):
-        patch_num = 1 + (self.get_height_patch_num(input_img_shape) * self.get_width_patch_num(input_img_shape))
-        return patch_num
-
-    def get_head_num(self):
-        if "dino" in self.model_name:
-            return 6 if "s" in self.model_name else 12
-        return 6 if "small" in self.model_name else 12
-
-    def get_embedding_dim(self):
-        if "dino" in self.model_name:
-            return 384 if "s" in self.model_name else 768
-        return 384 if "small" in self.model_name else 768
-
-    def get_queries_from_qkv(self, qkv, input_img_shape):
-        patch_num = self.get_patch_num(input_img_shape)
-        head_num = self.get_head_num()
-        embedding_dim = self.get_embedding_dim()
-        q = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[0]
-        return q
-
-    def get_keys_from_qkv(self, qkv, input_img_shape):
-        patch_num = self.get_patch_num(input_img_shape)
-        head_num = self.get_head_num()
-        embedding_dim = self.get_embedding_dim()
-        k = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[1]
-        return k
-
-    def get_values_from_qkv(self, qkv, input_img_shape):
-        patch_num = self.get_patch_num(input_img_shape)
-        head_num = self.get_head_num()
-        embedding_dim = self.get_embedding_dim()
-        v = qkv.reshape(patch_num, 3, head_num, embedding_dim // head_num).permute(1, 2, 0, 3)[2]
-        return v
-
-    def get_keys_from_input(self, input_img, layer_num):
-        qkv_features = self.get_qkv_feature_from_input(input_img)[layer_num]
-        keys = self.get_keys_from_qkv(qkv_features, input_img.shape)
-        return keys
-
-    def get_keys_self_sim_from_input(self, input_img, layer_num):
-        keys = self.get_keys_from_input(input_img, layer_num=layer_num)
-        h, t, d = keys.shape
-        concatenated_keys = keys.transpose(0, 1).reshape(t, h * d)
-        ssim_map = attn_cosine_sim(concatenated_keys[None, None, ...])
-        return ssim_map

options/base_option.py

@@ -1,47 +0,0 @@
-import argparse
-import models
-
-model_names = sorted(name for name in models.__dict__
-                     if name.islower() and not name.startswith("__")
-                     and callable(models.__dict__[name]))
-
-
-class BaseOptions():
-    def __init__(self):
-        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-        self.initialized = False
-
-    def initialize(self):
-        # experiment specifics
-        self.parser.add_argument('--name', type=str, default='ytmt_ucs_sirs',
-                                 help='name of the experiment. It decides where to store samples and models')
-        self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
-        self.parser.add_argument('--model', type=str, default='revcol', help='chooses which model to use.')
-        self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
-        self.parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
-        self.parser.add_argument('--resume_epoch', '-re', type=int, default=None,
-                                 help='checkpoint to use. (default: latest')
-        self.parser.add_argument('--seed', type=int, default=2018, help='random seed to use. Default=2018')
-        self.parser.add_argument('--supp_eval', action='store_true', help='supplementary evaluation')
-        self.parser.add_argument('--start_now', action='store_true', help='supplementary evaluation')
-        self.parser.add_argument('--testr', action='store_true', help='test for reflections')
-        self.parser.add_argument('--select', type=str, default=None)
-
-        # for setting input
-        self.parser.add_argument('--serial_batches', action='store_true',
-                                 help='if true, takes images in order to make batches, otherwise takes them randomly')
-        self.parser.add_argument('--nThreads', default=8, type=int, help='# threads for loading data')
-        self.parser.add_argument('--max_dataset_size', type=int, default=None,
-                                 help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
-
-        # for display
-        self.parser.add_argument('--no-log', action='store_true', help='disable tf logger?')
-        self.parser.add_argument('--no-verbose', action='store_true', help='disable verbose info?')
-        self.parser.add_argument('--display_winsize', type=int, default=256, help='display window size')
-        self.parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
-        self.parser.add_argument('--display_id', type=int, default=0,
-                                 help='window id of the web display (use 0 to disable visdom)')
-        self.parser.add_argument('--display_single_pane_ncols', type=int, default=0,
-                                 help='if positive, display all images in a single visdom web panel with certain number of images per row.')
-
-        self.initialized = True

options/net_options/base_options.py

@@ -1,71 +0,0 @@
-from options.base_option import BaseOptions as Base
-from util import util
-import os
-import torch
-import numpy as np
-import random
-
-class BaseOptions(Base):
-    def initialize(self):
-        Base.initialize(self)
-        # experiment specifics
-        self.parser.add_argument('--inet', type=str, default='ytmt_ucs', help='chooses which architecture to use for inet.')
-        self.parser.add_argument('--icnn_path', type=str, default=None, help='icnn checkpoint to use.')
-        self.parser.add_argument('--init_type', type=str, default='edsr', help='network initialization [normal|xavier|kaiming|orthogonal|uniform]')
-        # for network
-        self.parser.add_argument('--hyper', action='store_true', help='if true, augment input with vgg hypercolumn feature')
-        
-        self.initialized = True
-
-    def parse(self):
-        if not self.initialized:
-            self.initialize()
-        self.opt = self.parser.parse_args()
-        self.opt.isTrain = self.isTrain   # train or test
-
-        torch.backends.cudnn.deterministic = True
-        torch.manual_seed(self.opt.seed)
-        np.random.seed(self.opt.seed) # seed for every module
-        random.seed(self.opt.seed)
-
-        str_ids = self.opt.gpu_ids.split(',')
-        self.opt.gpu_ids = []
-        for str_id in str_ids:
-            id = int(str_id)
-            if id >= 0:
-                self.opt.gpu_ids.append(id)
-
-        # set gpu ids
-        if len(self.opt.gpu_ids) > 0:
-            torch.cuda.set_device(self.opt.gpu_ids[0])
-
-        args = vars(self.opt)
-
-        print('------------ Options -------------')
-        for k, v in sorted(args.items()):
-            print('%s: %s' % (str(k), str(v)))
-        print('-------------- End ----------------')
-
-        # save to the disk
-        self.opt.name = self.opt.name or '_'.join([self.opt.model])
-        expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
-        util.mkdirs(expr_dir)
-        file_name = os.path.join(expr_dir, 'opt.txt')
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write('------------ Options -------------\n')
-            for k, v in sorted(args.items()):
-                opt_file.write('%s: %s\n' % (str(k), str(v)))
-            opt_file.write('-------------- End ----------------\n')
-
-        if self.opt.debug:
-            self.opt.display_freq = 20
-            self.opt.print_freq = 20
-            self.opt.nEpochs = 40
-            self.opt.max_dataset_size = 100
-            self.opt.no_log = False
-            self.opt.nThreads = 0
-            self.opt.decay_iter = 0
-            self.opt.serial_batches = True
-            self.opt.no_flip = True
-        
-        return self.opt

options/net_options/train_options.py

@@ -1,75 +0,0 @@
-from .base_options import BaseOptions
-
-
-class TrainOptions(BaseOptions):
-    def initialize(self):
-        BaseOptions.initialize(self)
-        # for displays
-        self.parser.add_argument('--display_freq', type=int, default=100,
-                                 help='frequency of showing training results on screen')
-        self.parser.add_argument('--update_html_freq', type=int, default=1000,
-                                 help='frequency of saving training results to html')
-        self.parser.add_argument('--print_freq', type=int, default=100,
-                                 help='frequency of showing training results on console')
-        self.parser.add_argument('--eval_freq', type=int, default=1, help='frequency of evaluation')
-        self.parser.add_argument('--save_freq', type=int, default=1, help='frequency of save eval samples')
-        self.parser.add_argument('--no_html', action='store_true',
-                                 help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')
-        self.parser.add_argument('--save_epoch_freq', type=int, default=1,
-                                 help='frequency of saving checkpoints at the end of epochs')
-        self.parser.add_argument('--debug', action='store_true',
-                                 help='only do one epoch and displays at each iteration')
-        self.parser.add_argument('--finetune', action='store_true',
-                                 help='finetune the network using identity inputs and outputs')
-        self.parser.add_argument('--if_align', action='store_true',
-                                 help='if align 4x')
-
-        # self.parser.add_argument('--graph', action='store_true',
-        #                          help='print computation graph')
-        # for training (Note: in train_sirs.py, we mannually tune the training protocol, but you can also use following setting by modifying the code in errnet_model.py)
-        self.parser.add_argument('--nEpochs', '-n', type=int, default=60, help='# of epochs to run')
-        self.parser.add_argument('--lr', type=float, default=1e-4, help='initial learning rate for adam')
-        self.parser.add_argument('--wd', type=float, default=0, help='weight decay for adam')
-
-        self.parser.add_argument('--r_pixel_weight', '-rw', type=float, default=1.0, help='weight for r_pixel loss')
-
-        self.parser.add_argument('--low_sigma', type=float, default=2, help='min sigma in synthetic dataset')
-        self.parser.add_argument('--high_sigma', type=float, default=5, help='max sigma in synthetic dataset')
-        self.parser.add_argument('--low_gamma', type=float, default=1.3, help='max gamma in synthetic dataset')
-        self.parser.add_argument('--high_gamma', type=float, default=1.3, help='max gamma in synthetic dataset')
-
-        # data augmentation
-        self.parser.add_argument('--real20_size', type=int, default=420, help='scale images to compat size')
-        self.parser.add_argument('--batchSize', '-b', type=int, default=2, help='input batch size')
-        self.parser.add_argument('--loadSize', type=str, default='224,336,448', help='scale images to multiple size')
-        self.parser.add_argument('--fineSize', type=str, default='224,224', help='then crop to this size')
-        self.parser.add_argument('--no_flip', action='store_true',
-                                 help='if specified, do not flip the images for data augmentation')
-        self.parser.add_argument('--resize_or_crop', type=str, default='resize_and_crop',
-                                 help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]')
-        self.parser.add_argument('--debug_eval', action='store_true',
-                                 help='if specified, do not flip the images for data augmentation')
-        self.parser.add_argument('--graph', action='store_true', help='print graph')
-
-        # for discriminator
-        self.parser.add_argument('--which_model_D', type=str, default='disc_vgg', choices=['disc_vgg', 'disc_patch'])
-        self.parser.add_argument('--gan_type', type=str, default='rasgan',
-                                 help='gan/sgan : Vanilla GAN; rasgan : relativistic gan')
-        # loss weight
-        self.parser.add_argument('--unaligned_loss', type=str, default='vgg',
-                                 help='learning rate policy: vgg|mse|ctx|ctx_vgg')
-        self.parser.add_argument('--tv_type', type=str, default=None, choices=['ktv', 'mtv'])
-        self.parser.add_argument('--vgg_layer', type=int, default=31, help='vgg layer of unaligned loss')
-        self.parser.add_argument('--init_lr', type=float, default=1e-2, help='initial learning rate')
-        self.parser.add_argument('--fixed_lr', type=float, default=0, help='initial learning rate')
-        self.parser.add_argument('--lambda_gan', type=float, default=0.01, help='weight for gan loss')
-        self.parser.add_argument('--lambda_vgg', type=float, default=0.1, help='weight for vgg loss')
-        self.parser.add_argument('--weight_loss',type=float,default=0.25,help='weight fot overall loss')
-        self.parser.add_argument('--num_subnet',type=int,default=4,help='num_number of subnet')
-        self.parser.add_argument('--dataset',type=float,default=0.5,help='the setting of dataset')
-        self.parser.add_argument('--loss_col',type=int,default=4,help='numcol for loss')
-        self.parser.add_argument('--drop_path',type=float,default=0.6,help='drop_path')
-
-        
-
-        self.isTrain = True

pretrained/README.md

@@ -1,3 +0,0 @@
-# Pretrained models
-
-This folder is for pretrained models. 

script.py

@@ -1,64 +0,0 @@
-import torch
-
-# Load the original weights file
-original_weights = torch.load('/home/xteam/zhaohao/pycharmproject/YTMT/merge_stem_reg_014_00055524.pt')
-
-# Create a new weights dictionary
-# new_weights = {}
-
-# # Iterate through the original weights dictionary
-# for key, value in original_weights.items():
-#     # Check if the key contains 'projec_shit'
-#     if 'projback_shit' in key:
-#         # Replace 'projec_shit' with 'project_'
-#         new_key = key.replace('projback_shit', 'projback_')
-#         new_weights[new_key] = value
-#     else:
-#         # If the key doesn't contain 'projec_shit', keep it unchanged
-#         new_weights[key] = value
-#     if 'projback_shit_2' in key:
-#         # Replace 'projec_shit' with 'project_'
-#         new_key = key.replace('projback_shit_2', 'projback_2')
-#         new_weights[new_key] = value
-#     else:
-#         # If the key doesn't contain 'projec_shit', keep it unchanged
-#         new_weights[key] = value
-
-# # Save the modified weights
-# torch.save(new_weights, '/home/xteam/zhaohao/pycharmproject/RDNet/new_weights.pth')
-
-# print("Weights file has been updated.")
-
-# # 打印原始权重字典中的所有键,以检查确切的层名称
-# print("原始权重文件中的层名:")
-# for key in original_weights['icnn'].keys():
-#     print(key)
-
-# 创建一个新的权重字典
-new_weights = {'icnn': {}}
-
-# 遍历原始权重字典
-for key, value in original_weights['icnn'].items():
-    # 检查并替换包含 'projback_shit' 的键
-    if 'projback_shit_2' in key:
-        new_key = key.replace('projback_shit_2', 'projback_2')
-        new_weights['icnn'][new_key] = value
-    
-    # 检查并替换包含 'projback_shit_2' 的键
-    elif 'projback_shit' in key:
-        new_key = key.replace('projback_shit', 'projback_')
-        new_weights['icnn'][new_key] = value
-    else:
-        # 如果键不包含上述字符串,保持不变
-        new_weights['icnn'][key] = value
-
-# 打印新的权重字典中的所有键,以验证更改
-print("\n更新后的权重文件中的层名:")
-for key in new_weights['icnn'].keys():
-    print(key)
-
-# 保存修改后的权重
-torch.save(new_weights, '/home/xteam/zhaohao/pycharmproject/RDNet/new_weights_4.pth')
-
-print("\n权重文件已更新。")
-

test_sirs.py

@@ -1,60 +0,0 @@
-import os
-from os.path import join
-
-import torch.backends.cudnn as cudnn
-
-# import data.sirs_dataset as datasets
-import data.dataset_sir as datasets
-from data.image_folder import read_fns
-from engine import Engine
-from options.net_options.train_options import TrainOptions
-from tools import mutils
-
-opt = TrainOptions().parse()
-
-opt.isTrain = False
-cudnn.benchmark = True
-opt.no_log = True
-opt.display_id = 0
-opt.verbose = False
-datadir = os.path.join(os.path.expanduser('~'), '/opt/datasets/sirs')
-
-eval_dataset_real = datasets.DSRTestDataset(join(datadir, f'test/real20_{opt.real20_size}'),
-                                            fns=read_fns('data/real_test.txt'), if_align=opt.if_align)
-eval_dataset_solidobject = datasets.DSRTestDataset(join(datadir, 'test/SIR2/SolidObjectDataset'),
-                                                   if_align=opt.if_align)
-eval_dataset_postcard = datasets.DSRTestDataset(join(datadir, 'test/SIR2/PostcardDataset'), if_align=opt.if_align)
-eval_dataset_wild = datasets.DSRTestDataset(join(datadir, 'test/SIR2/WildSceneDataset'), if_align=opt.if_align)
-
-eval_dataloader_real = datasets.DataLoader(
-    eval_dataset_real, batch_size=1, shuffle=True,
-    num_workers=opt.nThreads, pin_memory=True)
-
-eval_dataloader_solidobject = datasets.DataLoader(
-    eval_dataset_solidobject, batch_size=1, shuffle=False,
-    num_workers=opt.nThreads, pin_memory=True)
-
-eval_dataloader_postcard = datasets.DataLoader(
-    eval_dataset_postcard, batch_size=1, shuffle=False,
-    num_workers=opt.nThreads, pin_memory=True)
-
-eval_dataloader_wild = datasets.DataLoader(
-    eval_dataset_wild, batch_size=1, shuffle=False,
-    num_workers=opt.nThreads, pin_memory=True)
-
-engine = Engine(opt, eval_dataset_real, eval_dataset_solidobject, eval_dataset_postcard, eval_dataloader_wild)
-
-"""Main Loop"""
-result_dir = os.path.join('./results', opt.name, mutils.get_formatted_time())
-
-res1 = engine.eval(eval_dataloader_real, dataset_name='testdata_real',
-                  savedir=join(result_dir, 'real20'), suffix='real20')
-
-res2 = engine.eval(eval_dataloader_solidobject, dataset_name='testdata_solidobject',
-                  savedir=join(result_dir, 'solidobject'), suffix='solidobject')
-res3 = engine.eval(eval_dataloader_postcard, dataset_name='testdata_postcard',
-                  savedir=join(result_dir, 'postcard'), suffix='postcard')
-
-res4 = engine.eval(eval_dataloader_wild, dataset_name='testdata_wild',
-                  savedir=join(result_dir, 'wild'), suffix='wild')
-