sync

.vscode/c_cpp_properties.json

@@ -11,7 +11,7 @@
             "cStandard": "gnu17",
             "cppStandard": "gnu++14",
             "intelliSenseMode": "gcc-x64"
-        },
+        }
     ],
     "version": 4
 }
\ No newline at end of file

README.md

@@ -25,4 +25,5 @@ OUTPUT_DIR = "/home/yejiannan/Project/LightField/outputE/lf_syn_full"
 这里要重新配置一下数据的位置
 
 Video generate:
-ffmpeg -r 50 -i view%04d.png -c:v libx264 -vf fps=50 -pix_fmt yuv420p ../mc_hint.mp4
+ffmpeg -y -r 30 -i view_%04d.png -c:v libx264 -vf fps=30 -pix_fmt yuv420p ../train.mp4
+trtexec --onnx=net@256x256x2.onnx --fp16 --saveEngine=net@256x256x2.trt --workspace=4096
\ No newline at end of file

my/foveation.py → components/foveation.py

@@ -2,7 +2,9 @@ import math
 import torch
 import torch.nn.functional as nn_f
 from typing import List, Tuple
-from . import util
+from utils import img
+from utils import view
+from utils import misc
 
 
 class Foveation(object):
@@ -17,7 +19,7 @@ class Foveation(object):
             self._gen_layer_blendmap(i)
             for i in range(self.n_layers - 1)
         ]  # blend maps of fovea layers
