'folder'

data/__init__.py

@@ -0,0 +1,43 @@
+'''create dataset and dataloader'''
+import logging
+import torch
+import torch.utils.data
+
+def create_dataloader(dataset, dataset_opt, opt=None, sampler=None):
+    phase = dataset_opt['phase']
+    if phase == 'train':
+        if opt['dist']:
+            world_size = torch.distributed.get_world_size()
+            num_workers = dataset_opt['n_workers']
+            assert dataset_opt['batch_size'] % world_size == 0
+            batch_size = dataset_opt['batch_size'] // world_size
+            shuffle = False
+        else:
+            num_workers = dataset_opt['n_workers'] * len(opt['gpu_ids'])
+            batch_size = dataset_opt['batch_size']
+            shuffle = True
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
+                                           num_workers=num_workers, sampler=sampler, drop_last=True,
+                                           pin_memory=False)
+    else:
+        return torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1,
+                                           pin_memory=True)
+
+
+def create_dataset(dataset_opt):
+    mode = dataset_opt['mode']
+    if mode == 'test':
+        from data.coco_test_dataset import imageTestDataset as D
+    elif mode == 'train':
+        from data.coco_dataset import CoCoDataset as D
+    elif mode == 'td':
+        from data.test_dataset_td import imageTestDataset as D
+    else:
+        raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode))
+    print(mode)
+    dataset = D(dataset_opt)
+
+    logger = logging.getLogger('base')
+    logger.info('Dataset [{:s} - {:s}] is created.'.format(dataset.__class__.__name__,
+                                                           dataset_opt['name']))
+    return dataset

data/coco_dataset.py

