Add files via upload

Author: zuru

dataset/__init__.py

@@ -0,0 +1 @@
+from .dataset_360D import *

dataset/dataset_360D.py

@@ -0,0 +1 @@
+from .image import *

exporters/__init__.py

@@ -0,0 +1,26 @@
+import torch
+import cv2
+import numpy
+
+def save_image(filename, tensor, scale=255.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        cv2.imwrite(filename.replace("#", str(n)), array)
+
+def save_depth(filename, tensor, scale=1000.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        array = numpy.uint16(array)
+        cv2.imwrite(filename.replace("#", str(n)), array)
+
+def save_data(filename, tensor, scale=1000.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        array = numpy.float32(array)
+        cv2.imwrite(filename.replace("#", str(n)), array)

exporters/image.py

@@ -0,0 +1,26 @@
+import os
+
+'''
+    Filesystem class
+    provides file control utilities like tensor saving etc.
+'''
+class Filesystem:
+    def __init__(self):
+        self.cwd = os.getcwd()
+        if os.path.isfile(self.cwd):
+            self.cwd = os.path.basename(self.cwd)
+    ''' 
+        Creates directory 
+        either by giving the absolute path to create
+        or the relative path w.r.t. the current working directory
+
+        \param path the path to create
+    '''
+    def mkdir(self, path):
+        if os.path.isabs(path):
+            if not os.path.exists(path):
+                os.mkdir(path)
+        else:
+            pathToCreate = os.path.join(self.cwd, path)
+            if not os.path.exists(pathToCreate):
+                os.mkdir(pathToCreate)

filesystem/file_utils.py

@@ -0,0 +1,13 @@
+from .resnet360 import *
+
+import sys
+
+def get_model(name, model_params):
+    if name == 'resnet_coord':
+        return ResNet360(            
+            # conv_type='standard', activation='elu', norm_type='none', \
+            conv_type='coord', activation='elu', norm_type='none', \
+            width=512,
+        )
+    else:
+        print("Could not find the requested model ({})".format(name), file=sys.stderr)

models/__init__.py

@@ -0,0 +1,111 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+'''
+    Code adapted from https://github.com/uber-research/coordconv 
+    accompanying the paper "An Intriguing Failing of Convolutional Neural Networks and the CoordConv Solution" (NeurIPS 2018)
+'''
+
+class AddCoords360(nn.Module):
+    def __init__(self, x_dim=64, y_dim=64, with_r=False):
+        super(AddCoords360, self).__init__()
+        self.x_dim = int(x_dim)
+        self.y_dim = int(y_dim)
+        self.with_r = with_r
+
+    def forward(self, input_tensor):
+        """
+        input_tensor: (batch, c, x_dim, y_dim)
+        """
+        batch_size_tensor = input_tensor.shape[0]
+
+        xx_ones = torch.ones([1, self.y_dim], dtype=torch.float32, device=input_tensor.device)
+        xx_ones = xx_ones.unsqueeze(-1)
+
+        xx_range = torch.arange(self.x_dim, dtype=torch.float32, device=input_tensor.device).unsqueeze(0)
+        xx_range = xx_range.unsqueeze(1)
+
+        xx_channel = torch.matmul(xx_ones, xx_range)
+        xx_channel = xx_channel.unsqueeze(-1)
+
+        yy_ones = torch.ones([1, self.x_dim], dtype=torch.float32, device=input_tensor.device)
+        yy_ones = yy_ones.unsqueeze(1)
+
+        yy_range = torch.arange(self.y_dim, dtype=torch.float32, device=input_tensor.device).unsqueeze(0)
+        yy_range = yy_range.unsqueeze(-1)
+
+        yy_channel = torch.matmul(yy_range, yy_ones)
+        yy_channel = yy_channel.unsqueeze(-1)
+
+        xx_channel = xx_channel.permute(0, 3, 2, 1)
+        yy_channel = yy_channel.permute(0, 3, 2, 1)
+
+        xx_channel = xx_channel.float() / (self.x_dim - 1)
+        yy_channel = yy_channel.float() / (self.y_dim - 1)
+
+        xx_channel = xx_channel * 2 - 1
+        yy_channel = yy_channel * 2 - 1
+
+        xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1)
+        yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1)
+
+        ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1)
+
+        if self.with_r:
+            rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
+            ret = torch.cat([ret, rr], dim=1)
+
+        return ret
+
+class CoordConv360(nn.Module):
+    """CoordConv layer as in the paper."""
+    def __init__(self, x_dim, y_dim, with_r, in_channels, out_channels, kernel_size, *args, **kwargs):
+        super(CoordConv360, self).__init__()
+        self.addcoords = AddCoords360(x_dim=x_dim, y_dim=y_dim, with_r=with_r)                
+        in_size = in_channels+2
+        if with_r:
+            in_size += 1            
+        self.conv = nn.Conv2d(in_size, out_channels, kernel_size, **kwargs)
+
+    def forward(self, input_tensor):
+        ret = self.addcoords(input_tensor)
+        ret = self.conv(ret)
+        return ret
+
+
+def create_conv(in_size, out_size, conv_type, padding=1, stride=1, kernel_size=3, width=512):
+    if conv_type == 'standard':
+        return nn.Conv2d(in_channels=in_size, out_channels=out_size, \
+            kernel_size=kernel_size, padding=padding, stride=stride)
+    elif conv_type == 'coord':
+        return CoordConv360(x_dim=width / 2.0, y_dim=width,\
+            with_r=False, kernel_size=kernel_size, stride=stride,\
+            in_channels=in_size, out_channels=out_size, padding=padding)    
+
+def create_activation(activation):
+    if activation == 'relu':
+        return nn.ReLU(inplace=True)
+    elif activation == 'elu':
+        return nn.ELU(inplace=True)
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+def create_normalization(out_size, norm_type):
+    if norm_type == 'batchnorm':
+        return nn.BatchNorm2d(out_size)
+    elif norm_type == 'groupnorm':
+        return nn.GroupNorm(out_size // 4, out_size)
+    elif norm_type == 'none':
+        return Identity()
+
+def create_downscale(out_size, down_mode):
+    if down_mode == 'pool':
+        return torch.nn.modules.MaxPool2d(2)
+    elif down_mode == 'downconv':
+        return nn.Conv2d(in_channels=out_size, out_channels=out_size, kernel_size=3,\
+            stride=2, padding=1, bias=False)
+    elif down_mode == 'gaussian':
+        print("Not implemented")

models/modules.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+
+import functools
+
+from .modules import *
+
+# adapted from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py 
+
+class ResNet360(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,
+        in_channels=3,
+        out_channels=1,
+        depth=5,
+        wf=32,
+        conv_type='coord',
+        padding='kernel',
+        norm_type='none',
+        activation='elu',
+        up_mode='upconv',
+        down_mode='downconv',
+        width=512,
+        use_dropout=False,
+        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(depth >= 0)
+        super(ResNet360, self).__init__()
+        model = (
+            [
+                create_conv(in_channels, wf, conv_type, \
+                    kernel_size=7, padding=3, stride=1, width=width),
+                create_normalization(wf, norm_type),
+                create_activation(activation)
+            ]
+        )
+
+        n_downsampling = 2
+        for i in range(n_downsampling): 
+            mult = 2 ** i
+            model += (
+                [
+                    create_conv(wf * mult, wf * mult * 2, conv_type, \
+                        kernel_size=3, stride=2, padding=1, width=width // (i+1)),
+                    create_normalization(wf * mult * 2, norm_type),
+                    create_activation(activation)
+                ]
+            )
+
+        mult = 2 ** n_downsampling
+        for i in range(depth):
+            model += [ResnetBlock(wf * mult, activation=activation, \
+                norm_type=norm_type, conv_type=conv_type, \
+                width=width // (2 ** n_downsampling))]
+
+        for i in range(n_downsampling):
+            mult = 2 ** (n_downsampling - i)
+            model += (
+                [
+                    nn.ConvTranspose2d(wf * mult, int(wf * mult / 2),
+                        kernel_size=3, stride=2,
+                        padding=1, output_padding=1),
+                    create_normalization(int(wf * mult / 2), norm_type),
+                    create_activation(activation)
+                ]
+            )
+        
+        model += [create_conv(wf, out_channels, conv_type, \
+            kernel_size=7, padding=3, width=width)]
+
+        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, norm_type, conv_type, activation, width):
+        """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__()
+        conv_block = []
+        conv_block +=(
+            [
+                create_conv(dim, dim, conv_type, width=width),
+                create_normalization(dim, norm_type),
+                create_activation(activation),
+            ]
+        )
+        conv_block +=(
+            [
+                create_conv(dim, dim, conv_type, width=width),
+                create_normalization(dim, norm_type),
+            ]
+        )
+
+        self.block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        """Forward function (with skip connections)"""
+        out = x + self.block(x)  # add skip connections
+        return out

models/resnet360.py

@@ -0,0 +1,4 @@
+from .grid import *
+from .cartesian import *
+from .derivatives import *
+from .weights import *

spherical/__init__.py

@@ -0,0 +1,50 @@
+import torch
+
+from .grid import *
+
+'''
+    Cartesian coordinates extraction from Spherical coordinates
+        z is forward axis
+        y is the up axis
+        x is the right axis
+        r is the radius (i.e. spherical depth)
+        phi is the longitude/azimuthial rotation angle (defined on the x-z plane)
+        theta is the latitude/elevation rotation angle (defined on the y-z plane)
+'''
+def coord_x(sgrid, depth):
+    return ( # r * sin(phi) * sin(theta) -> r * cos(phi) * -cos(theta) in our offsets
+        depth # this is due to the offsets as explained below
+        * torch.cos(phi(sgrid)) # long = x - 3 * pi / 2
+        * -1 * torch.cos(theta(sgrid)) # lat = y - pi / 2
+    )
+
+def coord_y(sgrid, depth):
+    return ( # r * cos(theta) -> r * sin(theta) in our offsets
+        depth # this is due to the offsets as explained below
+        * torch.sin(theta(sgrid)) # lat = y - pi / 2
+    )
+
+def coord_z(sgrid, depth):
+    return ( # r * cos(phi) * sin(theta) -> r * -sin(phi) * -cos(theta) in our offsets
+        depth # this is due to the offsets as explained above
+        * torch.sin(phi(sgrid)) # * -1
+        * torch.cos(theta(sgrid)) # * -1
+    ) # the -1s cancel out
+
+def coords_3d(sgrid, depth):
+    return torch.cat(
+        (
+            coord_x(sgrid, depth),
+            coord_y(sgrid, depth),
+            coord_z(sgrid, depth)
+        ), dim=1
+    )
+
+def xi(pcloud):
+    return pcloud[:, 0, :, :].unsqueeze(1)
+
+def yi(pcloud):
+    return pcloud[:, 1, :, :].unsqueeze(1)
+
+def zeta(pcloud):
+    return pcloud[:, 2, :, :].unsqueeze(1)