-        self.coords = util.MeshGrid(out_res).to(device=device)
+        self.coords = misc.meshgrid(*out_res).to(device=device)
 
     def to(self, device):
         self.eye_fovea_blend = [x.to(device=device)
@@ -32,13 +34,13 @@ class Foveation(object):
         Generate foveated retinal image by blending fovea layers
         **Note: current implementation only support two fovea layers**
 
-        :param layers ```List(Tensor(B, C, H'{l}, W'{l}))```: list of foveated layers
-        :return ```Tensor(B, C, H:out, W:out)```: foveated images
+        :param layers `List(Tensor(B, C, H'{l}, W'{l}))`: list of foveated layers
+        :return `Tensor(B, C, H:out, W:out)`: foveated images
         """
         output: torch.Tensor = nn_f.interpolate(layers[-1], self.out_res,
                                                 mode='bilinear', align_corners=False)
         if shifts != None:
-            output = util.horizontal_shift_image(output, shifts[-1])
+            output = img.horizontal_shift(output, shifts[-1])
         c = torch.tensor([
             fovea_center[0] + self.out_res[1] / 2,
             fovea_center[1] + self.out_res[0] / 2
@@ -49,7 +51,7 @@ class Foveation(object):
             R = self.get_layer_size_in_final_image(i) / 2
             grid = ((self.coords - c) / R)[None, ...]
             if shifts != None:
-                grid = util.horizontal_shift_image(grid, shifts[i], -2)
+                grid = img.horizontal_shift(grid, shifts[i], -2)
             blend = nn_f.grid_sample(self.eye_fovea_blend[i][None, None, ...], grid) # (1, 1, H:out, W:out)
             output.mul_(1 - blend).add_(nn_f.grid_sample(layers[i], grid) * blend)
         return output
@@ -61,8 +63,8 @@ class Foveation(object):
         :param i: index of layer
         :return: size of layer i in final image (in pixels)
         """
-        length_i = util.Fov2Length(self.fov_list[i])
-        length = util.Fov2Length(self.fov_list[-1])
+        length_i = view.fov2length(self.fov_list[i])
+        length = view.fov2length(self.fov_list[-1])
         k = length_i / length
         return int(math.ceil(self.out_res[0] * k))
 
@@ -71,11 +73,10 @@ class Foveation(object):
         Generate blend map for fovea layer i
 
         :param i: index of fovea layer
-        :return ```Tensor(H{i}, W{i})```: blend map
+        :return `Tensor(H{i}, W{i})`: blend map
         """
         size = self.get_layer_size_in_final_image(i)
         R = size / 2
-        p = util.MeshGrid((size, size)).to(
-            device=self.device)  # (size, size, 2)
+        p = misc.meshgrid(size, size).to(device=self.device)  # (size, size, 2)
         r = torch.norm(p - R, dim=2)  # (size, size, 2)
-        return util.SmoothStep(R, R * 0.6, r)
+        return misc.smooth_step(R, R * 0.6, r)

my/gen_final.py → components/gen_final.py

 import torch
-import math
-from torch import nn
 import torch.nn.functional as nn_f
 from typing import Any, List, Mapping, Tuple
-from . import view
+from torch import nn
+from utils import view
+from utils import misc
+from utils.perf import Perf
 from . import refine
-from . import util
 from .foveation import Foveation
-from .simple_perf import SimplePerf
 
 
 class GenFinal(object):
@@ -63,7 +62,7 @@ class GenFinal(object):
         periph_cam = self.layer_cams[2]
         trans_periph = mono_trans if mono_trans != None else trans
 
-        perf = SimplePerf(True, True) if perf_time else None
+        perf = Perf(True, True) if perf_time else None
 
         # *_rays_o, *_rays_d: (1, N, 3)
         fovea_rays_o, fovea_rays_d = fovea_cam.get_global_rays(trans, True)
@@ -73,22 +72,22 @@ class GenFinal(object):
         mid_periph_rays_o = torch.cat([mid_rays_o, periph_rays_o], 1)
         mid_periph_rays_d = torch.cat([mid_rays_d, periph_rays_d], 1)
         if perf_time:
-            perf.Checkpoint('Get rays')
+            perf.checkpoint('Get rays')
 
-        perf1 = SimplePerf(True, True) if perf_time else None
+        perf1 = Perf(True, True) if perf_time else None
 
         fovea_inferred, fovea_depth = self._infer(
             self.fovea_net, fovea_rays_o, fovea_rays_d, [fovea_cam.res], True)
 
         if perf_time:
-            perf1.Checkpoint('Infer fovea')
+            perf1.checkpoint('Infer fovea')
         mid_inferred, mid_depth, periph_inferred, periph_depth = self._infer(
             self.periph_net, mid_periph_rays_o, mid_periph_rays_d,
             [mid_cam.res, periph_cam.res], True)
 
         if perf_time:
-            perf1.Checkpoint('Infer mid & periph')
-            perf.Checkpoint('Infer')
+            perf1.checkpoint('Infer mid & periph')
+            perf.checkpoint('Infer')
 
         if mono_trans != None and shift == 0:  # do warp
             fovea_depth[torch.isnan(fovea_depth)] = 50
@@ -96,13 +95,13 @@ class GenFinal(object):
             periph_depth[torch.isnan(periph_depth)] = 50
 
             if warp_by_depth:
-                z_list = util.depth_sample((1, 50), 4, True)
+                z_list = misc.depth_sample((1, 50), 4, True)
                 mid_inferred = self._warp(trans, mono_trans, mid_cam,
                                           z_list, mid_inferred, mid_depth)
                 periph_inferred = self._warp(trans, mono_trans, periph_cam,
                                              z_list, periph_inferred, periph_depth)
                 if perf_time:
-                    perf.Checkpoint('Mono warp')
+                    perf.checkpoint('Mono warp')
             else:   
                 p = torch.tensor([[0, 0, torch.mean(fovea_depth)]],
                                 device=self.device)
@@ -111,7 +110,7 @@ class GenFinal(object):
                     p_, center_as_origin=True)[..., 0].item()
                 shift = round(shift)
                 if perf_time:
-                    perf.Checkpoint('Mono shift')
+                    perf.checkpoint('Mono shift')
 
         fovea_refined = refine.grad_aware_median(fovea_inferred, 3, 3, True)
         fovea_refined = refine.constrast_enhance(fovea_refined, 3, 0.2)
@@ -119,7 +118,7 @@ class GenFinal(object):
         periph_refined = refine.constrast_enhance(periph_inferred, 5, 0.2)
 
         if perf_time:
-            perf.Checkpoint('Refine')
+            perf.checkpoint('Refine')
 
         blended = self.foveation.synthesis([
             fovea_refined,
@@ -128,7 +127,7 @@ class GenFinal(object):
         ], (gaze[0], gaze[1]), [0, shift, shift] if shift != 0 else None)
 
         if perf_time:
-            perf.Checkpoint('Blend')
+            perf.checkpoint('Blend')
 
         if ret_raw:
             return {
@@ -195,9 +194,9 @@ class GenFinal(object):
         :param trans0: [description]
         :param cam: [description]
         :param z_list: [description]
-        :param image ```Tensor(B, C, H, W)```:
-        :param depthmap ```Tensor(B, H, W)```:
-        :return ```Tensor(B, C, H, W)```:
+        :param image `Tensor(B, C, H, W)`:
+        :param depthmap `Tensor(B, H, W)`:
+        :return `Tensor(B, C, H, W)`:
         """
         B = image.size(0)
         rays_d = cam.get_global_rays(trans, norm=False)[1]  # (1, H, W, 3)

my/refine.py → components/refine.py

@@ -2,9 +2,9 @@ import cv2
 import torch
 import numpy as np
 import torch.nn.functional as nn_f
-from . import view
-from . import util
-
+from utils import view
+from utils import img
+from utils.constants import *
 
 class GuideRefinement(object):
 
@@ -17,7 +17,7 @@ class GuideRefinement(object):
             for i in range(guides_image.size(0))
         ], 0)
         self.guides_diff = (guides_image - guides_inferred) / \