@@ -0,0 +1,90 @@
+'''
+Vimeo90K dataset
+support reading images from lmdb, image folder and memcached
+'''
+import logging
+import os
+import os.path as osp
+import pickle
+import random
+
+import cv2
+import lmdb
+import numpy as np
+import torch
+import torch.utils.data as data
+
+import data.util as util
+
+try:
+    import mc
+except ImportError:
+    pass
+logger = logging.getLogger('base')
+
+class CoCoDataset(data.Dataset):
+    def __init__(self, opt):
+        super(CoCoDataset, self).__init__()
+        self.opt = opt
+        # get train indexes
+        self.data_path = self.opt['data_path']
+        self.txt_path = self.opt['txt_path']
+        with open(self.txt_path) as f:
+            self.list_image = f.readlines()
+        self.list_image = [line.strip('\n') for line in self.list_image]
+        # temporal augmentation
+        self.interval_list = opt['interval_list']
+        self.random_reverse = opt['random_reverse']
+        logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
+            ','.join(str(x) for x in opt['interval_list']), self.random_reverse))
+        self.data_type = self.opt['data_type']
+        random.shuffle(self.list_image)
+        self.LR_input = True
+        self.num_image = self.opt['num_image']
+
+    def _ensure_memcached(self):
+        if self.mclient is None:
+            # specify the config files
+            server_list_config_file = None
+            client_config_file = None
+            self.mclient = mc.MemcachedClient.GetInstance(server_list_config_file,
+                                                          client_config_file)
+
+    def __getitem__(self, index):
+        GT_size = self.opt['GT_size']
+        image_name = self.list_image[index]
+        path_frame = os.path.join(self.data_path, image_name)
+        img_GT = util.read_img(None, osp.join(path_frame, path_frame))
+        index_h = random.randint(0, len(self.list_image) - 1)
+
+        # random crop
+        H, W, C = img_GT.shape
+        rnd_h = random.randint(0, max(0, H - GT_size))
+        rnd_w = random.randint(0, max(0, W - GT_size))
+        img_frames = img_GT[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_frames = img_frames[:, :, [2, 1, 0]]
+        img_frames = torch.from_numpy(np.ascontiguousarray(np.transpose(img_frames, (2, 0, 1)))).float().unsqueeze(0)
+
+        # process h_list
+        if index_h % 100 == 0:
+            path_frame_h = "../dataset/locwatermark/blue.png"
+        else:
+            image_name_h = self.list_image[index_h]
+            path_frame_h = os.path.join(self.data_path, image_name_h)
+        
+        frame_h = util.read_img(None, osp.join(path_frame_h, path_frame_h))
+        H1, W1, C1 = frame_h.shape
+        rnd_h = random.randint(0, max(0, H1 - GT_size))
+        rnd_w = random.randint(0, max(0, W1 - GT_size))
+        img_frames_h = frame_h[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
+        img_frames_h = img_frames_h[:, :, [2, 1, 0]]
+        img_frames_h = torch.from_numpy(np.ascontiguousarray(np.transpose(img_frames_h, (2, 0, 1)))).float().unsqueeze(0)
+
+        img_frames_h = torch.nn.functional.interpolate(img_frames_h, size=(512, 512), mode='nearest', align_corners=None).unsqueeze(0)
+        img_frames = torch.nn.functional.interpolate(img_frames, size=(512, 512), mode='nearest', align_corners=None)
+
+        return {'GT': img_frames, 'LQ': img_frames_h}
+
+    def __len__(self):
+        return len(self.list_image)

data/coco_test_dataset.py

@@ -0,0 +1,61 @@
+import os
+import os.path as osp
+import torch
+import torch.utils.data as data
+import data.util as util
+
+import random
+import numpy as np
+from PIL import Image
+
+class imageTestDataset(data.Dataset):
+
+    def __init__(self, opt):
+        super(imageTestDataset, self).__init__()
+        self.opt = opt
+        self.half_N_frames = opt['N_frames'] // 2
+        self.data_path = opt['data_path']
+        self.bit_path = opt['bit_path']
+        self.txt_path = self.opt['txt_path']
+        self.num_image = self.opt['num_image']
+        with open(self.txt_path) as f:
+            self.list_image = f.readlines()
+        self.list_image = [line.strip('\n') for line in self.list_image]
+        self.list_image.sort()
+        self.list_image = self.list_image
+        l = len(self.list_image) // (self.num_image + 1)
+        self.image_list_gt = self.list_image
+
+    def __getitem__(self, index):
+        path_GT = self.image_list_gt[index]
+        
+        img_GT = util.read_img(None, osp.join(self.data_path, path_GT))
+        img_GT = img_GT[:, :, [2, 1, 0]]
+        img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float().unsqueeze(0)
+        img_GT = torch.nn.functional.interpolate(img_GT, size=(512, 512), mode='nearest', align_corners=None)
+
+        T, C, W, H = img_GT.shape
+        list_h = []
+        R = 0
+        G = 0
+        B = 255
+        image = Image.new('RGB', (W, H), (R, G, B))
+        result = np.array(image) / 255.
+        expanded_matrix = np.expand_dims(result, axis=0) 
+        expanded_matrix = np.repeat(expanded_matrix, T, axis=0)
+        imgs_LQ = torch.from_numpy(np.ascontiguousarray(expanded_matrix)).float()
+        imgs_LQ = imgs_LQ.permute(0, 3, 1, 2)
+
+        imgs_LQ = torch.nn.functional.interpolate(imgs_LQ, size=(W, H), mode='nearest', align_corners=None)
+
+        list_h.append(imgs_LQ)
+
+        list_h = torch.stack(list_h, dim=0)
+
+        return {
+                'LQ': list_h,
+                'GT': img_GT
+            }
+    
+    def __len__(self):
+        return len(self.image_list_gt)  

data/data_sampler.py

@@ -0,0 +1,65 @@
+"""
+Modified from torch.utils.data.distributed.DistributedSampler
+Support enlarging the dataset for *iter-oriented* training, for saving time when restart the
+dataloader after each epoch
+"""
+import math
+import torch
+from torch.utils.data.sampler import Sampler
+import torch.distributed as dist
+
+
+class DistIterSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+
+    .. note::
+        Dataset is assumed to be of constant size.
+
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self, dataset, num_replicas=None, rank=None, ratio=100):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        indices = torch.randperm(self.total_size, generator=g).tolist()
+
+        dsize = len(self.dataset)
+        indices = [v % dsize for v in indices]
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

data/test_dataset_td.py

@@ -0,0 +1,63 @@
+import os
+import os.path as osp
+import torch
+import torch.utils.data as data
+import data.util as util
+
+import random
+import numpy as np
+from PIL import Image
+
+class imageTestDataset(data.Dataset):
+
+    def __init__(self, opt):
+        super(imageTestDataset, self).__init__()
+        self.opt = opt
+        self.half_N_frames = opt['N_frames'] // 2
+        self.data_path = opt['data_path']
+        self.bit_path = opt['bit_path']
+        self.txt_path = self.opt['txt_path']
+        self.num_image = self.opt['num_image']
+        with open(self.txt_path) as f:
+            self.list_image = f.readlines()
+        self.list_image = [line.strip('\n') for line in self.list_image]
+        # self.list_image = sorted(self.list_image)
+        l = len(self.list_image) // (self.num_image + 1)
+        self.image_list_gt = self.list_image
+        self.image_list_bit = self.list_image
+
+
+    def __getitem__(self, index):
+        path_GT = self.image_list_gt[index]
+        
+        img_GT = util.read_img(None, osp.join(self.data_path, path_GT))
+        img_GT = img_GT[:, :, [2, 1, 0]]
+        img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float().unsqueeze(0)
+        img_GT = torch.nn.functional.interpolate(img_GT, size=(512, 512), mode='nearest', align_corners=None)
+
+        T, C, W, H = img_GT.shape
+        list_h = []
+        R = 0
+        G = 0
+        B = 255
+        image = Image.new('RGB', (W, H), (R, G, B))
+        result = np.array(image) / 255.
+        expanded_matrix = np.expand_dims(result, axis=0) 
+        expanded_matrix = np.repeat(expanded_matrix, T, axis=0)
+        imgs_LQ = torch.from_numpy(np.ascontiguousarray(expanded_matrix)).float()
+        imgs_LQ = imgs_LQ.permute(0, 3, 1, 2)
+
+
+        imgs_LQ = torch.nn.functional.interpolate(imgs_LQ, size=(W, H), mode='nearest', align_corners=None)
+
+        list_h.append(imgs_LQ)
+
+        list_h = torch.stack(list_h, dim=0)
+
+        return {
+                'LQ': list_h,
+                'GT': img_GT
+            }
+    
+    def __len__(self):
+        return len(self.image_list_gt)  

data/util.py

@@ -0,0 +1,551 @@
+import os
+import math
+import pickle
+import random
+import numpy as np
+import glob
+import torch
+import cv2
+
+####################
+# Files & IO
+####################
+
+###################### get image path list ######################
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']
+
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def _get_paths_from_images(path):
+    '''get image path list from image folder'''
+    assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+    images = []
+    for dirpath, _, fnames in sorted(os.walk(path)):
+        for fname in sorted(fnames):
+            if is_image_file(fname):
+                img_path = os.path.join(dirpath, fname)
+                images.append(img_path)
+    assert images, '{:s} has no valid image file'.format(path)
+    return images
+
+
+def _get_paths_from_lmdb(dataroot):
+    '''get image path list from lmdb meta info'''
+    meta_info = pickle.load(open(os.path.join(dataroot, 'meta_info.pkl'), 'rb'))
+    paths = meta_info['keys']
+    sizes = meta_info['resolution']
+    if len(sizes) == 1:
+        sizes = sizes * len(paths)
+    return paths, sizes
+
+
+def get_image_paths(data_type, dataroot):
+    '''get image path list
+    support lmdb or image files'''
+    paths, sizes = None, None
+    if dataroot is not None:
+        if data_type == 'lmdb':
+            paths, sizes = _get_paths_from_lmdb(dataroot)
+        elif data_type == 'img':
+            paths = sorted(_get_paths_from_images(dataroot))
+        else:
+            raise NotImplementedError('data_type [{:s}] is not recognized.'.format(data_type))
+    return paths, sizes
+
+
+def glob_file_list(root):
+    return sorted(glob.glob(os.path.join(root, '*')))
+
+
+###################### read images ######################
+def _read_img_lmdb(env, key, size):
+    '''read image from lmdb with key (w/ and w/o fixed size)
+    size: (C, H, W) tuple'''
+    with env.begin(write=False) as txn:
+        buf = txn.get(key.encode('ascii'))
+    img_flat = np.frombuffer(buf, dtype=np.uint8)
+    C, H, W = size
+    img = img_flat.reshape(H, W, C)
+    return img
+
+
+def read_img(env, path, size=None):
+    '''read image by cv2 or from lmdb
+    return: Numpy float32, HWC, BGR, [0,1]'''
+    if env is None:  # img
+#         print(path)
+        #img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+        img = cv2.imread(path, cv2.IMREAD_COLOR)
+    else:
+        img = _read_img_lmdb(env, path, size)
+#     print(img.shape)
+#     if img is None:
+    # print(path)
+#     print(img.shape)
+    img = img.astype(np.float32) / 255.
+    if img.ndim == 2:
+        img = np.expand_dims(img, axis=2)
+    # some images have 4 channels
+    if img.shape[2] > 3:
+        img = img[:, :, :3]
+    return img
+
+
+def read_img_seq(path):
+    """Read a sequence of images from a given folder path
+    Args:
+        path (list/str): list of image paths/image folder path
+
+    Returns:
+        imgs (Tensor): size (T, C, H, W), RGB, [0, 1]
+    """
+    if type(path) is list:
+        img_path_l = path
+    else:
+        img_path_l = sorted(glob.glob(os.path.join(path, '*.png')))
+#     print(path)
+#     print(path,img_path_l)
+    img_l = [read_img(None, v) for v in img_path_l]
+    # stack to Torch tensor
+    imgs = np.stack(img_l, axis=0)
+    imgs = imgs[:, :, :, [2, 1, 0]]
+    imgs = torch.from_numpy(np.ascontiguousarray(np.transpose(imgs, (0, 3, 1, 2)))).float()
+    return imgs
+
+
+def index_generation(crt_i, max_n, N, padding='reflection'):
+    """Generate an index list for reading N frames from a sequence of images
+    Args:
+        crt_i (int): current center index
+        max_n (int): max number of the sequence of images (calculated from 1)
+        N (int): reading N frames
+        padding (str): padding mode, one of replicate | reflection | new_info | circle
+            Example: crt_i = 0, N = 5
+            replicate: [0, 0, 0, 1, 2]
+            reflection: [2, 1, 0, 1, 2]
+            new_info: [4, 3, 0, 1, 2]
+            circle: [3, 4, 0, 1, 2]
+
+    Returns:
+        return_l (list [int]): a list of indexes
+    """
+    max_n = max_n - 1
+    n_pad = N // 2
+    return_l = []
+
+    for i in range(crt_i - n_pad, crt_i + n_pad + 1):
+        if i < 0:
+            if padding == 'replicate':
+                add_idx = 0
+            elif padding == 'reflection':
+                add_idx = -i
+            elif padding == 'new_info':
+                add_idx = (crt_i + n_pad) + (-i)
+            elif padding == 'circle':
+                add_idx = N + i
+            else:
+                raise ValueError('Wrong padding mode')
+        elif i > max_n:
+            if padding == 'replicate':
+                add_idx = max_n
+            elif padding == 'reflection':
+                add_idx = max_n * 2 - i
+            elif padding == 'new_info':
+                add_idx = (crt_i - n_pad) - (i - max_n)
+            elif padding == 'circle':
+                add_idx = i - N
+            else:
+                raise ValueError('Wrong padding mode')
+        else:
+            add_idx = i
+        return_l.append(add_idx)
+    return return_l
+
+
+####################
+# image processing
+# process on numpy image
+####################
+
+
+def augment(img_list, hflip=True, rot=True):
+    # horizontal flip OR rotate
+    hflip = hflip and random.random() < 0.5
+    vflip = rot and random.random() < 0.5
+    rot90 = rot and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:
+            img = img[:, ::-1, :]
+        if vflip:
+            img = img[::-1, :, :]
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    return [_augment(img) for img in img_list]
+
+
+def augment_flow(img_list, flow_list, hflip=True, rot=True):
+    # horizontal flip OR rotate
+    hflip = hflip and random.random() < 0.5
+    vflip = rot and random.random() < 0.5
+    rot90 = rot and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:
+            img = img[:, ::-1, :]
+        if vflip:
+            img = img[::-1, :, :]
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    def _augment_flow(flow):
+        if hflip:
+            flow = flow[:, ::-1, :]
+            flow[:, :, 0] *= -1
+        if vflip:
+            flow = flow[::-1, :, :]
+            flow[:, :, 1] *= -1
+        if rot90:
+            flow = flow.transpose(1, 0, 2)
+            flow = flow[:, :, [1, 0]]
+        return flow
+
+    rlt_img_list = [_augment(img) for img in img_list]
+    rlt_flow_list = [_augment_flow(flow) for flow in flow_list]
+
+    return rlt_img_list, rlt_flow_list
+
+
+def channel_convert(in_c, tar_type, img_list):
+    # conversion among BGR, gray and y
+    if in_c == 3 and tar_type == 'gray':  # BGR to gray
+        gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in gray_list]
+    elif in_c == 3 and tar_type == 'y':  # BGR to y
+        y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
+        return [np.expand_dims(img, axis=2) for img in y_list]
+    elif in_c == 1 and tar_type == 'RGB':  # gray/y to BGR
+        return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
+    else:
+        return img_list
+
+
+def rgb2ycbcr(img, only_y=True):
+    '''same as matlab rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+                              [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def bgr2ycbcr(img, only_y=True):
+    '''bgr version of rgb2ycbcr
+    only_y: only return Y channel
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    if only_y:
+        rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
+    else:
+        rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+                              [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def ycbcr2rgb(img):
+    '''same as matlab ycbcr2rgb
+    Input:
+        uint8, [0, 255]
+        float, [0, 1]
+    '''
+    in_img_type = img.dtype
+    img.astype(np.float32)
+    if in_img_type != np.uint8:
+        img *= 255.
+    # convert
+    rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                          [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
+    if in_img_type == np.uint8:
+        rlt = rlt.round()
+    else:
+        rlt /= 255.
+    return rlt.astype(in_img_type)
+
+
+def modcrop(img_in, scale):
+    # img_in: Numpy, HWC or HW
+    img = np.copy(img_in)
+    if img.ndim == 2:
+        H, W = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r]
+    elif img.ndim == 3:
+        H, W, C = img.shape
+        H_r, W_r = H % scale, W % scale
+        img = img[:H - H_r, :W - W_r, :]
+    else:
+        raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
+    return img
+
+
+####################
+# Functions
+####################
+
+
+# matlab 'imresize' function, now only support 'bicubic'
+def cubic(x):
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5 * absx3 - 2.5 * absx2 + 1) * (
+        (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * ((
+            (absx > 1) * (absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    if (scale < 1) and (antialiasing):
+        # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5+scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    P = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
+        1, P).expand(out_length, P)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
+    # apply cubic kernel
+    if (scale < 1) and (antialiasing):
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, P)
+
+    # If a column in weights is all zero, get rid of it. only consider the first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, P - 2)
+        weights = weights.narrow(1, 1, P - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, P - 2)
+        weights = weights.narrow(1, 0, P - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+def imresize(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: CHW RGB [0,1]
+    # output: CHW RGB [0,1] w/o round
+
+    in_C, in_H, in_W = img.size()
+    _, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
+    img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:, :sym_len_Hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_He:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_C, out_H, in_W)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        out_1[0, i, :] = img_aug[0, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+        out_1[1, i, :] = img_aug[1, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+        out_1[2, i, :] = img_aug[2, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
+    out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_Ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_We:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_C, out_H, out_W)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        out_2[0, :, i] = out_1_aug[0, :, idx:idx + kernel_width].mv(weights_W[i])
+        out_2[1, :, i] = out_1_aug[1, :, idx:idx + kernel_width].mv(weights_W[i])
+        out_2[2, :, i] = out_1_aug[2, :, idx:idx + kernel_width].mv(weights_W[i])
+
+    return out_2
+
+
+def imresize_np(img, scale, antialiasing=True):
+    # Now the scale should be the same for H and W
+    # input: img: Numpy, HWC BGR [0,1]
+    # output: HWC BGR [0,1] w/o round
+    img = torch.from_numpy(img)
+
+    in_H, in_W, in_C = img.size()
+    _, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # Return the desired dimension order for performing the resize.  The
+    # strategy is to perform the resize first along the dimension with the
+    # smallest scale factor.
+    # Now we do not support this.
+
+    # get weights and indices
+    weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+        in_H, out_H, scale, kernel, kernel_width, antialiasing)
+    weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+        in_W, out_W, scale, kernel, kernel_width, antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
+    img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
+
+    sym_patch = img[:sym_len_Hs, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+    sym_patch = img[-sym_len_He:, :, :]
+    inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(0, inv_idx)
+    img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(out_H, in_W, in_C)
+    kernel_width = weights_H.size(1)
+    for i in range(out_H):
+        idx = int(indices_H[i][0])
+        out_1[i, :, 0] = img_aug[idx:idx + kernel_width, :, 0].transpose(0, 1).mv(weights_H[i])
+        out_1[i, :, 1] = img_aug[idx:idx + kernel_width, :, 1].transpose(0, 1).mv(weights_H[i])
+        out_1[i, :, 2] = img_aug[idx:idx + kernel_width, :, 2].transpose(0, 1).mv(weights_H[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
+    out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
+
+    sym_patch = out_1[:, :sym_len_Ws, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, -sym_len_We:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(out_H, out_W, in_C)
+    kernel_width = weights_W.size(1)
+    for i in range(out_W):
+        idx = int(indices_W[i][0])
+        out_2[:, i, 0] = out_1_aug[:, idx:idx + kernel_width, 0].mv(weights_W[i])
+        out_2[:, i, 1] = out_1_aug[:, idx:idx + kernel_width, 1].mv(weights_W[i])
+        out_2[:, i, 2] = out_1_aug[:, idx:idx + kernel_width, 2].mv(weights_W[i])
+
+    return out_2.numpy()
+
+
+if __name__ == '__main__':
+    # test imresize function
+    # read images
+    img = cv2.imread('test.png')
+    img = img * 1.0 / 255
+    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+    # imresize
+    scale = 1 / 4
+    import time
+    total_time = 0
+    for i in range(10):
+        start_time = time.time()
+        rlt = imresize(img, scale, antialiasing=True)
+        use_time = time.time() - start_time
+        total_time += use_time
+    print('average time: {}'.format(total_time / 10))
+
+    import torchvision.utils
+    torchvision.utils.save_image((rlt * 255).round() / 255, 'rlt.png', nrow=1, padding=0,
+                                 normalize=False)

models/IBSN.py

@@ -0,0 +1,738 @@
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+
+import models.networks as networks
+import models.lr_scheduler as lr_scheduler
+from .base_model import BaseModel
+from models.modules.loss import ReconstructionLoss, ReconstructionMsgLoss
+from models.modules.Quantization import Quantization
+from .modules.common import DWT,IWT
+from utils.jpegtest import JpegTest
+from utils.JPEG import DiffJPEG
+import utils.util as util
+
+
+import numpy as np
+import random
+import cv2
+import time
+
+logger = logging.getLogger('base')
+dwt=DWT()
+iwt=IWT()
+
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler
+from diffusers import StableDiffusionXLInpaintPipeline
+from diffusers.utils import load_image
+from diffusers import RePaintPipeline, RePaintScheduler
+
+class Model_VSN(BaseModel):
+    def __init__(self, opt):
+        super(Model_VSN, self).__init__(opt)
+
+        if opt['dist']:
+            self.rank = torch.distributed.get_rank()
+        else:
+            self.rank = -1  # non dist training
+
+        self.gop = opt['gop']
+        train_opt = opt['train']
+        test_opt = opt['test']
+        self.opt = opt
+        self.train_opt = train_opt
+        self.test_opt = test_opt
+        self.opt_net = opt['network_G']
+        self.center = self.gop // 2
+        self.num_image = opt['num_image']
+        self.mode = opt["mode"]
+        self.idxx = 0
+
+        self.netG = networks.define_G_v2(opt).to(self.device)
+        if opt['dist']:
+            self.netG = DistributedDataParallel(self.netG, device_ids=[torch.cuda.current_device()])
+        else:
+            self.netG = DataParallel(self.netG)
+        # print network
+        self.print_network()
+        self.load()
+
+        self.Quantization = Quantization()
+
+        if not self.opt['hide']:
+            file_path = "bit_sequence.txt"
+
+            data_list = []
+
+            with open(file_path, "r") as file:
+                for line in file:
+                    data = [int(bit) for bit in line.strip()]
+                    data_list.append(data)
+            
+            self.msg_list = data_list
+
+        if self.opt['sdinpaint']:
+            self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-2-inpainting",
+                torch_dtype=torch.float16,
+            ).to("cuda")
+        
+        if self.opt['controlnetinpaint']:
+            controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float32
+            ).to("cuda")
+            self.pipe_control = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32
+            ).to("cuda")
+        
+        if self.opt['sdxl']:
+            self.pipe_sdxl = StableDiffusionXLInpaintPipeline.from_pretrained(
+                "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+                torch_dtype=torch.float16,
+                variant="fp16",
+                use_safetensors=True,
+            ).to("cuda")
+        
+        if self.opt['repaint']:
+            self.scheduler = RePaintScheduler.from_pretrained("google/ddpm-ema-celebahq-256")
+            self.pipe_repaint = RePaintPipeline.from_pretrained("google/ddpm-ema-celebahq-256", scheduler=self.scheduler)
+            self.pipe_repaint = self.pipe_repaint.to("cuda")
+
+        if self.is_train:
+            self.netG.train()
+
+            # loss
+            self.Reconstruction_forw = ReconstructionLoss(losstype=self.train_opt['pixel_criterion_forw'])
+            self.Reconstruction_back = ReconstructionLoss(losstype=self.train_opt['pixel_criterion_back'])
+            self.Reconstruction_center = ReconstructionLoss(losstype="center")
+            self.Reconstruction_msg = ReconstructionMsgLoss(losstype=self.opt['losstype'])
+
+            # optimizers
+            wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
+            optim_params = []
+
+            if self.mode == "image":
+                for k, v in self.netG.named_parameters():
+                    if (k.startswith('module.irn') or k.startswith('module.pm')) and v.requires_grad: 
+                        optim_params.append(v)
+                    else:
+                        if self.rank <= 0:
+                            logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+            elif self.mode == "bit":
+                for k, v in self.netG.named_parameters():
+                    if (k.startswith('module.bitencoder') or k.startswith('module.bitdecoder')) and v.requires_grad:
+                        optim_params.append(v)
+                    else:
+                        if self.rank <= 0:
+                            logger.warning('Params [{:s}] will not optimize.'.format(k))
+
+
+            self.optimizer_G = torch.optim.Adam(optim_params, lr=train_opt['lr_G'],
+                                                weight_decay=wd_G,
+                                                betas=(train_opt['beta1'], train_opt['beta2']))
+            self.optimizers.append(self.optimizer_G)
+
+            # schedulers
+            if train_opt['lr_scheme'] == 'MultiStepLR':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.MultiStepLR_Restart(optimizer, train_opt['lr_steps'],
+                                                         restarts=train_opt['restarts'],
+                                                         weights=train_opt['restart_weights'],
+                                                         gamma=train_opt['lr_gamma'],
+                                                         clear_state=train_opt['clear_state']))
+            elif train_opt['lr_scheme'] == 'CosineAnnealingLR_Restart':
+                for optimizer in self.optimizers:
+                    self.schedulers.append(
+                        lr_scheduler.CosineAnnealingLR_Restart(
+                            optimizer, train_opt['T_period'], eta_min=train_opt['eta_min'],
+                            restarts=train_opt['restarts'], weights=train_opt['restart_weights']))
+            else:
+                raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
+
+            self.log_dict = OrderedDict()
+
+    def feed_data(self, data):
+        self.ref_L = data['LQ'].to(self.device)  
+        self.real_H = data['GT'].to(self.device)
+        self.mes = data['MES']
+
+    def init_hidden_state(self, z):
+        b, c, h, w = z.shape
+        h_t = []
+        c_t = []
+        for _ in range(self.opt_net['block_num_rbm']):
+            h_t.append(torch.zeros([b, c, h, w]).cuda())
+            c_t.append(torch.zeros([b, c, h, w]).cuda())
+        memory = torch.zeros([b, c, h, w]).cuda()
+
+        return h_t, c_t, memory
+
+    def loss_forward(self, out, y):
+        l_forw_fit = self.train_opt['lambda_fit_forw'] * self.Reconstruction_forw(out, y)
+        return l_forw_fit
+
+    def loss_back_rec(self, out, x):
+        l_back_rec = self.train_opt['lambda_rec_back'] * self.Reconstruction_back(out, x)
+        return l_back_rec
+    
+    def loss_back_rec_mul(self, out, x):
+        l_back_rec = self.train_opt['lambda_rec_back'] * self.Reconstruction_back(out, x)
+        return l_back_rec
+
+    def optimize_parameters(self, current_step):
+        self.optimizer_G.zero_grad()
+      
+        b, n, t, c, h, w = self.ref_L.shape
+        center = t // 2
+        intval = self.gop // 2
+
+        message = torch.Tensor(np.random.choice([-0.5, 0.5], (self.ref_L.shape[0], self.opt['message_length']))).to(self.device)
+
+        add_noise = self.opt['addnoise']
+        add_jpeg = self.opt['addjpeg']
+        add_possion = self.opt['addpossion']
+        add_sdinpaint = self.opt['sdinpaint']
+        degrade_shuffle = self.opt['degrade_shuffle']
+
+        self.host = self.real_H[:, center - intval:center + intval + 1]
+        self.secret = self.ref_L[:, :, center - intval:center + intval + 1]
+        self.output, container = self.netG(x=dwt(self.host.reshape(b, -1, h, w)), x_h=dwt(self.secret[:,0].reshape(b, -1, h, w)), message=message)
+
+        Gt_ref = self.real_H[:, center - intval:center + intval + 1].detach()
+
+        y_forw = container
+
+        l_forw_fit = self.loss_forward(y_forw, self.host[:,0])
+
+
+        if degrade_shuffle:
+            import random
+            choice = random.randint(0, 2)
+            
+            if choice == 0:
+                NL = float((np.random.randint(1, 16))/255)
+                noise = np.random.normal(0, NL, y_forw.shape)
+                torchnoise = torch.from_numpy(noise).cuda().float()
+                y_forw = y_forw + torchnoise
+
+            elif choice == 1:
+                NL = int(np.random.randint(70,95))
+                self.DiffJPEG = DiffJPEG(differentiable=True, quality=int(NL)).cuda()
+                y_forw = self.DiffJPEG(y_forw)
+            
+            elif choice == 2:
+                vals = 10**4
+                if random.random() < 0.5:
+                    noisy_img_tensor = torch.poisson(y_forw * vals) / vals
+                else:
+                    img_gray_tensor = torch.mean(y_forw, dim=0, keepdim=True)
+                    noisy_gray_tensor = torch.poisson(img_gray_tensor * vals) / vals
+                    noisy_img_tensor = y_forw + (noisy_gray_tensor - img_gray_tensor)
+
+                y_forw = torch.clamp(noisy_img_tensor, 0, 1)
+
+        else:
+
+            if add_noise:
+                NL = float((np.random.randint(1,16))/255)
+                noise = np.random.normal(0, NL, y_forw.shape)
+                torchnoise = torch.from_numpy(noise).cuda().float()
+                y_forw = y_forw + torchnoise
+
+            elif add_jpeg:
+                NL = int(np.random.randint(70,95))
+                self.DiffJPEG = DiffJPEG(differentiable=True, quality=int(NL)).cuda()
+                y_forw = self.DiffJPEG(y_forw)
+
+            elif add_possion:
+                vals = 10**4
+                if random.random() < 0.5:
+                    noisy_img_tensor = torch.poisson(y_forw * vals) / vals
+                else:
+                    img_gray_tensor = torch.mean(y_forw, dim=0, keepdim=True)
+                    noisy_gray_tensor = torch.poisson(img_gray_tensor * vals) / vals
+                    noisy_img_tensor = y_forw + (noisy_gray_tensor - img_gray_tensor)
+
+                y_forw = torch.clamp(noisy_img_tensor, 0, 1)
+
+        y = self.Quantization(y_forw)
+        all_zero = torch.zeros(message.shape).to(self.device)
+
+        if self.mode == "image":
+            out_x, out_x_h, out_z, recmessage = self.netG(x=y, message=all_zero, rev=True)
+            out_x = iwt(out_x)
+            out_x_h = [iwt(out_x_h_i) for out_x_h_i in out_x_h]
+
+            l_back_rec = self.loss_back_rec(out_x, self.host[:,0])
+            out_x_h = torch.stack(out_x_h, dim=1)
+
+            l_center_x = self.loss_back_rec(out_x_h[:, 0], self.secret[:,0].reshape(b, -1, h, w))
+
+            recmessage = torch.clamp(recmessage, -0.5, 0.5)
+
+            l_msg = self.Reconstruction_msg(message, recmessage)
+
+            loss = l_forw_fit*2 + l_back_rec + l_center_x*4
+
+            loss.backward()
+
+            if self.train_opt['lambda_center'] != 0:
+                self.log_dict['l_center_x'] = l_center_x.item()
+
+            # set log
+            self.log_dict['l_back_rec'] = l_back_rec.item()
+            self.log_dict['l_forw_fit'] = l_forw_fit.item()
+            self.log_dict['l_msg'] = l_msg.item()
+            
+            self.log_dict['l_h'] = (l_center_x*10).item()
+
+            # gradient clipping
+            if self.train_opt['gradient_clipping']:
+                nn.utils.clip_grad_norm_(self.netG.parameters(), self.train_opt['gradient_clipping'])
+
+            self.optimizer_G.step()
+
+        elif self.mode == "bit":
+            recmessage = self.netG(x=y, message=all_zero, rev=True)
+
+            recmessage = torch.clamp(recmessage, -0.5, 0.5)
+
+            l_msg = self.Reconstruction_msg(message, recmessage)
+            
+            lambda_msg = self.train_opt['lambda_msg']
+
+            loss = l_msg * lambda_msg + l_forw_fit
+
+            loss.backward()
+
+            # set log
+            self.log_dict['l_forw_fit'] = l_forw_fit.item()
+            self.log_dict['l_msg'] = l_msg.item()
+
+            # gradient clipping
+            if self.train_opt['gradient_clipping']:
+                nn.utils.clip_grad_norm_(self.netG.parameters(), self.train_opt['gradient_clipping'])
+
+            self.optimizer_G.step()
+
+    def test(self, image_id):
+        self.netG.eval()
+        add_noise = self.opt['addnoise']
+        add_jpeg = self.opt['addjpeg']
+        add_possion = self.opt['addpossion']
+        add_sdinpaint = self.opt['sdinpaint']
+        add_controlnet = self.opt['controlnetinpaint']
+        add_sdxl = self.opt['sdxl']
+        add_repaint = self.opt['repaint']
+        degrade_shuffle = self.opt['degrade_shuffle']
+
+        with torch.no_grad():
+            forw_L = []
+            forw_L_h = []
+            fake_H = []
+            fake_H_h = []
+            pred_z = []
+            recmsglist = []
+            msglist = []
+            b, t, c, h, w = self.real_H.shape
+            center = t // 2
+            intval = self.gop // 2
+            b, n, t, c, h, w = self.ref_L.shape
+            id=0
+            # forward downscaling
+            self.host = self.real_H[:, center - intval+id:center + intval + 1+id]
+            self.secret = self.ref_L[:, :, center - intval+id:center + intval + 1+id]
+            self.secret = [dwt(self.secret[:,i].reshape(b, -1, h, w)) for i in range(n)]
+
+            messagenp = np.random.choice([-0.5, 0.5], (self.ref_L.shape[0], self.opt['message_length']))
+
+            message = torch.Tensor(messagenp).to(self.device)
+
+            if self.opt['bitrecord']:
+                mymsg = message.clone()
+
+                mymsg[mymsg>0] = 1
+                mymsg[mymsg<0] = 0
+                mymsg = mymsg.squeeze(0).to(torch.int)
+
+                bit_list = mymsg.tolist()
+
+                bit_string = ''.join(map(str, bit_list))
+
+                file_name = "bit_sequence.txt"
+
+                with open(file_name, "a") as file:
+                    file.write(bit_string + "\n")
+
+            if self.opt['hide']:
+                self.output, container = self.netG(x=dwt(self.host.reshape(b, -1, h, w)), x_h=self.secret, message=message)
+                y_forw = container
+            else:
+                
+                message = torch.tensor(self.msg_list[image_id]).unsqueeze(0).cuda()
+                self.output = self.host
+                y_forw = self.output.squeeze(1)
+
+            if add_sdinpaint:
+                import random
+                from PIL import Image
+                prompt = ""
+
+                b, _, _, _ = y_forw.shape
+                
+                image_batch = y_forw.permute(0, 2, 3, 1).detach().cpu().numpy()
+                forw_list = []
+
+                for j in range(b):
+                    i = image_id + 1
+                    masksrc = "../dataset/valAGE-Set-Mask/"
+                    mask_image = Image.open(masksrc + str(i).zfill(4) + ".png").convert("L")
+                    mask_image = mask_image.resize((512, 512))
+                    h, w = mask_image.size
+                    
+                    image = image_batch[j, :, :, :]
+                    image_init = Image.fromarray((image * 255).astype(np.uint8), mode = "RGB")
+                    image_inpaint = self.pipe(prompt=prompt, image=image_init, mask_image=mask_image, height=w, width=h).images[0]
+                    image_inpaint = np.array(image_inpaint) / 255.
+                    mask_image = np.array(mask_image)
+                    mask_image = np.stack([mask_image] * 3, axis=-1) / 255.
+                    mask_image = mask_image.astype(np.uint8)
+                    image_fuse = image * (1 - mask_image) + image_inpaint * mask_image
+                    forw_list.append(torch.from_numpy(image_fuse).permute(2, 0, 1))
+                
+                y_forw = torch.stack(forw_list, dim=0).float().cuda()
+
+            if add_controlnet:
+                from diffusers.utils import load_image
+                from PIL import Image
+
+                b, _, _, _ = y_forw.shape
+                forw_list = []
+                
+                image_batch = y_forw.permute(0, 2, 3, 1).detach().cpu().numpy()
+                generator = torch.Generator(device="cuda").manual_seed(1)
+
+                for j in range(b):
+                    i = image_id + 1
+                    mask_path = "../dataset/valAGE-Set-Mask/" + str(i).zfill(4) + ".png"
+                    mask_image = load_image(mask_path)
+                    mask_image = mask_image.resize((512, 512))
+                    image_init = image_batch[j, :, :, :]
+                    image_init1 = Image.fromarray((image_init * 255).astype(np.uint8), mode = "RGB")
+                    image_mask = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
+
+                    assert image_init.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
+                    image_init[image_mask > 0.5] = -1.0  # set as masked pixel
+                    image = np.expand_dims(image_init, 0).transpose(0, 3, 1, 2)
+                    control_image = torch.from_numpy(image)
+
+                    # generate image
+                    image_inpaint = self.pipe_control(
+                        "",
+                        num_inference_steps=20,
+                        generator=generator,
+                        eta=1.0,
+                        image=image_init1,
+                        mask_image=image_mask,
+                        control_image=control_image,
+                    ).images[0]
+                    
+                    image_inpaint = np.array(image_inpaint) / 255.
+                    image_mask = np.stack([image_mask] * 3, axis=-1)
+                    image_mask = image_mask.astype(np.uint8)
+                    image_fuse = image_init * (1 - image_mask) + image_inpaint * image_mask
+                    forw_list.append(torch.from_numpy(image_fuse).permute(2, 0, 1))
+
+                y_forw = torch.stack(forw_list, dim=0).float().cuda()
+            
+            if add_sdxl:
+                import random
+                from PIL import Image
+                from diffusers.utils import load_image
+                prompt = ""
+
+                b, _, _, _ = y_forw.shape
+                
+                image_batch = y_forw.permute(0, 2, 3, 1).detach().cpu().numpy()
+                forw_list = []
+
+                for j in range(b):
+                    i = image_id + 1
+                    masksrc = "../dataset/valAGE-Set-Mask/"
+                    mask_image = load_image(masksrc + str(i).zfill(4) + ".png").convert("RGB")
+                    mask_image = mask_image.resize((512, 512))
+                    h, w = mask_image.size
+                    
+                    image = image_batch[j, :, :, :]
+                    image_init = Image.fromarray((image * 255).astype(np.uint8), mode = "RGB")
+                    image_inpaint = self.pipe_sdxl(
+                        prompt=prompt, image=image_init, mask_image=mask_image, num_inference_steps=50, strength=0.80, target_size=(512, 512)
+                    ).images[0]
+                    image_inpaint = image_inpaint.resize((512, 512))
+                    image_inpaint = np.array(image_inpaint) / 255.
+                    mask_image = np.array(mask_image) / 255.
+                    mask_image = mask_image.astype(np.uint8)
+                    image_fuse = image * (1 - mask_image) + image_inpaint * mask_image
+                    forw_list.append(torch.from_numpy(image_fuse).permute(2, 0, 1))
+                
+                y_forw = torch.stack(forw_list, dim=0).float().cuda()
+
+            
+            if add_repaint:
+                from PIL import Image
+                
+                b, _, _, _ = y_forw.shape
+                
+                image_batch = y_forw.permute(0, 2, 3, 1).detach().cpu().numpy()
+                forw_list = []
+
+                generator = torch.Generator(device="cuda").manual_seed(0)
+                for j in range(b):
+                    i = image_id + 1
+                    masksrc = "../dataset/valAGE-Set-Mask/" + str(i).zfill(4) + ".png"
+                    mask_image = Image.open(masksrc).convert("RGB")
+                    mask_image = mask_image.resize((256, 256))
+                    mask_image = Image.fromarray(255 - np.array(mask_image))
+                    image = image_batch[j, :, :, :]
+                    original_image = Image.fromarray((image * 255).astype(np.uint8), mode = "RGB")
+                    original_image = original_image.resize((256, 256))
+                    output = self.pipe_repaint(
+                        image=original_image,
+                        mask_image=mask_image,
+                        num_inference_steps=150,
+                        eta=0.0,
+                        jump_length=10,
+                        jump_n_sample=10,
+                        generator=generator,
+                    )
+                    image_inpaint = output.images[0]
+                    image_inpaint = image_inpaint.resize((512, 512))
+                    image_inpaint = np.array(image_inpaint) / 255.
+                    mask_image = mask_image.resize((512, 512))
+                    mask_image = np.array(mask_image) / 255.
+                    mask_image = mask_image.astype(np.uint8)
+                    image_fuse = image * mask_image + image_inpaint * (1 - mask_image)
+                    forw_list.append(torch.from_numpy(image_fuse).permute(2, 0, 1))
+                
+                y_forw = torch.stack(forw_list, dim=0).float().cuda()
+
+            if degrade_shuffle:
+                import random
+                choice = random.randint(0, 2)
+                
+                if choice == 0:
+                    NL = float((np.random.randint(1,5))/255)
+                    noise = np.random.normal(0, NL, y_forw.shape)
+                    torchnoise = torch.from_numpy(noise).cuda().float()
+                    y_forw = y_forw + torchnoise
+
+                elif choice == 1:
+                    NL = 90
+                    self.DiffJPEG = DiffJPEG(differentiable=True, quality=int(NL)).cuda()
+                    y_forw = self.DiffJPEG(y_forw)
+                
+                elif choice == 2:
+                    vals = 10**4
+                    if random.random() < 0.5:
+                        noisy_img_tensor = torch.poisson(y_forw * vals) / vals
+                    else:
+                        img_gray_tensor = torch.mean(y_forw, dim=0, keepdim=True)
+                        noisy_gray_tensor = torch.poisson(img_gray_tensor * vals) / vals
+                        noisy_img_tensor = y_forw + (noisy_gray_tensor - img_gray_tensor)
+
+                    y_forw = torch.clamp(noisy_img_tensor, 0, 1)
+
+            else:
+
+                if add_noise:
+                    NL = self.opt['noisesigma'] / 255.0
+                    noise = np.random.normal(0, NL, y_forw.shape)
+                    torchnoise = torch.from_numpy(noise).cuda().float()
+                    y_forw = y_forw + torchnoise
+
+                elif add_jpeg:
+                    Q = self.opt['jpegfactor']
+                    self.DiffJPEG = DiffJPEG(differentiable=True, quality=int(Q)).cuda()
+                    y_forw = self.DiffJPEG(y_forw)
+
+                elif add_possion:
+                    vals = 10**4
+                    if random.random() < 0.5:
+                        noisy_img_tensor = torch.poisson(y_forw * vals) / vals
+                    else:
+                        img_gray_tensor = torch.mean(y_forw, dim=0, keepdim=True)
+                        noisy_gray_tensor = torch.poisson(img_gray_tensor * vals) / vals
+                        noisy_img_tensor = y_forw + (noisy_gray_tensor - img_gray_tensor)
+
+                    y_forw = torch.clamp(noisy_img_tensor, 0, 1)
+
+            # backward upscaling
+            if self.opt['hide']:
+                y = self.Quantization(y_forw)
+            else:
+                y = y_forw
+
+            if self.mode == "image":
+                out_x, out_x_h, out_z, recmessage = self.netG(x=y, rev=True)
+                out_x = iwt(out_x)
+
+                out_x_h = [iwt(out_x_h_i) for out_x_h_i in out_x_h]
+                out_x = out_x.reshape(-1, self.gop, 3, h, w)
+                out_x_h = torch.stack(out_x_h, dim=1)
+                out_x_h = out_x_h.reshape(-1, 1, self.gop, 3, h, w)
+
+                forw_L.append(y_forw)
+                fake_H.append(out_x[:, self.gop//2])
+                fake_H_h.append(out_x_h[:,:, self.gop//2])
+                recmsglist.append(recmessage)
+                msglist.append(message)
+            
+            elif self.mode == "bit":
+                recmessage = self.netG(x=y, rev=True)
+                forw_L.append(y_forw)
+                recmsglist.append(recmessage)
+                msglist.append(message)
+
+        if self.mode == "image":
+            self.fake_H = torch.clamp(torch.stack(fake_H, dim=1),0,1)
+            self.fake_H_h = torch.clamp(torch.stack(fake_H_h, dim=2),0,1)
+
+        self.forw_L = torch.clamp(torch.stack(forw_L, dim=1),0,1)
+        remesg = torch.clamp(torch.stack(recmsglist, dim=0),-0.5,0.5)
+
+        if self.opt['hide']:
+            mesg = torch.clamp(torch.stack(msglist, dim=0),-0.5,0.5)
+        else:
+            mesg = torch.stack(msglist, dim=0)
+
+        self.recmessage = remesg.clone()
+        self.recmessage[remesg > 0] = 1
+        self.recmessage[remesg <= 0] = 0
+
+        self.message = mesg.clone()
+        self.message[mesg > 0] = 1
+        self.message[mesg <= 0] = 0
+
+        self.netG.train()
+
+
+    def image_hiding(self, ):
+        self.netG.eval()
+        with torch.no_grad():
+            b, t, c, h, w = self.real_H.shape
+            center = t // 2
+            intval = self.gop // 2
+            b, n, t, c, h, w = self.ref_L.shape
+            id=0
+            # forward downscaling
+            self.host = self.real_H[:, center - intval+id:center + intval + 1+id]
+            self.secret = self.ref_L[:, :, center - intval+id:center + intval + 1+id]
+            self.secret = [dwt(self.secret[:,i].reshape(b, -1, h, w)) for i in range(n)]
+
+            message = torch.Tensor(self.mes).to(self.device)
+
+            self.output, container = self.netG(x=dwt(self.host.reshape(b, -1, h, w)), x_h=self.secret, message=message)
+            y_forw = container
+
+            result = torch.clamp(y_forw,0,1)
+
+            lr_img = util.tensor2img(result)
+
+            return lr_img
+
+    def image_recovery(self, number):
+        self.netG.eval()
+        with torch.no_grad():
+            b, t, c, h, w = self.real_H.shape
+            center = t // 2
+            intval = self.gop // 2
+            b, n, t, c, h, w = self.ref_L.shape
+            id=0
+            # forward downscaling
+            self.host = self.real_H[:, center - intval+id:center + intval + 1+id]
+            self.secret = self.ref_L[:, :, center - intval+id:center + intval + 1+id]
+            template = self.secret.reshape(b, -1, h, w)
+            self.secret = [dwt(self.secret[:,i].reshape(b, -1, h, w)) for i in range(n)]
+
+            self.output = self.host
+            y_forw = self.output.squeeze(1)
+
+            y = self.Quantization(y_forw)
+
+            out_x, out_x_h, out_z, recmessage = self.netG(x=y, rev=True)
+            out_x = iwt(out_x)
+
+            out_x_h = [iwt(out_x_h_i) for out_x_h_i in out_x_h]
+            out_x = out_x.reshape(-1, self.gop, 3, h, w)
+            out_x_h = torch.stack(out_x_h, dim=1)
+            out_x_h = out_x_h.reshape(-1, 1, self.gop, 3, h, w)
+
+            rec_loc = out_x_h[:,:, self.gop//2]
+            # from PIL import Image
+            # tmp = util.tensor2img(rec_loc)
+            # save
+            residual = torch.abs(template - rec_loc)
+            binary_residual = (residual > number).float()
+            residual = util.tensor2img(binary_residual)
+            mask = np.sum(residual, axis=2)
+            # print(mask)
+
+            remesg = torch.clamp(recmessage,-0.5,0.5)
+            remesg[remesg > 0] = 1
+            remesg[remesg <= 0] = 0
+
+            return mask, remesg
+        
+    def get_current_log(self):
+        return self.log_dict
+
+    def get_current_visuals(self):
+        b, n, t, c, h, w = self.ref_L.shape
+        center = t // 2
+        intval = self.gop // 2
+        out_dict = OrderedDict()
+        LR_ref = self.ref_L[:, :, center - intval:center + intval + 1].detach()[0].float().cpu()
+        LR_ref = torch.chunk(LR_ref, self.num_image, dim=0)
+        out_dict['LR_ref'] = [image.squeeze(0) for image in LR_ref]
+        
+        if self.mode == "image":
+            out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
+            SR_h = self.fake_H_h.detach()[0].float().cpu()
+            SR_h = torch.chunk(SR_h, self.num_image, dim=0)
+            out_dict['SR_h'] = [image.squeeze(0) for image in SR_h]
+        
+        out_dict['LR'] = self.forw_L.detach()[0].float().cpu()
+        out_dict['GT'] = self.real_H[:, center - intval:center + intval + 1].detach()[0].float().cpu()
+        out_dict['message'] = self.message
+        out_dict['recmessage'] = self.recmessage
+
+        return out_dict
+
+    def print_network(self):
+        s, n = self.get_network_description(self.netG)
+        if isinstance(self.netG, nn.DataParallel) or isinstance(self.netG, DistributedDataParallel):
+            net_struc_str = '{} - {}'.format(self.netG.__class__.__name__,
+                                             self.netG.module.__class__.__name__)
+        else:
+            net_struc_str = '{}'.format(self.netG.__class__.__name__)
+        if self.rank <= 0:
+            logger.info('Network G structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
+            logger.info(s)
+
+    def load(self):
+        load_path_G = self.opt['path']['pretrain_model_G']
+        if load_path_G is not None:
+            logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
+    
+    def load_test(self,load_path_G):
+        self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
+
+    def save(self, iter_label):
+        self.save_network(self.netG, 'G', iter_label)

models/__init__.py

@@ -0,0 +1,11 @@
+import logging
+logger = logging.getLogger('base')
+
+def create_model(opt):
+    model = opt['model']
+    frame_num = opt['gop']
+    from .IBSN import Model_VSN as M
+
+    m = M(opt)
+    logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
+    return m

models/base_model.py

@@ -0,0 +1,119 @@
+import os
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+
+class BaseModel():
+    def __init__(self, opt):
+        self.opt = opt
+        self.device = torch.device('cuda' if opt['gpu_ids'] is not None else 'cpu')
+        self.is_train = opt['is_train']
+        self.schedulers = []
+        self.optimizers = []
+
+    def feed_data(self, data):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        pass
+
+    def get_current_losses(self):
+        pass
+
+    def print_network(self):
+        pass
+
+    def save(self, label):
+        pass
+
+    def load(self):
+        pass
+
+    def _set_lr(self, lr_groups_l):
+        ''' set learning rate for warmup,
+        lr_groups_l: list for lr_groups. each for a optimizer'''
+        for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
+            for param_group, lr in zip(optimizer.param_groups, lr_groups):
+                param_group['lr'] = lr
+
+    def _get_init_lr(self):
+        # get the initial lr, which is set by the scheduler
+        init_lr_groups_l = []
+        for optimizer in self.optimizers:
+            init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
+        return init_lr_groups_l
+
+    def update_learning_rate(self, cur_iter, warmup_iter=-1):
+        for scheduler in self.schedulers:
+            scheduler.step()
+        #### set up warm up learning rate
+        if cur_iter < warmup_iter:
+            # get initial lr for each group
+            init_lr_g_l = self._get_init_lr()
+            # modify warming-up learning rates
+            warm_up_lr_l = []
+            for init_lr_g in init_lr_g_l:
+                warm_up_lr_l.append([v / warmup_iter * cur_iter for v in init_lr_g])
+            # set learning rate
+            self._set_lr(warm_up_lr_l)
+
+    def get_current_learning_rate(self):
+        # return self.schedulers[0].get_lr()[0]
+        return self.optimizers[0].param_groups[0]['lr']
+
+    def get_network_description(self, network):
+        '''Get the string and total parameters of the network'''
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        s = str(network)
+        n = sum(map(lambda x: x.numel(), network.parameters()))
+        return s, n
+
+    def save_network(self, network, network_label, iter_label):
+        save_filename = '{}_{}.pth'.format(iter_label, network_label)
+        save_path = os.path.join(self.opt['path']['models'], save_filename)
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        state_dict = network.state_dict()
+        for key, param in state_dict.items():
+            state_dict[key] = param.cpu()
+        torch.save(state_dict, save_path)
+
+    def load_network(self, load_path, network, strict=True):
+        if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+            network = network.module
+        load_net = torch.load(load_path)
+        load_net_clean = OrderedDict()  # remove unnecessary 'module.'
+        for k, v in load_net.items():
+            if k.startswith('module.'):
+                load_net_clean[k[7:]] = v
+            else:
+                load_net_clean[k] = v
+        network.load_state_dict(load_net_clean, strict=strict)
+
+    def save_training_state(self, epoch, iter_step):
+        '''Saves training state during training, which will be used for resuming'''
+        state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []}
+        for s in self.schedulers:
+            state['schedulers'].append(s.state_dict())
+        for o in self.optimizers:
+            state['optimizers'].append(o.state_dict())
+        save_filename = '{}.state'.format(iter_step)
+        save_path = os.path.join(self.opt['path']['training_state'], save_filename)
+        torch.save(state, save_path)
+
+    def resume_training(self, resume_state):
+        '''Resume the optimizers and schedulers for training'''
+        resume_optimizers = resume_state['optimizers']
+        resume_schedulers = resume_state['schedulers']
+        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+        for i, o in enumerate(resume_optimizers):
+            self.optimizers[i].load_state_dict(o)
+        for i, s in enumerate(resume_schedulers):
+            self.schedulers[i].load_state_dict(s)

models/bitnetwork/ConvBlock.py

@@ -0,0 +1,38 @@
+import torch.nn as nn
+
+
+class ConvINRelu(nn.Module):
+	"""
+	A sequence of Convolution, Instance Normalization, and ReLU activation
+	"""
+
+	def __init__(self, channels_in, channels_out, stride):
+		super(ConvINRelu, self).__init__()
+
+		self.layers = nn.Sequential(
+			nn.Conv2d(channels_in, channels_out, 3, stride, padding=1),
+			nn.InstanceNorm2d(channels_out),
+			nn.ReLU(inplace=True)
+		)
+
+	def forward(self, x):
+		return self.layers(x)
+
+
+class ConvBlock(nn.Module):
+	'''
+	Network that composed by layers of ConvINRelu
+	'''
+
+	def __init__(self, in_channels, out_channels, blocks=1, stride=1):
+		super(ConvBlock, self).__init__()
+
+		layers = [ConvINRelu(in_channels, out_channels, stride)] if blocks != 0 else []
+		for _ in range(blocks - 1):
+			layer = ConvINRelu(out_channels, out_channels, 1)
+			layers.append(layer)
+
+		self.layers = nn.Sequential(*layers)
+
+	def forward(self, x):
+		return self.layers(x)

models/bitnetwork/DW_EncoderDecoder.py

@@ -0,0 +1,28 @@
+from . import *
+from .Encoder_U import DW_Encoder
+from .Decoder_U import DW_Decoder
+from .Noise import Noise
+from .Random_Noise import Random_Noise
+
+
+class DW_EncoderDecoder(nn.Module):
+	'''
+	A Sequential of Encoder_MP-Noise-Decoder
+	'''
+
+	def __init__(self, message_length, noise_layers_R, noise_layers_F, attention_encoder, attention_decoder):
+		super(DW_EncoderDecoder, self).__init__()
+		self.encoder = DW_Encoder(message_length, attention = attention_encoder)
+		self.noise = Random_Noise(noise_layers_R + noise_layers_F, len(noise_layers_R), len(noise_layers_F))
+		self.decoder_C = DW_Decoder(message_length, attention = attention_decoder)
+		self.decoder_RF = DW_Decoder(message_length, attention = attention_decoder)
+
+
+	def forward(self, image, message, mask):
+		encoded_image = self.encoder(image, message)
+		noised_image_C, noised_image_R, noised_image_F = self.noise([encoded_image, image, mask])
+		decoded_message_C = self.decoder_C(noised_image_C)
+		decoded_message_R = self.decoder_RF(noised_image_R)
+		decoded_message_F = self.decoder_RF(noised_image_F)
+		return encoded_image, noised_image_C, decoded_message_C, decoded_message_R, decoded_message_F
+

models/bitnetwork/Decoder_U.py

@@ -0,0 +1,87 @@
+from . import *
+
+
+class DW_Decoder(nn.Module):
+
+    def __init__(self, message_length, blocks=2, channels=64, attention=None):
+        super(DW_Decoder, self).__init__()
+
+        self.conv1 = ConvBlock(3, 16, blocks=blocks)
+        self.down1 = Down(16, 32, blocks=blocks)
+        self.down2 = Down(32, 64, blocks=blocks)
+        self.down3 = Down(64, 128, blocks=blocks)
+
+        self.down4 = Down(128, 256, blocks=blocks)
+
+        self.up3 = UP(256, 128)
+        self.att3 = ResBlock(128 * 2, 128, blocks=blocks, attention=attention)
+
+        self.up2 = UP(128, 64)
+        self.att2 = ResBlock(64 * 2, 64, blocks=blocks, attention=attention)
+
+        self.up1 = UP(64, 32)
+        self.att1 = ResBlock(32 * 2, 32, blocks=blocks, attention=attention)
+
+        self.up0 = UP(32, 16)
+        self.att0 = ResBlock(16 * 2, 16, blocks=blocks, attention=attention)
+
+        self.Conv_1x1 = nn.Conv2d(16, 1, kernel_size=1, stride=1, padding=0, bias=False)
+
+        self.message_layer = nn.Linear(message_length * message_length, message_length)
+        self.message_length = message_length
+
+
+    def forward(self, x):
+        d0 = self.conv1(x)
+        d1 = self.down1(d0)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+
+        d4 = self.down4(d3)
+
+        u3 = self.up3(d4)
+        u3 = torch.cat((d3, u3), dim=1)
+        u3 = self.att3(u3)
+
+        u2 = self.up2(u3)
+        u2 = torch.cat((d2, u2), dim=1)
+        u2 = self.att2(u2)
+
+        u1 = self.up1(u2)
+        u1 = torch.cat((d1, u1), dim=1)
+        u1 = self.att1(u1)
+
+        u0 = self.up0(u1)
+        u0 = torch.cat((d0, u0), dim=1)
+        u0 = self.att0(u0)
+
+        residual = self.Conv_1x1(u0)
+
+        message = F.interpolate(residual, size=(self.message_length, self.message_length),
+                                                           mode='nearest')
+        message = message.view(message.shape[0], -1)
+        message = self.message_layer(message)
+
+        return message
+
+
+class Down(nn.Module):
+    def __init__(self, in_channels, out_channels, blocks):
+        super(Down, self).__init__()
+        self.layer = torch.nn.Sequential(
+            ConvBlock(in_channels, in_channels, stride=2),
+            ConvBlock(in_channels, out_channels, blocks=blocks)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class UP(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UP, self).__init__()
+        self.conv = ConvBlock(in_channels, out_channels)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        return self.conv(x)

models/bitnetwork/Dual_Mark.py

@@ -0,0 +1,249 @@
+from .DW_EncoderDecoder import *
+from .Patch_Discriminator import Patch_Discriminator
+import torch
+import kornia.losses
+import lpips
+
+
+class Network:
+
+	def __init__(self, message_length, noise_layers_R, noise_layers_F, device, batch_size, lr, beta1, attention_encoder, attention_decoder, weight):
+		# device
+		self.device = device
+
+		# loss function
+		self.criterion_MSE = nn.MSELoss().to(device)
+		self.criterion_LPIPS = lpips.LPIPS().to(device)
+
+		# weight of encoder-decoder loss
+		self.encoder_weight = weight[0]
+		self.decoder_weight_C = weight[1]
+		self.decoder_weight_R = weight[2]
+		self.decoder_weight_F = weight[3]
+		self.discriminator_weight = weight[4]
+
+		# network
+		self.encoder_decoder = DW_EncoderDecoder(message_length, noise_layers_R, noise_layers_F, attention_encoder, attention_decoder).to(device)
+		self.discriminator = Patch_Discriminator().to(device)
+
+		self.encoder_decoder = torch.nn.DataParallel(self.encoder_decoder)
+		self.discriminator = torch.nn.DataParallel(self.discriminator)
+
+		# mark "cover" as 1, "encoded" as -1
+		self.label_cover = 1.0
+		self.label_encoded = - 1.0
+
+		for p in self.encoder_decoder.module.noise.parameters():
+			p.requires_grad = False
+
+		# optimizer
+		self.opt_encoder_decoder = torch.optim.Adam(
+			filter(lambda p: p.requires_grad, self.encoder_decoder.parameters()), lr=lr, betas=(beta1, 0.999))
+		self.opt_discriminator = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(beta1, 0.999))
+
+
+	def train(self, images: torch.Tensor, messages: torch.Tensor, masks: torch.Tensor):
+		self.encoder_decoder.train()
+		self.discriminator.train()
+
+		with torch.enable_grad():
+			# use device to compute
+			images, messages, masks = images.to(self.device), messages.to(self.device), masks.to(self.device)
+			encoded_images, noised_images, decoded_messages_C, decoded_messages_R, decoded_messages_F = self.encoder_decoder(images, messages, masks)
+
+			'''
+			train discriminator
+			'''
+			for p in self.discriminator.parameters():
+				p.requires_grad = True
+
+			self.opt_discriminator.zero_grad()
+
+			# RAW : target label for image should be "cover"(1)
+			d_label_cover = self.discriminator(images)
+			#d_cover_loss = self.criterion_MSE(d_label_cover, torch.ones_like(d_label_cover))
+			#d_cover_loss.backward()
+
+			# GAN : target label for encoded image should be "encoded"(0)
+			d_label_encoded = self.discriminator(encoded_images.detach())
+			#d_encoded_loss = self.criterion_MSE(d_label_encoded, torch.zeros_like(d_label_encoded))
+			#d_encoded_loss.backward()
+
+			d_loss = self.criterion_MSE(d_label_cover - torch.mean(d_label_encoded), self.label_cover * torch.ones_like(d_label_cover)) +\
+			         self.criterion_MSE(d_label_encoded - torch.mean(d_label_cover), self.label_encoded * torch.ones_like(d_label_encoded))
+			d_loss.backward()
+
+			self.opt_discriminator.step()
+
+			'''
+			train encoder and decoder
+			'''
+			# Make it a tiny bit faster
+			for p in self.discriminator.parameters():
+				p.requires_grad = False
+
+			self.opt_encoder_decoder.zero_grad()
+
+			# GAN : target label for encoded image should be "cover"(0)
+			g_label_cover = self.discriminator(images)
+			g_label_encoded = self.discriminator(encoded_images)
+			g_loss_on_discriminator = self.criterion_MSE(g_label_cover - torch.mean(g_label_encoded), self.label_encoded * torch.ones_like(g_label_cover)) +\
+									  self.criterion_MSE(g_label_encoded - torch.mean(g_label_cover), self.label_cover * torch.ones_like(g_label_encoded))
+
+			# RAW : the encoded image should be similar to cover image
+			g_loss_on_encoder_MSE = self.criterion_MSE(encoded_images, images)
+			g_loss_on_encoder_LPIPS = torch.mean(self.criterion_LPIPS(encoded_images, images))
+
+			# RESULT : the decoded message should be similar to the raw message /Dual
+			g_loss_on_decoder_C = self.criterion_MSE(decoded_messages_C, messages)
+			g_loss_on_decoder_R = self.criterion_MSE(decoded_messages_R, messages)
+			g_loss_on_decoder_F = self.criterion_MSE(decoded_messages_F, torch.zeros_like(messages))
+
+			# full loss
+			g_loss = self.discriminator_weight * g_loss_on_discriminator + self.encoder_weight * g_loss_on_encoder_MSE +\
+					 self.decoder_weight_C * g_loss_on_decoder_C + self.decoder_weight_R * g_loss_on_decoder_R + self.decoder_weight_F * g_loss_on_decoder_F
+
+			g_loss.backward()
+			self.opt_encoder_decoder.step()
+
+			# psnr
+			psnr = - kornia.losses.psnr_loss(encoded_images.detach(), images, 2)
+
+			# ssim
+			ssim = 1 - 2 * kornia.losses.ssim_loss(encoded_images.detach(), images, window_size=11, reduction="mean")
+
+		'''
+		decoded message error rate /Dual
+		'''
+		error_rate_C = self.decoded_message_error_rate_batch(messages, decoded_messages_C)
+		error_rate_R = self.decoded_message_error_rate_batch(messages, decoded_messages_R)
+		error_rate_F = self.decoded_message_error_rate_batch(messages, decoded_messages_F)
+
+		result = {
+			"g_loss": g_loss,
+			"error_rate_C": error_rate_C,
+			"error_rate_R": error_rate_R,
+			"error_rate_F": error_rate_F,
+			"psnr": psnr,
+			"ssim": ssim,
+			"g_loss_on_discriminator": g_loss_on_discriminator,
+			"g_loss_on_encoder_MSE": g_loss_on_encoder_MSE,
+			"g_loss_on_encoder_LPIPS": g_loss_on_encoder_LPIPS,
+			"g_loss_on_decoder_C": g_loss_on_decoder_C,
+			"g_loss_on_decoder_R": g_loss_on_decoder_R,
+			"g_loss_on_decoder_F": g_loss_on_decoder_F,
+			"d_loss": d_loss
+		}
+		return result
+
+
+	def validation(self, images: torch.Tensor, messages: torch.Tensor, masks: torch.Tensor):
+		self.encoder_decoder.eval()
+		self.encoder_decoder.module.noise.train()
+		self.discriminator.eval()
+
+		with torch.no_grad():
+			# use device to compute
+			images, messages, masks = images.to(self.device), messages.to(self.device), masks.to(self.device)
+			encoded_images, noised_images, decoded_messages_C, decoded_messages_R, decoded_messages_F = self.encoder_decoder(images, messages, masks)
+
+			'''
+			validate discriminator
+			'''
+			# RAW : target label for image should be "cover"(1)
+			d_label_cover = self.discriminator(images)
+			#d_cover_loss = self.criterion_MSE(d_label_cover, torch.ones_like(d_label_cover))
+
+			# GAN : target label for encoded image should be "encoded"(0)
+			d_label_encoded = self.discriminator(encoded_images.detach())
+			#d_encoded_loss = self.criterion_MSE(d_label_encoded, torch.zeros_like(d_label_encoded))
+
+			d_loss = self.criterion_MSE(d_label_cover - torch.mean(d_label_encoded), self.label_cover * torch.ones_like(d_label_cover)) +\
+			         self.criterion_MSE(d_label_encoded - torch.mean(d_label_cover), self.label_encoded * torch.ones_like(d_label_encoded))
+
+			'''
+			validate encoder and decoder
+			'''
+
+			# GAN : target label for encoded image should be "cover"(0)
+			g_label_cover = self.discriminator(images)
+			g_label_encoded = self.discriminator(encoded_images)
+			g_loss_on_discriminator = self.criterion_MSE(g_label_cover - torch.mean(g_label_encoded), self.label_encoded * torch.ones_like(g_label_cover)) +\
+									  self.criterion_MSE(g_label_encoded - torch.mean(g_label_cover), self.label_cover * torch.ones_like(g_label_encoded))
+
+			# RAW : the encoded image should be similar to cover image
+			g_loss_on_encoder_MSE = self.criterion_MSE(encoded_images, images)
+			g_loss_on_encoder_LPIPS = torch.mean(self.criterion_LPIPS(encoded_images, images))
+
+			# RESULT : the decoded message should be similar to the raw message /Dual
+			g_loss_on_decoder_C = self.criterion_MSE(decoded_messages_C, messages)
+			g_loss_on_decoder_R = self.criterion_MSE(decoded_messages_R, messages)
+			g_loss_on_decoder_F = self.criterion_MSE(decoded_messages_F, torch.zeros_like(messages))
+
+			# full loss
+			# unstable g_loss_on_discriminator is not used during validation
+
+			g_loss = 0 * g_loss_on_discriminator + self.encoder_weight * g_loss_on_encoder_LPIPS +\
+					 self.decoder_weight_C * g_loss_on_decoder_C + self.decoder_weight_R * g_loss_on_decoder_R + self.decoder_weight_F * g_loss_on_decoder_F
+
+
+			# psnr
+			psnr = - kornia.losses.psnr_loss(encoded_images.detach(), images, 2)
+
+			# ssim
+			ssim = 1 - 2 * kornia.losses.ssim_loss(encoded_images.detach(), images, window_size=11, reduction="mean")
+
+		'''
+		decoded message error rate /Dual
+		'''
+		error_rate_C = self.decoded_message_error_rate_batch(messages, decoded_messages_C)
+		error_rate_R = self.decoded_message_error_rate_batch(messages, decoded_messages_R)
+		error_rate_F = self.decoded_message_error_rate_batch(messages, decoded_messages_F)
+
+		result = {
+			"g_loss": g_loss,
+			"error_rate_C": error_rate_C,
+			"error_rate_R": error_rate_R,
+			"error_rate_F": error_rate_F,
+			"psnr": psnr,
+			"ssim": ssim,
+			"g_loss_on_discriminator": g_loss_on_discriminator,
+			"g_loss_on_encoder_MSE": g_loss_on_encoder_MSE,
+			"g_loss_on_encoder_LPIPS": g_loss_on_encoder_LPIPS,
+			"g_loss_on_decoder_C": g_loss_on_decoder_C,
+			"g_loss_on_decoder_R": g_loss_on_decoder_R,
+			"g_loss_on_decoder_F": g_loss_on_decoder_F,
+			"d_loss": d_loss
+		}
+
+		return result, (images, encoded_images, noised_images)
+
+	def decoded_message_error_rate(self, message, decoded_message):
+		length = message.shape[0]
+
+		message = message.gt(0)
+		decoded_message = decoded_message.gt(0)
+		error_rate = float(sum(message != decoded_message)) / length
+		return error_rate
+
+	def decoded_message_error_rate_batch(self, messages, decoded_messages):
+		error_rate = 0.0
+		batch_size = len(messages)
+		for i in range(batch_size):
+			error_rate += self.decoded_message_error_rate(messages[i], decoded_messages[i])
+		error_rate /= batch_size
+		return error_rate
+
+	def save_model(self, path_encoder_decoder: str, path_discriminator: str):
+		torch.save(self.encoder_decoder.module.state_dict(), path_encoder_decoder)
+		torch.save(self.discriminator.module.state_dict(), path_discriminator)
+
+	def load_model(self, path_encoder_decoder: str, path_discriminator: str):
+		self.load_model_ed(path_encoder_decoder)
+		self.load_model_dis(path_discriminator)
+
+	def load_model_ed(self, path_encoder_decoder: str):
+		self.encoder_decoder.module.load_state_dict(torch.load(path_encoder_decoder), strict=False)
+
+	def load_model_dis(self, path_discriminator: str):
+		self.discriminator.module.load_state_dict(torch.load(path_discriminator))

models/bitnetwork/Encoder_U.py

@@ -0,0 +1,125 @@
+from . import *
+
+class DW_Encoder(nn.Module):
+
+    def __init__(self, message_length, blocks=2, channels=64, attention=None):
+        super(DW_Encoder, self).__init__()
+
+        self.conv1 = ConvBlock(3, 16, blocks=blocks)
+        self.down1 = Down(16, 32, blocks=blocks)
+        self.down2 = Down(32, 64, blocks=blocks)
+        self.down3 = Down(64, 128, blocks=blocks)
+
+        self.down4 = Down(128, 256, blocks=blocks)
+
+        self.up3 = UP(256, 128)
+        self.linear3 = nn.Linear(message_length, message_length * message_length)
+        self.Conv_message3 = ConvBlock(1, channels, blocks=blocks)
+        self.att3 = ResBlock(128 * 2 + channels, 128, blocks=blocks, attention=attention)
+
+        self.up2 = UP(128, 64)
+        self.linear2 = nn.Linear(message_length, message_length * message_length)
+        self.Conv_message2 = ConvBlock(1, channels, blocks=blocks)
+        self.att2 = ResBlock(64 * 2 + channels, 64, blocks=blocks, attention=attention)
+
+        self.up1 = UP(64, 32)
+        self.linear1 = nn.Linear(message_length, message_length * message_length)
+        self.Conv_message1 = ConvBlock(1, channels, blocks=blocks)
+        self.att1 = ResBlock(32 * 2 + channels, 32, blocks=blocks, attention=attention)
+
+        self.up0 = UP(32, 16)
+        self.linear0 = nn.Linear(message_length, message_length * message_length)
+        self.Conv_message0 = ConvBlock(1, channels, blocks=blocks)
+        self.att0 = ResBlock(16 * 2 + channels, 16, blocks=blocks, attention=attention)
+
+        self.Conv_1x1 = nn.Conv2d(16 + 3, 3, kernel_size=1, stride=1, padding=0)
+
+        self.message_length = message_length
+
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+        ])
+
+
+    def forward(self, x, watermark):
+        d0 = self.conv1(x)
+        d1 = self.down1(d0)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+
+        d4 = self.down4(d3)
+
+        u3 = self.up3(d4)
+        expanded_message = self.linear3(watermark)
+        expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)
+        expanded_message = F.interpolate(expanded_message, size=(d3.shape[2], d3.shape[3]),
+                                                           mode='nearest')
+        expanded_message = self.Conv_message3(expanded_message)
+        u3 = torch.cat((d3, u3, expanded_message), dim=1)
+        u3 = self.att3(u3)
+
+        u2 = self.up2(u3)
+        expanded_message = self.linear2(watermark)
+        expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)
+        expanded_message = F.interpolate(expanded_message, size=(d2.shape[2], d2.shape[3]),
+                                                           mode='nearest')
+        expanded_message = self.Conv_message2(expanded_message)
+        u2 = torch.cat((d2, u2, expanded_message), dim=1)
+        u2 = self.att2(u2)
+
+        u1 = self.up1(u2)
+        expanded_message = self.linear1(watermark)
+        expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)
+        expanded_message = F.interpolate(expanded_message, size=(d1.shape[2], d1.shape[3]),
+                                                           mode='nearest')
+        expanded_message = self.Conv_message1(expanded_message)
+        u1 = torch.cat((d1, u1, expanded_message), dim=1)
+        u1 = self.att1(u1)
+
+        u0 = self.up0(u1)
+        expanded_message = self.linear0(watermark)
+        expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)
+        expanded_message = F.interpolate(expanded_message, size=(d0.shape[2], d0.shape[3]),
+                                                           mode='nearest')
+        expanded_message = self.Conv_message0(expanded_message)
+        u0 = torch.cat((d0, u0, expanded_message), dim=1)
+        u0 = self.att0(u0)
+
+        image = self.Conv_1x1(torch.cat((x, u0), dim=1))
+
+        forward_image = image.clone().detach()
+        '''read_image = torch.zeros_like(forward_image)
+
+        for index in range(forward_image.shape[0]):
+            single_image = ((forward_image[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(single_image)
+            read = np.array(im, dtype=np.uint8)
+            read_image[index] = self.transform(read).unsqueeze(0).to(image.device)
+
+        gap = read_image - forward_image'''
+        gap = forward_image.clamp(-1, 1) - forward_image
+
+        return image + gap
+
+
+class Down(nn.Module):
+    def __init__(self, in_channels, out_channels, blocks):
+        super(Down, self).__init__()
+        self.layer = torch.nn.Sequential(
+            ConvBlock(in_channels, in_channels, stride=2),
+            ConvBlock(in_channels, out_channels, blocks=blocks)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class UP(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UP, self).__init__()
+        self.conv = ConvBlock(in_channels, out_channels)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        return self.conv(x)

models/bitnetwork/Random_Noise.py

@@ -0,0 +1,59 @@
+from . import *
+from .noise_layers import *
+
+
+class Random_Noise(nn.Module):
+
+    def __init__(self, layers, len_layers_R, len_layers_F):
+        super(Random_Noise, self).__init__()
+        for i in range(len(layers)):
+            layers[i] = eval(layers[i])
+        self.noise = nn.Sequential(*layers)
+        self.len_layers_R = len_layers_R
+        self.len_layers_F = len_layers_F
+        print(self.noise)
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+        ])
+
+    def forward(self, image_cover_mask):
+        image, cover_image, mask = image_cover_mask[0], image_cover_mask[1], image_cover_mask[2]
+        forward_image = image.clone().detach()
+        forward_cover_image = cover_image.clone().detach()
+        forward_mask = mask.clone().detach()
+        noised_image_C = torch.zeros_like(forward_image)
+        noised_image_R = torch.zeros_like(forward_image)
+        noised_image_F = torch.zeros_like(forward_image)
+
+        for index in range(forward_image.shape[0]):
+            random_noise_layer_C = np.random.choice(self.noise, 1)[0]
+            random_noise_layer_R = np.random.choice(self.noise[0:self.len_layers_R], 1)[0]
+            random_noise_layer_F = np.random.choice(self.noise[self.len_layers_R:self.len_layers_R + self.len_layers_F], 1)[0]
+            noised_image_C[index] = random_noise_layer_C([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
+            noised_image_R[index] = random_noise_layer_R([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
+            noised_image_F[index] = random_noise_layer_F([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
+
+            '''single_image = ((noised_image_C[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(single_image)
+            read = np.array(im, dtype=np.uint8)
+            noised_image_C[index] = self.transform(read).unsqueeze(0).to(image.device)
+
+            single_image = ((noised_image_R[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(single_image)
+            read = np.array(im, dtype=np.uint8)
+            noised_image_R[index] = self.transform(read).unsqueeze(0).to(image.device)
+
+            single_image = ((noised_image_F[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(single_image)
+            read = np.array(im, dtype=np.uint8)
+            noised_image_F[index] = self.transform(read).unsqueeze(0).to(image.device)
+
+        noised_image_gap_C = noised_image_C - forward_image
+        noised_image_gap_R = noised_image_R - forward_image
+        noised_image_gap_F = noised_image_F - forward_image'''
+        noised_image_gap_C = noised_image_C.clamp(-1, 1) - forward_image
+        noised_image_gap_R = noised_image_R.clamp(-1, 1) - forward_image
+        noised_image_gap_F = noised_image_F.clamp(-1, 1) - forward_image
+
+        return image + noised_image_gap_C, image + noised_image_gap_R, image + noised_image_gap_F

models/bitnetwork/ResBlock.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SEAttention(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=8):
+        super(SEAttention, self).__init__()
+        self.se = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels // reduction, kernel_size=1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_channels // reduction, out_channels=out_channels, kernel_size=1, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        x = self.se(x) * x
+        return x
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=8):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.max_pool = nn.AdaptiveMaxPool2d((1, 1))
+
+        self.fc = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels // reduction, kernel_size=1, bias=False),
+                                nn.ReLU(inplace=True),
+                                nn.Conv2d(in_channels=out_channels // reduction, out_channels=out_channels, kernel_size=1, bias=False))
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super(SpatialAttention, self).__init__()
+
+        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv1(x)
+        return self.sigmoid(x)
+
+
+class CBAMAttention(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=8):
+        super(CBAMAttention, self).__init__()
+        self.ca = ChannelAttention(in_channels=in_channels, out_channels=out_channels, reduction=reduction)
+        self.sa = SpatialAttention()
+
+    def forward(self, x):
+        x = self.ca(x) * x
+        x = self.sa(x) * x
+        return x
+
+
+class h_sigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(h_sigmoid, self).__init__()
+        self.relu = nn.ReLU6(inplace=inplace)
+
+    def forward(self, x):
+        return self.relu(x + 3) / 6
+
+
+class h_swish(nn.Module):
+    def __init__(self, inplace=True):
+        super(h_swish, self).__init__()
+        self.sigmoid = h_sigmoid(inplace=inplace)
+
+    def forward(self, x):
+        return x * self.sigmoid(x)
+
+
+class CoordAttention(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction=8):
+        super(CoordAttention, self).__init__()
+        self.pool_w, self.pool_h = nn.AdaptiveAvgPool2d((1, None)), nn.AdaptiveAvgPool2d((None, 1))
+        temp_c = max(8, in_channels // reduction)
+        self.conv1 = nn.Conv2d(in_channels, temp_c, kernel_size=1, stride=1, padding=0)
+
+        self.bn1 = nn.InstanceNorm2d(temp_c)
+        self.act1 = h_swish() # nn.SiLU() # nn.Hardswish() # nn.SiLU()
+
+        self.conv2 = nn.Conv2d(temp_c, out_channels, kernel_size=1, stride=1, padding=0)
+        self.conv3 = nn.Conv2d(temp_c, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        short = x
+        n, c, H, W = x.shape
+        x_h, x_w = self.pool_h(x), self.pool_w(x).permute(0, 1, 3, 2)
+        x_cat = torch.cat([x_h, x_w], dim=2)
+        out = self.act1(self.bn1(self.conv1(x_cat)))
+        x_h, x_w = torch.split(out, [H, W], dim=2)
+        x_w = x_w.permute(0, 1, 3, 2)
+        out_h = torch.sigmoid(self.conv2(x_h))
+        out_w = torch.sigmoid(self.conv3(x_w))
+        return short * out_w * out_h
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction, stride, attention=None):
+        super(BasicBlock, self).__init__()
+
+        self.change = None
+        if (in_channels != out_channels or stride != 1):
+            self.change = nn.Sequential(
+                nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, padding=0,
+                          stride=stride, bias=False),
+                nn.InstanceNorm2d(out_channels)
+            )
+
+        self.left = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1,
+                      stride=stride, bias=False),
+            nn.InstanceNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1, bias=False),
+            nn.InstanceNorm2d(out_channels)
+        )
+
+        if attention == 'se':
+            print('SEAttention')
+            self.attention = SEAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        elif attention == 'cbam':
+            print('CBAMAttention')
+            self.attention = CBAMAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        elif attention == 'coord':
+            print('CoordAttention')
+            self.attention = CoordAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        else:
+            print('None Attention')
+            self.attention = nn.Identity()
+
+    def forward(self, x):
+        identity = x
+        x = self.left(x)
+        x = self.attention(x)
+
+        if self.change is not None:
+            identity = self.change(identity)
+
+        x += identity
+        x = F.relu(x)
+        return x
+
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, reduction, stride, attention=None):
+        super(BottleneckBlock, self).__init__()
+
+        self.change = None
+        if (in_channels != out_channels or stride != 1):
+            self.change = nn.Sequential(
+                nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, padding=0,
+                          stride=stride, bias=False),
+                nn.InstanceNorm2d(out_channels)
+            )
+
+        self.left = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1,
+                      stride=stride, padding=0, bias=False),
+            nn.InstanceNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1, bias=False),
+            nn.InstanceNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=1, padding=0, bias=False),
+            nn.InstanceNorm2d(out_channels)
+        )
+
+        if attention == 'se':
+            print('SEAttention')
+            self.attention = SEAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        elif attention == 'cbam':
+            print('CBAMAttention')
+            self.attention = CBAMAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        elif attention == 'coord':
+            print('CoordAttention')
+            self.attention = CoordAttention(in_channels=out_channels, out_channels=out_channels, reduction=reduction)
+        else:
+            print('None Attention')
+            self.attention = nn.Identity()
+
+    def forward(self, x):
+        identity = x
+        x = self.left(x)
+        x = self.attention(x)
+
+        if self.change is not None:
+            identity = self.change(identity)
+
+        x += identity
+        x = F.relu(x)
+        return x
+
+
+class ResBlock(nn.Module):
+
+    def __init__(self, in_channels, out_channels, blocks=1, block_type="BottleneckBlock", reduction=8, stride=1, attention=None):
+        super(ResBlock, self).__init__()
+
+        layers = [eval(block_type)(in_channels, out_channels, reduction, stride, attention=attention)] if blocks != 0 else []
+        for _ in range(blocks - 1):
+            layer = eval(block_type)(out_channels, out_channels, reduction, 1, attention=attention)
+            layers.append(layer)
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+

