hfspace gradio demo

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,2 @@
----
-title: BayesCap
-emoji: 🦀
-colorFrom: purple
-colorTo: green
-sdk: gradio
-sdk_version: 3.0.24
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BayesCap
+Bayesian Identity Cap for Calibrated Uncertainty in Pretrained Neural Networks

app.py

@@ -0,0 +1,119 @@
+import numpy as np
+import random
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode as IMode
+
+from PIL import Image
+
+from ds import *
+from losses import *
+from networks_SRGAN import *
+from utils import *
+
+device = 'cuda'
+
+
+NetG = Generator()
+model_parameters = filter(lambda p: True, NetG.parameters())
+params = sum([np.prod(p.size()) for p in model_parameters])
+print("Number of Parameters:", params)
+NetC = BayesCap(in_channels=3, out_channels=3)
+
+ensure_checkpoint_exists('BayesCap_SRGAN.pth')
+NetG.load_state_dict(torch.load('BayesCap_SRGAN.pth', map_location=device))
+NetG.to(device)
+NetG.eval()
+
+ensure_checkpoint_exists('BayesCap_ckpt.pth')
+NetC.load_state_dict(torch.load('BayesCap_ckpt.pth', map_location=device))
+NetC.to(device)
+NetC.eval()
+
+def tensor01_to_pil(xt):
+    r = transforms.ToPILImage(mode='RGB')(xt.squeeze())
+    return r
+
+
+def predict(img):
+    """
+    img: image
+    """
+    image_size = (256,256)
+    upscale_factor = 4
+    lr_transforms = transforms.Resize((image_size[0]//upscale_factor, image_size[1]//upscale_factor), interpolation=IMode.BICUBIC, antialias=True)
+    # lr_transforms = transforms.Resize((128, 128), interpolation=IMode.BICUBIC, antialias=True)
+    
+    img = Image.fromarray(np.array(img))
+    img = lr_transforms(img)
+    lr_tensor = utils.image2tensor(img, range_norm=False, half=False)
+    
+    device = 'cuda'
+    dtype = torch.cuda.FloatTensor
+    xLR = lr_tensor.to(device).unsqueeze(0)
+    xLR = xLR.type(dtype)
+    # pass them through the network
+    with torch.no_grad():
+        xSR = NetG(xLR)
+        xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+        
+    a_map = (1/(xSRC_alpha[0] + 1e-5)).to('cpu').data
+    b_map = xSRC_beta[0].to('cpu').data
+    u_map = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2)))) 
+    
+    
+    x_LR = tensor01_to_pil(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    
+    x_mean = tensor01_to_pil(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    
+    #im = Image.fromarray(np.uint8(cm.gist_earth(myarray)*255))
+    
+    a_map = torch.clamp(a_map, min=0, max=0.1)
+    a_map = (a_map - a_map.min())/(a_map.max() - a_map.min())
+    x_alpha = Image.fromarray(np.uint8(cm.inferno(a_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    b_map = torch.clamp(b_map, min=0.45, max=0.75)
+    b_map = (b_map - b_map.min())/(b_map.max() - b_map.min())
+    x_beta = Image.fromarray(np.uint8(cm.cividis(b_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    u_map = torch.clamp(u_map, min=0, max=0.15)
+    u_map = (u_map - u_map.min())/(u_map.max() - u_map.min())
+    x_uncer = Image.fromarray(np.uint8(cm.hot(u_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    return x_LR, x_mean, x_alpha, x_beta, x_uncer
+
+import gradio as gr
+
+title = "BayesCap"
+description = "BayesCap: Bayesian Identity Cap for Calibrated Uncertainty in Frozen Neural Networks (ECCV 2022)"
+article = "<p style='text-align: center'> BayesCap: Bayesian Identity Cap for Calibrated Uncertainty in Frozen Neural Networks| <a href='https://github.com/ExplainableML/BayesCap'>Github Repo</a></p>"
+
+
+gr.Interface(
+    fn=predict, 
+    inputs=gr.inputs.Image(type='pil', label="Orignal"), 
+    outputs=[
+        gr.outputs.Image(type='pil', label="Low-res"), 
+        gr.outputs.Image(type='pil', label="Super-res"), 
+        gr.outputs.Image(type='pil', label="Alpha"), 
+        gr.outputs.Image(type='pil', label="Beta"), 
+        gr.outputs.Image(type='pil', label="Uncertainty")
+     ],
+    title=title,
+    description=description,
+    article=article,
+     examples=[
+        ["./demo_examples/baby.png"],
+        ["./demo_examples/bird.png"],
+        ["./demo_examples/butterfly.png"],
+        ["./demo_examples/head.png"],
+        ["./demo_examples/woman.png"],
+    ]
+).launch(share=True)

ds.py

@@ -0,0 +1,485 @@
+from __future__ import absolute_import, division, print_function
+
+import random
+import copy
+import io
+import os
+import numpy as np
+from PIL import Image
+import skimage.transform
+from collections import Counter
+
+
+import torch
+import torch.utils.data as data
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode as IMode
+
+import utils
+
+class ImgDset(Dataset):
+    """Customize the data set loading function and prepare low/high resolution image data in advance.
+
+    Args:
+        dataroot         (str): Training data set address
+        image_size       (int): High resolution image size
+        upscale_factor   (int): Image magnification
+        mode             (str): Data set loading method, the training data set is for data enhancement,
+                             and the verification data set is not for data enhancement
+
+    """
+
+    def __init__(self, dataroot: str, image_size: int, upscale_factor: int, mode: str) -> None:
+        super(ImgDset, self).__init__()
+        self.filenames = [os.path.join(dataroot, x) for x in os.listdir(dataroot)]
+
+        if mode == "train":
+            self.hr_transforms = transforms.Compose([
+                transforms.RandomCrop(image_size),
+                transforms.RandomRotation(90),
+                transforms.RandomHorizontalFlip(0.5),
+            ])
+        else:
+            self.hr_transforms = transforms.Resize(image_size)
+
+        self.lr_transforms = transforms.Resize((image_size[0]//upscale_factor, image_size[1]//upscale_factor), interpolation=IMode.BICUBIC, antialias=True)
+
+    def __getitem__(self, batch_index: int) -> [Tensor, Tensor]:
+        # Read a batch of image data
+        image = Image.open(self.filenames[batch_index])
+
+        # Transform image
+        hr_image = self.hr_transforms(image)
+        lr_image = self.lr_transforms(hr_image)
+
+        # Convert image data into Tensor stream format (PyTorch).
+        # Note: The range of input and output is between [0, 1]
+        lr_tensor = utils.image2tensor(lr_image, range_norm=False, half=False)
+        hr_tensor = utils.image2tensor(hr_image, range_norm=False, half=False)
+
+        return lr_tensor, hr_tensor
+
+    def __len__(self) -> int:
+        return len(self.filenames)
+    
+
+class PairedImages_w_nameList(Dataset):
+    '''
+    can act as supervised or un-supervised based on flists
+    '''
+    def __init__(self, flist1, flist2, transform1=None, transform2=None, do_aug=False):
+        self.flist1 = flist1
+        self.flist2 = flist2
+        self.transform1 = transform1
+        self.transform2 = transform2
+        self.do_aug = do_aug
+    def __getitem__(self, index):
+        impath1 = self.flist1[index]
+        img1 = Image.open(impath1).convert('RGB')
+        impath2 = self.flist2[index]
+        img2 = Image.open(impath2).convert('RGB')
+        
+        img1 = utils.image2tensor(img1, range_norm=False, half=False)
+        img2 = utils.image2tensor(img2, range_norm=False, half=False)
+        
+        if self.transform1 is not None:
+            img1 = self.transform1(img1)
+        if self.transform2 is not None:
+            img2 = self.transform2(img2)
+        
+        return img1, img2
+    def __len__(self):
+        return len(self.flist1)
+    
+class PairedImages_w_nameList_npy(Dataset):
+    '''
+    can act as supervised or un-supervised based on flists
+    '''
+    def __init__(self, flist1, flist2, transform1=None, transform2=None, do_aug=False):
+        self.flist1 = flist1
+        self.flist2 = flist2
+        self.transform1 = transform1
+        self.transform2 = transform2
+        self.do_aug = do_aug
+    def __getitem__(self, index):
+        impath1 = self.flist1[index]
+        img1 = np.load(impath1)
+        impath2 = self.flist2[index]
+        img2 = np.load(impath2)
+        
+        if self.transform1 is not None:
+            img1 = self.transform1(img1)
+        if self.transform2 is not None:
+            img2 = self.transform2(img2)
+        
+        return img1, img2
+    def __len__(self):
+        return len(self.flist1)
+
+# def call_paired():
+#     root1='./GOPRO_3840FPS_AVG_3-21/train/blur/'
+#     root2='./GOPRO_3840FPS_AVG_3-21/train/sharp/'
+
+#     flist1=glob.glob(root1+'/*/*.png')
+#     flist2=glob.glob(root2+'/*/*.png')
+
+#     dset = PairedImages_w_nameList(root1,root2,flist1,flist2)
+
+#### KITTI depth
+
+def load_velodyne_points(filename):
+    """Load 3D point cloud from KITTI file format
+    (adapted from https://github.com/hunse/kitti)
+    """
+    points = np.fromfile(filename, dtype=np.float32).reshape(-1, 4)
+    points[:, 3] = 1.0  # homogeneous
+    return points
+
+
+def read_calib_file(path):
+    """Read KITTI calibration file
+    (from https://github.com/hunse/kitti)
+    """
+    float_chars = set("0123456789.e+- ")
+    data = {}
+    with open(path, 'r') as f:
+        for line in f.readlines():
+            key, value = line.split(':', 1)
+            value = value.strip()
+            data[key] = value
+            if float_chars.issuperset(value):
+                # try to cast to float array
+                try:
+                    data[key] = np.array(list(map(float, value.split(' '))))
+                except ValueError:
+                    # casting error: data[key] already eq. value, so pass
+                    pass
+
+    return data
+
+
+def sub2ind(matrixSize, rowSub, colSub):
+    """Convert row, col matrix subscripts to linear indices
+    """
+    m, n = matrixSize
+    return rowSub * (n-1) + colSub - 1
+
+
+def generate_depth_map(calib_dir, velo_filename, cam=2, vel_depth=False):
+    """Generate a depth map from velodyne data
+    """
+    # load calibration files
+    cam2cam = read_calib_file(os.path.join(calib_dir, 'calib_cam_to_cam.txt'))
+    velo2cam = read_calib_file(os.path.join(calib_dir, 'calib_velo_to_cam.txt'))
+    velo2cam = np.hstack((velo2cam['R'].reshape(3, 3), velo2cam['T'][..., np.newaxis]))
+    velo2cam = np.vstack((velo2cam, np.array([0, 0, 0, 1.0])))
+
+    # get image shape
+    im_shape = cam2cam["S_rect_02"][::-1].astype(np.int32)
+
+    # compute projection matrix velodyne->image plane
+    R_cam2rect = np.eye(4)
+    R_cam2rect[:3, :3] = cam2cam['R_rect_00'].reshape(3, 3)
+    P_rect = cam2cam['P_rect_0'+str(cam)].reshape(3, 4)
+    P_velo2im = np.dot(np.dot(P_rect, R_cam2rect), velo2cam)
+
+    # load velodyne points and remove all behind image plane (approximation)
+    # each row of the velodyne data is forward, left, up, reflectance
+    velo = load_velodyne_points(velo_filename)
+    velo = velo[velo[:, 0] >= 0, :]
+
+    # project the points to the camera
+    velo_pts_im = np.dot(P_velo2im, velo.T).T
+    velo_pts_im[:, :2] = velo_pts_im[:, :2] / velo_pts_im[:, 2][..., np.newaxis]
+
+    if vel_depth:
+        velo_pts_im[:, 2] = velo[:, 0]
+
+    # check if in bounds
+    # use minus 1 to get the exact same value as KITTI matlab code
+    velo_pts_im[:, 0] = np.round(velo_pts_im[:, 0]) - 1
+    velo_pts_im[:, 1] = np.round(velo_pts_im[:, 1]) - 1
+    val_inds = (velo_pts_im[:, 0] >= 0) & (velo_pts_im[:, 1] >= 0)
+    val_inds = val_inds & (velo_pts_im[:, 0] < im_shape[1]) & (velo_pts_im[:, 1] < im_shape[0])
+    velo_pts_im = velo_pts_im[val_inds, :]
+
+    # project to image
+    depth = np.zeros((im_shape[:2]))
+    depth[velo_pts_im[:, 1].astype(np.int), velo_pts_im[:, 0].astype(np.int)] = velo_pts_im[:, 2]
+
+    # find the duplicate points and choose the closest depth
+    inds = sub2ind(depth.shape, velo_pts_im[:, 1], velo_pts_im[:, 0])
+    dupe_inds = [item for item, count in Counter(inds).items() if count > 1]
+    for dd in dupe_inds:
+        pts = np.where(inds == dd)[0]
+        x_loc = int(velo_pts_im[pts[0], 0])
+        y_loc = int(velo_pts_im[pts[0], 1])
+        depth[y_loc, x_loc] = velo_pts_im[pts, 2].min()
+    depth[depth < 0] = 0
+
+    return depth
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning
+    # (https://github.com/python-pillow/Pillow/issues/835)
+    with open(path, 'rb') as f:
+        with Image.open(f) as img:
+            return img.convert('RGB')
+
+
+class MonoDataset(data.Dataset):
+    """Superclass for monocular dataloaders
+
+    Args:
+        data_path
+        filenames
+        height
+        width
+        frame_idxs
+        num_scales
+        is_train
+        img_ext
+    """
+    def __init__(self,
+                 data_path,
+                 filenames,
+                 height,
+                 width,
+                 frame_idxs,
+                 num_scales,
+                 is_train=False,
+                 img_ext='.jpg'):
+        super(MonoDataset, self).__init__()
+
+        self.data_path = data_path
+        self.filenames = filenames
+        self.height = height
+        self.width = width
+        self.num_scales = num_scales
+        self.interp = Image.ANTIALIAS
+
+        self.frame_idxs = frame_idxs
+
+        self.is_train = is_train
+        self.img_ext = img_ext
+
+        self.loader = pil_loader
+        self.to_tensor = transforms.ToTensor()
+
+        # We need to specify augmentations differently in newer versions of torchvision.
+        # We first try the newer tuple version; if this fails we fall back to scalars
+        try:
+            self.brightness = (0.8, 1.2)
+            self.contrast = (0.8, 1.2)
+            self.saturation = (0.8, 1.2)
+            self.hue = (-0.1, 0.1)
+            transforms.ColorJitter.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+        except TypeError:
+            self.brightness = 0.2
+            self.contrast = 0.2
+            self.saturation = 0.2
+            self.hue = 0.1
+
+        self.resize = {}
+        for i in range(self.num_scales):
+            s = 2 ** i
+            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
+                                               interpolation=self.interp)
+
+        self.load_depth = self.check_depth()
+
+    def preprocess(self, inputs, color_aug):
+        """Resize colour images to the required scales and augment if required
+
+        We create the color_aug object in advance and apply the same augmentation to all
+        images in this item. This ensures that all images input to the pose network receive the
+        same augmentation.
+        """
+        for k in list(inputs):
+            frame = inputs[k]
+            if "color" in k:
+                n, im, i = k
+                for i in range(self.num_scales):
+                    inputs[(n, im, i)] = self.resize[i](inputs[(n, im, i - 1)])
+
+        for k in list(inputs):
+            f = inputs[k]
+            if "color" in k:
+                n, im, i = k
+                inputs[(n, im, i)] = self.to_tensor(f)
+                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f))
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        """Returns a single training item from the dataset as a dictionary.
+
+        Values correspond to torch tensors.
+        Keys in the dictionary are either strings or tuples:
+
+            ("color", <frame_id>, <scale>)          for raw colour images,
+            ("color_aug", <frame_id>, <scale>)      for augmented colour images,
+            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
+            "stereo_T"                              for camera extrinsics, and
+            "depth_gt"                              for ground truth depth maps.
+
+        <frame_id> is either:
+            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
+        or
+            "s" for the opposite image in the stereo pair.
+
+        <scale> is an integer representing the scale of the image relative to the fullsize image:
+            -1      images at native resolution as loaded from disk
+            0       images resized to (self.width,      self.height     )
+            1       images resized to (self.width // 2, self.height // 2)
+            2       images resized to (self.width // 4, self.height // 4)
+            3       images resized to (self.width // 8, self.height // 8)
+        """
+        inputs = {}
+
+        do_color_aug = self.is_train and random.random() > 0.5
+        do_flip = self.is_train and random.random() > 0.5
+
+        line = self.filenames[index].split()
+        folder = line[0]
+
+        if len(line) == 3:
+            frame_index = int(line[1])
+        else:
+            frame_index = 0
+
+        if len(line) == 3:
+            side = line[2]
+        else:
+            side = None
+
+        for i in self.frame_idxs:
+            if i == "s":
+                other_side = {"r": "l", "l": "r"}[side]
+                inputs[("color", i, -1)] = self.get_color(folder, frame_index, other_side, do_flip)
+            else:
+                inputs[("color", i, -1)] = self.get_color(folder, frame_index + i, side, do_flip)
+
+        # adjusting intrinsics to match each scale in the pyramid
+        for scale in range(self.num_scales):
+            K = self.K.copy()
+
+            K[0, :] *= self.width // (2 ** scale)
+            K[1, :] *= self.height // (2 ** scale)
+
+            inv_K = np.linalg.pinv(K)
+
+            inputs[("K", scale)] = torch.from_numpy(K)
+            inputs[("inv_K", scale)] = torch.from_numpy(inv_K)
+
+        if do_color_aug:
+            color_aug = transforms.ColorJitter.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+        else:
+            color_aug = (lambda x: x)
+
+        self.preprocess(inputs, color_aug)
+
+        for i in self.frame_idxs:
+            del inputs[("color", i, -1)]
+            del inputs[("color_aug", i, -1)]
+
+        if self.load_depth:
+            depth_gt = self.get_depth(folder, frame_index, side, do_flip)
+            inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
+            inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))
+
+        if "s" in self.frame_idxs:
+            stereo_T = np.eye(4, dtype=np.float32)
+            baseline_sign = -1 if do_flip else 1
+            side_sign = -1 if side == "l" else 1
+            stereo_T[0, 3] = side_sign * baseline_sign * 0.1
+
+            inputs["stereo_T"] = torch.from_numpy(stereo_T)
+
+        return inputs
+
+    def get_color(self, folder, frame_index, side, do_flip):
+        raise NotImplementedError
+
+    def check_depth(self):
+        raise NotImplementedError
+
+    def get_depth(self, folder, frame_index, side, do_flip):
+        raise NotImplementedError
+
+class KITTIDataset(MonoDataset):
+    """Superclass for different types of KITTI dataset loaders
+    """
+    def __init__(self, *args, **kwargs):
+        super(KITTIDataset, self).__init__(*args, **kwargs)
+
+        # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size.
+        # To normalize you need to scale the first row by 1 / image_width and the second row
+        # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered.
+        # If your principal point is far from the center you might need to disable the horizontal
+        # flip augmentation.
+        self.K = np.array([[0.58, 0, 0.5, 0],
+                           [0, 1.92, 0.5, 0],
+                           [0, 0, 1, 0],
+                           [0, 0, 0, 1]], dtype=np.float32)
+
+        self.full_res_shape = (1242, 375)
+        self.side_map = {"2": 2, "3": 3, "l": 2, "r": 3}
+
+    def check_depth(self):
+        line = self.filenames[0].split()
+        scene_name = line[0]
+        frame_index = int(line[1])
+
+        velo_filename = os.path.join(
+            self.data_path,
+            scene_name,
+            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))
+
+        return os.path.isfile(velo_filename)
+
+    def get_color(self, folder, frame_index, side, do_flip):
+        color = self.loader(self.get_image_path(folder, frame_index, side))
+
+        if do_flip:
+            color = color.transpose(Image.FLIP_LEFT_RIGHT)
+
+        return color
+
+
+class KITTIDepthDataset(KITTIDataset):
+    """KITTI dataset which uses the updated ground truth depth maps
+    """
+    def __init__(self, *args, **kwargs):
+        super(KITTIDepthDataset, self).__init__(*args, **kwargs)
+
+    def get_image_path(self, folder, frame_index, side):
+        f_str = "{:010d}{}".format(frame_index, self.img_ext)
+        image_path = os.path.join(
+            self.data_path,
+            folder,
+            "image_0{}/data".format(self.side_map[side]),
+            f_str)
+        return image_path
+
+    def get_depth(self, folder, frame_index, side, do_flip):
+        f_str = "{:010d}.png".format(frame_index)
+        depth_path = os.path.join(
+            self.data_path,
+            folder,
+            "proj_depth/groundtruth/image_0{}".format(self.side_map[side]),
+            f_str)
+
+        depth_gt = Image.open(depth_path)
+        depth_gt = depth_gt.resize(self.full_res_shape, Image.NEAREST)
+        depth_gt = np.array(depth_gt).astype(np.float32) / 256
+
+        if do_flip:
+            depth_gt = np.fliplr(depth_gt)
+
+        return depth_gt

losses.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch import Tensor
+
+class ContentLoss(nn.Module):
+	"""Constructs a content loss function based on the VGG19 network.
+	Using high-level feature mapping layers from the latter layers will focus more on the texture content of the image.
+
+	Paper reference list:
+		-`Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network <https://arxiv.org/pdf/1609.04802.pdf>` paper.
+		-`ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks                    <https://arxiv.org/pdf/1809.00219.pdf>` paper.
+		-`Perceptual Extreme Super Resolution Network with Receptive Field Block               <https://arxiv.org/pdf/2005.12597.pdf>` paper.
+
+	 """
+
+	def __init__(self) -> None:
+		super(ContentLoss, self).__init__()
+		# Load the VGG19 model trained on the ImageNet dataset.
+		vgg19 = models.vgg19(pretrained=True).eval()
+		# Extract the thirty-sixth layer output in the VGG19 model as the content loss.
+		self.feature_extractor = nn.Sequential(*list(vgg19.features.children())[:36])
+		# Freeze model parameters.
+		for parameters in self.feature_extractor.parameters():
+			parameters.requires_grad = False
+
+		# The preprocessing method of the input data. This is the VGG model preprocessing method of the ImageNet dataset.
+		self.register_buffer("mean", torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+		self.register_buffer("std", torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+	def forward(self, sr: Tensor, hr: Tensor) -> Tensor:
+		# Standardized operations
+		sr = sr.sub(self.mean).div(self.std)
+		hr = hr.sub(self.mean).div(self.std)
+
+		# Find the feature map difference between the two images
+		loss = F.l1_loss(self.feature_extractor(sr), self.feature_extractor(hr))
+
+		return loss
+
+
+class GenGaussLoss(nn.Module):
+	def __init__(
+		self, reduction='mean',
+		alpha_eps = 1e-4, beta_eps=1e-4,
+		resi_min = 1e-4, resi_max=1e3
+	) -> None:
+		super(GenGaussLoss, self).__init__()
+		self.reduction = reduction
+		self.alpha_eps = alpha_eps
+		self.beta_eps = beta_eps
+		self.resi_min = resi_min
+		self.resi_max = resi_max
+	
+	def forward(
+		self, 
+		mean: Tensor, one_over_alpha: Tensor, beta: Tensor, target: Tensor
+	):
+		one_over_alpha1 = one_over_alpha + self.alpha_eps
+		beta1 = beta + self.beta_eps
+
+		resi = torch.abs(mean - target)
+		# resi = torch.pow(resi*one_over_alpha1, beta1).clamp(min=self.resi_min, max=self.resi_max)
+		resi = (resi*one_over_alpha1*beta1).clamp(min=self.resi_min, max=self.resi_max)
+		## check if resi has nans
+		if torch.sum(resi != resi) > 0:
+			print('resi has nans!!')
+			return None
+		
+		log_one_over_alpha = torch.log(one_over_alpha1)
+		log_beta = torch.log(beta1)
+		lgamma_beta = torch.lgamma(torch.pow(beta1, -1))
+		
+		if torch.sum(log_one_over_alpha != log_one_over_alpha) > 0:
+			print('log_one_over_alpha has nan')
+		if torch.sum(lgamma_beta != lgamma_beta) > 0:
+			print('lgamma_beta has nan')
+		if torch.sum(log_beta != log_beta) > 0:
+			print('log_beta has nan')
+		
+		l = resi - log_one_over_alpha + lgamma_beta - log_beta
+
+		if self.reduction == 'mean':
+			return l.mean()
+		elif self.reduction == 'sum':
+			return l.sum()
+		else:
+			print('Reduction not supported')
+			return None
+
+class TempCombLoss(nn.Module):
+	def __init__(
+		self, reduction='mean',
+		alpha_eps = 1e-4, beta_eps=1e-4,
+		resi_min = 1e-4, resi_max=1e3
+	) -> None:
+		super(TempCombLoss, self).__init__()
+		self.reduction = reduction
+		self.alpha_eps = alpha_eps
+		self.beta_eps = beta_eps
+		self.resi_min = resi_min
+		self.resi_max = resi_max
+
+		self.L_GenGauss = GenGaussLoss(
+			reduction=self.reduction,
+			alpha_eps=self.alpha_eps, beta_eps=self.beta_eps, 
+			resi_min=self.resi_min, resi_max=self.resi_max
+		)
+		self.L_l1 = nn.L1Loss(reduction=self.reduction)
+	
+	def forward(
+		self,
+		mean: Tensor, one_over_alpha: Tensor, beta: Tensor, target: Tensor,
+		T1: float, T2: float
+	):
+		l1 = self.L_l1(mean, target)
+		l2 = self.L_GenGauss(mean, one_over_alpha, beta, target)
+		l = T1*l1 + T2*l2
+
+		return l
+
+
+# x1 = torch.randn(4,3,32,32)
+# x2 = torch.rand(4,3,32,32)
+# x3 = torch.rand(4,3,32,32)
+# x4 = torch.randn(4,3,32,32)
+
+# L = GenGaussLoss(alpha_eps=1e-4, beta_eps=1e-4, resi_min=1e-4, resi_max=1e3)
+# L2 =  TempCombLoss(alpha_eps=1e-4, beta_eps=1e-4, resi_min=1e-4, resi_max=1e3)
+# print(L(x1, x2, x3, x4), L2(x1, x2, x3, x4, 1e0, 1e-2))

networks_SRGAN.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch import Tensor
+
+# __all__ = [
+#     "ResidualConvBlock",
+#     "Discriminator", "Generator",
+# ]
+
+
+class ResidualConvBlock(nn.Module):
+	"""Implements residual conv function.
+
+	Args:
+		channels (int): Number of channels in the input image.
+	"""
+
+	def __init__(self, channels: int) -> None:
+		super(ResidualConvBlock, self).__init__()
+		self.rcb = nn.Sequential(
+			nn.Conv2d(channels, channels, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(channels),
+			nn.PReLU(),
+			nn.Conv2d(channels, channels, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(channels),
+		)
+
+	def forward(self, x: Tensor) -> Tensor:
+		identity = x
+
+		out = self.rcb(x)
+		out = torch.add(out, identity)
+
+		return out
+
+
+class Discriminator(nn.Module):
+	def __init__(self) -> None:
+		super(Discriminator, self).__init__()
+		self.features = nn.Sequential(
+			# input size. (3) x 96 x 96
+			nn.Conv2d(3, 64, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.LeakyReLU(0.2, True),
+			# state size. (64) x 48 x 48
+			nn.Conv2d(64, 64, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(64),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(128),
+			nn.LeakyReLU(0.2, True),
+			# state size. (128) x 24 x 24
+			nn.Conv2d(128, 128, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(128),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(128, 256, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(256),
+			nn.LeakyReLU(0.2, True),
+			# state size. (256) x 12 x 12
+			nn.Conv2d(256, 256, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(256),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(256, 512, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(512),
+			nn.LeakyReLU(0.2, True),
+			# state size. (512) x 6 x 6
+			nn.Conv2d(512, 512, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(512),
+			nn.LeakyReLU(0.2, True),
+		)
+
+		self.classifier = nn.Sequential(
+			nn.Linear(512 * 6 * 6, 1024),
+			nn.LeakyReLU(0.2, True),
+			nn.Linear(1024, 1),
+		)
+
+	def forward(self, x: Tensor) -> Tensor:
+		out = self.features(x)
+		out = torch.flatten(out, 1)
+		out = self.classifier(out)
+
+		return out
+
+
+class Generator(nn.Module):
+	def __init__(self) -> None:
+		super(Generator, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(3, 64, (9, 9), (1, 1), (4, 4)),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(64),
+		)
+
+		# Upscale conv block.
+		self.upsampling = nn.Sequential(
+			nn.Conv2d(64, 256, (3, 3), (1, 1), (1, 1)),
+			nn.PixelShuffle(2),
+			nn.PReLU(),
+			nn.Conv2d(64, 256, (3, 3), (1, 1), (1, 1)),
+			nn.PixelShuffle(2),
+			nn.PReLU(),
+		)
+
+		# Output layer.
+		self.conv_block3 = nn.Conv2d(64, 3, (9, 9), (1, 1), (4, 4))
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor, dop=None) -> Tensor:
+		if not dop:
+			return self._forward_impl(x)
+		else:
+			return self._forward_w_dop_impl(x, dop)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = torch.add(out1, out2)
+		out = self.upsampling(out)
+		out = self.conv_block3(out)
+
+		return out
+	
+	def _forward_w_dop_impl(self, x: Tensor, dop) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = F.dropout2d(self.conv_block2(out), p=dop)
+		out = torch.add(out1, out2)
+		out = self.upsampling(out)
+		out = self.conv_block3(out)
+
+		return out
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)
+
+
+#### BayesCap
+class BayesCap(nn.Module):
+	def __init__(self, in_channels=3, out_channels=3) -> None:
+		super(BayesCap, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(
+				in_channels, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=3, stride=1, padding=1, bias=False
+			),
+			nn.BatchNorm2d(64),
+		)
+
+		# Output layer.
+		self.conv_block3_mu = nn.Conv2d(
+			64, out_channels=out_channels, 
+			kernel_size=9, stride=1, padding=4
+		)
+		self.conv_block3_alpha = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+		self.conv_block3_beta = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor) -> Tensor:
+		return self._forward_impl(x)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = out1 + out2
+		out_mu = self.conv_block3_mu(out)
+		out_alpha = self.conv_block3_alpha(out)
+		out_beta = self.conv_block3_beta(out)
+		return out_mu, out_alpha, out_beta
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)
+                
+                
+class BayesCap_noID(nn.Module):
+	def __init__(self, in_channels=3, out_channels=3) -> None:
+		super(BayesCap_noID, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(
+				in_channels, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=3, stride=1, padding=1, bias=False
+			),
+			nn.BatchNorm2d(64),
+		)
+
+		# Output layer.
+		# self.conv_block3_mu = nn.Conv2d(
+		# 	64, out_channels=out_channels, 
+		# 	kernel_size=9, stride=1, padding=4
+		# )
+		self.conv_block3_alpha = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+		self.conv_block3_beta = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor) -> Tensor:
+		return self._forward_impl(x)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = out1 + out2
+		# out_mu = self.conv_block3_mu(out)
+		out_alpha = self.conv_block3_alpha(out)
+		out_beta = self.conv_block3_beta(out)
+		return out_alpha, out_beta
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)

networks_T1toT2.py

@@ -0,0 +1,477 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+
+### components
+class ResConv(nn.Module):
+    """
+    Residual convolutional block, where
+    convolutional block consists: (convolution => [BN] => ReLU) * 3
+    residual connection adds the input to the output
+    """
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.double_conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+    def forward(self, x):
+        x_in = self.double_conv1(x)
+        x1 = self.double_conv(x)
+        return self.double_conv(x) + x_in
+
+class Down(nn.Module):
+    """Downscaling with maxpool then Resconv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            ResConv(in_channels, out_channels)
+        )
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+class Up(nn.Module):
+	"""Upscaling then double conv"""
+	def __init__(self, in_channels, out_channels, bilinear=True):
+		super().__init__()
+		# if bilinear, use the normal convolutions to reduce the number of channels
+		if bilinear:
+			self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+			self.conv = ResConv(in_channels, out_channels, in_channels // 2)
+		else:
+			self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
+			self.conv = ResConv(in_channels, out_channels)
+	def forward(self, x1, x2):
+		x1 = self.up(x1)
+		# input is CHW
+		diffY = x2.size()[2] - x1.size()[2]
+		diffX = x2.size()[3] - x1.size()[3]
+		x1 = F.pad(
+			x1, 
+			[
+				diffX // 2, diffX - diffX // 2,
+				diffY // 2, diffY - diffY // 2
+			]
+		)
+		# if you have padding issues, see
+		# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+		# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+		x = torch.cat([x2, x1], dim=1)
+		return self.conv(x)
+
+class OutConv(nn.Module):
+	def __init__(self, in_channels, out_channels):
+		super(OutConv, self).__init__()
+		self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+	def forward(self, x):
+		# return F.relu(self.conv(x))
+		return self.conv(x)
+
+##### The composite networks
+class UNet(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		self.outc = OutConv(64, out_channels)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y = self.outc(x)
+		return y
+
+class CasUNet(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x, dop=None):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				if dop is not None:
+					y = F.dropout2d(self.unet_list[i](y), p=dop)
+				else:
+					y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		return y
+
+class CasUNet_2head(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet_2head, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			if i != self.n_unet-1:
+				self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+			else:
+				self.unet_list.append(UNet_2head(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		y_mean, y_sigma = y[0], y[1]
+		return y_mean, y_sigma
+
+class CasUNet_3head(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet_3head, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			if i != self.n_unet-1:
+				self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+			else:
+				self.unet_list.append(UNet_3head(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		y_mean, y_alpha, y_beta = y[0], y[1], y[2]
+		return y_mean, y_alpha, y_beta
+
+class UNet_2head(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet_2head, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		#per pixel multiple channels may exist
+		self.out_mean = OutConv(64, out_channels)
+		#variance will always be a single number for a pixel
+		self.out_var = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+		)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y_mean, y_var = self.out_mean(x), self.out_var(x)
+		return y_mean, y_var
+
+class UNet_3head(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet_3head, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		#per pixel multiple channels may exist
+		self.out_mean = OutConv(64, out_channels)
+		#variance will always be a single number for a pixel
+		self.out_alpha = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+			nn.ReLU()
+		)
+		self.out_beta = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+			nn.ReLU()
+		)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y_mean, y_alpha, y_beta = self.out_mean(x), \
+		self.out_alpha(x), self.out_beta(x)
+		return y_mean, y_alpha, y_beta
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+        conv_block = [  
+			nn.ReflectionPad2d(1),
+			nn.Conv2d(in_features, in_features, 3),
+			nn.InstanceNorm2d(in_features),
+			nn.ReLU(inplace=True),
+			nn.ReflectionPad2d(1),
+			nn.Conv2d(in_features, in_features, 3),
+			nn.InstanceNorm2d(in_features)
+		]
+        self.conv_block = nn.Sequential(*conv_block)
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9):
+        super(Generator, self).__init__()
+        # Initial convolution block       
+        model = [
+			nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, 7),
+            nn.InstanceNorm2d(64), nn.ReLU(inplace=True)
+		]
+        # Downsampling
+        in_features = 64
+        out_features = in_features*2
+        for _ in range(2):
+            model += [  
+				nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True) 
+			]
+            in_features = out_features
+            out_features = in_features*2
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model += [ResidualBlock(in_features)]
+        # Upsampling
+        out_features = in_features//2
+        for _ in range(2):
+            model += [  
+				nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+				nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True)
+			]
+            in_features = out_features
+            out_features = in_features//2
+        # Output layer
+        model += [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7), nn.Tanh()]
+        self.model = nn.Sequential(*model)
+    def forward(self, x):
+        return self.model(x)
+    
+    
+class ResnetGenerator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        """Construct a Resnet-based generator
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2 ** n_downsampling
+        for i in range(n_blocks):       # add ResNet blocks
+
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=1,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class ResnetBlock(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Construct a convolutional block.
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)  # add skip connections
+        return out
+
+### discriminator
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator"""
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.model = nn.Sequential(*sequence)
+    def forward(self, input):
+        """Standard forward."""
+        return self.model(input)

requirements.txt

@@ -0,0 +1,334 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=conda_forge
+_openmp_mutex=4.5=2_kmp_llvm
+aiohttp=3.8.1=pypi_0
+aiosignal=1.2.0=pypi_0
+albumentations=1.2.0=pyhd8ed1ab_0
+alsa-lib=1.2.6.1=h7f98852_0
+analytics-python=1.4.0=pypi_0
+anyio=3.6.1=pypi_0
+aom=3.3.0=h27087fc_1
+argon2-cffi=21.3.0=pypi_0
+argon2-cffi-bindings=21.2.0=pypi_0
+asttokens=2.0.5=pypi_0
+async-timeout=4.0.2=pypi_0
+attr=2.5.1=h166bdaf_0
+attrs=21.4.0=pypi_0
+babel=2.10.1=pypi_0
+backcall=0.2.0=pypi_0
+backoff=1.10.0=pypi_0
+bcrypt=3.2.2=pypi_0
+beautifulsoup4=4.11.1=pypi_0
+blas=1.0=mkl
+bleach=5.0.0=pypi_0
+blosc=1.21.1=h83bc5f7_3
+brotli=1.0.9=h166bdaf_7
+brotli-bin=1.0.9=h166bdaf_7
+brotlipy=0.7.0=py310h7f8727e_1002
+brunsli=0.1=h9c3ff4c_0
+bzip2=1.0.8=h7b6447c_0
+c-ares=1.18.1=h7f98852_0
+c-blosc2=2.2.0=h7a311fb_0
+ca-certificates=2022.6.15=ha878542_0
+cairo=1.16.0=ha61ee94_1011
+certifi=2022.6.15=py310hff52083_0
+cffi=1.15.0=py310hd667e15_1
+cfitsio=4.1.0=hd9d235c_0
+charls=2.3.4=h9c3ff4c_0
+charset-normalizer=2.0.4=pyhd3eb1b0_0
+click=8.1.3=pypi_0
+cloudpickle=2.1.0=pyhd8ed1ab_0
+cryptography=37.0.1=py310h9ce1e76_0
+cudatoolkit=10.2.89=hfd86e86_1
+cycler=0.11.0=pypi_0
+cytoolz=0.11.2=py310h5764c6d_2
+dask-core=2022.7.0=pyhd8ed1ab_0
+dbus=1.13.6=h5008d03_3
+debugpy=1.6.0=pypi_0
+decorator=5.1.1=pypi_0
+defusedxml=0.7.1=pypi_0
+entrypoints=0.4=pypi_0
+executing=0.8.3=pypi_0
+expat=2.4.8=h27087fc_0
+fastapi=0.78.0=pypi_0
+fastjsonschema=2.15.3=pypi_0
+ffmpeg=4.4.2=habc3f16_0
+ffmpy=0.3.0=pypi_0
+fftw=3.3.10=nompi_h77c792f_102
+fire=0.4.0=pypi_0
+font-ttf-dejavu-sans-mono=2.37=hab24e00_0
+font-ttf-inconsolata=3.000=h77eed37_0
+font-ttf-source-code-pro=2.038=h77eed37_0
+font-ttf-ubuntu=0.83=hab24e00_0
+fontconfig=2.14.0=h8e229c2_0
+fonts-conda-ecosystem=1=0
+fonts-conda-forge=1=0
+fonttools=4.33.3=pypi_0
+freeglut=3.2.2=h9c3ff4c_1
+freetype=2.11.0=h70c0345_0
+frozenlist=1.3.0=pypi_0
+fsspec=2022.5.0=pyhd8ed1ab_0
+ftfy=6.1.1=pypi_0
+gettext=0.19.8.1=h73d1719_1008
+giflib=5.2.1=h7b6447c_0
+glib=2.70.2=h780b84a_4
+glib-tools=2.70.2=h780b84a_4
+gmp=6.2.1=h295c915_3
+gnutls=3.7.6=hb5d6004_1
+gradio=3.0.24=pypi_0
+graphite2=1.3.13=h58526e2_1001
+gst-plugins-base=1.20.3=hf6a322e_0
+gstreamer=1.20.3=hd4edc92_0
+h11=0.12.0=pypi_0
+harfbuzz=4.4.1=hf9f4e7c_0
+hdf5=1.12.1=nompi_h2386368_104
+httpcore=0.15.0=pypi_0
+httpx=0.23.0=pypi_0
+icu=70.1=h27087fc_0
+idna=3.3=pyhd3eb1b0_0
+imagecodecs=2022.2.22=py310h3ac3b6e_6
+imageio=2.19.3=pyhcf75d05_0
+intel-openmp=2021.4.0=h06a4308_3561
+ipykernel=6.13.0=pypi_0
+ipython=8.4.0=pypi_0
+ipython-genutils=0.2.0=pypi_0
+jack=1.9.18=h8c3723f_1002
+jasper=2.0.33=ha77e612_0
+jedi=0.18.1=pypi_0
+jinja2=3.1.2=pypi_0
+joblib=1.1.0=pyhd8ed1ab_0
+jpeg=9e=h7f8727e_0
+json5=0.9.8=pypi_0
+jsonschema=4.6.0=pypi_0
+jupyter-client=7.3.1=pypi_0
+jupyter-core=4.10.0=pypi_0
+jupyter-server=1.17.0=pypi_0
+jupyterlab=3.4.2=pypi_0
+jupyterlab-pygments=0.2.2=pypi_0
+jupyterlab-server=2.14.0=pypi_0
+jxrlib=1.1=h7f98852_2
+keyutils=1.6.1=h166bdaf_0
+kiwisolver=1.4.2=pypi_0
+kornia=0.6.5=pypi_0
+krb5=1.19.3=h3790be6_0
+lame=3.100=h7b6447c_0
+lcms2=2.12=h3be6417_0
+ld_impl_linux-64=2.38=h1181459_1
+lerc=3.0=h9c3ff4c_0
+libaec=1.0.6=h9c3ff4c_0
+libavif=0.10.1=h166bdaf_0
+libblas=3.9.0=12_linux64_mkl
+libbrotlicommon=1.0.9=h166bdaf_7
+libbrotlidec=1.0.9=h166bdaf_7
+libbrotlienc=1.0.9=h166bdaf_7
+libcap=2.64=ha37c62d_0
+libcblas=3.9.0=12_linux64_mkl
+libclang=14.0.6=default_h2e3cab8_0
+libclang13=14.0.6=default_h3a83d3e_0
+libcups=2.3.3=hf5a7f15_1
+libcurl=7.83.1=h7bff187_0
+libdb=6.2.32=h9c3ff4c_0
+libdeflate=1.12=h166bdaf_0
+libdrm=2.4.112=h166bdaf_0
+libedit=3.1.20191231=he28a2e2_2
+libev=4.33=h516909a_1
+libevent=2.1.10=h9b69904_4
+libffi=3.4.2=h7f98852_5
+libflac=1.3.4=h27087fc_0
+libgcc-ng=12.1.0=h8d9b700_16
+libgfortran-ng=12.1.0=h69a702a_16
+libgfortran5=12.1.0=hdcd56e2_16
+libglib=2.70.2=h174f98d_4
+libglu=9.0.0=he1b5a44_1001
+libiconv=1.16=h7f8727e_2
+libidn2=2.3.2=h7f8727e_0
+liblapack=3.9.0=12_linux64_mkl
+liblapacke=3.9.0=12_linux64_mkl
+libllvm14=14.0.6=he0ac6c6_0
+libnghttp2=1.47.0=h727a467_0
+libnsl=2.0.0=h7f98852_0
+libogg=1.3.4=h7f98852_1
+libopencv=4.5.5=py310hcb97b83_13
+libopus=1.3.1=h7f98852_1
+libpciaccess=0.16=h516909a_0
+libpng=1.6.37=hbc83047_0
+libpq=14.4=hd77ab85_0
+libprotobuf=3.20.1=h6239696_0
+libsndfile=1.0.31=h9c3ff4c_1
+libssh2=1.10.0=ha56f1ee_2
+libstdcxx-ng=12.1.0=ha89aaad_16
+libtasn1=4.16.0=h27cfd23_0
+libtiff=4.4.0=hc85c160_1
+libtool=2.4.6=h9c3ff4c_1008
+libudev1=249=h166bdaf_4
+libunistring=0.9.10=h27cfd23_0
+libuuid=2.32.1=h7f98852_1000
+libuv=1.40.0=h7b6447c_0
+libva=2.15.0=h166bdaf_0
+libvorbis=1.3.7=h9c3ff4c_0
+libvpx=1.11.0=h9c3ff4c_3
+libwebp=1.2.2=h55f646e_0
+libwebp-base=1.2.2=h7f8727e_0
+libxcb=1.13=h7f98852_1004
+libxkbcommon=1.0.3=he3ba5ed_0
+libxml2=2.9.14=h22db469_3
+libzlib=1.2.12=h166bdaf_1
+libzopfli=1.0.3=h9c3ff4c_0
+linkify-it-py=1.0.3=pypi_0
+llvm-openmp=14.0.4=he0ac6c6_0
+locket=1.0.0=pyhd8ed1ab_0
+lz4-c=1.9.3=h295c915_1
+markdown-it-py=2.1.0=pypi_0
+markupsafe=2.1.1=pypi_0
+matplotlib=3.5.2=pypi_0
+matplotlib-inline=0.1.3=pypi_0
+mdit-py-plugins=0.3.0=pypi_0
+mdurl=0.1.1=pypi_0
+mistune=0.8.4=pypi_0
+mkl=2021.4.0=h06a4308_640
+mkl-service=2.4.0=py310h7f8727e_0
+mkl_fft=1.3.1=py310hd6ae3a3_0
+mkl_random=1.2.2=py310h00e6091_0
+mltk=0.0.5=pypi_0
+monotonic=1.6=pypi_0
+multidict=6.0.2=pypi_0
+munch=2.5.0=pypi_0
+mysql-common=8.0.29=haf5c9bc_1
+mysql-libs=8.0.29=h28c427c_1
+nbclassic=0.3.7=pypi_0
+nbclient=0.6.4=pypi_0
+nbconvert=6.5.0=pypi_0
+nbformat=5.4.0=pypi_0
+ncurses=6.3=h7f8727e_2
+nest-asyncio=1.5.5=pypi_0
+nettle=3.7.3=hbbd107a_1
+networkx=2.8.4=pyhd8ed1ab_0
+nltk=3.7=pypi_0
+notebook=6.4.11=pypi_0
+notebook-shim=0.1.0=pypi_0
+nspr=4.32=h9c3ff4c_1
+nss=3.78=h2350873_0
+ntk=1.1.3=pypi_0
+numpy=1.22.3=py310hfa59a62_0
+numpy-base=1.22.3=py310h9585f30_0
+opencv=4.5.5=py310hff52083_13
+opencv-python=4.6.0.66=pypi_0
+openh264=2.1.1=h4ff587b_0
+openjpeg=2.4.0=hb52868f_1
+openssl=1.1.1q=h166bdaf_0
+orjson=3.7.7=pypi_0
+packaging=21.3=pyhd8ed1ab_0
+pandas=1.4.2=pypi_0
+pandocfilters=1.5.0=pypi_0
+paramiko=2.11.0=pypi_0
+parso=0.8.3=pypi_0
+partd=1.2.0=pyhd8ed1ab_0
+pcre=8.45=h9c3ff4c_0
+pexpect=4.8.0=pypi_0
+pickleshare=0.7.5=pypi_0
+pillow=9.0.1=py310h22f2fdc_0
+pip=21.2.4=py310h06a4308_0
+pixman=0.40.0=h36c2ea0_0
+portaudio=19.6.0=h57a0ea0_5
+prometheus-client=0.14.1=pypi_0
+prompt-toolkit=3.0.29=pypi_0
+psutil=5.9.1=pypi_0
+pthread-stubs=0.4=h36c2ea0_1001
+ptyprocess=0.7.0=pypi_0
+pulseaudio=14.0=h7f54b18_8
+pure-eval=0.2.2=pypi_0
+py-opencv=4.5.5=py310hfdc917e_13
+pycocotools=2.0.4=pypi_0
+pycparser=2.21=pyhd3eb1b0_0
+pycryptodome=3.15.0=pypi_0
+pydantic=1.9.1=pypi_0
+pydub=0.25.1=pypi_0
+pygments=2.12.0=pypi_0
+pynacl=1.5.0=pypi_0
+pyopenssl=22.0.0=pyhd3eb1b0_0
+pyparsing=3.0.9=pyhd8ed1ab_0
+pyrsistent=0.18.1=pypi_0
+pysocks=1.7.1=py310h06a4308_0
+python=3.10.5=h582c2e5_0_cpython
+python-dateutil=2.8.2=pypi_0
+python-multipart=0.0.5=pypi_0
+python_abi=3.10=2_cp310
+pytorch=1.11.0=py3.10_cuda10.2_cudnn7.6.5_0
+pytorch-mutex=1.0=cuda
+pytz=2022.1=pypi_0
+pywavelets=1.3.0=py310hde88566_1
+pyyaml=6.0=py310h5764c6d_4
+pyzmq=23.1.0=pypi_0
+qt-main=5.15.4=ha5833f6_2
+qudida=0.0.4=pyhd8ed1ab_0
+readline=8.1.2=h7f8727e_1
+regex=2022.6.2=pypi_0
+requests=2.27.1=pyhd3eb1b0_0
+rfc3986=1.5.0=pypi_0
+scikit-image=0.19.3=py310h769672d_0
+scikit-learn=1.1.1=py310hffb9edd_0
+scipy=1.8.1=py310h7612f91_0
+seaborn=0.11.2=pypi_0
+send2trash=1.8.0=pypi_0
+setuptools=61.2.0=py310h06a4308_0
+six=1.16.0=pyhd3eb1b0_1
+snappy=1.1.9=hbd366e4_1
+sniffio=1.2.0=pypi_0
+soupsieve=2.3.2.post1=pypi_0
+sqlite=3.39.0=h4ff8645_0
+stack-data=0.2.0=pypi_0
+starlette=0.19.1=pypi_0
+svt-av1=1.1.0=h27087fc_1
+termcolor=1.1.0=pypi_0
+terminado=0.15.0=pypi_0
+threadpoolctl=3.1.0=pyh8a188c0_0
+tifffile=2022.5.4=pyhd8ed1ab_0
+tinycss2=1.1.1=pypi_0
+tk=8.6.12=h1ccaba5_0
+toolz=0.11.2=pyhd8ed1ab_0
+torchaudio=0.11.0=py310_cu102
+torchvision=0.12.0=py310_cu102
+tornado=6.1=pypi_0
+tqdm=4.64.0=pypi_0
+traitlets=5.2.2.post1=pypi_0
+typing-extensions=4.1.1=hd3eb1b0_0
+typing_extensions=4.1.1=pyh06a4308_0
+tzdata=2022a=hda174b7_0
+uc-micro-py=1.0.1=pypi_0
+urllib3=1.26.9=py310h06a4308_0
+uvicorn=0.18.2=pypi_0
+wcwidth=0.2.5=pypi_0
+webencodings=0.5.1=pypi_0
+websocket-client=1.3.2=pypi_0
+wheel=0.37.1=pyhd3eb1b0_0
+x264=1!161.3030=h7f98852_1
+x265=3.5=h924138e_3
+xcb-util=0.4.0=h166bdaf_0
+xcb-util-image=0.4.0=h166bdaf_0
+xcb-util-keysyms=0.4.0=h166bdaf_0
+xcb-util-renderutil=0.3.9=h166bdaf_0
+xcb-util-wm=0.4.1=h166bdaf_0
+xorg-fixesproto=5.0=h7f98852_1002
+xorg-inputproto=2.3.2=h7f98852_1002
+xorg-kbproto=1.0.7=h7f98852_1002
+xorg-libice=1.0.10=h7f98852_0
+xorg-libsm=1.2.3=hd9c2040_1000
+xorg-libx11=1.7.2=h7f98852_0
+xorg-libxau=1.0.9=h7f98852_0
+xorg-libxdmcp=1.1.3=h7f98852_0
+xorg-libxext=1.3.4=h7f98852_1
+xorg-libxfixes=5.0.3=h7f98852_1004
+xorg-libxi=1.7.10=h7f98852_0
+xorg-libxrender=0.9.10=h7f98852_1003
+xorg-renderproto=0.11.1=h7f98852_1002
+xorg-xextproto=7.3.0=h7f98852_1002
+xorg-xproto=7.0.31=h7f98852_1007
+xz=5.2.5=h7f8727e_1
+yaml=0.2.5=h7f98852_2
+yarl=1.7.2=pypi_0
+zfp=0.5.5=h9c3ff4c_8
+zlib=1.2.12=h166bdaf_1
+zlib-ng=2.0.6=h166bdaf_0
+zstd=1.5.2=ha4553b6_0

src/app.py

@@ -0,0 +1,115 @@
+import numpy as np
+import random
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode as IMode
+
+from PIL import Image
+
+from ds import *
+from losses import *
+from networks_SRGAN import *
+from utils import *
+
+
+NetG = Generator()
+model_parameters = filter(lambda p: True, NetG.parameters())
+params = sum([np.prod(p.size()) for p in model_parameters])
+print("Number of Parameters:",params)
+NetC = BayesCap(in_channels=3, out_channels=3)
+
+
+NetG = Generator()
+NetG.load_state_dict(torch.load('../ckpt/srgan-ImageNet-bc347d67.pth', map_location='cuda:0'))
+NetG.to('cuda')
+NetG.eval()
+
+NetC = BayesCap(in_channels=3, out_channels=3)
+NetC.load_state_dict(torch.load('../ckpt/BayesCap_SRGAN_best.pth', map_location='cuda:0'))
+NetC.to('cuda')
+NetC.eval()
+
+def tensor01_to_pil(xt):
+    r = transforms.ToPILImage(mode='RGB')(xt.squeeze())
+    return r
+
+
+def predict(img):
+    """
+    img: image
+    """
+    image_size = (256,256)
+    upscale_factor = 4
+    lr_transforms = transforms.Resize((image_size[0]//upscale_factor, image_size[1]//upscale_factor), interpolation=IMode.BICUBIC, antialias=True)
+    # lr_transforms = transforms.Resize((128, 128), interpolation=IMode.BICUBIC, antialias=True)
+    
+    img = Image.fromarray(np.array(img))
+    img = lr_transforms(img)
+    lr_tensor = utils.image2tensor(img, range_norm=False, half=False)
+    
+    device = 'cuda'
+    dtype = torch.cuda.FloatTensor
+    xLR = lr_tensor.to(device).unsqueeze(0)
+    xLR = xLR.type(dtype)
+    # pass them through the network
+    with torch.no_grad():
+        xSR = NetG(xLR)
+        xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+        
+    a_map = (1/(xSRC_alpha[0] + 1e-5)).to('cpu').data
+    b_map = xSRC_beta[0].to('cpu').data
+    u_map = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2)))) 
+    
+    
+    x_LR = tensor01_to_pil(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    
+    x_mean = tensor01_to_pil(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    
+    #im = Image.fromarray(np.uint8(cm.gist_earth(myarray)*255))
+    
+    a_map = torch.clamp(a_map, min=0, max=0.1)
+    a_map = (a_map - a_map.min())/(a_map.max() - a_map.min())
+    x_alpha = Image.fromarray(np.uint8(cm.inferno(a_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    b_map = torch.clamp(b_map, min=0.45, max=0.75)
+    b_map = (b_map - b_map.min())/(b_map.max() - b_map.min())
+    x_beta = Image.fromarray(np.uint8(cm.cividis(b_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    u_map = torch.clamp(u_map, min=0, max=0.15)
+    u_map = (u_map - u_map.min())/(u_map.max() - u_map.min())
+    x_uncer = Image.fromarray(np.uint8(cm.hot(u_map.transpose(0,2).transpose(0,1).squeeze())*255))
+    
+    return x_LR, x_mean, x_alpha, x_beta, x_uncer
+
+import gradio as gr
+
+title = "BayesCap"
+description = "BayesCap: Bayesian Identity Cap for Calibrated Uncertainty in Frozen Neural Networks (ECCV 2022)"
+article = "<p style='text-align: center'> BayesCap: Bayesian Identity Cap for Calibrated Uncertainty in Frozen Neural Networks| <a href='https://github.com/ExplainableML/BayesCap'>Github Repo</a></p>"
+
+
+gr.Interface(
+    fn=predict, 
+    inputs=gr.inputs.Image(type='pil', label="Orignal"), 
+    outputs=[
+        gr.outputs.Image(type='pil', label="Low-res"), 
+        gr.outputs.Image(type='pil', label="Super-res"), 
+        gr.outputs.Image(type='pil', label="Alpha"), 
+        gr.outputs.Image(type='pil', label="Beta"), 
+        gr.outputs.Image(type='pil', label="Uncertainty")
+     ],
+    title=title,
+    description=description,
+    article=article,
+     examples=[
+        ["../demo_examples/baby.png"],
+        ["../demo_examples/bird.png"]
+    ]
+).launch(share=True)

src/ds.py

@@ -0,0 +1,485 @@
+from __future__ import absolute_import, division, print_function
+
+import random
+import copy
+import io
+import os
+import numpy as np
+from PIL import Image
+import skimage.transform
+from collections import Counter
+
+
+import torch
+import torch.utils.data as data
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode as IMode
+
+import utils
+
+class ImgDset(Dataset):
+    """Customize the data set loading function and prepare low/high resolution image data in advance.
+
+    Args:
+        dataroot         (str): Training data set address
+        image_size       (int): High resolution image size
+        upscale_factor   (int): Image magnification
+        mode             (str): Data set loading method, the training data set is for data enhancement,
+                             and the verification data set is not for data enhancement
+
+    """
+
+    def __init__(self, dataroot: str, image_size: int, upscale_factor: int, mode: str) -> None:
+        super(ImgDset, self).__init__()
+        self.filenames = [os.path.join(dataroot, x) for x in os.listdir(dataroot)]
+
+        if mode == "train":
+            self.hr_transforms = transforms.Compose([
+                transforms.RandomCrop(image_size),
+                transforms.RandomRotation(90),
+                transforms.RandomHorizontalFlip(0.5),
+            ])
+        else:
+            self.hr_transforms = transforms.Resize(image_size)
+
+        self.lr_transforms = transforms.Resize((image_size[0]//upscale_factor, image_size[1]//upscale_factor), interpolation=IMode.BICUBIC, antialias=True)
+
+    def __getitem__(self, batch_index: int) -> [Tensor, Tensor]:
+        # Read a batch of image data
+        image = Image.open(self.filenames[batch_index])
+
+        # Transform image
+        hr_image = self.hr_transforms(image)
+        lr_image = self.lr_transforms(hr_image)
+
+        # Convert image data into Tensor stream format (PyTorch).
+        # Note: The range of input and output is between [0, 1]
+        lr_tensor = utils.image2tensor(lr_image, range_norm=False, half=False)
+        hr_tensor = utils.image2tensor(hr_image, range_norm=False, half=False)
+
+        return lr_tensor, hr_tensor
+
+    def __len__(self) -> int:
+        return len(self.filenames)
+    
+
+class PairedImages_w_nameList(Dataset):
+    '''
+    can act as supervised or un-supervised based on flists
+    '''
+    def __init__(self, flist1, flist2, transform1=None, transform2=None, do_aug=False):
+        self.flist1 = flist1
+        self.flist2 = flist2
+        self.transform1 = transform1
+        self.transform2 = transform2
+        self.do_aug = do_aug
+    def __getitem__(self, index):
+        impath1 = self.flist1[index]
+        img1 = Image.open(impath1).convert('RGB')
+        impath2 = self.flist2[index]
+        img2 = Image.open(impath2).convert('RGB')
+        
+        img1 = utils.image2tensor(img1, range_norm=False, half=False)
+        img2 = utils.image2tensor(img2, range_norm=False, half=False)
+        
+        if self.transform1 is not None:
+            img1 = self.transform1(img1)
+        if self.transform2 is not None:
+            img2 = self.transform2(img2)
+        
+        return img1, img2
+    def __len__(self):
+        return len(self.flist1)
+    
+class PairedImages_w_nameList_npy(Dataset):
+    '''
+    can act as supervised or un-supervised based on flists
+    '''
+    def __init__(self, flist1, flist2, transform1=None, transform2=None, do_aug=False):
+        self.flist1 = flist1
+        self.flist2 = flist2
+        self.transform1 = transform1
+        self.transform2 = transform2
+        self.do_aug = do_aug
+    def __getitem__(self, index):
+        impath1 = self.flist1[index]
+        img1 = np.load(impath1)
+        impath2 = self.flist2[index]
+        img2 = np.load(impath2)
+        
+        if self.transform1 is not None:
+            img1 = self.transform1(img1)
+        if self.transform2 is not None:
+            img2 = self.transform2(img2)
+        
+        return img1, img2
+    def __len__(self):
+        return len(self.flist1)
+
+# def call_paired():
+#     root1='./GOPRO_3840FPS_AVG_3-21/train/blur/'
+#     root2='./GOPRO_3840FPS_AVG_3-21/train/sharp/'
+
+#     flist1=glob.glob(root1+'/*/*.png')
+#     flist2=glob.glob(root2+'/*/*.png')
+
+#     dset = PairedImages_w_nameList(root1,root2,flist1,flist2)
+
+#### KITTI depth
+
+def load_velodyne_points(filename):
+    """Load 3D point cloud from KITTI file format
+    (adapted from https://github.com/hunse/kitti)
+    """
+    points = np.fromfile(filename, dtype=np.float32).reshape(-1, 4)
+    points[:, 3] = 1.0  # homogeneous
+    return points
+
+
+def read_calib_file(path):
+    """Read KITTI calibration file
+    (from https://github.com/hunse/kitti)
+    """
+    float_chars = set("0123456789.e+- ")
+    data = {}
+    with open(path, 'r') as f:
+        for line in f.readlines():
+            key, value = line.split(':', 1)
+            value = value.strip()
+            data[key] = value
+            if float_chars.issuperset(value):
+                # try to cast to float array
+                try:
+                    data[key] = np.array(list(map(float, value.split(' '))))
+                except ValueError:
+                    # casting error: data[key] already eq. value, so pass
+                    pass
+
+    return data
+
+
+def sub2ind(matrixSize, rowSub, colSub):
+    """Convert row, col matrix subscripts to linear indices
+    """
+    m, n = matrixSize
+    return rowSub * (n-1) + colSub - 1
+
+
+def generate_depth_map(calib_dir, velo_filename, cam=2, vel_depth=False):
+    """Generate a depth map from velodyne data
+    """
+    # load calibration files
+    cam2cam = read_calib_file(os.path.join(calib_dir, 'calib_cam_to_cam.txt'))
+    velo2cam = read_calib_file(os.path.join(calib_dir, 'calib_velo_to_cam.txt'))
+    velo2cam = np.hstack((velo2cam['R'].reshape(3, 3), velo2cam['T'][..., np.newaxis]))
+    velo2cam = np.vstack((velo2cam, np.array([0, 0, 0, 1.0])))
+
+    # get image shape
+    im_shape = cam2cam["S_rect_02"][::-1].astype(np.int32)
+
+    # compute projection matrix velodyne->image plane
+    R_cam2rect = np.eye(4)
+    R_cam2rect[:3, :3] = cam2cam['R_rect_00'].reshape(3, 3)
+    P_rect = cam2cam['P_rect_0'+str(cam)].reshape(3, 4)
+    P_velo2im = np.dot(np.dot(P_rect, R_cam2rect), velo2cam)
+
+    # load velodyne points and remove all behind image plane (approximation)
+    # each row of the velodyne data is forward, left, up, reflectance
+    velo = load_velodyne_points(velo_filename)
+    velo = velo[velo[:, 0] >= 0, :]
+
+    # project the points to the camera
+    velo_pts_im = np.dot(P_velo2im, velo.T).T
+    velo_pts_im[:, :2] = velo_pts_im[:, :2] / velo_pts_im[:, 2][..., np.newaxis]
+
+    if vel_depth:
+        velo_pts_im[:, 2] = velo[:, 0]
+
+    # check if in bounds
+    # use minus 1 to get the exact same value as KITTI matlab code
+    velo_pts_im[:, 0] = np.round(velo_pts_im[:, 0]) - 1
+    velo_pts_im[:, 1] = np.round(velo_pts_im[:, 1]) - 1
+    val_inds = (velo_pts_im[:, 0] >= 0) & (velo_pts_im[:, 1] >= 0)
+    val_inds = val_inds & (velo_pts_im[:, 0] < im_shape[1]) & (velo_pts_im[:, 1] < im_shape[0])
+    velo_pts_im = velo_pts_im[val_inds, :]
+
+    # project to image
+    depth = np.zeros((im_shape[:2]))
+    depth[velo_pts_im[:, 1].astype(np.int), velo_pts_im[:, 0].astype(np.int)] = velo_pts_im[:, 2]
+
+    # find the duplicate points and choose the closest depth
+    inds = sub2ind(depth.shape, velo_pts_im[:, 1], velo_pts_im[:, 0])
+    dupe_inds = [item for item, count in Counter(inds).items() if count > 1]
+    for dd in dupe_inds:
+        pts = np.where(inds == dd)[0]
+        x_loc = int(velo_pts_im[pts[0], 0])
+        y_loc = int(velo_pts_im[pts[0], 1])
+        depth[y_loc, x_loc] = velo_pts_im[pts, 2].min()
+    depth[depth < 0] = 0
+
+    return depth
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning
+    # (https://github.com/python-pillow/Pillow/issues/835)
+    with open(path, 'rb') as f:
+        with Image.open(f) as img:
+            return img.convert('RGB')
+
+
+class MonoDataset(data.Dataset):
+    """Superclass for monocular dataloaders
+
+    Args:
+        data_path
+        filenames
+        height
+        width
+        frame_idxs
+        num_scales
+        is_train
+        img_ext
+    """
+    def __init__(self,
+                 data_path,
+                 filenames,
+                 height,
+                 width,
+                 frame_idxs,
+                 num_scales,
+                 is_train=False,
+                 img_ext='.jpg'):
+        super(MonoDataset, self).__init__()
+
+        self.data_path = data_path
+        self.filenames = filenames
+        self.height = height
+        self.width = width
+        self.num_scales = num_scales
+        self.interp = Image.ANTIALIAS
+
+        self.frame_idxs = frame_idxs
+
+        self.is_train = is_train
+        self.img_ext = img_ext
+
+        self.loader = pil_loader
+        self.to_tensor = transforms.ToTensor()
+
+        # We need to specify augmentations differently in newer versions of torchvision.
+        # We first try the newer tuple version; if this fails we fall back to scalars
+        try:
+            self.brightness = (0.8, 1.2)
+            self.contrast = (0.8, 1.2)
+            self.saturation = (0.8, 1.2)
+            self.hue = (-0.1, 0.1)
+            transforms.ColorJitter.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+        except TypeError:
+            self.brightness = 0.2
+            self.contrast = 0.2
+            self.saturation = 0.2
+            self.hue = 0.1
+
+        self.resize = {}
+        for i in range(self.num_scales):
+            s = 2 ** i
+            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
+                                               interpolation=self.interp)
+
+        self.load_depth = self.check_depth()
+
+    def preprocess(self, inputs, color_aug):
+        """Resize colour images to the required scales and augment if required
+
+        We create the color_aug object in advance and apply the same augmentation to all
+        images in this item. This ensures that all images input to the pose network receive the
+        same augmentation.
+        """
+        for k in list(inputs):
+            frame = inputs[k]
+            if "color" in k:
+                n, im, i = k
+                for i in range(self.num_scales):
+                    inputs[(n, im, i)] = self.resize[i](inputs[(n, im, i - 1)])
+
+        for k in list(inputs):
+            f = inputs[k]
+            if "color" in k:
+                n, im, i = k
+                inputs[(n, im, i)] = self.to_tensor(f)
+                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f))
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        """Returns a single training item from the dataset as a dictionary.
+
+        Values correspond to torch tensors.
+        Keys in the dictionary are either strings or tuples:
+
+            ("color", <frame_id>, <scale>)          for raw colour images,
+            ("color_aug", <frame_id>, <scale>)      for augmented colour images,
+            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
+            "stereo_T"                              for camera extrinsics, and
+            "depth_gt"                              for ground truth depth maps.
+
+        <frame_id> is either:
+            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
+        or
+            "s" for the opposite image in the stereo pair.
+
+        <scale> is an integer representing the scale of the image relative to the fullsize image:
+            -1      images at native resolution as loaded from disk
+            0       images resized to (self.width,      self.height     )
+            1       images resized to (self.width // 2, self.height // 2)
+            2       images resized to (self.width // 4, self.height // 4)
+            3       images resized to (self.width // 8, self.height // 8)
+        """
+        inputs = {}
+
+        do_color_aug = self.is_train and random.random() > 0.5
+        do_flip = self.is_train and random.random() > 0.5
+
+        line = self.filenames[index].split()
+        folder = line[0]
+
+        if len(line) == 3:
+            frame_index = int(line[1])
+        else:
+            frame_index = 0
+
+        if len(line) == 3:
+            side = line[2]
+        else:
+            side = None
+
+        for i in self.frame_idxs:
+            if i == "s":
+                other_side = {"r": "l", "l": "r"}[side]
+                inputs[("color", i, -1)] = self.get_color(folder, frame_index, other_side, do_flip)
+            else:
+                inputs[("color", i, -1)] = self.get_color(folder, frame_index + i, side, do_flip)
+
+        # adjusting intrinsics to match each scale in the pyramid
+        for scale in range(self.num_scales):
+            K = self.K.copy()
+
+            K[0, :] *= self.width // (2 ** scale)
+            K[1, :] *= self.height // (2 ** scale)
+
+            inv_K = np.linalg.pinv(K)
+
+            inputs[("K", scale)] = torch.from_numpy(K)
+            inputs[("inv_K", scale)] = torch.from_numpy(inv_K)
+
+        if do_color_aug:
+            color_aug = transforms.ColorJitter.get_params(
+                self.brightness, self.contrast, self.saturation, self.hue)
+        else:
+            color_aug = (lambda x: x)
+
+        self.preprocess(inputs, color_aug)
+
+        for i in self.frame_idxs:
+            del inputs[("color", i, -1)]
+            del inputs[("color_aug", i, -1)]
+
+        if self.load_depth:
+            depth_gt = self.get_depth(folder, frame_index, side, do_flip)
+            inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
+            inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))
+
+        if "s" in self.frame_idxs:
+            stereo_T = np.eye(4, dtype=np.float32)
+            baseline_sign = -1 if do_flip else 1
+            side_sign = -1 if side == "l" else 1
+            stereo_T[0, 3] = side_sign * baseline_sign * 0.1
+
+            inputs["stereo_T"] = torch.from_numpy(stereo_T)
+
+        return inputs
+
+    def get_color(self, folder, frame_index, side, do_flip):
+        raise NotImplementedError
+
+    def check_depth(self):
+        raise NotImplementedError
+
+    def get_depth(self, folder, frame_index, side, do_flip):
+        raise NotImplementedError
+
+class KITTIDataset(MonoDataset):
+    """Superclass for different types of KITTI dataset loaders
+    """
+    def __init__(self, *args, **kwargs):
+        super(KITTIDataset, self).__init__(*args, **kwargs)
+
+        # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size.
+        # To normalize you need to scale the first row by 1 / image_width and the second row
+        # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered.
+        # If your principal point is far from the center you might need to disable the horizontal
+        # flip augmentation.
+        self.K = np.array([[0.58, 0, 0.5, 0],
+                           [0, 1.92, 0.5, 0],
+                           [0, 0, 1, 0],
+                           [0, 0, 0, 1]], dtype=np.float32)
+
+        self.full_res_shape = (1242, 375)
+        self.side_map = {"2": 2, "3": 3, "l": 2, "r": 3}
+
+    def check_depth(self):
+        line = self.filenames[0].split()
+        scene_name = line[0]
+        frame_index = int(line[1])
+
+        velo_filename = os.path.join(
+            self.data_path,
+            scene_name,
+            "velodyne_points/data/{:010d}.bin".format(int(frame_index)))
+
+        return os.path.isfile(velo_filename)
+
+    def get_color(self, folder, frame_index, side, do_flip):
+        color = self.loader(self.get_image_path(folder, frame_index, side))
+
+        if do_flip:
+            color = color.transpose(Image.FLIP_LEFT_RIGHT)
+
+        return color
+
+
+class KITTIDepthDataset(KITTIDataset):
+    """KITTI dataset which uses the updated ground truth depth maps
+    """
+    def __init__(self, *args, **kwargs):
+        super(KITTIDepthDataset, self).__init__(*args, **kwargs)
+
+    def get_image_path(self, folder, frame_index, side):
+        f_str = "{:010d}{}".format(frame_index, self.img_ext)
+        image_path = os.path.join(
+            self.data_path,
+            folder,
+            "image_0{}/data".format(self.side_map[side]),
+            f_str)
+        return image_path
+
+    def get_depth(self, folder, frame_index, side, do_flip):
+        f_str = "{:010d}.png".format(frame_index)
+        depth_path = os.path.join(
+            self.data_path,
+            folder,
+            "proj_depth/groundtruth/image_0{}".format(self.side_map[side]),
+            f_str)
+
+        depth_gt = Image.open(depth_path)
+        depth_gt = depth_gt.resize(self.full_res_shape, Image.NEAREST)
+        depth_gt = np.array(depth_gt).astype(np.float32) / 256
+
+        if do_flip:
+            depth_gt = np.fliplr(depth_gt)
+
+        return depth_gt

src/flagged/log.csv

@@ -0,0 +1,2 @@
+'Orignal','Low-res','Super-res','Alpha','Beta','Uncertainty','flag','username','timestamp'
+'Orignal/0.png','Low-res/0.png','Super-res/0.png','Alpha/0.png','Beta/0.png','Uncertainty/0.png','','','2022-07-09 14:01:12.964411'

src/losses.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch import Tensor
+
+class ContentLoss(nn.Module):
+	"""Constructs a content loss function based on the VGG19 network.
+	Using high-level feature mapping layers from the latter layers will focus more on the texture content of the image.
+
+	Paper reference list:
+		-`Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network <https://arxiv.org/pdf/1609.04802.pdf>` paper.
+		-`ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks                    <https://arxiv.org/pdf/1809.00219.pdf>` paper.
+		-`Perceptual Extreme Super Resolution Network with Receptive Field Block               <https://arxiv.org/pdf/2005.12597.pdf>` paper.
+
+	 """
+
+	def __init__(self) -> None:
+		super(ContentLoss, self).__init__()
+		# Load the VGG19 model trained on the ImageNet dataset.
+		vgg19 = models.vgg19(pretrained=True).eval()
+		# Extract the thirty-sixth layer output in the VGG19 model as the content loss.
+		self.feature_extractor = nn.Sequential(*list(vgg19.features.children())[:36])
+		# Freeze model parameters.
+		for parameters in self.feature_extractor.parameters():
+			parameters.requires_grad = False
+
+		# The preprocessing method of the input data. This is the VGG model preprocessing method of the ImageNet dataset.
+		self.register_buffer("mean", torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+		self.register_buffer("std", torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+	def forward(self, sr: Tensor, hr: Tensor) -> Tensor:
+		# Standardized operations
+		sr = sr.sub(self.mean).div(self.std)
+		hr = hr.sub(self.mean).div(self.std)
+
+		# Find the feature map difference between the two images
+		loss = F.l1_loss(self.feature_extractor(sr), self.feature_extractor(hr))
+
+		return loss
+
+
+class GenGaussLoss(nn.Module):
+	def __init__(
+		self, reduction='mean',
+		alpha_eps = 1e-4, beta_eps=1e-4,
+		resi_min = 1e-4, resi_max=1e3
+	) -> None:
+		super(GenGaussLoss, self).__init__()
+		self.reduction = reduction
+		self.alpha_eps = alpha_eps
+		self.beta_eps = beta_eps
+		self.resi_min = resi_min
+		self.resi_max = resi_max
+	
+	def forward(
+		self, 
+		mean: Tensor, one_over_alpha: Tensor, beta: Tensor, target: Tensor
+	):
+		one_over_alpha1 = one_over_alpha + self.alpha_eps
+		beta1 = beta + self.beta_eps
+
+		resi = torch.abs(mean - target)
+		# resi = torch.pow(resi*one_over_alpha1, beta1).clamp(min=self.resi_min, max=self.resi_max)
+		resi = (resi*one_over_alpha1*beta1).clamp(min=self.resi_min, max=self.resi_max)
+		## check if resi has nans
+		if torch.sum(resi != resi) > 0:
+			print('resi has nans!!')
+			return None
+		
+		log_one_over_alpha = torch.log(one_over_alpha1)
+		log_beta = torch.log(beta1)
+		lgamma_beta = torch.lgamma(torch.pow(beta1, -1))
+		
+		if torch.sum(log_one_over_alpha != log_one_over_alpha) > 0:
+			print('log_one_over_alpha has nan')
+		if torch.sum(lgamma_beta != lgamma_beta) > 0:
+			print('lgamma_beta has nan')
+		if torch.sum(log_beta != log_beta) > 0:
+			print('log_beta has nan')
+		
+		l = resi - log_one_over_alpha + lgamma_beta - log_beta
+
+		if self.reduction == 'mean':
+			return l.mean()
+		elif self.reduction == 'sum':
+			return l.sum()
+		else:
+			print('Reduction not supported')
+			return None
+
+class TempCombLoss(nn.Module):
+	def __init__(
+		self, reduction='mean',
+		alpha_eps = 1e-4, beta_eps=1e-4,
+		resi_min = 1e-4, resi_max=1e3
+	) -> None:
+		super(TempCombLoss, self).__init__()
+		self.reduction = reduction
+		self.alpha_eps = alpha_eps
+		self.beta_eps = beta_eps
+		self.resi_min = resi_min
+		self.resi_max = resi_max
+
+		self.L_GenGauss = GenGaussLoss(
+			reduction=self.reduction,
+			alpha_eps=self.alpha_eps, beta_eps=self.beta_eps, 
+			resi_min=self.resi_min, resi_max=self.resi_max
+		)
+		self.L_l1 = nn.L1Loss(reduction=self.reduction)
+	
+	def forward(
+		self,
+		mean: Tensor, one_over_alpha: Tensor, beta: Tensor, target: Tensor,
+		T1: float, T2: float
+	):
+		l1 = self.L_l1(mean, target)
+		l2 = self.L_GenGauss(mean, one_over_alpha, beta, target)
+		l = T1*l1 + T2*l2
+
+		return l
+
+
+# x1 = torch.randn(4,3,32,32)
+# x2 = torch.rand(4,3,32,32)
+# x3 = torch.rand(4,3,32,32)
+# x4 = torch.randn(4,3,32,32)
+
+# L = GenGaussLoss(alpha_eps=1e-4, beta_eps=1e-4, resi_min=1e-4, resi_max=1e3)
+# L2 =  TempCombLoss(alpha_eps=1e-4, beta_eps=1e-4, resi_min=1e-4, resi_max=1e3)
+# print(L(x1, x2, x3, x4), L2(x1, x2, x3, x4, 1e0, 1e-2))

src/networks_SRGAN.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+from torch import Tensor
+
+# __all__ = [
+#     "ResidualConvBlock",
+#     "Discriminator", "Generator",
+# ]
+
+
+class ResidualConvBlock(nn.Module):
+	"""Implements residual conv function.
+
+	Args:
+		channels (int): Number of channels in the input image.
+	"""
+
+	def __init__(self, channels: int) -> None:
+		super(ResidualConvBlock, self).__init__()
+		self.rcb = nn.Sequential(
+			nn.Conv2d(channels, channels, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(channels),
+			nn.PReLU(),
+			nn.Conv2d(channels, channels, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(channels),
+		)
+
+	def forward(self, x: Tensor) -> Tensor:
+		identity = x
+
+		out = self.rcb(x)
+		out = torch.add(out, identity)
+
+		return out
+
+
+class Discriminator(nn.Module):
+	def __init__(self) -> None:
+		super(Discriminator, self).__init__()
+		self.features = nn.Sequential(
+			# input size. (3) x 96 x 96
+			nn.Conv2d(3, 64, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.LeakyReLU(0.2, True),
+			# state size. (64) x 48 x 48
+			nn.Conv2d(64, 64, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(64),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(128),
+			nn.LeakyReLU(0.2, True),
+			# state size. (128) x 24 x 24
+			nn.Conv2d(128, 128, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(128),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(128, 256, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(256),
+			nn.LeakyReLU(0.2, True),
+			# state size. (256) x 12 x 12
+			nn.Conv2d(256, 256, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(256),
+			nn.LeakyReLU(0.2, True),
+			nn.Conv2d(256, 512, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(512),
+			nn.LeakyReLU(0.2, True),
+			# state size. (512) x 6 x 6
+			nn.Conv2d(512, 512, (3, 3), (2, 2), (1, 1), bias=False),
+			nn.BatchNorm2d(512),
+			nn.LeakyReLU(0.2, True),
+		)
+
+		self.classifier = nn.Sequential(
+			nn.Linear(512 * 6 * 6, 1024),
+			nn.LeakyReLU(0.2, True),
+			nn.Linear(1024, 1),
+		)
+
+	def forward(self, x: Tensor) -> Tensor:
+		out = self.features(x)
+		out = torch.flatten(out, 1)
+		out = self.classifier(out)
+
+		return out
+
+
+class Generator(nn.Module):
+	def __init__(self) -> None:
+		super(Generator, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(3, 64, (9, 9), (1, 1), (4, 4)),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1), bias=False),
+			nn.BatchNorm2d(64),
+		)
+
+		# Upscale conv block.
+		self.upsampling = nn.Sequential(
+			nn.Conv2d(64, 256, (3, 3), (1, 1), (1, 1)),
+			nn.PixelShuffle(2),
+			nn.PReLU(),
+			nn.Conv2d(64, 256, (3, 3), (1, 1), (1, 1)),
+			nn.PixelShuffle(2),
+			nn.PReLU(),
+		)
+
+		# Output layer.
+		self.conv_block3 = nn.Conv2d(64, 3, (9, 9), (1, 1), (4, 4))
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor, dop=None) -> Tensor:
+		if not dop:
+			return self._forward_impl(x)
+		else:
+			return self._forward_w_dop_impl(x, dop)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = torch.add(out1, out2)
+		out = self.upsampling(out)
+		out = self.conv_block3(out)
+
+		return out
+	
+	def _forward_w_dop_impl(self, x: Tensor, dop) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = F.dropout2d(self.conv_block2(out), p=dop)
+		out = torch.add(out1, out2)
+		out = self.upsampling(out)
+		out = self.conv_block3(out)
+
+		return out
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)
+
+
+#### BayesCap
+class BayesCap(nn.Module):
+	def __init__(self, in_channels=3, out_channels=3) -> None:
+		super(BayesCap, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(
+				in_channels, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=3, stride=1, padding=1, bias=False
+			),
+			nn.BatchNorm2d(64),
+		)
+
+		# Output layer.
+		self.conv_block3_mu = nn.Conv2d(
+			64, out_channels=out_channels, 
+			kernel_size=9, stride=1, padding=4
+		)
+		self.conv_block3_alpha = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+		self.conv_block3_beta = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor) -> Tensor:
+		return self._forward_impl(x)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = out1 + out2
+		out_mu = self.conv_block3_mu(out)
+		out_alpha = self.conv_block3_alpha(out)
+		out_beta = self.conv_block3_beta(out)
+		return out_mu, out_alpha, out_beta
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)
+                
+                
+class BayesCap_noID(nn.Module):
+	def __init__(self, in_channels=3, out_channels=3) -> None:
+		super(BayesCap_noID, self).__init__()
+		# First conv layer.
+		self.conv_block1 = nn.Sequential(
+			nn.Conv2d(
+				in_channels, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+		)
+
+		# Features trunk blocks.
+		trunk = []
+		for _ in range(16):
+			trunk.append(ResidualConvBlock(64))
+		self.trunk = nn.Sequential(*trunk)
+
+		# Second conv layer.
+		self.conv_block2 = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=3, stride=1, padding=1, bias=False
+			),
+			nn.BatchNorm2d(64),
+		)
+
+		# Output layer.
+		# self.conv_block3_mu = nn.Conv2d(
+		# 	64, out_channels=out_channels, 
+		# 	kernel_size=9, stride=1, padding=4
+		# )
+		self.conv_block3_alpha = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+		self.conv_block3_beta = nn.Sequential(
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 64, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.PReLU(),
+			nn.Conv2d(
+				64, 1, 
+				kernel_size=9, stride=1, padding=4
+			),
+			nn.ReLU(),
+		)
+
+		# Initialize neural network weights.
+		self._initialize_weights()
+
+	def forward(self, x: Tensor) -> Tensor:
+		return self._forward_impl(x)
+
+	# Support torch.script function.
+	def _forward_impl(self, x: Tensor) -> Tensor:
+		out1 = self.conv_block1(x)
+		out = self.trunk(out1)
+		out2 = self.conv_block2(out)
+		out = out1 + out2
+		# out_mu = self.conv_block3_mu(out)
+		out_alpha = self.conv_block3_alpha(out)
+		out_beta = self.conv_block3_beta(out)
+		return out_alpha, out_beta
+
+	def _initialize_weights(self) -> None:
+		for module in self.modules():
+			if isinstance(module, nn.Conv2d):
+				nn.init.kaiming_normal_(module.weight)
+				if module.bias is not None:
+					nn.init.constant_(module.bias, 0)
+			elif isinstance(module, nn.BatchNorm2d):
+				nn.init.constant_(module.weight, 1)

src/networks_T1toT2.py

@@ -0,0 +1,477 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+
+### components
+class ResConv(nn.Module):
+    """
+    Residual convolutional block, where
+    convolutional block consists: (convolution => [BN] => ReLU) * 3
+    residual connection adds the input to the output
+    """
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        self.double_conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+    def forward(self, x):
+        x_in = self.double_conv1(x)
+        x1 = self.double_conv(x)
+        return self.double_conv(x) + x_in
+
+class Down(nn.Module):
+    """Downscaling with maxpool then Resconv"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            ResConv(in_channels, out_channels)
+        )
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+class Up(nn.Module):
+	"""Upscaling then double conv"""
+	def __init__(self, in_channels, out_channels, bilinear=True):
+		super().__init__()
+		# if bilinear, use the normal convolutions to reduce the number of channels
+		if bilinear:
+			self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+			self.conv = ResConv(in_channels, out_channels, in_channels // 2)
+		else:
+			self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
+			self.conv = ResConv(in_channels, out_channels)
+	def forward(self, x1, x2):
+		x1 = self.up(x1)
+		# input is CHW
+		diffY = x2.size()[2] - x1.size()[2]
+		diffX = x2.size()[3] - x1.size()[3]
+		x1 = F.pad(
+			x1, 
+			[
+				diffX // 2, diffX - diffX // 2,
+				diffY // 2, diffY - diffY // 2
+			]
+		)
+		# if you have padding issues, see
+		# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
+		# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
+		x = torch.cat([x2, x1], dim=1)
+		return self.conv(x)
+
+class OutConv(nn.Module):
+	def __init__(self, in_channels, out_channels):
+		super(OutConv, self).__init__()
+		self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+	def forward(self, x):
+		# return F.relu(self.conv(x))
+		return self.conv(x)
+
+##### The composite networks
+class UNet(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		self.outc = OutConv(64, out_channels)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y = self.outc(x)
+		return y
+
+class CasUNet(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x, dop=None):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				if dop is not None:
+					y = F.dropout2d(self.unet_list[i](y), p=dop)
+				else:
+					y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		return y
+
+class CasUNet_2head(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet_2head, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			if i != self.n_unet-1:
+				self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+			else:
+				self.unet_list.append(UNet_2head(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		y_mean, y_sigma = y[0], y[1]
+		return y_mean, y_sigma
+
+class CasUNet_3head(nn.Module):
+	def __init__(self, n_unet, io_channels, bilinear=True):
+		super(CasUNet_3head, self).__init__()
+		self.n_unet = n_unet
+		self.io_channels = io_channels
+		self.bilinear = bilinear
+		####
+		self.unet_list = nn.ModuleList()
+		for i in range(self.n_unet):
+			if i != self.n_unet-1:
+				self.unet_list.append(UNet(self.io_channels, self.io_channels, self.bilinear))
+			else:
+				self.unet_list.append(UNet_3head(self.io_channels, self.io_channels, self.bilinear))
+	def forward(self, x):
+		y = x
+		for i in range(self.n_unet):
+			if i==0:
+				y = self.unet_list[i](y)
+			else:
+				y = self.unet_list[i](y+x)
+		y_mean, y_alpha, y_beta = y[0], y[1], y[2]
+		return y_mean, y_alpha, y_beta
+
+class UNet_2head(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet_2head, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		#per pixel multiple channels may exist
+		self.out_mean = OutConv(64, out_channels)
+		#variance will always be a single number for a pixel
+		self.out_var = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+		)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y_mean, y_var = self.out_mean(x), self.out_var(x)
+		return y_mean, y_var
+
+class UNet_3head(nn.Module):
+	def __init__(self, n_channels, out_channels, bilinear=True):
+		super(UNet_3head, self).__init__()
+		self.n_channels = n_channels
+		self.out_channels = out_channels
+		self.bilinear = bilinear
+		####
+		self.inc = ResConv(n_channels, 64)
+		self.down1 = Down(64, 128)
+		self.down2 = Down(128, 256)
+		self.down3 = Down(256, 512)
+		factor = 2 if bilinear else 1
+		self.down4 = Down(512, 1024 // factor)
+		self.up1 = Up(1024, 512 // factor, bilinear)
+		self.up2 = Up(512, 256 // factor, bilinear)
+		self.up3 = Up(256, 128 // factor, bilinear)
+		self.up4 = Up(128, 64, bilinear)
+		#per pixel multiple channels may exist
+		self.out_mean = OutConv(64, out_channels)
+		#variance will always be a single number for a pixel
+		self.out_alpha = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+			nn.ReLU()
+		)
+		self.out_beta = nn.Sequential(
+			OutConv(64, 128),
+			OutConv(128, 1),
+			nn.ReLU()
+		)
+	def forward(self, x):
+		x1 = self.inc(x)
+		x2 = self.down1(x1)
+		x3 = self.down2(x2)
+		x4 = self.down3(x3)
+		x5 = self.down4(x4)
+		x = self.up1(x5, x4)
+		x = self.up2(x, x3)
+		x = self.up3(x, x2)
+		x = self.up4(x, x1)
+		y_mean, y_alpha, y_beta = self.out_mean(x), \
+		self.out_alpha(x), self.out_beta(x)
+		return y_mean, y_alpha, y_beta
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+        conv_block = [  
+			nn.ReflectionPad2d(1),
+			nn.Conv2d(in_features, in_features, 3),
+			nn.InstanceNorm2d(in_features),
+			nn.ReLU(inplace=True),
+			nn.ReflectionPad2d(1),
+			nn.Conv2d(in_features, in_features, 3),
+			nn.InstanceNorm2d(in_features)
+		]
+        self.conv_block = nn.Sequential(*conv_block)
+    def forward(self, x):
+        return x + self.conv_block(x)
+
+class Generator(nn.Module):
+    def __init__(self, input_nc, output_nc, n_residual_blocks=9):
+        super(Generator, self).__init__()
+        # Initial convolution block       
+        model = [
+			nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, 7),
+            nn.InstanceNorm2d(64), nn.ReLU(inplace=True)
+		]
+        # Downsampling
+        in_features = 64
+        out_features = in_features*2
+        for _ in range(2):
+            model += [  
+				nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True) 
+			]
+            in_features = out_features
+            out_features = in_features*2
+        # Residual blocks
+        for _ in range(n_residual_blocks):
+            model += [ResidualBlock(in_features)]
+        # Upsampling
+        out_features = in_features//2
+        for _ in range(2):
+            model += [  
+				nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+				nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True)
+			]
+            in_features = out_features
+            out_features = in_features//2
+        # Output layer
+        model += [nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, 7), nn.Tanh()]
+        self.model = nn.Sequential(*model)
+    def forward(self, x):
+        return self.model(x)
+    
+    
+class ResnetGenerator(nn.Module):
+    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
+    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
+    """
+
+    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
+        """Construct a Resnet-based generator
+        Parameters:
+            input_nc (int)      -- the number of channels in input images
+            output_nc (int)     -- the number of channels in output images
+            ngf (int)           -- the number of filters in the last conv layer
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers
+            n_blocks (int)      -- the number of ResNet blocks
+            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero
+        """
+        assert(n_blocks >= 0)
+        super(ResnetGenerator, self).__init__()
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+
+        model = [nn.ReflectionPad2d(3),
+                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
+                 norm_layer(ngf),
+                 nn.ReLU(True)]
+
+        n_downsampling = 2
+        for i in range(n_downsampling):  # add downsampling layers
+            mult = 2 ** i
+            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
+                      norm_layer(ngf * mult * 2),
+                      nn.ReLU(True)]
+
+        mult = 2 ** n_downsampling
+        for i in range(n_blocks):       # add ResNet blocks
+
+            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
+
+        for i in range(n_downsampling):  # add upsampling layers
+            mult = 2 ** (n_downsampling - i)
+            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
+                                         kernel_size=3, stride=2,
+                                         padding=1, output_padding=1,
+                                         bias=use_bias),
+                      norm_layer(int(ngf * mult / 2)),
+                      nn.ReLU(True)]
+        model += [nn.ReflectionPad2d(3)]
+        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
+        model += [nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        """Standard forward"""
+        return self.model(input)
+
+
+class ResnetBlock(nn.Module):
+    """Define a Resnet block"""
+
+    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Initialize the Resnet block
+        A resnet block is a conv block with skip connections
+        We construct a conv block with build_conv_block function,
+        and implement skip connections in <forward> function.
+        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
+        """
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
+        """Construct a convolutional block.
+        Parameters:
+            dim (int)           -- the number of channels in the conv layer.
+            padding_type (str)  -- the name of padding layer: reflect | replicate | zero
+            norm_layer          -- normalization layer
+            use_dropout (bool)  -- if use dropout layers.
+            use_bias (bool)     -- if the conv layer uses bias or not
+        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
+        """
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' % padding_type)
+        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.conv_block(x)  # add skip connections
+        return out
+
+### discriminator
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator"""
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.model = nn.Sequential(*sequence)
+    def forward(self, input):
+        """Standard forward."""
+        return self.model(input)

src/utils.py

@@ -0,0 +1,1273 @@
+import random
+from typing import Any, Optional
+import numpy as np
+import os
+import cv2
+from glob import glob
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+import kornia
+import matplotlib.pyplot as plt
+import seaborn as sns
+import albumentations as albu
+import functools
+import math
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+import torchvision as tv
+import torchvision.models as models
+from torchvision import transforms
+from torchvision.transforms import functional as F
+from losses import TempCombLoss
+
+########### DeblurGAN function
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+def _array_to_batch(x):
+        x = np.transpose(x, (2, 0, 1))
+        x = np.expand_dims(x, 0)
+        return torch.from_numpy(x)
+
+def get_normalize():
+    normalize = albu.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    normalize = albu.Compose([normalize], additional_targets={'target': 'image'})
+
+    def process(a, b):
+        r = normalize(image=a, target=b)
+        return r['image'], r['target']
+
+    return process
+
+def preprocess(x: np.ndarray, mask: Optional[np.ndarray]):
+    x, _ = get_normalize()(x, x)
+    if mask is None:
+        mask = np.ones_like(x, dtype=np.float32)
+    else:
+        mask = np.round(mask.astype('float32') / 255)
+
+    h, w, _ = x.shape
+    block_size = 32
+    min_height = (h // block_size + 1) * block_size
+    min_width = (w // block_size + 1) * block_size
+
+    pad_params = {'mode': 'constant',
+                    'constant_values': 0,
+                    'pad_width': ((0, min_height - h), (0, min_width - w), (0, 0))
+                    }
+    x = np.pad(x, **pad_params)
+    mask = np.pad(mask, **pad_params)
+
+    return map(_array_to_batch, (x, mask)), h, w
+
+def postprocess(x: torch.Tensor) -> np.ndarray:
+    x, = x
+    x = x.detach().cpu().float().numpy()
+    x = (np.transpose(x, (1, 2, 0)) + 1) / 2.0 * 255.0
+    return x.astype('uint8')
+
+def sorted_glob(pattern):
+    return sorted(glob(pattern))
+###########
+
+def normalize(image: np.ndarray) -> np.ndarray:
+	"""Normalize the ``OpenCV.imread`` or ``skimage.io.imread`` data.
+	Args:
+		image (np.ndarray): The image data read by ``OpenCV.imread`` or ``skimage.io.imread``.
+	Returns:
+		Normalized image data. Data range [0, 1].
+	"""
+	return image.astype(np.float64) / 255.0
+
+
+def unnormalize(image: np.ndarray) -> np.ndarray:
+	"""Un-normalize the ``OpenCV.imread`` or ``skimage.io.imread`` data.
+	Args:
+		image (np.ndarray): The image data read by ``OpenCV.imread`` or ``skimage.io.imread``.
+	Returns:
+		Denormalized image data. Data range [0, 255].
+	"""
+	return image.astype(np.float64) * 255.0
+
+
+def image2tensor(image: np.ndarray, range_norm: bool, half: bool) -> torch.Tensor:
+	"""Convert ``PIL.Image`` to Tensor.
+	Args:
+		image (np.ndarray): The image data read by ``PIL.Image``
+		range_norm (bool): Scale [0, 1] data to between [-1, 1]
+		half (bool): Whether to convert torch.float32 similarly to torch.half type.
+	Returns:
+		Normalized image data
+	Examples:
+		>>> image = Image.open("image.bmp")
+		>>> tensor_image = image2tensor(image, range_norm=False, half=False)
+	"""
+	tensor = F.to_tensor(image)
+
+	if range_norm:
+		tensor = tensor.mul_(2.0).sub_(1.0)
+	if half:
+		tensor = tensor.half()
+
+	return tensor
+
+
+def tensor2image(tensor: torch.Tensor, range_norm: bool, half: bool) -> Any:
+	"""Converts ``torch.Tensor`` to ``PIL.Image``.
+	Args:
+		tensor (torch.Tensor): The image that needs to be converted to ``PIL.Image``
+		range_norm (bool): Scale [-1, 1] data to between [0, 1]
+		half (bool): Whether to convert torch.float32 similarly to torch.half type.
+	Returns:
+		Convert image data to support PIL library
+	Examples:
+		>>> tensor = torch.randn([1, 3, 128, 128])
+		>>> image = tensor2image(tensor, range_norm=False, half=False)
+	"""
+	if range_norm:
+		tensor = tensor.add_(1.0).div_(2.0)
+	if half:
+		tensor = tensor.half()
+
+	image = tensor.squeeze_(0).permute(1, 2, 0).mul_(255).clamp_(0, 255).cpu().numpy().astype("uint8")
+
+	return image
+
+
+def convert_rgb_to_y(image: Any) -> Any:
+	"""Convert RGB image or tensor image data to YCbCr(Y) format.
+	Args:
+		image: RGB image data read by ``PIL.Image''.
+	Returns:
+		Y image array data.
+	"""
+	if type(image) == np.ndarray:
+		return 16. + (64.738 * image[:, :, 0] + 129.057 * image[:, :, 1] + 25.064 * image[:, :, 2]) / 256.
+	elif type(image) == torch.Tensor:
+		if len(image.shape) == 4:
+			image = image.squeeze_(0)
+		return 16. + (64.738 * image[0, :, :] + 129.057 * image[1, :, :] + 25.064 * image[2, :, :]) / 256.
+	else:
+		raise Exception("Unknown Type", type(image))
+
+
+def convert_rgb_to_ycbcr(image: Any) -> Any:
+	"""Convert RGB image or tensor image data to YCbCr format.
+	Args:
+		image: RGB image data read by ``PIL.Image''.
+	Returns:
+		YCbCr image array data.
+	"""
+	if type(image) == np.ndarray:
+		y = 16. + (64.738 * image[:, :, 0] + 129.057 * image[:, :, 1] + 25.064 * image[:, :, 2]) / 256.
+		cb = 128. + (-37.945 * image[:, :, 0] - 74.494 * image[:, :, 1] + 112.439 * image[:, :, 2]) / 256.
+		cr = 128. + (112.439 * image[:, :, 0] - 94.154 * image[:, :, 1] - 18.285 * image[:, :, 2]) / 256.
+		return np.array([y, cb, cr]).transpose([1, 2, 0])
+	elif type(image) == torch.Tensor:
+		if len(image.shape) == 4:
+			image = image.squeeze(0)
+		y = 16. + (64.738 * image[0, :, :] + 129.057 * image[1, :, :] + 25.064 * image[2, :, :]) / 256.
+		cb = 128. + (-37.945 * image[0, :, :] - 74.494 * image[1, :, :] + 112.439 * image[2, :, :]) / 256.
+		cr = 128. + (112.439 * image[0, :, :] - 94.154 * image[1, :, :] - 18.285 * image[2, :, :]) / 256.
+		return torch.cat([y, cb, cr], 0).permute(1, 2, 0)
+	else:
+		raise Exception("Unknown Type", type(image))
+
+
+def convert_ycbcr_to_rgb(image: Any) -> Any:
+	"""Convert YCbCr format image to RGB format.
+	Args:
+	   image: YCbCr image data read by ``PIL.Image''.
+	Returns:
+		RGB image array data.
+	"""
+	if type(image) == np.ndarray:
+		r = 298.082 * image[:, :, 0] / 256. + 408.583 * image[:, :, 2] / 256. - 222.921
+		g = 298.082 * image[:, :, 0] / 256. - 100.291 * image[:, :, 1] / 256. - 208.120 * image[:, :, 2] / 256. + 135.576
+		b = 298.082 * image[:, :, 0] / 256. + 516.412 * image[:, :, 1] / 256. - 276.836
+		return np.array([r, g, b]).transpose([1, 2, 0])
+	elif type(image) == torch.Tensor:
+		if len(image.shape) == 4:
+			image = image.squeeze(0)
+		r = 298.082 * image[0, :, :] / 256. + 408.583 * image[2, :, :] / 256. - 222.921
+		g = 298.082 * image[0, :, :] / 256. - 100.291 * image[1, :, :] / 256. - 208.120 * image[2, :, :] / 256. + 135.576
+		b = 298.082 * image[0, :, :] / 256. + 516.412 * image[1, :, :] / 256. - 276.836
+		return torch.cat([r, g, b], 0).permute(1, 2, 0)
+	else:
+		raise Exception("Unknown Type", type(image))
+
+
+def center_crop(lr: Any, hr: Any, image_size: int, upscale_factor: int) -> [Any, Any]:
+	"""Cut ``PIL.Image`` in the center area of the image.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+		image_size (int): The size of the captured image area. It should be the size of the high-resolution image.
+		upscale_factor (int): magnification factor.
+	Returns:
+		Randomly cropped low-resolution images and high-resolution images.
+	"""
+	w, h = hr.size
+
+	left = (w - image_size) // 2
+	top = (h - image_size) // 2
+	right = left + image_size
+	bottom = top + image_size
+
+	lr = lr.crop((left // upscale_factor,
+				  top // upscale_factor,
+				  right // upscale_factor,
+				  bottom // upscale_factor))
+	hr = hr.crop((left, top, right, bottom))
+
+	return lr, hr
+
+
+def random_crop(lr: Any, hr: Any, image_size: int, upscale_factor: int) -> [Any, Any]:
+	"""Will ``PIL.Image`` randomly capture the specified area of the image.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+		image_size (int): The size of the captured image area. It should be the size of the high-resolution image.
+		upscale_factor (int): magnification factor.
+	Returns:
+		Randomly cropped low-resolution images and high-resolution images.
+	"""
+	w, h = hr.size
+	left = torch.randint(0, w - image_size + 1, size=(1,)).item()
+	top = torch.randint(0, h - image_size + 1, size=(1,)).item()
+	right = left + image_size
+	bottom = top + image_size
+
+	lr = lr.crop((left // upscale_factor,
+				  top // upscale_factor,
+				  right // upscale_factor,
+				  bottom // upscale_factor))
+	hr = hr.crop((left, top, right, bottom))
+
+	return lr, hr
+
+
+def random_rotate(lr: Any, hr: Any, angle: int) -> [Any, Any]:
+	"""Will ``PIL.Image`` randomly rotate the image.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+		angle (int): rotation angle, clockwise and counterclockwise rotation.
+	Returns:
+		Randomly rotated low-resolution images and high-resolution images.
+	"""
+	angle = random.choice((+angle, -angle))
+	lr = F.rotate(lr, angle)
+	hr = F.rotate(hr, angle)
+
+	return lr, hr
+
+
+def random_horizontally_flip(lr: Any, hr: Any, p=0.5) -> [Any, Any]:
+	"""Flip the ``PIL.Image`` image horizontally randomly.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+		p (optional, float): rollover probability. (Default: 0.5)
+	Returns:
+		Low-resolution image and high-resolution image after random horizontal flip.
+	"""
+	if torch.rand(1).item() > p:
+		lr = F.hflip(lr)
+		hr = F.hflip(hr)
+
+	return lr, hr
+
+
+def random_vertically_flip(lr: Any, hr: Any, p=0.5) -> [Any, Any]:
+	"""Turn the ``PIL.Image`` image upside down randomly.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+		p (optional, float): rollover probability. (Default: 0.5)
+	Returns:
+		Randomly rotated up and down low-resolution images and high-resolution images.
+	"""
+	if torch.rand(1).item() > p:
+		lr = F.vflip(lr)
+		hr = F.vflip(hr)
+
+	return lr, hr
+
+
+def random_adjust_brightness(lr: Any, hr: Any) -> [Any, Any]:
+	"""Set ``PIL.Image`` to randomly adjust the image brightness.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+	Returns:
+		Low-resolution image and high-resolution image with randomly adjusted brightness.
+	"""
+	# Randomly adjust the brightness gain range.
+	factor = random.uniform(0.5, 2)
+	lr = F.adjust_brightness(lr, factor)
+	hr = F.adjust_brightness(hr, factor)
+
+	return lr, hr
+
+
+def random_adjust_contrast(lr: Any, hr: Any) -> [Any, Any]:
+	"""Set ``PIL.Image`` to randomly adjust the image contrast.
+	Args:
+		lr: Low-resolution image data read by ``PIL.Image``.
+		hr: High-resolution image data read by ``PIL.Image``.
+	Returns:
+		Low-resolution image and high-resolution image with randomly adjusted contrast.
+	"""
+	# Randomly adjust the contrast gain range.
+	factor = random.uniform(0.5, 2)
+	lr = F.adjust_contrast(lr, factor)
+	hr = F.adjust_contrast(hr, factor)
+
+	return lr, hr
+
+#### metrics to compute -- assumes single images, i.e., tensor of 3 dims
+def img_mae(x1, x2):
+	m = torch.abs(x1-x2).mean()
+	return m
+
+def img_mse(x1, x2):
+	m = torch.pow(torch.abs(x1-x2),2).mean()
+	return m
+
+def img_psnr(x1, x2):
+	m = kornia.metrics.psnr(x1, x2, 1)
+	return m
+
+def img_ssim(x1, x2):
+	m = kornia.metrics.ssim(x1.unsqueeze(0), x2.unsqueeze(0), 5)
+	m = m.mean()
+	return m
+
+def show_SR_w_uncer(xLR, xHR, xSR, xSRvar, elim=(0,0.01), ulim=(0,0.15)):
+	'''
+	xLR/SR/HR: 3xHxW
+	xSRvar: 1xHxW
+	'''
+	plt.figure(figsize=(30,10))
+	
+	plt.subplot(1,5,1)
+	plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+	plt.axis('off')
+	
+	plt.subplot(1,5,2)
+	plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+	plt.axis('off')
+	
+	plt.subplot(1,5,3)
+	plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+	plt.axis('off')
+	
+	plt.subplot(1,5,4)
+	error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0) 
+	print('error', error_map.min(), error_map.max())
+	plt.imshow(error_map.transpose(0,2).transpose(0,1), cmap='jet')
+	plt.clim(elim[0], elim[1])
+	plt.axis('off')
+
+	plt.subplot(1,5,5)
+	print('uncer', xSRvar.min(), xSRvar.max())
+	plt.imshow(xSRvar.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+	plt.clim(ulim[0], ulim[1])
+	plt.axis('off')
+
+	plt.subplots_adjust(wspace=0, hspace=0)
+	plt.show()
+
+def show_SR_w_err(xLR, xHR, xSR, elim=(0,0.01), task=None, xMask=None):
+	'''
+	xLR/SR/HR: 3xHxW
+	'''
+	plt.figure(figsize=(30,10))
+	
+	if task != 'm':
+		plt.subplot(1,4,1)
+		plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+		plt.axis('off')
+		
+		plt.subplot(1,4,2)
+		plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+		plt.axis('off')
+		
+		plt.subplot(1,4,3)
+		plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+		plt.axis('off')
+	else:
+		plt.subplot(1,4,1)
+		plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+		plt.clim(0,0.9)
+		plt.axis('off')
+		
+		plt.subplot(1,4,2)
+		plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+		plt.clim(0,0.9)
+		plt.axis('off')
+		
+		plt.subplot(1,4,3)
+		plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+		plt.clim(0,0.9)
+		plt.axis('off')
+	
+	plt.subplot(1,4,4)
+	if task == 'inpainting':
+		error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0)*xMask.to('cpu').data
+	else:
+		error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0) 
+	print('error', error_map.min(), error_map.max())
+	plt.imshow(error_map.transpose(0,2).transpose(0,1), cmap='jet')
+	plt.clim(elim[0], elim[1])
+	plt.axis('off')
+
+	plt.subplots_adjust(wspace=0, hspace=0)
+	plt.show()
+
+def show_uncer4(xSRvar1, xSRvar2, xSRvar3, xSRvar4, ulim=(0,0.15)):
+	'''
+	xSRvar: 1xHxW
+	'''
+	plt.figure(figsize=(30,10))
+	
+	plt.subplot(1,4,1)
+	print('uncer', xSRvar1.min(), xSRvar1.max())
+	plt.imshow(xSRvar1.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+	plt.clim(ulim[0], ulim[1])
+	plt.axis('off')
+
+	plt.subplot(1,4,2)
+	print('uncer', xSRvar2.min(), xSRvar2.max())
+	plt.imshow(xSRvar2.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+	plt.clim(ulim[0], ulim[1])
+	plt.axis('off')
+
+	plt.subplot(1,4,3)
+	print('uncer', xSRvar3.min(), xSRvar3.max())
+	plt.imshow(xSRvar3.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+	plt.clim(ulim[0], ulim[1])
+	plt.axis('off')
+
+	plt.subplot(1,4,4)
+	print('uncer', xSRvar4.min(), xSRvar4.max())
+	plt.imshow(xSRvar4.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+	plt.clim(ulim[0], ulim[1])
+	plt.axis('off')
+
+	plt.subplots_adjust(wspace=0, hspace=0)
+	plt.show()
+
+def get_UCE(list_err, list_yout_var, num_bins=100):
+	err_min = np.min(list_err)
+	err_max = np.max(list_err)
+	err_len = (err_max-err_min)/num_bins
+	num_points = len(list_err)
+	
+	bin_stats = {}
+	for i in range(num_bins):
+		bin_stats[i] = {
+			'start_idx': err_min + i*err_len,
+			'end_idx': err_min + (i+1)*err_len,
+			'num_points': 0,
+			'mean_err': 0,
+			'mean_var': 0,
+		}
+	
+	for e,v in zip(list_err, list_yout_var):
+		for i in range(num_bins):
+			if e>=bin_stats[i]['start_idx'] and e<bin_stats[i]['end_idx']:
+				bin_stats[i]['num_points'] += 1
+				bin_stats[i]['mean_err'] += e
+				bin_stats[i]['mean_var'] += v
+	
+	uce = 0
+	eps = 1e-8
+	for i in range(num_bins):
+		bin_stats[i]['mean_err'] /= bin_stats[i]['num_points'] + eps
+		bin_stats[i]['mean_var'] /= bin_stats[i]['num_points'] + eps
+		bin_stats[i]['uce_bin'] = (bin_stats[i]['num_points']/num_points) \
+			*(np.abs(bin_stats[i]['mean_err'] - bin_stats[i]['mean_var']))
+		uce += bin_stats[i]['uce_bin']
+	
+	list_x, list_y = [], []
+	for i in range(num_bins):
+		if bin_stats[i]['num_points']>0:
+			list_x.append(bin_stats[i]['mean_err'])
+			list_y.append(bin_stats[i]['mean_var'])
+	
+	# sns.set_style('darkgrid')
+	# sns.scatterplot(x=list_x, y=list_y)
+	# sns.regplot(x=list_x, y=list_y, order=1)
+	# plt.xlabel('MSE', fontsize=34)
+	# plt.ylabel('Uncertainty', fontsize=34)
+	# plt.plot(list_x, list_x, color='r')
+	# plt.xlim(np.min(list_x), np.max(list_x))
+	# plt.ylim(np.min(list_err), np.max(list_x))
+	# plt.show()
+
+	return bin_stats, uce
+
+##################### training BayesCap
+def train_BayesCap(
+	NetC,
+	NetG,
+	train_loader,
+	eval_loader,
+	Cri = TempCombLoss(),
+	device='cuda',
+	dtype=torch.cuda.FloatTensor(),
+	init_lr=1e-4,
+	num_epochs=100,
+	eval_every=1,
+	ckpt_path='../ckpt/BayesCap',
+	T1=1e0,
+	T2=5e-2,
+	task=None,
+):
+	NetC.to(device)
+	NetC.train()
+	NetG.to(device)
+	NetG.eval()
+	optimizer = torch.optim.Adam(list(NetC.parameters()), lr=init_lr)
+	optim_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)
+
+	score = -1e8
+	all_loss = []
+	for eph in range(num_epochs):
+		eph_loss = 0
+		with tqdm(train_loader, unit='batch') as tepoch:
+			for (idx, batch) in enumerate(tepoch):
+				if idx>2000:
+					break
+				tepoch.set_description('Epoch {}'.format(eph))
+				##
+				xLR, xHR = batch[0].to(device), batch[1].to(device)
+				xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+				if task == 'inpainting':
+					xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+					xMask = xMask.to(device).type(dtype)
+				# pass them through the network
+				with torch.no_grad():
+					if task == 'inpainting':
+						_, xSR1 = NetG(xLR, xMask)
+					elif task == 'depth':
+						xSR1 = NetG(xLR)[("disp", 0)]
+					else:
+						xSR1 = NetG(xLR)
+				# with torch.autograd.set_detect_anomaly(True):
+				xSR = xSR1.clone()
+				xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+				# print(xSRC_alpha) 
+				optimizer.zero_grad()
+				if task == 'depth':
+					loss = Cri(xSRC_mu, xSRC_alpha, xSRC_beta, xSR, T1=T1, T2=T2)
+				else:
+					loss = Cri(xSRC_mu, xSRC_alpha, xSRC_beta, xHR, T1=T1, T2=T2)
+				# print(loss)
+				loss.backward()
+				optimizer.step()
+				##
+				eph_loss += loss.item()
+				tepoch.set_postfix(loss=loss.item())
+			eph_loss /= len(train_loader)
+			all_loss.append(eph_loss)
+			print('Avg. loss: {}'.format(eph_loss))
+		# evaluate and save the models
+		torch.save(NetC.state_dict(), ckpt_path+'_last.pth')
+		if eph%eval_every == 0:
+			curr_score = eval_BayesCap(
+				NetC,
+				NetG,
+				eval_loader,
+				device=device,
+				dtype=dtype,
+				task=task,
+			)
+			print('current score: {} | Last best score: {}'.format(curr_score, score))
+			if curr_score >= score:
+				score = curr_score
+				torch.save(NetC.state_dict(), ckpt_path+'_best.pth')
+	optim_scheduler.step()
+
+#### get different uncertainty maps
+def get_uncer_BayesCap(
+	NetC,
+	NetG,
+	xin,
+	task=None,
+	xMask=None,
+):
+	with torch.no_grad():
+		if task == 'inpainting':
+			_, xSR = NetG(xin, xMask)
+		else:
+			xSR = NetG(xin)
+		xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+	a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+	b_map = xSRC_beta.to('cpu').data
+	xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+
+	return xSRvar
+
+def get_uncer_TTDAp(
+	NetG,
+	xin,
+	p_mag=0.05,
+	num_runs=50,
+	task=None,
+	xMask=None,
+):
+	list_xSR = []
+	with torch.no_grad():
+		for z in range(num_runs):
+			if task == 'inpainting':
+				_, xSRz = NetG(xin+p_mag*xin.max()*torch.randn_like(xin), xMask)
+			else:
+				xSRz = NetG(xin+p_mag*xin.max()*torch.randn_like(xin))
+			list_xSR.append(xSRz)
+	xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+	xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+	return xSRvar
+
+def get_uncer_DO(
+	NetG,
+	xin,
+	dop=0.2,
+	num_runs=50,
+	task=None,
+	xMask=None,
+):
+	list_xSR = []
+	with torch.no_grad():
+		for z in range(num_runs):
+			if task == 'inpainting':
+				_, xSRz = NetG(xin, xMask, dop=dop)
+			else:
+				xSRz = NetG(xin, dop=dop)
+			list_xSR.append(xSRz)
+	xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+	xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+	return xSRvar
+
+################### Different eval functions
+
+def eval_BayesCap(
+	NetC,
+	NetG,
+	eval_loader,
+	device='cuda',
+	dtype=torch.cuda.FloatTensor,
+	task=None,
+	xMask=None,
+):
+	NetC.to(device)
+	NetC.eval()
+	NetG.to(device)
+	NetG.eval()
+	
+	mean_ssim = 0
+	mean_psnr = 0
+	mean_mse = 0
+	mean_mae = 0
+	num_imgs = 0
+	list_error = []
+	list_var = []
+	with tqdm(eval_loader, unit='batch') as tepoch:
+		for (idx, batch) in enumerate(tepoch):
+			tepoch.set_description('Validating ...')
+			##
+			xLR, xHR = batch[0].to(device), batch[1].to(device)
+			xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+			if task == 'inpainting':
+				if xMask==None:
+					xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+					xMask = xMask.to(device).type(dtype)
+				else:
+					xMask = xMask.to(device).type(dtype)
+			# pass them through the network
+			with torch.no_grad():
+				if task == 'inpainting':
+					_, xSR = NetG(xLR, xMask)
+				elif task == 'depth':
+					xSR = NetG(xLR)[("disp", 0)]
+				else:
+					xSR = NetG(xLR)
+				xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+			a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+			b_map = xSRC_beta.to('cpu').data
+			xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+			n_batch = xSRC_mu.shape[0]
+			if task == 'depth':
+				xHR = xSR
+			for j in range(n_batch):
+				num_imgs += 1
+				mean_ssim += img_ssim(xSRC_mu[j], xHR[j])
+				mean_psnr += img_psnr(xSRC_mu[j], xHR[j])
+				mean_mse += img_mse(xSRC_mu[j], xHR[j])
+				mean_mae += img_mae(xSRC_mu[j], xHR[j])
+
+				show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+				
+				error_map = torch.mean(torch.pow(torch.abs(xSR[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+				var_map =  xSRvar[j].to('cpu').data.reshape(-1)
+				list_error.extend(list(error_map.numpy()))
+				list_var.extend(list(var_map.numpy()))
+			##
+		mean_ssim /= num_imgs
+		mean_psnr /= num_imgs
+		mean_mse /= num_imgs
+		mean_mae /= num_imgs
+		print(
+			'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+			(
+				mean_ssim, mean_psnr, mean_mse, mean_mae 
+			)
+		)
+		# print(len(list_error), len(list_var))
+		# print('UCE: ', get_UCE(list_error[::10], list_var[::10], num_bins=500)[1])
+		# print('C.Coeff: ', np.corrcoef(np.array(list_error[::10]), np.array(list_var[::10])))
+	return mean_ssim
+
+def eval_TTDA_p(
+	NetG,
+	eval_loader,
+	device='cuda',
+	dtype=torch.cuda.FloatTensor,
+	p_mag=0.05,
+	num_runs=50,
+	task = None,
+	xMask = None,
+):
+	NetG.to(device)
+	NetG.eval()
+	
+	mean_ssim = 0
+	mean_psnr = 0
+	mean_mse = 0
+	mean_mae = 0
+	num_imgs = 0
+	with tqdm(eval_loader, unit='batch') as tepoch:
+		for (idx, batch) in enumerate(tepoch):
+			tepoch.set_description('Validating ...')
+			##
+			xLR, xHR = batch[0].to(device), batch[1].to(device)
+			xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+			# pass them through the network
+			list_xSR = []
+			with torch.no_grad():
+				if task=='inpainting':
+					_, xSR = NetG(xLR, xMask)
+				else:
+					xSR = NetG(xLR)
+				for z in range(num_runs):
+					xSRz = NetG(xLR+p_mag*xLR.max()*torch.randn_like(xLR))	
+					list_xSR.append(xSRz)
+			xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+			xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+			n_batch = xSR.shape[0]
+			for j in range(n_batch):
+				num_imgs += 1
+				mean_ssim += img_ssim(xSR[j], xHR[j])
+				mean_psnr += img_psnr(xSR[j], xHR[j])
+				mean_mse += img_mse(xSR[j], xHR[j])
+				mean_mae += img_mae(xSR[j], xHR[j])
+
+				show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+				
+		mean_ssim /= num_imgs
+		mean_psnr /= num_imgs
+		mean_mse /= num_imgs
+		mean_mae /= num_imgs
+		print(
+			'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+			(
+				mean_ssim, mean_psnr, mean_mse, mean_mae 
+			)
+		)
+		
+	return mean_ssim
+
+def eval_DO(
+	NetG,
+	eval_loader,
+	device='cuda',
+	dtype=torch.cuda.FloatTensor,
+	dop=0.2,
+	num_runs=50,
+	task=None,
+	xMask=None,
+):
+	NetG.to(device)
+	NetG.eval()
+	
+	mean_ssim = 0
+	mean_psnr = 0
+	mean_mse = 0
+	mean_mae = 0
+	num_imgs = 0
+	with tqdm(eval_loader, unit='batch') as tepoch:
+		for (idx, batch) in enumerate(tepoch):
+			tepoch.set_description('Validating ...')
+			##
+			xLR, xHR = batch[0].to(device), batch[1].to(device)
+			xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+			# pass them through the network
+			list_xSR = []
+			with torch.no_grad():
+				if task == 'inpainting':
+					_, xSR = NetG(xLR, xMask)
+				else:
+					xSR = NetG(xLR)
+				for z in range(num_runs):
+					xSRz = NetG(xLR, dop=dop)	
+					list_xSR.append(xSRz)
+			xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+			xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+			n_batch = xSR.shape[0]
+			for j in range(n_batch):
+				num_imgs += 1
+				mean_ssim += img_ssim(xSR[j], xHR[j])
+				mean_psnr += img_psnr(xSR[j], xHR[j])
+				mean_mse += img_mse(xSR[j], xHR[j])
+				mean_mae += img_mae(xSR[j], xHR[j])
+
+				show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+						##
+		mean_ssim /= num_imgs
+		mean_psnr /= num_imgs
+		mean_mse /= num_imgs
+		mean_mae /= num_imgs
+		print(
+			'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+			(
+				mean_ssim, mean_psnr, mean_mse, mean_mae 
+			)
+		)
+	
+	return mean_ssim
+
+
+############### compare all function
+def compare_all(
+	NetC,
+	NetG,
+	eval_loader,
+	p_mag = 0.05,
+	dop = 0.2,
+	num_runs = 100,
+	device='cuda',
+	dtype=torch.cuda.FloatTensor,
+	task=None,
+):
+	NetC.to(device)
+	NetC.eval()
+	NetG.to(device)
+	NetG.eval()
+	
+	with tqdm(eval_loader, unit='batch') as tepoch:
+		for (idx, batch) in enumerate(tepoch):
+			tepoch.set_description('Comparing ...')
+			##
+			xLR, xHR = batch[0].to(device), batch[1].to(device)
+			xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+			if task == 'inpainting':
+				xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+				xMask = xMask.to(device).type(dtype)
+			# pass them through the network
+			with torch.no_grad():
+				if task == 'inpainting':
+					_, xSR = NetG(xLR, xMask)
+				else:
+					xSR = NetG(xLR)
+				xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+
+			if task == 'inpainting':
+				xSRvar1 = get_uncer_TTDAp(NetG, xLR, p_mag=p_mag, num_runs=num_runs, task='inpainting', xMask=xMask)
+				xSRvar2 = get_uncer_DO(NetG, xLR, dop=dop, num_runs=num_runs, task='inpainting', xMask=xMask)
+				xSRvar3 = get_uncer_BayesCap(NetC, NetG, xLR, task='inpainting', xMask=xMask)
+			else:
+				xSRvar1 = get_uncer_TTDAp(NetG, xLR, p_mag=p_mag, num_runs=num_runs)
+				xSRvar2 = get_uncer_DO(NetG, xLR, dop=dop, num_runs=num_runs)
+				xSRvar3 = get_uncer_BayesCap(NetC, NetG, xLR)
+
+			print('bdg', xSRvar1.shape, xSRvar2.shape, xSRvar3.shape)
+
+			n_batch = xSR.shape[0]
+			for j in range(n_batch):
+				if task=='s':
+					show_SR_w_err(xLR[j], xHR[j], xSR[j])
+					show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42))
+					show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 1.5), xSRvar3[j])
+				if task=='d':
+					show_SR_w_err(xLR[j], xHR[j], 0.5*xSR[j]+0.5*xHR[j])
+					show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42))
+					show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 0.8), xSRvar3[j])
+				if task=='inpainting':
+					show_SR_w_err(xLR[j]*(1-xMask[j]), xHR[j], xSR[j], elim=(0,0.25), task='inpainting', xMask=xMask[j])
+					show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.45), torch.pow(xSRvar1[j], 0.4))
+					show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 0.8), xSRvar3[j])
+				if task=='m':
+					show_SR_w_err(xLR[j], xHR[j], xSR[j], elim=(0,0.04), task='m')
+					show_uncer4(0.4*xSRvar1[j]+0.6*xSRvar2[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42), ulim=(0.02,0.15))
+					show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 1.5), xSRvar3[j], ulim=(0.02,0.15))
+
+
+################# Degrading Identity 
+def degrage_BayesCap_p(
+	NetC,
+	NetG,
+	eval_loader,
+	device='cuda',
+	dtype=torch.cuda.FloatTensor,
+	num_runs=50,
+):
+	NetC.to(device)
+	NetC.eval()
+	NetG.to(device)
+	NetG.eval()
+	
+	p_mag_list = [0, 0.05, 0.1, 0.15, 0.2]
+	list_s = []
+	list_p = []
+	list_u1 = []
+	list_u2 = []
+	list_c = []
+	for p_mag in p_mag_list:
+		mean_ssim = 0
+		mean_psnr = 0
+		mean_mse = 0
+		mean_mae = 0
+		num_imgs = 0
+		list_error = []
+		list_error2 = []
+		list_var = []
+
+		with tqdm(eval_loader, unit='batch') as tepoch:
+			for (idx, batch) in enumerate(tepoch):
+				tepoch.set_description('Validating ...')
+				##
+				xLR, xHR = batch[0].to(device), batch[1].to(device)
+				xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+				# pass them through the network
+				with torch.no_grad():
+					xSR = NetG(xLR)
+					xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR + p_mag*xSR.max()*torch.randn_like(xSR))
+				a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+				b_map = xSRC_beta.to('cpu').data
+				xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+				n_batch = xSRC_mu.shape[0]
+				for j in range(n_batch):
+					num_imgs += 1
+					mean_ssim += img_ssim(xSRC_mu[j], xSR[j])
+					mean_psnr += img_psnr(xSRC_mu[j], xSR[j])
+					mean_mse += img_mse(xSRC_mu[j], xSR[j])
+					mean_mae += img_mae(xSRC_mu[j], xSR[j])
+
+					error_map = torch.mean(torch.pow(torch.abs(xSR[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+					error_map2 = torch.mean(torch.pow(torch.abs(xSRC_mu[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+					var_map =  xSRvar[j].to('cpu').data.reshape(-1)
+					list_error.extend(list(error_map.numpy()))
+					list_error2.extend(list(error_map2.numpy()))
+					list_var.extend(list(var_map.numpy()))
+				##
+			mean_ssim /= num_imgs
+			mean_psnr /= num_imgs
+			mean_mse /= num_imgs
+			mean_mae /= num_imgs
+			print(
+				'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+				(
+					mean_ssim, mean_psnr, mean_mse, mean_mae 
+				)
+			)
+			uce1 = get_UCE(list_error[::100], list_var[::100], num_bins=200)[1]
+			uce2 = get_UCE(list_error2[::100], list_var[::100], num_bins=200)[1]
+			print('UCE1: ', uce1)
+			print('UCE2: ', uce2)
+			list_s.append(mean_ssim.item())
+			list_p.append(mean_psnr.item())
+			list_u1.append(uce1)
+			list_u2.append(uce2)
+	
+	plt.plot(list_s)
+	plt.show()
+	plt.plot(list_p)
+	plt.show()
+
+	plt.plot(list_u1, label='wrt SR output')
+	plt.plot(list_u2, label='wrt BayesCap output')
+	plt.legend()
+	plt.show()
+
+	sns.set_style('darkgrid')
+	fig,ax = plt.subplots()
+	# make a plot
+	ax.plot(p_mag_list, list_s, color="red", marker="o")
+	# set x-axis label
+	ax.set_xlabel("Reducing faithfulness of BayesCap Reconstruction",fontsize=10)
+	# set y-axis label
+	ax.set_ylabel("SSIM btwn BayesCap and SRGAN outputs", color="red",fontsize=10)
+
+	# twin object for two different y-axis on the sample plot
+	ax2=ax.twinx()
+	# make a plot with different y-axis using second axis object
+	ax2.plot(p_mag_list, list_u1, color="blue", marker="o", label='UCE wrt to error btwn SRGAN output and GT')
+	ax2.plot(p_mag_list, list_u2, color="orange", marker="o", label='UCE wrt to error btwn BayesCap output and GT')
+	ax2.set_ylabel("UCE", color="green", fontsize=10)
+	plt.legend(fontsize=10)
+	plt.tight_layout()
+	plt.show()
+
+################# DeepFill_v2
+   
+# ----------------------------------------
+#             PATH processing
+# ----------------------------------------
+def text_readlines(filename):
+    # Try to read a txt file and return a list.Return [] if there was a mistake.
+    try:
+        file = open(filename, 'r')
+    except IOError:
+        error = []
+        return error
+    content = file.readlines()
+    # This for loop deletes the EOF (like \n)
+    for i in range(len(content)):
+        content[i] = content[i][:len(content[i])-1]
+    file.close()
+    return content
+
+def savetxt(name, loss_log):
+    np_loss_log = np.array(loss_log)
+    np.savetxt(name, np_loss_log)
+
+def get_files(path):
+    # read a folder, return the complete path
+    ret = []
+    for root, dirs, files in os.walk(path):
+        for filespath in files:
+            ret.append(os.path.join(root, filespath))
+    return ret
+
+def get_names(path):
+    # read a folder, return the image name
+    ret = []
+    for root, dirs, files in os.walk(path):
+        for filespath in files:
+            ret.append(filespath)
+    return ret
+
+def text_save(content, filename, mode = 'a'):
+    # save a list to a txt
+    # Try to save a list variable in txt file.
+    file = open(filename, mode)
+    for i in range(len(content)):
+        file.write(str(content[i]) + '\n')
+    file.close()
+
+def check_path(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    
+# ----------------------------------------
+#    Validation and Sample at training
+# ----------------------------------------
+def save_sample_png(sample_folder, sample_name, img_list, name_list, pixel_max_cnt = 255):
+    # Save image one-by-one
+    for i in range(len(img_list)):
+        img = img_list[i]
+        # Recover normalization: * 255 because last layer is sigmoid activated
+        img = img * 255
+        # Process img_copy and do not destroy the data of img
+        img_copy = img.clone().data.permute(0, 2, 3, 1)[0, :, :, :].cpu().numpy()
+        img_copy = np.clip(img_copy, 0, pixel_max_cnt)
+        img_copy = img_copy.astype(np.uint8)
+        img_copy = cv2.cvtColor(img_copy, cv2.COLOR_RGB2BGR)
+        # Save to certain path
+        save_img_name = sample_name + '_' + name_list[i] + '.jpg'
+        save_img_path = os.path.join(sample_folder, save_img_name)
+        cv2.imwrite(save_img_path, img_copy)
+
+def psnr(pred, target, pixel_max_cnt = 255):
+    mse = torch.mul(target - pred, target - pred)
+    rmse_avg = (torch.mean(mse).item()) ** 0.5
+    p = 20 * np.log10(pixel_max_cnt / rmse_avg)
+    return p
+
+def grey_psnr(pred, target, pixel_max_cnt = 255):
+    pred = torch.sum(pred, dim = 0)
+    target = torch.sum(target, dim = 0)
+    mse = torch.mul(target - pred, target - pred)
+    rmse_avg = (torch.mean(mse).item()) ** 0.5
+    p = 20 * np.log10(pixel_max_cnt * 3 / rmse_avg)
+    return p
+
+def ssim(pred, target):
+    pred = pred.clone().data.permute(0, 2, 3, 1).cpu().numpy()
+    target = target.clone().data.permute(0, 2, 3, 1).cpu().numpy()
+    target = target[0]
+    pred = pred[0]
+    ssim = skimage.measure.compare_ssim(target, pred, multichannel = True)
+    return ssim
+
+## for contextual attention
+
+def extract_image_patches(images, ksizes, strides, rates, padding='same'):
+    """
+    Extract patches from images and put them in the C output dimension.
+    :param padding:
+    :param images: [batch, channels, in_rows, in_cols]. A 4-D Tensor with shape
+    :param ksizes: [ksize_rows, ksize_cols]. The size of the sliding window for
+     each dimension of images
+    :param strides: [stride_rows, stride_cols]
+    :param rates: [dilation_rows, dilation_cols]
+    :return: A Tensor
+    """
+    assert len(images.size()) == 4
+    assert padding in ['same', 'valid']
+    batch_size, channel, height, width = images.size()
+
+    if padding == 'same':
+        images = same_padding(images, ksizes, strides, rates)
+    elif padding == 'valid':
+        pass
+    else:
+        raise NotImplementedError('Unsupported padding type: {}.\
+                Only "same" or "valid" are supported.'.format(padding))
+
+    unfold = torch.nn.Unfold(kernel_size=ksizes,
+                             dilation=rates,
+                             padding=0,
+                             stride=strides)
+    patches = unfold(images)
+    return patches  # [N, C*k*k, L], L is the total number of such blocks
+
+def same_padding(images, ksizes, strides, rates):
+    assert len(images.size()) == 4
+    batch_size, channel, rows, cols = images.size()
+    out_rows = (rows + strides[0] - 1) // strides[0]
+    out_cols = (cols + strides[1] - 1) // strides[1]
+    effective_k_row = (ksizes[0] - 1) * rates[0] + 1
+    effective_k_col = (ksizes[1] - 1) * rates[1] + 1
+    padding_rows = max(0, (out_rows-1)*strides[0]+effective_k_row-rows)
+    padding_cols = max(0, (out_cols-1)*strides[1]+effective_k_col-cols)
+    # Pad the input
+    padding_top = int(padding_rows / 2.)
+    padding_left = int(padding_cols / 2.)
+    padding_bottom = padding_rows - padding_top
+    padding_right = padding_cols - padding_left
+    paddings = (padding_left, padding_right, padding_top, padding_bottom)
+    images = torch.nn.ZeroPad2d(paddings)(images)
+    return images
+
+def reduce_mean(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.mean(x, dim=i, keepdim=keepdim)
+    return x
+
+
+def reduce_std(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.std(x, dim=i, keepdim=keepdim)
+    return x
+
+
+def reduce_sum(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.sum(x, dim=i, keepdim=keepdim)
+    return x
+
+def random_mask(num_batch=1, mask_shape=(256,256)):
+    list_mask = []
+    for _ in range(num_batch):
+        # rectangle mask
+        image_height = mask_shape[0]
+        image_width = mask_shape[1]
+        max_delta_height = image_height//8
+        max_delta_width = image_width//8
+        height = image_height//4
+        width = image_width//4
+        max_t = image_height - height
+        max_l = image_width - width
+        t = random.randint(0, max_t)
+        l = random.randint(0, max_l)
+        # bbox = (t, l, height, width)
+        h = random.randint(0, max_delta_height//2)
+        w = random.randint(0, max_delta_width//2)
+        mask = torch.zeros((1, 1, image_height, image_width))
+        mask[:, :, t+h:t+height-h, l+w:l+width-w] = 1
+        rect_mask = mask
+
+        # brush mask
+        min_num_vertex = 4
+        max_num_vertex = 12
+        mean_angle = 2 * math.pi / 5
+        angle_range = 2 * math.pi / 15
+        min_width = 12
+        max_width = 40
+        H, W = image_height, image_width
+        average_radius = math.sqrt(H*H+W*W) / 8
+        mask = Image.new('L', (W, H), 0)
+
+        for _ in range(np.random.randint(1, 4)):
+            num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
+            angle_min = mean_angle - np.random.uniform(0, angle_range)
+            angle_max = mean_angle + np.random.uniform(0, angle_range)
+            angles = []
+            vertex = []
+            for i in range(num_vertex):
+                if i % 2 == 0:
+                    angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
+                else:
+                    angles.append(np.random.uniform(angle_min, angle_max))
+
+            h, w = mask.size
+            vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
+            for i in range(num_vertex):
+                r = np.clip(
+                    np.random.normal(loc=average_radius, scale=average_radius//2),
+                    0, 2*average_radius)
+                new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
+                new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
+                vertex.append((int(new_x), int(new_y)))
+
+            draw = ImageDraw.Draw(mask)
+            width = int(np.random.uniform(min_width, max_width))
+            draw.line(vertex, fill=255, width=width)
+            for v in vertex:
+                draw.ellipse((v[0] - width//2,
+                              v[1] - width//2,
+                              v[0] + width//2,
+                              v[1] + width//2),
+                             fill=255)
+
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_LEFT_RIGHT)
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_TOP_BOTTOM)
+
+        mask = transforms.ToTensor()(mask)
+        mask = mask.reshape((1, 1, H, W))
+        brush_mask = mask
+
+        mask = torch.cat([rect_mask, brush_mask], dim=1).max(dim=1, keepdim=True)[0]
+        list_mask.append(mask)
+    mask = torch.cat(list_mask, dim=0)
+    return mask

utils.py

@@ -0,0 +1,1304 @@
+import random
+from typing import Any, Optional
+import numpy as np
+import os
+import cv2
+from glob import glob
+from PIL import Image, ImageDraw
+from tqdm import tqdm
+import kornia
+import matplotlib.pyplot as plt
+import seaborn as sns
+import albumentations as albu
+import functools
+import math
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+import torchvision as tv
+import torchvision.models as models
+from torchvision import transforms
+from torchvision.transforms import functional as F
+from losses import TempCombLoss
+
+
+######## for loading checkpoint from googledrive
+google_drive_paths = {
+    "BayesCap_SRGAN.pth": "https://drive.google.com/uc?id=1d_5j1f8-vN79htZTfRUqP1ddHZIYsNvL",
+    "BayesCap_ckpt.pth": "https://drive.google.com/uc?id=1Vg1r6gKgQ1J3M51n6BeKXYS8auT9NhA9",
+}
+
+def ensure_checkpoint_exists(model_weights_filename):
+    if not os.path.isfile(model_weights_filename) and (
+        model_weights_filename in google_drive_paths
+    ):
+        gdrive_url = google_drive_paths[model_weights_filename]
+        try:
+            from gdown import download as drive_download
+
+            drive_download(gdrive_url, model_weights_filename, quiet=False)
+        except ModuleNotFoundError:
+            print(
+                "gdown module not found.",
+                "pip3 install gdown or, manually download the checkpoint file:",
+                gdrive_url
+            )
+
+    if not os.path.isfile(model_weights_filename) and (
+        model_weights_filename not in google_drive_paths
+    ):
+        print(
+            model_weights_filename,
+            " not found, you may need to manually download the model weights."
+        )
+
+########### DeblurGAN function
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=True)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
+    return norm_layer
+
+def _array_to_batch(x):
+        x = np.transpose(x, (2, 0, 1))
+        x = np.expand_dims(x, 0)
+        return torch.from_numpy(x)
+
+def get_normalize():
+    normalize = albu.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+    normalize = albu.Compose([normalize], additional_targets={'target': 'image'})
+
+    def process(a, b):
+        r = normalize(image=a, target=b)
+        return r['image'], r['target']
+
+    return process
+
+def preprocess(x: np.ndarray, mask: Optional[np.ndarray]):
+    x, _ = get_normalize()(x, x)
+    if mask is None:
+        mask = np.ones_like(x, dtype=np.float32)
+    else:
+        mask = np.round(mask.astype('float32') / 255)
+
+    h, w, _ = x.shape
+    block_size = 32
+    min_height = (h // block_size + 1) * block_size
+    min_width = (w // block_size + 1) * block_size
+
+    pad_params = {'mode': 'constant',
+                    'constant_values': 0,
+                    'pad_width': ((0, min_height - h), (0, min_width - w), (0, 0))
+                    }
+    x = np.pad(x, **pad_params)
+    mask = np.pad(mask, **pad_params)
+
+    return map(_array_to_batch, (x, mask)), h, w
+
+def postprocess(x: torch.Tensor) -> np.ndarray:
+    x, = x
+    x = x.detach().cpu().float().numpy()
+    x = (np.transpose(x, (1, 2, 0)) + 1) / 2.0 * 255.0
+    return x.astype('uint8')
+
+def sorted_glob(pattern):
+    return sorted(glob(pattern))
+###########
+
+def normalize(image: np.ndarray) -> np.ndarray:
+    """Normalize the ``OpenCV.imread`` or ``skimage.io.imread`` data.
+    Args:
+        image (np.ndarray): The image data read by ``OpenCV.imread`` or ``skimage.io.imread``.
+    Returns:
+        Normalized image data. Data range [0, 1].
+    """
+    return image.astype(np.float64) / 255.0
+
+
+def unnormalize(image: np.ndarray) -> np.ndarray:
+    """Un-normalize the ``OpenCV.imread`` or ``skimage.io.imread`` data.
+    Args:
+        image (np.ndarray): The image data read by ``OpenCV.imread`` or ``skimage.io.imread``.
+    Returns:
+        Denormalized image data. Data range [0, 255].
+    """
+    return image.astype(np.float64) * 255.0
+
+
+def image2tensor(image: np.ndarray, range_norm: bool, half: bool) -> torch.Tensor:
+    """Convert ``PIL.Image`` to Tensor.
+    Args:
+        image (np.ndarray): The image data read by ``PIL.Image``
+        range_norm (bool): Scale [0, 1] data to between [-1, 1]
+        half (bool): Whether to convert torch.float32 similarly to torch.half type.
+    Returns:
+        Normalized image data
+    Examples:
+        >>> image = Image.open("image.bmp")
+        >>> tensor_image = image2tensor(image, range_norm=False, half=False)
+    """
+    tensor = F.to_tensor(image)
+
+    if range_norm:
+        tensor = tensor.mul_(2.0).sub_(1.0)
+    if half:
+        tensor = tensor.half()
+
+    return tensor
+
+
+def tensor2image(tensor: torch.Tensor, range_norm: bool, half: bool) -> Any:
+    """Converts ``torch.Tensor`` to ``PIL.Image``.
+    Args:
+        tensor (torch.Tensor): The image that needs to be converted to ``PIL.Image``
+        range_norm (bool): Scale [-1, 1] data to between [0, 1]
+        half (bool): Whether to convert torch.float32 similarly to torch.half type.
+    Returns:
+        Convert image data to support PIL library
+    Examples:
+        >>> tensor = torch.randn([1, 3, 128, 128])
+        >>> image = tensor2image(tensor, range_norm=False, half=False)
+    """
+    if range_norm:
+        tensor = tensor.add_(1.0).div_(2.0)
+    if half:
+        tensor = tensor.half()
+
+    image = tensor.squeeze_(0).permute(1, 2, 0).mul_(255).clamp_(0, 255).cpu().numpy().astype("uint8")
+
+    return image
+
+
+def convert_rgb_to_y(image: Any) -> Any:
+    """Convert RGB image or tensor image data to YCbCr(Y) format.
+    Args:
+        image: RGB image data read by ``PIL.Image''.
+    Returns:
+        Y image array data.
+    """
+    if type(image) == np.ndarray:
+        return 16. + (64.738 * image[:, :, 0] + 129.057 * image[:, :, 1] + 25.064 * image[:, :, 2]) / 256.
+    elif type(image) == torch.Tensor:
+        if len(image.shape) == 4:
+            image = image.squeeze_(0)
+        return 16. + (64.738 * image[0, :, :] + 129.057 * image[1, :, :] + 25.064 * image[2, :, :]) / 256.
+    else:
+        raise Exception("Unknown Type", type(image))
+
+
+def convert_rgb_to_ycbcr(image: Any) -> Any:
+    """Convert RGB image or tensor image data to YCbCr format.
+    Args:
+        image: RGB image data read by ``PIL.Image''.
+    Returns:
+        YCbCr image array data.
+    """
+    if type(image) == np.ndarray:
+        y = 16. + (64.738 * image[:, :, 0] + 129.057 * image[:, :, 1] + 25.064 * image[:, :, 2]) / 256.
+        cb = 128. + (-37.945 * image[:, :, 0] - 74.494 * image[:, :, 1] + 112.439 * image[:, :, 2]) / 256.
+        cr = 128. + (112.439 * image[:, :, 0] - 94.154 * image[:, :, 1] - 18.285 * image[:, :, 2]) / 256.
+        return np.array([y, cb, cr]).transpose([1, 2, 0])
+    elif type(image) == torch.Tensor:
+        if len(image.shape) == 4:
+            image = image.squeeze(0)
+        y = 16. + (64.738 * image[0, :, :] + 129.057 * image[1, :, :] + 25.064 * image[2, :, :]) / 256.
+        cb = 128. + (-37.945 * image[0, :, :] - 74.494 * image[1, :, :] + 112.439 * image[2, :, :]) / 256.
+        cr = 128. + (112.439 * image[0, :, :] - 94.154 * image[1, :, :] - 18.285 * image[2, :, :]) / 256.
+        return torch.cat([y, cb, cr], 0).permute(1, 2, 0)
+    else:
+        raise Exception("Unknown Type", type(image))
+
+
+def convert_ycbcr_to_rgb(image: Any) -> Any:
+    """Convert YCbCr format image to RGB format.
+    Args:
+       image: YCbCr image data read by ``PIL.Image''.
+    Returns:
+        RGB image array data.
+    """
+    if type(image) == np.ndarray:
+        r = 298.082 * image[:, :, 0] / 256. + 408.583 * image[:, :, 2] / 256. - 222.921
+        g = 298.082 * image[:, :, 0] / 256. - 100.291 * image[:, :, 1] / 256. - 208.120 * image[:, :, 2] / 256. + 135.576
+        b = 298.082 * image[:, :, 0] / 256. + 516.412 * image[:, :, 1] / 256. - 276.836
+        return np.array([r, g, b]).transpose([1, 2, 0])
+    elif type(image) == torch.Tensor:
+        if len(image.shape) == 4:
+            image = image.squeeze(0)
+        r = 298.082 * image[0, :, :] / 256. + 408.583 * image[2, :, :] / 256. - 222.921
+        g = 298.082 * image[0, :, :] / 256. - 100.291 * image[1, :, :] / 256. - 208.120 * image[2, :, :] / 256. + 135.576
+        b = 298.082 * image[0, :, :] / 256. + 516.412 * image[1, :, :] / 256. - 276.836
+        return torch.cat([r, g, b], 0).permute(1, 2, 0)
+    else:
+        raise Exception("Unknown Type", type(image))
+
+
+def center_crop(lr: Any, hr: Any, image_size: int, upscale_factor: int) -> [Any, Any]:
+    """Cut ``PIL.Image`` in the center area of the image.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+        image_size (int): The size of the captured image area. It should be the size of the high-resolution image.
+        upscale_factor (int): magnification factor.
+    Returns:
+        Randomly cropped low-resolution images and high-resolution images.
+    """
+    w, h = hr.size
+
+    left = (w - image_size) // 2
+    top = (h - image_size) // 2
+    right = left + image_size
+    bottom = top + image_size
+
+    lr = lr.crop((left // upscale_factor,
+                  top // upscale_factor,
+                  right // upscale_factor,
+                  bottom // upscale_factor))
+    hr = hr.crop((left, top, right, bottom))
+
+    return lr, hr
+
+
+def random_crop(lr: Any, hr: Any, image_size: int, upscale_factor: int) -> [Any, Any]:
+    """Will ``PIL.Image`` randomly capture the specified area of the image.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+        image_size (int): The size of the captured image area. It should be the size of the high-resolution image.
+        upscale_factor (int): magnification factor.
+    Returns:
+        Randomly cropped low-resolution images and high-resolution images.
+    """
+    w, h = hr.size
+    left = torch.randint(0, w - image_size + 1, size=(1,)).item()
+    top = torch.randint(0, h - image_size + 1, size=(1,)).item()
+    right = left + image_size
+    bottom = top + image_size
+
+    lr = lr.crop((left // upscale_factor,
+                  top // upscale_factor,
+                  right // upscale_factor,
+                  bottom // upscale_factor))
+    hr = hr.crop((left, top, right, bottom))
+
+    return lr, hr
+
+
+def random_rotate(lr: Any, hr: Any, angle: int) -> [Any, Any]:
+    """Will ``PIL.Image`` randomly rotate the image.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+        angle (int): rotation angle, clockwise and counterclockwise rotation.
+    Returns:
+        Randomly rotated low-resolution images and high-resolution images.
+    """
+    angle = random.choice((+angle, -angle))
+    lr = F.rotate(lr, angle)
+    hr = F.rotate(hr, angle)
+
+    return lr, hr
+
+
+def random_horizontally_flip(lr: Any, hr: Any, p=0.5) -> [Any, Any]:
+    """Flip the ``PIL.Image`` image horizontally randomly.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+        p (optional, float): rollover probability. (Default: 0.5)
+    Returns:
+        Low-resolution image and high-resolution image after random horizontal flip.
+    """
+    if torch.rand(1).item() > p:
+        lr = F.hflip(lr)
+        hr = F.hflip(hr)
+
+    return lr, hr
+
+
+def random_vertically_flip(lr: Any, hr: Any, p=0.5) -> [Any, Any]:
+    """Turn the ``PIL.Image`` image upside down randomly.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+        p (optional, float): rollover probability. (Default: 0.5)
+    Returns:
+        Randomly rotated up and down low-resolution images and high-resolution images.
+    """
+    if torch.rand(1).item() > p:
+        lr = F.vflip(lr)
+        hr = F.vflip(hr)
+
+    return lr, hr
+
+
+def random_adjust_brightness(lr: Any, hr: Any) -> [Any, Any]:
+    """Set ``PIL.Image`` to randomly adjust the image brightness.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+    Returns:
+        Low-resolution image and high-resolution image with randomly adjusted brightness.
+    """
+    # Randomly adjust the brightness gain range.
+    factor = random.uniform(0.5, 2)
+    lr = F.adjust_brightness(lr, factor)
+    hr = F.adjust_brightness(hr, factor)
+
+    return lr, hr
+
+
+def random_adjust_contrast(lr: Any, hr: Any) -> [Any, Any]:
+    """Set ``PIL.Image`` to randomly adjust the image contrast.
+    Args:
+        lr: Low-resolution image data read by ``PIL.Image``.
+        hr: High-resolution image data read by ``PIL.Image``.
+    Returns:
+        Low-resolution image and high-resolution image with randomly adjusted contrast.
+    """
+    # Randomly adjust the contrast gain range.
+    factor = random.uniform(0.5, 2)
+    lr = F.adjust_contrast(lr, factor)
+    hr = F.adjust_contrast(hr, factor)
+
+    return lr, hr
+
+#### metrics to compute -- assumes single images, i.e., tensor of 3 dims
+def img_mae(x1, x2):
+    m = torch.abs(x1-x2).mean()
+    return m
+
+def img_mse(x1, x2):
+    m = torch.pow(torch.abs(x1-x2),2).mean()
+    return m
+
+def img_psnr(x1, x2):
+    m = kornia.metrics.psnr(x1, x2, 1)
+    return m
+
+def img_ssim(x1, x2):
+    m = kornia.metrics.ssim(x1.unsqueeze(0), x2.unsqueeze(0), 5)
+    m = m.mean()
+    return m
+
+def show_SR_w_uncer(xLR, xHR, xSR, xSRvar, elim=(0,0.01), ulim=(0,0.15)):
+    '''
+    xLR/SR/HR: 3xHxW
+    xSRvar: 1xHxW
+    '''
+    plt.figure(figsize=(30,10))
+    
+    plt.subplot(1,5,1)
+    plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    plt.axis('off')
+    
+    plt.subplot(1,5,2)
+    plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    plt.axis('off')
+    
+    plt.subplot(1,5,3)
+    plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+    plt.axis('off')
+    
+    plt.subplot(1,5,4)
+    error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0) 
+    print('error', error_map.min(), error_map.max())
+    plt.imshow(error_map.transpose(0,2).transpose(0,1), cmap='jet')
+    plt.clim(elim[0], elim[1])
+    plt.axis('off')
+
+    plt.subplot(1,5,5)
+    print('uncer', xSRvar.min(), xSRvar.max())
+    plt.imshow(xSRvar.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+    plt.clim(ulim[0], ulim[1])
+    plt.axis('off')
+
+    plt.subplots_adjust(wspace=0, hspace=0)
+    plt.show()
+
+def show_SR_w_err(xLR, xHR, xSR, elim=(0,0.01), task=None, xMask=None):
+    '''
+    xLR/SR/HR: 3xHxW
+    '''
+    plt.figure(figsize=(30,10))
+    
+    if task != 'm':
+        plt.subplot(1,4,1)
+        plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+        plt.axis('off')
+        
+        plt.subplot(1,4,2)
+        plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+        plt.axis('off')
+        
+        plt.subplot(1,4,3)
+        plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1))
+        plt.axis('off')
+    else:
+        plt.subplot(1,4,1)
+        plt.imshow(xLR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+        plt.clim(0,0.9)
+        plt.axis('off')
+        
+        plt.subplot(1,4,2)
+        plt.imshow(xHR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+        plt.clim(0,0.9)
+        plt.axis('off')
+        
+        plt.subplot(1,4,3)
+        plt.imshow(xSR.to('cpu').data.clip(0,1).transpose(0,2).transpose(0,1), cmap='gray')
+        plt.clim(0,0.9)
+        plt.axis('off')
+    
+    plt.subplot(1,4,4)
+    if task == 'inpainting':
+        error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0)*xMask.to('cpu').data
+    else:
+        error_map = torch.mean(torch.pow(torch.abs(xSR-xHR),2), dim=0).to('cpu').data.unsqueeze(0) 
+    print('error', error_map.min(), error_map.max())
+    plt.imshow(error_map.transpose(0,2).transpose(0,1), cmap='jet')
+    plt.clim(elim[0], elim[1])
+    plt.axis('off')
+
+    plt.subplots_adjust(wspace=0, hspace=0)
+    plt.show()
+
+def show_uncer4(xSRvar1, xSRvar2, xSRvar3, xSRvar4, ulim=(0,0.15)):
+    '''
+    xSRvar: 1xHxW
+    '''
+    plt.figure(figsize=(30,10))
+    
+    plt.subplot(1,4,1)
+    print('uncer', xSRvar1.min(), xSRvar1.max())
+    plt.imshow(xSRvar1.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+    plt.clim(ulim[0], ulim[1])
+    plt.axis('off')
+
+    plt.subplot(1,4,2)
+    print('uncer', xSRvar2.min(), xSRvar2.max())
+    plt.imshow(xSRvar2.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+    plt.clim(ulim[0], ulim[1])
+    plt.axis('off')
+
+    plt.subplot(1,4,3)
+    print('uncer', xSRvar3.min(), xSRvar3.max())
+    plt.imshow(xSRvar3.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+    plt.clim(ulim[0], ulim[1])
+    plt.axis('off')
+
+    plt.subplot(1,4,4)
+    print('uncer', xSRvar4.min(), xSRvar4.max())
+    plt.imshow(xSRvar4.to('cpu').data.transpose(0,2).transpose(0,1), cmap='hot')
+    plt.clim(ulim[0], ulim[1])
+    plt.axis('off')
+
+    plt.subplots_adjust(wspace=0, hspace=0)
+    plt.show()
+
+def get_UCE(list_err, list_yout_var, num_bins=100):
+    err_min = np.min(list_err)
+    err_max = np.max(list_err)
+    err_len = (err_max-err_min)/num_bins
+    num_points = len(list_err)
+    
+    bin_stats = {}
+    for i in range(num_bins):
+        bin_stats[i] = {
+            'start_idx': err_min + i*err_len,
+            'end_idx': err_min + (i+1)*err_len,
+            'num_points': 0,
+            'mean_err': 0,
+            'mean_var': 0,
+        }
+    
+    for e,v in zip(list_err, list_yout_var):
+        for i in range(num_bins):
+            if e>=bin_stats[i]['start_idx'] and e<bin_stats[i]['end_idx']:
+                bin_stats[i]['num_points'] += 1
+                bin_stats[i]['mean_err'] += e
+                bin_stats[i]['mean_var'] += v
+    
+    uce = 0
+    eps = 1e-8
+    for i in range(num_bins):
+        bin_stats[i]['mean_err'] /= bin_stats[i]['num_points'] + eps
+        bin_stats[i]['mean_var'] /= bin_stats[i]['num_points'] + eps
+        bin_stats[i]['uce_bin'] = (bin_stats[i]['num_points']/num_points) \
+            *(np.abs(bin_stats[i]['mean_err'] - bin_stats[i]['mean_var']))
+        uce += bin_stats[i]['uce_bin']
+    
+    list_x, list_y = [], []
+    for i in range(num_bins):
+        if bin_stats[i]['num_points']>0:
+            list_x.append(bin_stats[i]['mean_err'])
+            list_y.append(bin_stats[i]['mean_var'])
+    
+    # sns.set_style('darkgrid')
+    # sns.scatterplot(x=list_x, y=list_y)
+    # sns.regplot(x=list_x, y=list_y, order=1)
+    # plt.xlabel('MSE', fontsize=34)
+    # plt.ylabel('Uncertainty', fontsize=34)
+    # plt.plot(list_x, list_x, color='r')
+    # plt.xlim(np.min(list_x), np.max(list_x))
+    # plt.ylim(np.min(list_err), np.max(list_x))
+    # plt.show()
+
+    return bin_stats, uce
+
+##################### training BayesCap
+def train_BayesCap(
+    NetC,
+    NetG,
+    train_loader,
+    eval_loader,
+    Cri = TempCombLoss(),
+    device='cuda',
+    dtype=torch.cuda.FloatTensor(),
+    init_lr=1e-4,
+    num_epochs=100,
+    eval_every=1,
+    ckpt_path='../ckpt/BayesCap',
+    T1=1e0,
+    T2=5e-2,
+    task=None,
+):
+    NetC.to(device)
+    NetC.train()
+    NetG.to(device)
+    NetG.eval()
+    optimizer = torch.optim.Adam(list(NetC.parameters()), lr=init_lr)
+    optim_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs)
+
+    score = -1e8
+    all_loss = []
+    for eph in range(num_epochs):
+        eph_loss = 0
+        with tqdm(train_loader, unit='batch') as tepoch:
+            for (idx, batch) in enumerate(tepoch):
+                if idx>2000:
+                    break
+                tepoch.set_description('Epoch {}'.format(eph))
+                ##
+                xLR, xHR = batch[0].to(device), batch[1].to(device)
+                xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+                if task == 'inpainting':
+                    xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+                    xMask = xMask.to(device).type(dtype)
+                # pass them through the network
+                with torch.no_grad():
+                    if task == 'inpainting':
+                        _, xSR1 = NetG(xLR, xMask)
+                    elif task == 'depth':
+                        xSR1 = NetG(xLR)[("disp", 0)]
+                    else:
+                        xSR1 = NetG(xLR)
+                # with torch.autograd.set_detect_anomaly(True):
+                xSR = xSR1.clone()
+                xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+                # print(xSRC_alpha) 
+                optimizer.zero_grad()
+                if task == 'depth':
+                    loss = Cri(xSRC_mu, xSRC_alpha, xSRC_beta, xSR, T1=T1, T2=T2)
+                else:
+                    loss = Cri(xSRC_mu, xSRC_alpha, xSRC_beta, xHR, T1=T1, T2=T2)
+                # print(loss)
+                loss.backward()
+                optimizer.step()
+                ##
+                eph_loss += loss.item()
+                tepoch.set_postfix(loss=loss.item())
+            eph_loss /= len(train_loader)
+            all_loss.append(eph_loss)
+            print('Avg. loss: {}'.format(eph_loss))
+        # evaluate and save the models
+        torch.save(NetC.state_dict(), ckpt_path+'_last.pth')
+        if eph%eval_every == 0:
+            curr_score = eval_BayesCap(
+                NetC,
+                NetG,
+                eval_loader,
+                device=device,
+                dtype=dtype,
+                task=task,
+            )
+            print('current score: {} | Last best score: {}'.format(curr_score, score))
+            if curr_score >= score:
+                score = curr_score
+                torch.save(NetC.state_dict(), ckpt_path+'_best.pth')
+    optim_scheduler.step()
+
+#### get different uncertainty maps
+def get_uncer_BayesCap(
+    NetC,
+    NetG,
+    xin,
+    task=None,
+    xMask=None,
+):
+    with torch.no_grad():
+        if task == 'inpainting':
+            _, xSR = NetG(xin, xMask)
+        else:
+            xSR = NetG(xin)
+        xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+    a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+    b_map = xSRC_beta.to('cpu').data
+    xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+
+    return xSRvar
+
+def get_uncer_TTDAp(
+    NetG,
+    xin,
+    p_mag=0.05,
+    num_runs=50,
+    task=None,
+    xMask=None,
+):
+    list_xSR = []
+    with torch.no_grad():
+        for z in range(num_runs):
+            if task == 'inpainting':
+                _, xSRz = NetG(xin+p_mag*xin.max()*torch.randn_like(xin), xMask)
+            else:
+                xSRz = NetG(xin+p_mag*xin.max()*torch.randn_like(xin))
+            list_xSR.append(xSRz)
+    xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+    xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+    return xSRvar
+
+def get_uncer_DO(
+    NetG,
+    xin,
+    dop=0.2,
+    num_runs=50,
+    task=None,
+    xMask=None,
+):
+    list_xSR = []
+    with torch.no_grad():
+        for z in range(num_runs):
+            if task == 'inpainting':
+                _, xSRz = NetG(xin, xMask, dop=dop)
+            else:
+                xSRz = NetG(xin, dop=dop)
+            list_xSR.append(xSRz)
+    xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+    xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+    return xSRvar
+
+################### Different eval functions
+
+def eval_BayesCap(
+    NetC,
+    NetG,
+    eval_loader,
+    device='cuda',
+    dtype=torch.cuda.FloatTensor,
+    task=None,
+    xMask=None,
+):
+    NetC.to(device)
+    NetC.eval()
+    NetG.to(device)
+    NetG.eval()
+    
+    mean_ssim = 0
+    mean_psnr = 0
+    mean_mse = 0
+    mean_mae = 0
+    num_imgs = 0
+    list_error = []
+    list_var = []
+    with tqdm(eval_loader, unit='batch') as tepoch:
+        for (idx, batch) in enumerate(tepoch):
+            tepoch.set_description('Validating ...')
+            ##
+            xLR, xHR = batch[0].to(device), batch[1].to(device)
+            xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+            if task == 'inpainting':
+                if xMask==None:
+                    xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+                    xMask = xMask.to(device).type(dtype)
+                else:
+                    xMask = xMask.to(device).type(dtype)
+            # pass them through the network
+            with torch.no_grad():
+                if task == 'inpainting':
+                    _, xSR = NetG(xLR, xMask)
+                elif task == 'depth':
+                    xSR = NetG(xLR)[("disp", 0)]
+                else:
+                    xSR = NetG(xLR)
+                xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+            a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+            b_map = xSRC_beta.to('cpu').data
+            xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+            n_batch = xSRC_mu.shape[0]
+            if task == 'depth':
+                xHR = xSR
+            for j in range(n_batch):
+                num_imgs += 1
+                mean_ssim += img_ssim(xSRC_mu[j], xHR[j])
+                mean_psnr += img_psnr(xSRC_mu[j], xHR[j])
+                mean_mse += img_mse(xSRC_mu[j], xHR[j])
+                mean_mae += img_mae(xSRC_mu[j], xHR[j])
+
+                show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+                
+                error_map = torch.mean(torch.pow(torch.abs(xSR[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+                var_map =  xSRvar[j].to('cpu').data.reshape(-1)
+                list_error.extend(list(error_map.numpy()))
+                list_var.extend(list(var_map.numpy()))
+            ##
+        mean_ssim /= num_imgs
+        mean_psnr /= num_imgs
+        mean_mse /= num_imgs
+        mean_mae /= num_imgs
+        print(
+            'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+            (
+                mean_ssim, mean_psnr, mean_mse, mean_mae 
+            )
+        )
+        # print(len(list_error), len(list_var))
+        # print('UCE: ', get_UCE(list_error[::10], list_var[::10], num_bins=500)[1])
+        # print('C.Coeff: ', np.corrcoef(np.array(list_error[::10]), np.array(list_var[::10])))
+    return mean_ssim
+
+def eval_TTDA_p(
+    NetG,
+    eval_loader,
+    device='cuda',
+    dtype=torch.cuda.FloatTensor,
+    p_mag=0.05,
+    num_runs=50,
+    task = None,
+    xMask = None,
+):
+    NetG.to(device)
+    NetG.eval()
+    
+    mean_ssim = 0
+    mean_psnr = 0
+    mean_mse = 0
+    mean_mae = 0
+    num_imgs = 0
+    with tqdm(eval_loader, unit='batch') as tepoch:
+        for (idx, batch) in enumerate(tepoch):
+            tepoch.set_description('Validating ...')
+            ##
+            xLR, xHR = batch[0].to(device), batch[1].to(device)
+            xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+            # pass them through the network
+            list_xSR = []
+            with torch.no_grad():
+                if task=='inpainting':
+                    _, xSR = NetG(xLR, xMask)
+                else:
+                    xSR = NetG(xLR)
+                for z in range(num_runs):
+                    xSRz = NetG(xLR+p_mag*xLR.max()*torch.randn_like(xLR))	
+                    list_xSR.append(xSRz)
+            xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+            xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+            n_batch = xSR.shape[0]
+            for j in range(n_batch):
+                num_imgs += 1
+                mean_ssim += img_ssim(xSR[j], xHR[j])
+                mean_psnr += img_psnr(xSR[j], xHR[j])
+                mean_mse += img_mse(xSR[j], xHR[j])
+                mean_mae += img_mae(xSR[j], xHR[j])
+
+                show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+                
+        mean_ssim /= num_imgs
+        mean_psnr /= num_imgs
+        mean_mse /= num_imgs
+        mean_mae /= num_imgs
+        print(
+            'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+            (
+                mean_ssim, mean_psnr, mean_mse, mean_mae 
+            )
+        )
+        
+    return mean_ssim
+
+def eval_DO(
+    NetG,
+    eval_loader,
+    device='cuda',
+    dtype=torch.cuda.FloatTensor,
+    dop=0.2,
+    num_runs=50,
+    task=None,
+    xMask=None,
+):
+    NetG.to(device)
+    NetG.eval()
+    
+    mean_ssim = 0
+    mean_psnr = 0
+    mean_mse = 0
+    mean_mae = 0
+    num_imgs = 0
+    with tqdm(eval_loader, unit='batch') as tepoch:
+        for (idx, batch) in enumerate(tepoch):
+            tepoch.set_description('Validating ...')
+            ##
+            xLR, xHR = batch[0].to(device), batch[1].to(device)
+            xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+            # pass them through the network
+            list_xSR = []
+            with torch.no_grad():
+                if task == 'inpainting':
+                    _, xSR = NetG(xLR, xMask)
+                else:
+                    xSR = NetG(xLR)
+                for z in range(num_runs):
+                    xSRz = NetG(xLR, dop=dop)	
+                    list_xSR.append(xSRz)
+            xSRmean = torch.mean(torch.cat(list_xSR, dim=0), dim=0).unsqueeze(0)
+            xSRvar = torch.mean(torch.var(torch.cat(list_xSR, dim=0), dim=0), dim=0).unsqueeze(0).unsqueeze(1)
+            n_batch = xSR.shape[0]
+            for j in range(n_batch):
+                num_imgs += 1
+                mean_ssim += img_ssim(xSR[j], xHR[j])
+                mean_psnr += img_psnr(xSR[j], xHR[j])
+                mean_mse += img_mse(xSR[j], xHR[j])
+                mean_mae += img_mae(xSR[j], xHR[j])
+
+                show_SR_w_uncer(xLR[j], xHR[j], xSR[j], xSRvar[j])
+                        ##
+        mean_ssim /= num_imgs
+        mean_psnr /= num_imgs
+        mean_mse /= num_imgs
+        mean_mae /= num_imgs
+        print(
+            'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+            (
+                mean_ssim, mean_psnr, mean_mse, mean_mae 
+            )
+        )
+    
+    return mean_ssim
+
+
+############### compare all function
+def compare_all(
+    NetC,
+    NetG,
+    eval_loader,
+    p_mag = 0.05,
+    dop = 0.2,
+    num_runs = 100,
+    device='cuda',
+    dtype=torch.cuda.FloatTensor,
+    task=None,
+):
+    NetC.to(device)
+    NetC.eval()
+    NetG.to(device)
+    NetG.eval()
+    
+    with tqdm(eval_loader, unit='batch') as tepoch:
+        for (idx, batch) in enumerate(tepoch):
+            tepoch.set_description('Comparing ...')
+            ##
+            xLR, xHR = batch[0].to(device), batch[1].to(device)
+            xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+            if task == 'inpainting':
+                xMask = random_mask(xLR.shape[0], (xLR.shape[2], xLR.shape[3]))
+                xMask = xMask.to(device).type(dtype)
+            # pass them through the network
+            with torch.no_grad():
+                if task == 'inpainting':
+                    _, xSR = NetG(xLR, xMask)
+                else:
+                    xSR = NetG(xLR)
+                xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR)
+
+            if task == 'inpainting':
+                xSRvar1 = get_uncer_TTDAp(NetG, xLR, p_mag=p_mag, num_runs=num_runs, task='inpainting', xMask=xMask)
+                xSRvar2 = get_uncer_DO(NetG, xLR, dop=dop, num_runs=num_runs, task='inpainting', xMask=xMask)
+                xSRvar3 = get_uncer_BayesCap(NetC, NetG, xLR, task='inpainting', xMask=xMask)
+            else:
+                xSRvar1 = get_uncer_TTDAp(NetG, xLR, p_mag=p_mag, num_runs=num_runs)
+                xSRvar2 = get_uncer_DO(NetG, xLR, dop=dop, num_runs=num_runs)
+                xSRvar3 = get_uncer_BayesCap(NetC, NetG, xLR)
+
+            print('bdg', xSRvar1.shape, xSRvar2.shape, xSRvar3.shape)
+
+            n_batch = xSR.shape[0]
+            for j in range(n_batch):
+                if task=='s':
+                    show_SR_w_err(xLR[j], xHR[j], xSR[j])
+                    show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42))
+                    show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 1.5), xSRvar3[j])
+                if task=='d':
+                    show_SR_w_err(xLR[j], xHR[j], 0.5*xSR[j]+0.5*xHR[j])
+                    show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42))
+                    show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 0.8), xSRvar3[j])
+                if task=='inpainting':
+                    show_SR_w_err(xLR[j]*(1-xMask[j]), xHR[j], xSR[j], elim=(0,0.25), task='inpainting', xMask=xMask[j])
+                    show_uncer4(xSRvar1[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.45), torch.pow(xSRvar1[j], 0.4))
+                    show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 0.8), xSRvar3[j])
+                if task=='m':
+                    show_SR_w_err(xLR[j], xHR[j], xSR[j], elim=(0,0.04), task='m')
+                    show_uncer4(0.4*xSRvar1[j]+0.6*xSRvar2[j], torch.sqrt(xSRvar1[j]), torch.pow(xSRvar1[j], 0.48), torch.pow(xSRvar1[j], 0.42), ulim=(0.02,0.15))
+                    show_uncer4(xSRvar2[j], torch.sqrt(xSRvar2[j]), torch.pow(xSRvar3[j], 1.5), xSRvar3[j], ulim=(0.02,0.15))
+
+
+################# Degrading Identity 
+def degrage_BayesCap_p(
+    NetC,
+    NetG,
+    eval_loader,
+    device='cuda',
+    dtype=torch.cuda.FloatTensor,
+    num_runs=50,
+):
+    NetC.to(device)
+    NetC.eval()
+    NetG.to(device)
+    NetG.eval()
+    
+    p_mag_list = [0, 0.05, 0.1, 0.15, 0.2]
+    list_s = []
+    list_p = []
+    list_u1 = []
+    list_u2 = []
+    list_c = []
+    for p_mag in p_mag_list:
+        mean_ssim = 0
+        mean_psnr = 0
+        mean_mse = 0
+        mean_mae = 0
+        num_imgs = 0
+        list_error = []
+        list_error2 = []
+        list_var = []
+
+        with tqdm(eval_loader, unit='batch') as tepoch:
+            for (idx, batch) in enumerate(tepoch):
+                tepoch.set_description('Validating ...')
+                ##
+                xLR, xHR = batch[0].to(device), batch[1].to(device)
+                xLR, xHR = xLR.type(dtype), xHR.type(dtype)
+                # pass them through the network
+                with torch.no_grad():
+                    xSR = NetG(xLR)
+                    xSRC_mu, xSRC_alpha, xSRC_beta = NetC(xSR + p_mag*xSR.max()*torch.randn_like(xSR))
+                a_map = (1/(xSRC_alpha + 1e-5)).to('cpu').data
+                b_map = xSRC_beta.to('cpu').data
+                xSRvar = (a_map**2)*(torch.exp(torch.lgamma(3/(b_map + 1e-2)))/torch.exp(torch.lgamma(1/(b_map + 1e-2))))
+                n_batch = xSRC_mu.shape[0]
+                for j in range(n_batch):
+                    num_imgs += 1
+                    mean_ssim += img_ssim(xSRC_mu[j], xSR[j])
+                    mean_psnr += img_psnr(xSRC_mu[j], xSR[j])
+                    mean_mse += img_mse(xSRC_mu[j], xSR[j])
+                    mean_mae += img_mae(xSRC_mu[j], xSR[j])
+
+                    error_map = torch.mean(torch.pow(torch.abs(xSR[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+                    error_map2 = torch.mean(torch.pow(torch.abs(xSRC_mu[j]-xHR[j]),2), dim=0).to('cpu').data.reshape(-1)
+                    var_map =  xSRvar[j].to('cpu').data.reshape(-1)
+                    list_error.extend(list(error_map.numpy()))
+                    list_error2.extend(list(error_map2.numpy()))
+                    list_var.extend(list(var_map.numpy()))
+                ##
+            mean_ssim /= num_imgs
+            mean_psnr /= num_imgs
+            mean_mse /= num_imgs
+            mean_mae /= num_imgs
+            print(
+                'Avg. SSIM: {} | Avg. PSNR: {} | Avg. MSE: {} | Avg. MAE: {}'.format
+                (
+                    mean_ssim, mean_psnr, mean_mse, mean_mae 
+                )
+            )
+            uce1 = get_UCE(list_error[::100], list_var[::100], num_bins=200)[1]
+            uce2 = get_UCE(list_error2[::100], list_var[::100], num_bins=200)[1]
+            print('UCE1: ', uce1)
+            print('UCE2: ', uce2)
+            list_s.append(mean_ssim.item())
+            list_p.append(mean_psnr.item())
+            list_u1.append(uce1)
+            list_u2.append(uce2)
+    
+    plt.plot(list_s)
+    plt.show()
+    plt.plot(list_p)
+    plt.show()
+
+    plt.plot(list_u1, label='wrt SR output')
+    plt.plot(list_u2, label='wrt BayesCap output')
+    plt.legend()
+    plt.show()
+
+    sns.set_style('darkgrid')
+    fig,ax = plt.subplots()
+    # make a plot
+    ax.plot(p_mag_list, list_s, color="red", marker="o")
+    # set x-axis label
+    ax.set_xlabel("Reducing faithfulness of BayesCap Reconstruction",fontsize=10)
+    # set y-axis label
+    ax.set_ylabel("SSIM btwn BayesCap and SRGAN outputs", color="red",fontsize=10)
+
+    # twin object for two different y-axis on the sample plot
+    ax2=ax.twinx()
+    # make a plot with different y-axis using second axis object
+    ax2.plot(p_mag_list, list_u1, color="blue", marker="o", label='UCE wrt to error btwn SRGAN output and GT')
+    ax2.plot(p_mag_list, list_u2, color="orange", marker="o", label='UCE wrt to error btwn BayesCap output and GT')
+    ax2.set_ylabel("UCE", color="green", fontsize=10)
+    plt.legend(fontsize=10)
+    plt.tight_layout()
+    plt.show()
+
+################# DeepFill_v2
+   
+# ----------------------------------------
+#             PATH processing
+# ----------------------------------------
+def text_readlines(filename):
+    # Try to read a txt file and return a list.Return [] if there was a mistake.
+    try:
+        file = open(filename, 'r')
+    except IOError:
+        error = []
+        return error
+    content = file.readlines()
+    # This for loop deletes the EOF (like \n)
+    for i in range(len(content)):
+        content[i] = content[i][:len(content[i])-1]
+    file.close()
+    return content
+
+def savetxt(name, loss_log):
+    np_loss_log = np.array(loss_log)
+    np.savetxt(name, np_loss_log)
+
+def get_files(path):
+    # read a folder, return the complete path
+    ret = []
+    for root, dirs, files in os.walk(path):
+        for filespath in files:
+            ret.append(os.path.join(root, filespath))
+    return ret
+
+def get_names(path):
+    # read a folder, return the image name
+    ret = []
+    for root, dirs, files in os.walk(path):
+        for filespath in files:
+            ret.append(filespath)
+    return ret
+
+def text_save(content, filename, mode = 'a'):
+    # save a list to a txt
+    # Try to save a list variable in txt file.
+    file = open(filename, mode)
+    for i in range(len(content)):
+        file.write(str(content[i]) + '\n')
+    file.close()
+
+def check_path(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+    
+# ----------------------------------------
+#    Validation and Sample at training
+# ----------------------------------------
+def save_sample_png(sample_folder, sample_name, img_list, name_list, pixel_max_cnt = 255):
+    # Save image one-by-one
+    for i in range(len(img_list)):
+        img = img_list[i]
+        # Recover normalization: * 255 because last layer is sigmoid activated
+        img = img * 255
+        # Process img_copy and do not destroy the data of img
+        img_copy = img.clone().data.permute(0, 2, 3, 1)[0, :, :, :].cpu().numpy()
+        img_copy = np.clip(img_copy, 0, pixel_max_cnt)
+        img_copy = img_copy.astype(np.uint8)
+        img_copy = cv2.cvtColor(img_copy, cv2.COLOR_RGB2BGR)
+        # Save to certain path
+        save_img_name = sample_name + '_' + name_list[i] + '.jpg'
+        save_img_path = os.path.join(sample_folder, save_img_name)
+        cv2.imwrite(save_img_path, img_copy)
+
+def psnr(pred, target, pixel_max_cnt = 255):
+    mse = torch.mul(target - pred, target - pred)
+    rmse_avg = (torch.mean(mse).item()) ** 0.5
+    p = 20 * np.log10(pixel_max_cnt / rmse_avg)
+    return p
+
+def grey_psnr(pred, target, pixel_max_cnt = 255):
+    pred = torch.sum(pred, dim = 0)
+    target = torch.sum(target, dim = 0)
+    mse = torch.mul(target - pred, target - pred)
+    rmse_avg = (torch.mean(mse).item()) ** 0.5
+    p = 20 * np.log10(pixel_max_cnt * 3 / rmse_avg)
+    return p
+
+def ssim(pred, target):
+    pred = pred.clone().data.permute(0, 2, 3, 1).cpu().numpy()
+    target = target.clone().data.permute(0, 2, 3, 1).cpu().numpy()
+    target = target[0]
+    pred = pred[0]
+    ssim = skimage.measure.compare_ssim(target, pred, multichannel = True)
+    return ssim
+
+## for contextual attention
+
+def extract_image_patches(images, ksizes, strides, rates, padding='same'):
+    """
+    Extract patches from images and put them in the C output dimension.
+    :param padding:
+    :param images: [batch, channels, in_rows, in_cols]. A 4-D Tensor with shape
+    :param ksizes: [ksize_rows, ksize_cols]. The size of the sliding window for
+     each dimension of images
+    :param strides: [stride_rows, stride_cols]
+    :param rates: [dilation_rows, dilation_cols]
+    :return: A Tensor
+    """
+    assert len(images.size()) == 4
+    assert padding in ['same', 'valid']
+    batch_size, channel, height, width = images.size()
+
+    if padding == 'same':
+        images = same_padding(images, ksizes, strides, rates)
+    elif padding == 'valid':
+        pass
+    else:
+        raise NotImplementedError('Unsupported padding type: {}.\
+                Only "same" or "valid" are supported.'.format(padding))
+
+    unfold = torch.nn.Unfold(kernel_size=ksizes,
+                             dilation=rates,
+                             padding=0,
+                             stride=strides)
+    patches = unfold(images)
+    return patches  # [N, C*k*k, L], L is the total number of such blocks
+
+def same_padding(images, ksizes, strides, rates):
+    assert len(images.size()) == 4
+    batch_size, channel, rows, cols = images.size()
+    out_rows = (rows + strides[0] - 1) // strides[0]
+    out_cols = (cols + strides[1] - 1) // strides[1]
+    effective_k_row = (ksizes[0] - 1) * rates[0] + 1
+    effective_k_col = (ksizes[1] - 1) * rates[1] + 1
+    padding_rows = max(0, (out_rows-1)*strides[0]+effective_k_row-rows)
+    padding_cols = max(0, (out_cols-1)*strides[1]+effective_k_col-cols)
+    # Pad the input
+    padding_top = int(padding_rows / 2.)
+    padding_left = int(padding_cols / 2.)
+    padding_bottom = padding_rows - padding_top
+    padding_right = padding_cols - padding_left
+    paddings = (padding_left, padding_right, padding_top, padding_bottom)
+    images = torch.nn.ZeroPad2d(paddings)(images)
+    return images
+
+def reduce_mean(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.mean(x, dim=i, keepdim=keepdim)
+    return x
+
+
+def reduce_std(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.std(x, dim=i, keepdim=keepdim)
+    return x
+
+
+def reduce_sum(x, axis=None, keepdim=False):
+    if not axis:
+        axis = range(len(x.shape))
+    for i in sorted(axis, reverse=True):
+        x = torch.sum(x, dim=i, keepdim=keepdim)
+    return x
+
+def random_mask(num_batch=1, mask_shape=(256,256)):
+    list_mask = []
+    for _ in range(num_batch):
+        # rectangle mask
+        image_height = mask_shape[0]
+        image_width = mask_shape[1]
+        max_delta_height = image_height//8
+        max_delta_width = image_width//8
+        height = image_height//4
+        width = image_width//4
+        max_t = image_height - height
+        max_l = image_width - width
+        t = random.randint(0, max_t)
+        l = random.randint(0, max_l)
+        # bbox = (t, l, height, width)
+        h = random.randint(0, max_delta_height//2)
+        w = random.randint(0, max_delta_width//2)
+        mask = torch.zeros((1, 1, image_height, image_width))
+        mask[:, :, t+h:t+height-h, l+w:l+width-w] = 1
+        rect_mask = mask
+
+        # brush mask
+        min_num_vertex = 4
+        max_num_vertex = 12
+        mean_angle = 2 * math.pi / 5
+        angle_range = 2 * math.pi / 15
+        min_width = 12
+        max_width = 40
+        H, W = image_height, image_width
+        average_radius = math.sqrt(H*H+W*W) / 8
+        mask = Image.new('L', (W, H), 0)
+
+        for _ in range(np.random.randint(1, 4)):
+            num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
+            angle_min = mean_angle - np.random.uniform(0, angle_range)
+            angle_max = mean_angle + np.random.uniform(0, angle_range)
+            angles = []
+            vertex = []
+            for i in range(num_vertex):
+                if i % 2 == 0:
+                    angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
+                else:
+                    angles.append(np.random.uniform(angle_min, angle_max))
+
+            h, w = mask.size
+            vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
+            for i in range(num_vertex):
+                r = np.clip(
+                    np.random.normal(loc=average_radius, scale=average_radius//2),
+                    0, 2*average_radius)
+                new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
+                new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
+                vertex.append((int(new_x), int(new_y)))
+
+            draw = ImageDraw.Draw(mask)
+            width = int(np.random.uniform(min_width, max_width))
+            draw.line(vertex, fill=255, width=width)
+            for v in vertex:
+                draw.ellipse((v[0] - width//2,
+                              v[1] - width//2,
+                              v[0] + width//2,
+                              v[1] + width//2),
+                             fill=255)
+
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_LEFT_RIGHT)
+        if np.random.normal() > 0:
+            mask.transpose(Image.FLIP_TOP_BOTTOM)
+
+        mask = transforms.ToTensor()(mask)
+        mask = mask.reshape((1, 1, H, W))
+        brush_mask = mask
+
+        mask = torch.cat([rect_mask, brush_mask], dim=1).max(dim=1, keepdim=True)[0]
+        list_mask.append(mask)
+    mask = torch.cat(list_mask, dim=0)
+    return mask