BayesCap demo to EuroCrypt

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,26 @@
 ---
 title: BayesCap
-emoji: 👁
-colorFrom: gray
-colorTo: pink
+emoji: 🔥
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 3.0.24
 app_file: app.py
 pinned: false
-license: cc
 ---
+# Configuration
+`title`: _string_  
+Display title for the Space
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,152 @@
+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 = 'cpu'
+if device == 'cuda':
+    dtype = torch.cuda.FloatTensor
+else:
+    dtype = torch.FloatTensor
+
+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 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 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)
+    # to retain aspect ratio
+    lr_transforms = transforms.Resize(image_size[0]//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 = image2tensor(img, range_norm=False, half=False)
+    
+    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: Bayesian Identity Cap for Calibrated Uncertainty in Frozen Neural Networks"
+
+abstract="<b>Abstract.</b> High-quality calibrated uncertainty estimates are crucial for numerous real-world applications, especially for deep learning-based deployed ML systems. While Bayesian deep learning techniques allow uncertainty estimation, training them with large-scale datasets is an expensive process that does not always yield models competitive with non-Bayesian counterparts. Moreover, many of the high-performing deep learning models that are already trained and deployed are non-Bayesian in nature and do not provide uncertainty estimates. To address these issues, we propose BayesCap that learns a Bayesian identity mapping for the frozen model, allowing uncertainty estimation. BayesCap is a memory-efficient method that can be trained on a small fraction of the original dataset, enhancing pretrained non-Bayesian computer vision models by providing calibrated uncertainty estimates for the predictions without (i) hampering the performance of the model and (ii) the need for expensive retraining the model from scratch. The proposed method is agnostic to various architectures and tasks. We show the efficacy of our method on a wide variety of tasks with a diverse set of architectures, including image super-resolution, deblurring, inpainting, and crucial application such as medical image translation. Moreover, we apply the derived uncertainty estimates to detect out-of-distribution samples in critical scenarios like depth estimation in autonomous driving. Code is available <a href='https://github.com/ExplainableML/BayesCap'>here</a>. <br> <br>"
+
+method = "In this demo, we show an application of BayesCap on top of SRGAN for the task of super resolution. BayesCap estimates the per-pixel uncertainty of a pretrained computer vision model like SRGAN (used for super-resolution). BayesCap takes the ouput of the pretrained model (in this case SRGAN), and predicts the per-pixel distribution parameters for the output, that can be used to quantify the per-pixel uncertainty. In our work, we model the per-pixel output as a <a href='https://en.wikipedia.org/wiki/Generalized_normal_distribution'>Generalized Gaussian distribution</a> that is parameterized by 3 parameters the mean, scale (alpha), and the shape (beta). As a result our model predicts these three parameters as shown below. From these 3 parameters one can compute the uncertainty as shown in <a href='https://en.wikipedia.org/wiki/Generalized_normal_distribution'>this article</a>. <br><br>"
+
+closing = "For more details, please find the <a href='https://arxiv.org/'>ECCV 2022 paper here</a>."
+
+description =  abstract + method + closing
+
+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-resolution image (input to SRGAN)"), 
+        gr.outputs.Image(type='pil', label="Super-resolved image (output of SRGAN)"), 
+        gr.outputs.Image(type='pil', label="Alpha parameter map characterizing per-pixel distribution (output of BayesCap)"), 
+        gr.outputs.Image(type='pil', label="Beta parameter map characterizing per-pixel distribution (output of BayesCap)"), 
+        gr.outputs.Image(type='pil', label="Per-pixel uncertainty map (derived using outputs of BayesCap)")
+     ],
+    title=title,
+    description=description,
+    article=article,
+     examples=[
+        ["./demo_examples/tue.jpeg"],
+        ["./demo_examples/baby.png"],
+        ["./demo_examples/bird.png"],
+        ["./demo_examples/butterfly.png"],
+        ["./demo_examples/head.png"],
+        ["./demo_examples/woman.png"],
+    ]
+).launch()

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,158 @@
+aiohttp==3.8.1
+aiosignal==1.2.0
+albumentations
+analytics-python==1.4.0
+anyio==3.6.1
+argon2-cffi==21.3.0
+argon2-cffi-bindings==21.2.0
+asttokens==2.0.5
+async-timeout==4.0.2
+attrs==21.4.0
+Babel==2.10.1
+backcall==0.2.0
+backoff==1.10.0
+bcrypt==3.2.2
+beautifulsoup4==4.11.1
+bleach==5.0.0
+brotlipy==0.7.0
+certifi
+cffi
+charset-normalizer
+click==8.1.3
+cloudpickle
+cryptography
+cycler==0.11.0
+cytoolz==0.11.2
+dask
+debugpy==1.6.0
+decorator==5.1.1
+defusedxml==0.7.1
+entrypoints==0.4
+executing==0.8.3
+fastapi==0.78.0
+fastjsonschema==2.15.3
+ffmpy==0.3.0
+filelock==3.7.1
+fire==0.4.0
+fonttools==4.33.3
+frozenlist==1.3.0
+fsspec
+ftfy==6.1.1
+gdown==4.5.1
+gradio==3.0.24
+h11==0.12.0
+httpcore==0.15.0
+httpx==0.23.0
+idna
+imagecodecs
+imageio
+ipykernel==6.13.0
+ipython==8.4.0
+ipython-genutils==0.2.0
+jedi==0.18.1
+Jinja2==3.1.2
+joblib
+json5==0.9.8
+jsonschema==4.6.0
+jupyter-client==7.3.1
+jupyter-core==4.10.0
+jupyter-server==1.17.0
+jupyterlab==3.4.2
+jupyterlab-pygments==0.2.2
+jupyterlab-server==2.14.0
+kiwisolver==1.4.2
+kornia==0.6.5
+linkify-it-py==1.0.3
+locket
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.5.2
+matplotlib-inline==0.1.3
+mdit-py-plugins==0.3.0
+mdurl==0.1.1
+mistune==0.8.4
+mkl-fft==1.3.1
+mkl-random
+mkl-service==2.4.0
+mltk==0.0.5
+monotonic==1.6
+multidict==6.0.2
+munch==2.5.0
+nbclassic==0.3.7
+nbclient==0.6.4
+nbconvert==6.5.0
+nbformat==5.4.0
+nest-asyncio==1.5.5
+networkx
+nltk==3.7
+notebook==6.4.11
+notebook-shim==0.1.0
+ntk==1.1.3
+numpy
+opencv-python==4.6.0.66
+orjson==3.7.7
+packaging
+pandas==1.4.2
+pandocfilters==1.5.0
+paramiko==2.11.0
+parso==0.8.3
+partd
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==9.0.1
+prometheus-client==0.14.1
+prompt-toolkit==3.0.29
+psutil==5.9.1
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pycocotools==2.0.4
+pycparser
+pycryptodome==3.15.0
+pydantic==1.9.1
+pydub==0.25.1
+Pygments==2.12.0
+PyNaCl==1.5.0
+pyOpenSSL
+pyparsing
+pyrsistent==0.18.1
+PySocks
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytz==2022.1
+PyWavelets
+PyYAML
+pyzmq==23.1.0
+qudida
+regex==2022.6.2
+requests
+rfc3986==1.5.0
+scikit-image
+scikit-learn
+scipy
+seaborn==0.11.2
+Send2Trash==1.8.0
+six
+sniffio==1.2.0
+soupsieve==2.3.2.post1
+stack-data==0.2.0
+starlette==0.19.1
+termcolor==1.1.0
+terminado==0.15.0
+threadpoolctl
+tifffile
+tinycss2==1.1.1
+toolz
+torch==1.11.0
+torchaudio==0.11.0
+torchvision==0.12.0
+tornado==6.1
+tqdm==4.64.0
+traitlets==5.2.2.post1
+typing_extensions
+uc-micro-py==1.0.1
+urllib3
+uvicorn==0.18.2
+wcwidth==0.2.5
+webencodings==0.5.1
+websocket-client==1.3.2
+yarl==1.7.2

src/README.md

@@ -0,0 +1,26 @@
+---
+title: BayesCap
+emoji: 🔥
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+app_file: app.py
+pinned: false
+---
+# Configuration
+`title`: _string_  
+Display title for the Space
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

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,117 @@
+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."
+        )
+
+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