models/bitnetwork/__init__.py

@@ -0,0 +1,9 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+# import kornia.losses
+from PIL import Image
+from torchvision import transforms
+from .ResBlock import *
+from .ConvBlock import *

models/discrim.py

@@ -0,0 +1,169 @@
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn.utils import spectral_norm
+
+
+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):
+        # downsample
+        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)
+
+        return out
+    
+
+class GANLoss(nn.Module):
+    """Define GAN loss.
+
+    Args:
+        gan_type (str): Support 'vanilla', 'lsgan', 'wgan', 'hinge'.
+        real_label_val (float): The value for real label. Default: 1.0.
+        fake_label_val (float): The value for fake label. Default: 0.0.
+        loss_weight (float): Loss weight. Default: 1.0.
+            Note that loss_weight is only for generators; and it is always 1.0
+            for discriminators.
+    """
+
+    def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0, loss_weight=1.0):
+        super(GANLoss, self).__init__()
+        self.gan_type = gan_type
+        self.loss_weight = loss_weight
+        self.real_label_val = real_label_val
+        self.fake_label_val = fake_label_val
+
+        if self.gan_type == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif self.gan_type == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif self.gan_type == 'wgan':
+            self.loss = self._wgan_loss
+        elif self.gan_type == 'wgan_softplus':
+            self.loss = self._wgan_softplus_loss
+        elif self.gan_type == 'hinge':
+            self.loss = nn.ReLU()
+        else:
+            raise NotImplementedError(f'GAN type {self.gan_type} is not implemented.')
+
+    def _wgan_loss(self, input, target):
+        """wgan loss.
+
+        Args:
+            input (Tensor): Input tensor.
+            target (bool): Target label.
+
+        Returns:
+            Tensor: wgan loss.
+        """
+        return -input.mean() if target else input.mean()
+
+    def _wgan_softplus_loss(self, input, target):
+        """wgan loss with soft plus. softplus is a smooth approximation to the
+        ReLU function.
+
+        In StyleGAN2, it is called:
+            Logistic loss for discriminator;
+            Non-saturating loss for generator.
+
+        Args:
+            input (Tensor): Input tensor.
+            target (bool): Target label.
+
+        Returns:
+            Tensor: wgan loss.
+        """
+        return F.softplus(-input).mean() if target else F.softplus(input).mean()
+
+    def get_target_label(self, input, target_is_real):
+        """Get target label.
+
+        Args:
+            input (Tensor): Input tensor.
+            target_is_real (bool): Whether the target is real or fake.
+
+        Returns:
+            (bool | Tensor): Target tensor. Return bool for wgan, otherwise,
+                return Tensor.
+        """
+
+        if self.gan_type in ['wgan', 'wgan_softplus']:
+            return target_is_real
+        target_val = (self.real_label_val if target_is_real else self.fake_label_val)
+        return input.new_ones(input.size()) * target_val
+
+    def forward(self, input, target_is_real, is_disc=False):
+        """
+        Args:
+            input (Tensor): The input for the loss module, i.e., the network
+                prediction.
+            target_is_real (bool): Whether the targe is real or fake.
+            is_disc (bool): Whether the loss for discriminators or not.
+                Default: False.
+
+        Returns:
+            Tensor: GAN loss value.
+        """
+        target_label = self.get_target_label(input, target_is_real)
+        if self.gan_type == 'hinge':
+            if is_disc:  # for discriminators in hinge-gan
+                input = -input if target_is_real else input
+                loss = self.loss(1 + input).mean()
+            else:  # for generators in hinge-gan
+                loss = -input.mean()
+        else:  # other gan types
+            loss = self.loss(input, target_label)
+
+        # loss_weight is always 1.0 for discriminators
+        return loss if is_disc else loss * self.loss_weight