-            (guides_inferred + 1e-5)
+            (guides_inferred + TINY_FLOAT)
         self.guides_view = guides_view
         self.guides_cam = guides_cam
 
@@ -62,10 +62,10 @@ def constrast_enhance(image, sigma, fe):
 
 
 def morph_close(image: torch.Tensor):
-    image_ = util.Tensor2MatImg(image)
+    image_ = img.torch2np(image)
     kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
     image_ = cv2.morphologyEx(image_, cv2.MORPH_CLOSE, kernel)
-    return util.MatImg2Tensor(image_).to(image.device)
+    return img.np2torch(image_).to(image.device)
 
 
 def get_grad(image: torch.Tensor, k=1, do_morph_close=False):
@@ -105,9 +105,9 @@ def getGaussianKernel(ksize, sigma=0):
 
 def grad_aware_median(image: torch.Tensor, median_kernel_size: int, grad_k: float,
                       grad_do_morph_close: bool):
-    image_ = util.Tensor2MatImg(image)
+    image_ = img.torch2np(image)
     blur = cv2.medianBlur(image_, median_kernel_size)
-    blur = util.MatImg2Tensor(blur).to(image.device)
+    blur = img.np2torch(blur).to(image.device)
     grad = get_grad(image, grad_k, grad_do_morph_close)
     return image * grad + blur * (1 - grad)
 

configs/bgnet.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'bgnet'
+    config.N_ENCODE_DIM = 10
+    config.FC_PARAMS.update({
+        'nf': 128,
+        'n_layers': 4
+    })

configs/cnerf.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'cnerf128'
+    config.N_ENCODE_DIM = 10
+    config.DEPTH_REF = True
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.4, 6),
+        'n_samples': 8,
+        'perturb_sample': False
+    })

configs/us_periph.py → configs/dnerfabins.py

-from my import color_mode
-
 def update_config(config):
-    # Dataset settings
-    config.COLOR = color_mode.RGB
-
     # Net parameters
-    config.NET_TYPE = 'nmsl'
+    config.NET_TYPE = 'dnerfa'
     config.N_ENCODE_DIM = 10
+    config.DEPTH_REF = True
     config.FC_PARAMS.update({
-        'nf': 64,
-        'n_layers': 4
+        'nf': 256,
+        'n_layers': 8
     })
     config.SAMPLE_PARAMS.update({
-        'depth_range': (1, 50),
-        'n_samples': 16
+        'depth_range': (0.4, 6),
+        'n_samples': 8,
+        'perturb_sample': False
     })
\ No newline at end of file

configs/fovea.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'snerffast4'
+    config.N_ENCODE_DIM = 10
+    #config.N_DIR_ENCODE = 4
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.7, 10),
+        'n_samples': 64,
+        'perturb_sample': False
+    })