models/lr_scheduler.py

@@ -0,0 +1,142 @@
+import math
+from collections import Counter
+from collections import defaultdict
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class MultiStepLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, milestones, restarts=None, weights=None, gamma=0.1,
+                 clear_state=False, last_epoch=-1):
+        self.milestones = Counter(milestones)
+        self.gamma = gamma
+        self.clear_state = clear_state
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(MultiStepLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch in self.restarts:
+            if self.clear_state:
+                self.optimizer.state = defaultdict(dict)
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        if self.last_epoch not in self.milestones:
+            return [group['lr'] for group in self.optimizer.param_groups]
+        return [
+            group['lr'] * self.gamma**self.milestones[self.last_epoch]
+            for group in self.optimizer.param_groups
+        ]
+
+
+class CosineAnnealingLR_Restart(_LRScheduler):
+    def __init__(self, optimizer, T_period, restarts=None, weights=None, eta_min=0, last_epoch=-1):
+        self.T_period = T_period
+        self.T_max = self.T_period[0]  # current T period
+        self.eta_min = eta_min
+        self.restarts = restarts if restarts else [0]
+        self.restart_weights = weights if weights else [1]
+        self.last_restart = 0
+        assert len(self.restarts) == len(
+            self.restart_weights), 'restarts and their weights do not match.'
+        super(CosineAnnealingLR_Restart, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch == 0:
+            return self.base_lrs
+        elif self.last_epoch in self.restarts:
+            self.last_restart = self.last_epoch
+            self.T_max = self.T_period[self.restarts.index(self.last_epoch) + 1]
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        elif (self.last_epoch - self.last_restart - 1 - self.T_max) % (2 * self.T_max) == 0:
+            return [
+                group['lr'] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2
+                for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            ]
+        return [(1 + math.cos(math.pi * (self.last_epoch - self.last_restart) / self.T_max)) /
+                (1 + math.cos(math.pi * ((self.last_epoch - self.last_restart) - 1) / self.T_max)) *
+                (group['lr'] - self.eta_min) + self.eta_min
+                for group in self.optimizer.param_groups]
+
+
+if __name__ == "__main__":
+    optimizer = torch.optim.Adam([torch.zeros(3, 64, 3, 3)], lr=2e-4, weight_decay=0,
+                                 betas=(0.9, 0.99))
+    ##############################
+    # MultiStepLR_Restart
+    ##############################
+    ## Original
+    lr_steps = [200000, 400000, 600000, 800000]
+    restarts = None
+    restart_weights = None
+
+    ## two
+    lr_steps = [100000, 200000, 300000, 400000, 490000, 600000, 700000, 800000, 900000, 990000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    lr_steps = [
+        50000, 100000, 150000, 200000, 240000, 300000, 350000, 400000, 450000, 490000, 550000,
+        600000, 650000, 700000, 740000, 800000, 850000, 900000, 950000, 990000
+    ]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = MultiStepLR_Restart(optimizer, lr_steps, restarts, restart_weights, gamma=0.5,
+                                    clear_state=False)
+
+    ##############################
+    # Cosine Annealing Restart
+    ##############################
+    ## two
+    T_period = [500000, 500000]
+    restarts = [500000]
+    restart_weights = [1]
+
+    ## four
+    T_period = [250000, 250000, 250000, 250000]
+    restarts = [250000, 500000, 750000]
+    restart_weights = [1, 1, 1]
+
+    scheduler = CosineAnnealingLR_Restart(optimizer, T_period, eta_min=1e-7, restarts=restarts,
+                                          weights=restart_weights)
+
+    ##############################
+    # Draw figure
+    ##############################
+    N_iter = 1000000
+    lr_l = list(range(N_iter))
+    for i in range(N_iter):
+        scheduler.step()
+        current_lr = optimizer.param_groups[0]['lr']
+        lr_l[i] = current_lr
+
+    import matplotlib as mpl
+    from matplotlib import pyplot as plt
+    import matplotlib.ticker as mtick
+    mpl.style.use('default')
+    import seaborn
+    seaborn.set(style='whitegrid')
+    seaborn.set_context('paper')
+
+    plt.figure(1)
+    plt.subplot(111)
+    plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0))
+    plt.title('Title', fontsize=16, color='k')
+    plt.plot(list(range(N_iter)), lr_l, linewidth=1.5, label='learning rate scheme')
+    legend = plt.legend(loc='upper right', shadow=False)
+    ax = plt.gca()
+    labels = ax.get_xticks().tolist()
+    for k, v in enumerate(labels):
+        labels[k] = str(int(v / 1000)) + 'K'
+    ax.set_xticklabels(labels)
+    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))
+
+    ax.set_ylabel('Learning rate')
+    ax.set_xlabel('Iteration')
+    fig = plt.gcf()
+    plt.show()

models/modules/Inv_arch.py

@@ -0,0 +1,584 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .module_util import initialize_weights_xavier
+from torch.nn import init
+from .common import DWT,IWT
+import cv2
+from basicsr.archs.arch_util import flow_warp
+from models.modules.Subnet_constructor import subnet
+import numpy as np
+
+from pdb import set_trace as stx
+import numbers
+
+from einops import rearrange
+from models.bitnetwork.Encoder_U import DW_Encoder
+from models.bitnetwork.Decoder_U import DW_Decoder
+
+
+## Layer Norm
+def to_3d(x):
+    return rearrange(x, 'b c h w -> b (h w) c')
+
+
+def to_4d(x, h, w):
+    return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+
+
+class BiasFree_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(BiasFree_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return x / torch.sqrt(sigma + 1e-5) * self.weight
+
+
+class WithBias_LayerNorm(nn.Module):
+    def __init__(self, normalized_shape):
+        super(WithBias_LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        normalized_shape = torch.Size(normalized_shape)
+
+        assert len(normalized_shape) == 1
+
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.normalized_shape = normalized_shape
+
+    def forward(self, x):
+        mu = x.mean(-1, keepdim=True)
+        sigma = x.var(-1, keepdim=True, unbiased=False)
+        return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, LayerNorm_type):
+        super(LayerNorm, self).__init__()
+        if LayerNorm_type == 'BiasFree':
+            self.body = BiasFree_LayerNorm(dim)
+        else:
+            self.body = WithBias_LayerNorm(dim)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        return to_4d(self.body(to_3d(x)), h, w)
+
+
+##########################################################################
+## Gated-Dconv Feed-Forward Network (GDFN)
+class FeedForward(nn.Module):
+    def __init__(self, dim, ffn_expansion_factor, bias):
+        super(FeedForward, self).__init__()
+
+        hidden_features = int(dim * ffn_expansion_factor)
+
+        self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
+
+        self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1,
+                                groups=hidden_features * 2, bias=bias)
+
+        self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        x = self.project_in(x)
+        x1, x2 = self.dwconv(x).chunk(2, dim=1)
+        x = F.gelu(x1) * x2
+        x = self.project_out(x)
+        return x
+
+
+##########################################################################
+## Multi-DConv Head Transposed Self-Attention (MDTA)
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads, bias):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
+
+        self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
+        self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias)
+        self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        qkv = self.qkv_dwconv(self.qkv(x))
+        q, k, v = qkv.chunk(3, dim=1)
+
+        q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+        v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
+
+        q = torch.nn.functional.normalize(q, dim=-1)
+        k = torch.nn.functional.normalize(k, dim=-1)
+
+        attn = (q @ k.transpose(-2, -1)) * self.temperature
+        attn = attn.softmax(dim=-1)
+
+        out = (attn @ v)
+
+        out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
+
+        out = self.project_out(out)
+        return out
+
+
+##########################################################################
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads=4, ffn_expansion_factor=4, bias=False, LayerNorm_type="withbias"):
+        super(TransformerBlock, self).__init__()
+
+        self.norm1 = LayerNorm(dim, LayerNorm_type)
+        self.attn = Attention(dim, num_heads, bias)
+        self.norm2 = LayerNorm(dim, LayerNorm_type)
+        self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+
+        return x
+
+dwt=DWT()
+iwt=IWT()
+
+class LayerNormFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, eps):
+        ctx.eps = eps
+        N, C, H, W = x.size()
+        mu = x.mean(1, keepdim=True)
+        var = (x - mu).pow(2).mean(1, keepdim=True)
+        y = (x - mu) / (var + eps).sqrt()
+        ctx.save_for_backward(y, var, weight)
+        y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        eps = ctx.eps
+
+        N, C, H, W = grad_output.size()
+        y, var, weight = ctx.saved_variables
+        g = grad_output * weight.view(1, C, 1, 1)
+        mean_g = g.mean(dim=1, keepdim=True)
+
+        mean_gy = (g * y).mean(dim=1, keepdim=True)
+        gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
+        return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum(
+            dim=0), None
+
+class LayerNorm2d(nn.Module):
+
+    def __init__(self, channels, eps=1e-6):
+        super(LayerNorm2d, self).__init__()
+        self.register_parameter('weight', nn.Parameter(torch.ones(channels)))
+        self.register_parameter('bias', nn.Parameter(torch.zeros(channels)))
+        self.eps = eps
+
+    def forward(self, x):
+        return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+class SimpleGate(nn.Module):
+    def forward(self, x):
+        x1, x2 = x.chunk(2, dim=1)
+        return x1 * x2
+
+class NAFBlock(nn.Module):
+    def __init__(self, c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.):
+        super().__init__()
+        dw_channel = c * DW_Expand
+        self.conv1 = nn.Conv2d(in_channels=c, out_channels=dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        self.conv2 = nn.Conv2d(in_channels=dw_channel, out_channels=dw_channel, kernel_size=3, padding=1, stride=1, groups=dw_channel,
+                               bias=True)
+        self.conv3 = nn.Conv2d(in_channels=dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        
+        # Simplified Channel Attention
+        self.sca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels=dw_channel // 2, out_channels=dw_channel // 2, kernel_size=1, padding=0, stride=1,
+                      groups=1, bias=True),
+        )
+
+        # SimpleGate
+        self.sg = SimpleGate()
+
+        ffn_channel = FFN_Expand * c
+        self.conv4 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        self.conv5 = nn.Conv2d(in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+
+        self.norm1 = LayerNorm2d(c)
+        self.norm2 = LayerNorm2d(c)
+
+        self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
+        self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
+
+        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+
+    def forward(self, inp):
+        x = inp
+
+        x = self.norm1(x)
+
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.sg(x)
+        x = x * self.sca(x)
+        x = self.conv3(x)
+
+        x = self.dropout1(x)
+
+        y = inp + x * self.beta
+
+        x = self.conv4(self.norm2(y))
+        x = self.sg(x)
+        x = self.conv5(x)
+
+        x = self.dropout2(x)
+
+        return y + x * self.gamma
+
+def thops_mean(tensor, dim=None, keepdim=False):
+    if dim is None:
+        # mean all dim
+        return torch.mean(tensor)
+    else:
+        if isinstance(dim, int):
+            dim = [dim]
+        dim = sorted(dim)
+        for d in dim:
+            tensor = tensor.mean(dim=d, keepdim=True)
+        if not keepdim:
+            for i, d in enumerate(dim):
+                tensor.squeeze_(d-i)
+        return tensor
+
+
+class ResidualBlockNoBN(nn.Module):
+    def __init__(self, nf=64, model='MIMO-VRN'):
+        super(ResidualBlockNoBN, self).__init__()
+        self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        # honestly, there's no significant difference between ReLU and leaky ReLU in terms of performance here
+        # but this is how we trained the model in the first place and what we reported in the paper
+        if model == 'LSTM-VRN':
+            self.relu = nn.ReLU(inplace=True)
+        elif model == 'MIMO-VRN':
+            self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        initialize_weights_xavier([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = self.relu(self.conv1(x))
+        out = self.conv2(out)
+        return identity + out
+
+
+class InvBlock(nn.Module):
+    def __init__(self, subnet_constructor, subnet_constructor_v2, channel_num_ho, channel_num_hi, groups, clamp=1.):
+        super(InvBlock, self).__init__()
+        self.split_len1 = channel_num_ho  # channel_split_num
+        self.split_len2 = channel_num_hi  # channel_num - channel_split_num
+        self.clamp = clamp
+
+        self.F = subnet_constructor_v2(self.split_len2, self.split_len1, groups=groups)
+        self.NF = NAFBlock(self.split_len2)
+        if groups == 1: 
+            self.G = subnet_constructor(self.split_len1, self.split_len2, groups=groups)
+            self.NG = NAFBlock(self.split_len1)
+            self.H = subnet_constructor(self.split_len1, self.split_len2, groups=groups)
+            self.NH = NAFBlock(self.split_len1)
+        else:
+            self.G = subnet_constructor(self.split_len1, self.split_len2)
+            self.NG = NAFBlock(self.split_len1)
+            self.H = subnet_constructor(self.split_len1, self.split_len2)
+            self.NH = NAFBlock(self.split_len1)
+
+    def forward(self, x1, x2, rev=False):
+        if not rev:
+            y1 = x1 + self.NF(self.F(x2))
+            self.s = self.clamp * (torch.sigmoid(self.NH(self.H(y1))) * 2 - 1)
+            y2 = [x2i.mul(torch.exp(self.s)) + self.NG(self.G(y1)) for x2i in x2]
+        else:
+            self.s = self.clamp * (torch.sigmoid(self.NH(self.H(x1))) * 2 - 1)
+            y2 = [(x2i - self.NG(self.G(x1))).div(torch.exp(self.s)) for x2i in x2]
+            y1 = x1 - self.NF(self.F(y2))
+
+        return y1, y2  # torch.cat((y1, y2), 1)
+
+    def jacobian(self, x, rev=False):
+        if not rev:
+            jac = torch.sum(self.s)
+        else:
+            jac = -torch.sum(self.s)
+
+        return jac / x.shape[0]
+
+class InvNN(nn.Module):
+    def __init__(self, channel_in_ho=3, channel_in_hi=3, subnet_constructor=None, subnet_constructor_v2=None, block_num=[], down_num=2, groups=None):
+        super(InvNN, self).__init__()
+        operations = []
+
+        current_channel_ho = channel_in_ho
+        current_channel_hi = channel_in_hi
+        for i in range(down_num):
+            for j in range(block_num[i]):
+                b = InvBlock(subnet_constructor, subnet_constructor_v2, current_channel_ho, current_channel_hi, groups=groups)
+                operations.append(b)
+
+        self.operations = nn.ModuleList(operations)
+
+    def forward(self, x, x_h, rev=False, cal_jacobian=False):
+        # 		out = x
+        jacobian = 0
+
+        if not rev:
+            for op in self.operations:
+                x, x_h = op.forward(x, x_h, rev)
+                if cal_jacobian:
+                    jacobian += op.jacobian(x, rev)
+        else:
+            for op in reversed(self.operations):
+                x, x_h = op.forward(x, x_h, rev)
+                if cal_jacobian:
+                    jacobian += op.jacobian(x, rev)
+
+        if cal_jacobian:
+            return x, x_h, jacobian
+        else:
+            return x, x_h
+
+class PredictiveModuleMIMO(nn.Module):
+    def __init__(self, channel_in, nf, block_num_rbm=8, block_num_trans=4):
+        super(PredictiveModuleMIMO, self).__init__()
+        self.conv_in = nn.Conv2d(channel_in, nf, 3, 1, 1, bias=True)
+        res_block = []
+        trans_block = []
+        for i in range(block_num_rbm):
+            res_block.append(ResidualBlockNoBN(nf))
+        for j in range(block_num_trans):
+            trans_block.append(TransformerBlock(nf))
+
+        self.res_block = nn.Sequential(*res_block)
+        self.transformer_block = nn.Sequential(*trans_block)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.res_block(x)
+        res = self.transformer_block(x) + x
+
+        return res
+
+class ConvRelu(nn.Module):
+    def __init__(self, channels_in, channels_out, stride=1, init_zero=False):
+        super(ConvRelu, self).__init__()
+        self.init_zero = init_zero
+        if self.init_zero:
+            self.layers = nn.Conv2d(channels_in, channels_out, 3, stride, padding=1)
+
+        else:
+            self.layers = nn.Sequential(
+                nn.Conv2d(channels_in, channels_out, 3, stride, padding=1),
+                nn.LeakyReLU(inplace=True)
+            )
+
+    def forward(self, x):
+        return self.layers(x)
+    
+class PredictiveModuleBit(nn.Module):
+    def __init__(self, channel_in, nf, block_num_rbm=4, block_num_trans=2):
+        super(PredictiveModuleBit, self).__init__()
+        self.conv_in = nn.Conv2d(channel_in, nf, 3, 1, 1, bias=True)
+        res_block = []
+        trans_block = []
+        for i in range(block_num_rbm):
+            res_block.append(ResidualBlockNoBN(nf))
+        for j in range(block_num_trans):
+            trans_block.append(TransformerBlock(nf))
+        
+        blocks = 4
+        layers = [ConvRelu(nf, 1, 2)]
+        for _ in range(blocks - 1):
+            layer = ConvRelu(1, 1, 2)
+            layers.append(layer)
+        self.layers = nn.Sequential(*layers)
+
+        self.res_block = nn.Sequential(*res_block)
+        self.transformer_block = nn.Sequential(*trans_block)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.res_block(x)
+        res = self.transformer_block(x) + x
+        res = self.layers(res)
+
+        return res
+
+
+##---------- Prompt Gen Module -----------------------
+class PromptGenBlock(nn.Module):
+    def __init__(self,prompt_dim=12,prompt_len=3,prompt_size = 36,lin_dim = 12):
+        super(PromptGenBlock,self).__init__()
+        self.prompt_param = nn.Parameter(torch.rand(1,prompt_len,prompt_dim,prompt_size,prompt_size))
+        self.linear_layer = nn.Linear(lin_dim,prompt_len)
+        self.conv3x3 = nn.Conv2d(prompt_dim,prompt_dim,kernel_size=3,stride=1,padding=1,bias=False)
+        
+
+    def forward(self,x):
+        B,C,H,W = x.shape
+        emb = x.mean(dim=(-2,-1))
+        prompt_weights = F.softmax(self.linear_layer(emb),dim=1)
+        prompt = prompt_weights.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) * self.prompt_param.unsqueeze(0).repeat(B,1,1,1,1,1).squeeze(1)
+        prompt = torch.sum(prompt,dim=1)
+        prompt = F.interpolate(prompt,(H,W),mode="bilinear")
+        prompt = self.conv3x3(prompt)
+
+        return prompt
+
+class PredictiveModuleMIMO_prompt(nn.Module):
+    def __init__(self, channel_in, nf, prompt_len=3, block_num_rbm=8, block_num_trans=4):
+        super(PredictiveModuleMIMO_prompt, self).__init__()
+        self.conv_in = nn.Conv2d(channel_in, nf, 3, 1, 1, bias=True)
+        res_block = []
+        trans_block = []
+        for i in range(block_num_rbm):
+            res_block.append(ResidualBlockNoBN(nf))
+        for j in range(block_num_trans):
+            trans_block.append(TransformerBlock(nf))
+
+        self.res_block = nn.Sequential(*res_block)
+        self.transformer_block = nn.Sequential(*trans_block)
+        self.prompt = PromptGenBlock(prompt_dim=nf,prompt_len=prompt_len,prompt_size = 36,lin_dim = nf)
+        self.fuse = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.res_block(x)
+        res = self.transformer_block(x) + x
+        prompt = self.prompt(res)
+
+        result = self.fuse(torch.cat([res, prompt], dim=1))
+
+        return result
+
+def gauss_noise(shape):
+    noise = torch.zeros(shape).cuda()
+    for i in range(noise.shape[0]):
+        noise[i] = torch.randn(noise[i].shape).cuda()
+
+    return noise
+
+def gauss_noise_mul(shape):
+    noise = torch.randn(shape).cuda()
+
+    return noise
+
+class PredictiveModuleBit_prompt(nn.Module):
+    def __init__(self, channel_in, nf, prompt_length, block_num_rbm=4, block_num_trans=2):
+        super(PredictiveModuleBit_prompt, self).__init__()
+        self.conv_in = nn.Conv2d(channel_in, nf, 3, 1, 1, bias=True)
+        res_block = []
+        trans_block = []
+        for i in range(block_num_rbm):
+            res_block.append(ResidualBlockNoBN(nf))
+        for j in range(block_num_trans):
+            trans_block.append(TransformerBlock(nf))
+        
+        blocks = 4
+        layers = [ConvRelu(nf, 1, 2)]
+        for _ in range(blocks - 1):
+            layer = ConvRelu(1, 1, 2)
+            layers.append(layer)
+        self.layers = nn.Sequential(*layers)
+
+        self.res_block = nn.Sequential(*res_block)
+        self.transformer_block = nn.Sequential(*trans_block)
+        self.prompt = PromptGenBlock(prompt_dim=nf,prompt_len=prompt_length,prompt_size = 36,lin_dim = nf)
+        self.fuse = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.res_block(x)
+        res = self.transformer_block(x) + x
+        prompt = self.prompt(res)
+        res = self.fuse(torch.cat([res, prompt], dim=1))
+        res = self.layers(res)
+
+        return res
+
+class VSN(nn.Module):
+    def __init__(self, opt, subnet_constructor=None, subnet_constructor_v2=None, down_num=2):
+        super(VSN, self).__init__()
+        self.model = opt['model']
+        self.mode = opt['mode']
+        opt_net = opt['network_G']
+        self.num_image = opt['num_image']
+        self.gop = opt['gop']
+        self.channel_in = opt_net['in_nc'] * self.gop
+        self.channel_out = opt_net['out_nc'] * self.gop
+        self.channel_in_hi = opt_net['in_nc'] * self.gop
+        self.channel_in_ho = opt_net['in_nc'] * self.gop
+        self.message_len = opt['message_length']
+
+        self.block_num = opt_net['block_num']
+        self.block_num_rbm = opt_net['block_num_rbm']
+        self.block_num_trans = opt_net['block_num_trans']
+        self.nf = self.channel_in_hi 
+        
+        self.bitencoder = DW_Encoder(self.message_len, attention = "se")
+        self.bitdecoder = DW_Decoder(self.message_len, attention = "se")
+        self.irn = InvNN(self.channel_in_ho, self.channel_in_hi, subnet_constructor, subnet_constructor_v2, self.block_num, down_num, groups=self.num_image)
+
+        if opt['prompt']:
+            self.pm = PredictiveModuleMIMO_prompt(self.channel_in_ho, self.nf* self.num_image, opt['prompt_len'], block_num_rbm=self.block_num_rbm, block_num_trans=self.block_num_trans)
+        else:
+            self.pm = PredictiveModuleMIMO(self.channel_in_ho, self.nf* self.num_image, opt['prompt_len'], block_num_rbm=self.block_num_rbm, block_num_trans=self.block_num_trans)
+            self.BitPM = PredictiveModuleBit(3, 4, block_num_rbm=4, block_num_trans=2)
+
+
+    def forward(self, x, x_h=None, message=None, rev=False, hs=[], direction='f'):
+        if not rev:
+            if self.mode == "image":
+                out_y, out_y_h = self.irn(x, x_h, rev)
+                out_y = iwt(out_y)
+                encoded_image = self.bitencoder(out_y, message)          
+                return out_y, encoded_image
+            
+            elif self.mode == "bit":
+                out_y = iwt(x)
+                encoded_image = self.bitencoder(out_y, message)            
+                return out_y, encoded_image
+
+        else:
+            if self.mode == "image":
+                recmessage = self.bitdecoder(x)
+
+                x = dwt(x)
+                out_z = self.pm(x).unsqueeze(1)
+                out_z_new = out_z.view(-1, self.num_image, self.channel_in, x.shape[-2], x.shape[-1])
+                out_z_new = [out_z_new[:,i] for i in range(self.num_image)]
+                out_x, out_x_h = self.irn(x, out_z_new, rev)
+
+                return out_x, out_x_h, out_z, recmessage
+            
+            elif self.mode == "bit":
+                recmessage = self.bitdecoder(x)
+                return recmessage
+

models/modules/Quantization.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+
+class Quant(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input):
+        input = torch.clamp(input, 0, 1)
+        output = (input * 255.).round() / 255.
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+class Quantization(nn.Module):
+    def __init__(self):
+        super(Quantization, self).__init__()
+
+    def forward(self, input):
+        return Quant.apply(input)

models/modules/Subnet_constructor.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import models.modules.module_util as mutil
+from basicsr.archs.arch_util import flow_warp, ResidualBlockNoBN
+from models.modules.module_util import initialize_weights_xavier
+
+class DenseBlock(nn.Module):
+    def __init__(self, channel_in, channel_out, init='xavier', gc=32, bias=True):
+        super(DenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channel_in, gc, 3, 1, 1, bias=bias)
+        self.conv2 = nn.Conv2d(channel_in + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = nn.Conv2d(channel_in + 2 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = nn.Conv2d(channel_in + 3 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = nn.Conv2d(channel_in + 4 * gc, channel_out, 3, 1, 1, bias=bias)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.H = None
+
+        if init == 'xavier':
+            mutil.initialize_weights_xavier([self.conv1, self.conv2, self.conv3, self.conv4], 0.1)
+        else:
+            mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4], 0.1)
+        mutil.initialize_weights(self.conv5, 0)
+
+    def forward(self, x):
+        if isinstance(x, list):
+            x = x[0]
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+
+        return x5
+
+class DenseBlock_v2(nn.Module):
+    def __init__(self, channel_in, channel_out, groups, init='xavier', gc=32, bias=True):
+        super(DenseBlock_v2, self).__init__()
+        self.conv1 = nn.Conv2d(channel_in, gc, 3, 1, 1, bias=bias)
+        self.conv2 = nn.Conv2d(channel_in + gc, gc, 3, 1, 1, bias=bias)
+        self.conv3 = nn.Conv2d(channel_in + 2 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv4 = nn.Conv2d(channel_in + 3 * gc, gc, 3, 1, 1, bias=bias)
+        self.conv5 = nn.Conv2d(channel_in + 4 * gc, channel_out, 3, 1, 1, bias=bias)
+        self.conv_final = nn.Conv2d(channel_out*groups, channel_out, 3, 1, 1, bias=bias)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        if init == 'xavier':
+            mutil.initialize_weights_xavier([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+        else:
+            mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+        mutil.initialize_weights(self.conv_final, 0)
+
+    def forward(self, x):
+        res = []
+        for xi in x:
+            x1 = self.lrelu(self.conv1(xi))
+            x2 = self.lrelu(self.conv2(torch.cat((xi, x1), 1)))
+            x3 = self.lrelu(self.conv3(torch.cat((xi, x1, x2), 1)))
+            x4 = self.lrelu(self.conv4(torch.cat((xi, x1, x2, x3), 1)))
+            x5 = self.lrelu(self.conv5(torch.cat((xi, x1, x2, x3, x4), 1)))
+            res.append(x5)
+        res = torch.cat(res, dim=1)
+        res = self.conv_final(res)
+
+        return res
+
+def subnet(net_structure, init='xavier'):
+    def constructor(channel_in, channel_out, groups=None):
+        if net_structure == 'DBNet':
+            if init == 'xavier':
+                return DenseBlock(channel_in, channel_out, init)
+            elif init == 'xavier_v2':
+                return DenseBlock_v2(channel_in, channel_out, groups, 'xavier')
+            else:
+                return DenseBlock(channel_in, channel_out)
+        else:
+            return None
+
+    return constructor