configs/msl2fast.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'msl2fast'
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.5, 5),
+        'n_samples': 128,
+        'perturb_sample': False
+    })

configs/fovea_rgb.py → configs/msl_fovea.py

-from my import color_mode
-
 def update_config(config):
-    # Dataset settings
-    config.COLOR = color_mode.RGB
-
     # Net parameters
     config.NET_TYPE = 'msl'
-    config.N_ENCODE_DIM = 10
     config.FC_PARAMS.update({
         'nf': 128,
         'n_layers': 4
@@ -15,5 +9,3 @@ def update_config(config):
         'depth_range': (1, 50),
         'n_samples': 32
     })
-
-

configs/mslfast.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'mslfast'
+    config.N_ENCODE_DIM = 6
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.5, 5),
+        'n_samples': 128,
+        'perturb_sample': False
+    })

configs/mslray.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'mslray'
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.5, 5),
+        'n_samples': 128,
+        'perturb_sample': False
+    })

configs/nerf.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'nerf'
+    config.N_ENCODE_DIM = 10
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (1, 50),
+        'n_samples': 64,
+        'perturb_sample': True
+    })
+    config.NERF_FINE_NET_PARAMS.update({
+        'enable': True,
+        'nf': 256,
+        'n_layers': 8,
+        'additional_samples': 128
+    })
\ No newline at end of file

configs/periph_rgb.py → configs/nerf_periph.py

-from my import color_mode
-
 def update_config(config):
-    # Dataset settings
-    config.COLOR = color_mode.RGB
-
     # Net parameters
-    config.NET_TYPE = 'msl'
+    config.NET_TYPE = 'nerf'
     config.N_ENCODE_DIM = 10
     config.FC_PARAMS.update({
-        'nf': 128,
+        'nf': 96,
         'n_layers': 4
     })
     config.SAMPLE_PARAMS.update({
         'depth_range': (1, 50),
-        'n_samples': 8
+        'n_samples': 16,
+        'perturb_sample': False
     })
\ No newline at end of file

configs/nerfsimple.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'nerf'
+    config.N_ENCODE_DIM = 10
+    #config.N_DIR_ENCODE = 4
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8,
+        'skips': [4]
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.7, 10),
+        'n_samples': 128,
+        'perturb_sample': False
+    })
\ No newline at end of file

configs/us_fovea.py → configs/nmsl_fovea.py

-from my import color_mode
-
 def update_config(config):
-    # Dataset settings
-    config.COLOR = color_mode.RGB
-
     # Net parameters
     config.NET_TYPE = 'nmsl'
-    config.N_ENCODE_DIM = 10
     config.FC_PARAMS.update({
         'nf': 128,
         'n_layers': 4

configs/fovea_nmsl4.py → configs/nnerf.py

-from my import color_mode
-
 def update_config(config):
-    # Dataset settings
-    config.COLOR = color_mode.RGB
-
     # Net parameters
-    config.NET_TYPE = 'nmsl4'
+    config.NET_TYPE = 'nnerf4'
     config.N_ENCODE_DIM = 10
     config.FC_PARAMS.update({
         'nf': 128,
@@ -13,7 +8,6 @@ def update_config(config):
     })
     config.SAMPLE_PARAMS.update({
         'depth_range': (1, 50),
-        'n_samples': 32
+        'n_samples': 32,
+        'perturb_sample': True
     })
-
-

configs/oracle.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'oracle'
+    config.N_ENCODE_DIM = 0
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8,
+        'activation': 'selu',
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.4, 6),
+        'n_samples': 128,
+        'perturb_sample': False
+    })

configs/part0.py

+def update_config(config):
+    # Net parameters
+    config.NET_TYPE = 'nerf'
+    config.N_ENCODE_DIM = 10
+    config.N_DIR_ENCODE = 4
+    config.FC_PARAMS.update({
+        'nf': 256,
+        'n_layers': 8,
+        'skips': [4]
+    })
+    config.SAMPLE_PARAMS.update({
+        'depth_range': (0.3, 7),
+        'n_samples': 256,
+        'perturb_sample': True
+    })