├── Data └── scene.yaml ├── LICENSE ├── MyRender ├── CloudProjection │ ├── helper_math.h │ ├── pcpr_cuda.cpp │ ├── point_render.cu │ └── point_render.cuh └── setup.py ├── READ ├── criterions │ ├── __pycache__ │ │ ├── vgg_loss.cpython-36.pyc │ │ └── vgg_loss.cpython-38.pyc │ └── vgg_loss.py ├── datasets │ ├── common.py │ ├── dynamic.py │ └── splitter.py ├── gl │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── camera.cpython-36.pyc │ │ ├── camera.cpython-38.pyc │ │ ├── dataset.cpython-36.pyc │ │ ├── dataset.cpython-38.pyc │ │ ├── nn.cpython-36.pyc │ │ ├── nn.cpython-38.pyc │ │ ├── programs.cpython-36.pyc │ │ ├── programs.cpython-38.pyc │ │ ├── render.cpython-36.pyc │ │ ├── render.cpython-38.pyc │ │ ├── utils.cpython-36.pyc │ │ └── utils.cpython-38.pyc │ ├── camera.py │ ├── dataset.py │ ├── nn.py │ ├── programs.py │ ├── render.py │ ├── utils.py │ └── viewer │ │ └── camera.py ├── models │ ├── __pycache__ │ │ ├── common.cpython-36.pyc │ │ ├── common.cpython-38.pyc │ │ ├── compose.cpython-36.pyc │ │ ├── compose.cpython-38.pyc │ │ ├── conv.cpython-36.pyc │ │ ├── conv.cpython-38.pyc │ │ ├── texture.cpython-36.pyc │ │ ├── texture.cpython-38.pyc │ │ ├── unet.cpython-36.pyc │ │ └── unet.cpython-38.pyc │ ├── compose.py │ ├── conv.py │ ├── texture.py │ └── unet.py ├── pipelines │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── ogl.cpython-36.pyc │ │ ├── ogl.cpython-38.pyc │ │ ├── pipeline.cpython-36.pyc │ │ └── pipeline.cpython-38.pyc │ ├── ogl.py │ └── pipeline.py └── utils │ ├── __pycache__ │ ├── arguments.cpython-36.pyc │ ├── arguments.cpython-38.pyc │ ├── perform.cpython-36.pyc │ ├── perform.cpython-38.pyc │ ├── train.cpython-36.pyc │ └── train.cpython-38.pyc │ ├── arguments.py │ ├── perform.py │ └── train.py ├── README.md ├── configs ├── paths_example.yaml └── train_example.yaml ├── downloads └── kitti6.yaml ├── image ├── NovelView.jpg ├── Scene_Editing.jpg ├── Scene_Stitching.jpg └── main.jpg ├── requirement.sh ├── src ├── Data │ └── scene.yaml ├── MyRender │ ├── CloudProjection │ │ ├── helper_math.h │ │ ├── pcpr_cuda.cpp │ │ ├── point_render.cu │ │ └── point_render.cuh │ └── setup.py ├── READ │ ├── criterions │ │ ├── __pycache__ │ │ │ ├── vgg_loss.cpython-36.pyc │ │ │ └── vgg_loss.cpython-39.pyc │ │ └── vgg_loss.py │ ├── datasets │ │ ├── __pycache__ │ │ │ ├── common.cpython-36.pyc │ │ │ ├── common.cpython-39.pyc │ │ │ ├── dynamic.cpython-36.pyc │ │ │ ├── dynamic.cpython-39.pyc │ │ │ ├── splitter.cpython-36.pyc │ │ │ └── splitter.cpython-39.pyc │ │ ├── common.py │ │ ├── dynamic.py │ │ └── splitter.py │ ├── gl │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── camera.cpython-39.pyc │ │ │ ├── dataset.cpython-36.pyc │ │ │ ├── dataset.cpython-39.pyc │ │ │ ├── myrender.cpython-36.pyc │ │ │ ├── myrender.cpython-39.pyc │ │ │ ├── nn.cpython-39.pyc │ │ │ ├── programs.cpython-36.pyc │ │ │ ├── programs.cpython-39.pyc │ │ │ ├── render.cpython-36.pyc │ │ │ ├── render.cpython-39.pyc │ │ │ ├── utils.cpython-36.pyc │ │ │ └── utils.cpython-39.pyc │ │ ├── camera.py │ │ ├── dataset.py │ │ ├── myrender.py │ │ ├── nn.py │ │ ├── programs.py │ │ ├── render.py │ │ ├── utils.py │ │ └── viewer │ │ │ └── camera.py │ ├── models │ │ ├── __pycache__ │ │ │ ├── compose.cpython-36.pyc │ │ │ ├── compose.cpython-39.pyc │ │ │ ├── conv.cpython-36.pyc │ │ │ ├── conv.cpython-39.pyc │ │ │ ├── texture.cpython-36.pyc │ │ │ ├── texture.cpython-39.pyc │ │ │ ├── unet.cpython-36.pyc │ │ │ └── unet.cpython-39.pyc │ │ ├── app_encoder.py │ │ ├── compose.py │ │ ├── conv.py │ │ ├── texture.py │ │ └── unet.py │ ├── pipelines │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-36.pyc │ │ │ ├── __init__.cpython-39.pyc │ │ │ ├── ogl.cpython-36.pyc │ │ │ ├── ogl.cpython-39.pyc │ │ │ ├── pipeline.cpython-36.pyc │ │ │ └── pipeline.cpython-39.pyc │ │ ├── ogl.py │ │ └── pipeline.py │ └── utils │ │ ├── __pycache__ │ │ ├── arguments.cpython-36.pyc │ │ ├── arguments.cpython-39.pyc │ │ ├── perform.cpython-36.pyc │ │ ├── perform.cpython-39.pyc │ │ ├── train.cpython-36.pyc │ │ └── train.cpython-39.pyc │ │ ├── arguments.py │ │ ├── perform.py │ │ └── train.py ├── configs │ ├── paths_example.yaml │ └── train_example.yaml ├── docker │ └── Dockerfile ├── downloads │ └── kitti6.yaml ├── scripts │ └── train_data.sh └── train.py ├── train.py └── viewer.py /Data/scene.yaml: -------------------------------------------------------------------------------- 1 | viewport_size: [1216, 368] 2 | intrinsic_matrix: camera.xml 3 | view_matrix: camera.xml 4 | pointcloud: pointcloud.ply 5 | 6 | -------------------------------------------------------------------------------- /MyRender/CloudProjection/pcpr_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "point_render.cuh" 4 | #include 5 | 6 | 7 | // CUDA forward declarations 8 | 9 | std::vector pcpr_cuda_forward( 10 | torch::Tensor in_points, //(num_points,3) 11 | torch::Tensor , 12 | int tar_width, int tar_height, int block_size 13 | ); 14 | 15 | // C++ interface 16 | 17 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 18 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 19 | #define CHECK_FLOAT(x) AT_ASSERTM(x.type().scalarType()==torch::ScalarType::Float, #x " must be a float tensor") 20 | #define CHECK_Int(x) AT_ASSERTM(x.type().scalarType()==torch::ScalarType::Int, #x " must be a Int tensor") 21 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x); 22 | 23 | std::vector pcpr_cuda_forward( 24 | torch::Tensor in_points, //(num_points,3) 25 | torch::Tensor total_m, 26 | int tar_width, int tar_height, int block_size 27 | ) 28 | { 29 | in_points = in_points.to(torch::kCUDA); 30 | total_m = total_m.to(torch::kCUDA); 31 | 32 | CHECK_INPUT(in_points); CHECK_FLOAT(in_points); 33 | CHECK_INPUT(total_m); CHECK_FLOAT(total_m); 34 | AT_ASSERTM(total_m.sizes().size()==3, "batch_size check"); 35 | 36 | return GPU_PCPR(in_points, total_m, tar_width, tar_height, block_size); 37 | } 38 | 39 | 40 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 41 | m.def("forward", &pcpr_cuda_forward, "PCPR forward (CUDA)"); 42 | } 43 | 44 | -------------------------------------------------------------------------------- /MyRender/CloudProjection/point_render.cu: -------------------------------------------------------------------------------- 1 | #include "point_render.cuh" 2 | #include 3 | 4 | #include "helper_math.h" 5 | 6 | #include "cooperative_groups.h" 7 | 8 | struct Matrix4x4 9 | { 10 | public: 11 | float4 col[4]; 12 | __device__ __forceinline__ 13 | Matrix4x4() 14 | { 15 | col[0] = col[1] = col[2] = col[3] = make_float4(0, 0, 0, 0); 16 | } 17 | __device__ __forceinline__ 18 | Matrix4x4(float4 a, float4 b, float4 c, float4 d) 19 | { 20 | col[0].x = a.x; 21 | col[0].y = a.y; 22 | col[0].z = a.z; 23 | col[0].w = a.w; 24 | 25 | col[1].x = b.x; 26 | col[1].y = b.y; 27 | col[1].z = b.z; 28 | col[1].w = b.w; 29 | 30 | col[2].x = c.x; 31 | col[2].y = c.y; 32 | col[2].z = c.z; 33 | col[2].w = c.w; 34 | 35 | col[3].x = d.x; 36 | col[3].y = d.y; 37 | col[3].z = d.z; 38 | col[3].w = d.w; 39 | } 40 | 41 | __device__ __forceinline__ 42 | Matrix4x4 transpose() const 43 | { 44 | Matrix4x4 res; 45 | 46 | res.col[0].x = col[0].x; 47 | res.col[0].y = col[1].x; 48 | res.col[0].z = col[2].x; 49 | res.col[0].w = col[3].x; 50 | 51 | res.col[1].x = col[0].y; 52 | res.col[1].y = col[1].y; 53 | res.col[1].z = col[2].y; 54 | res.col[1].w = col[3].y; 55 | 56 | res.col[2].x = col[0].z; 57 | res.col[2].y = col[1].z; 58 | res.col[2].z = col[2].z; 59 | res.col[2].w = col[3].z; 60 | 61 | res.col[3].x = col[0].w; 62 | res.col[3].y = col[1].w; 63 | res.col[3].z = col[2].w; 64 | res.col[3].w = col[3].w; 65 | return res; 66 | 67 | } 68 | __device__ __forceinline__ 69 | Matrix4x4 inv() const 70 | { 71 | Matrix4x4 res; 72 | res.col[0].x = col[0].x; 73 | res.col[0].y = col[1].x; 74 | res.col[0].z = col[2].x; 75 | res.col[0].w = 0; 76 | 77 | res.col[1].x = col[0].y; 78 | res.col[1].y = col[1].y; 79 | res.col[1].z = col[2].y; 80 | res.col[1].w = 0; 81 | 82 | res.col[2].x = col[0].z; 83 | res.col[2].y = col[1].z; 84 | res.col[2].z = col[2].z; 85 | res.col[2].w = 0; 86 | 87 | res.col[3].x = -dot(col[0], col[3]); 88 | res.col[3].y = -dot(col[1], col[3]); 89 | res.col[3].z = -dot(col[2], col[3]); 90 | res.col[3].w = 1; 91 | return res; 92 | } 93 | }; 94 | 95 | 96 | typedef struct CamMatrix 97 | { 98 | float4 m[4]; 99 | __device__ __forceinline__ 100 | Matrix4x4 getRT() const 101 | { 102 | return Matrix4x4(m[0],m[1],m[2],m[3]); 103 | } 104 | 105 | }; 106 | 107 | namespace math 108 | { 109 | __device__ __forceinline__ 110 | float4 MatrixMul(const Matrix4x4& mat, float4& p) 111 | { 112 | Matrix4x4 res = mat; 113 | float4 ans; 114 | ans.x = dot(res.col[0], p); 115 | ans.y = dot(res.col[1], p); 116 | ans.z = dot(res.col[2], p); 117 | ans.w = dot(res.col[3], p); 118 | 119 | ans = ans / ans.w; 120 | return ans; 121 | } 122 | } 123 | 124 | 125 | __global__ 126 | void DepthProject(float3 * point_clouds, int num_points, int batch_id, 127 | CamMatrix* total_m, int tar_width, int tar_height, 128 | int* mutex_map, float* out_depth, float* out_index) 129 | { 130 | // int ids = blockDim.x * blockIdx.x + threadIdx.x; // index of point 131 | cooperative_groups::grid_group grid = cooperative_groups::this_grid(); 132 | for(int ids = grid.thread_rank(); ids num_points) return; 135 | float4 p = make_float4(point_clouds[ids], 1.0); 136 | 137 | float4 camp = math::MatrixMul(total_m->getRT(), p); 138 | 139 | if(camp.x<-1 || camp.x>1 || camp.y<-1 || camp.y>1 || camp.z<-1 || camp.z>1) return; 140 | 141 | float u = tar_width*(camp.x+1)*0.5; 142 | float v = tar_height*(1-camp.y)*0.5; 143 | float tdepth = ((camp.z+1)*0.5); 144 | 145 | int xx = int(u); 146 | int yy = int(v); 147 | if (xx < 0 || xx >= tar_width || yy < 0 || yy >= tar_height) return; 148 | int ind = batch_id * tar_width * tar_height + yy * tar_width + xx ; 149 | bool isSet = false; 150 | do 151 | { 152 | if ((isSet = atomicCAS(mutex_map + ind, 0, 1)) == false) 153 | { 154 | // critical section goes here 155 | if (out_depth[ind] > tdepth || out_depth[ind]==0) 156 | { 157 | out_depth[ind] = tdepth; 158 | out_index[ind] = (float)ids ; // 0 denote empty 159 | } 160 | } 161 | if (isSet) 162 | { 163 | mutex_map[ind] = 0; 164 | } 165 | } while (!isSet); 166 | } 167 | } 168 | 169 | std::vector GPU_PCPR( 170 | torch::Tensor in_points, //(num_points,3) 171 | torch::Tensor total_m, int tar_width, int tar_height, int block_size) 172 | { 173 | const auto num_points = in_points.size(0); 174 | const auto batch_size = total_m.size(0); 175 | 176 | torch::Tensor out_index = torch::zeros({batch_size, tar_height, tar_width}).to(torch::kCUDA); 177 | torch::Tensor out_depth = torch::zeros({batch_size, tar_height, tar_width}).to(torch::kCUDA); 178 | 179 | int default_block_size = block_size; 180 | int block_num = (num_points+default_block_size-1) / default_block_size; 181 | 182 | int *mutex_map; 183 | cudaMalloc(&mutex_map, sizeof(int) * batch_size * tar_width * tar_height); 184 | cudaMemset(mutex_map, 0, sizeof(int) * batch_size * tar_width * tar_height); 185 | 186 | for(int b=0; b> > ( 188 | (float3*)in_points.data(), num_points, b, 189 | (CamMatrix*)total_m[b].data(), 190 | tar_width, tar_height, mutex_map, 191 | out_depth.data(), out_index.data()); 192 | } 193 | cudaFree(mutex_map); 194 | cudaDeviceSynchronize(); 195 | 196 | out_index = out_index.to(torch::kCPU); 197 | out_depth = out_depth.to(torch::kCPU); 198 | 199 | return {out_index, out_depth}; 200 | } 201 | 202 | -------------------------------------------------------------------------------- /MyRender/CloudProjection/point_render.cuh: -------------------------------------------------------------------------------- 1 | #include "cuda_runtime.h" 2 | #include "device_launch_parameters.h" 3 | #include 4 | 5 | std::vector GPU_PCPR( 6 | torch::Tensor in_points, //(num_points,3) 7 | torch::Tensor total_m, 8 | int tar_width, int tar_height, int block_size); // (tar_heigh ,tar_width) 9 | -------------------------------------------------------------------------------- /MyRender/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='pcpr', 6 | version="0.1", 7 | ext_modules=[ 8 | CUDAExtension('pcpr', [ 9 | 'CloudProjection/pcpr_cuda.cpp', 10 | 'CloudProjection/point_render.cu' 11 | ] 12 | ) 13 | ], 14 | cmdclass={ 15 | 'build_ext': BuildExtension 16 | }) -------------------------------------------------------------------------------- /READ/criterions/__pycache__/vgg_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/criterions/__pycache__/vgg_loss.cpython-36.pyc -------------------------------------------------------------------------------- /READ/criterions/__pycache__/vgg_loss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/criterions/__pycache__/vgg_loss.cpython-38.pyc -------------------------------------------------------------------------------- /READ/criterions/vgg_loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | import torch.nn as nn 3 | import torchvision 4 | import torch 5 | from collections import OrderedDict 6 | from os.path import expanduser, join 7 | import os 8 | 9 | from READ.models.conv import PartialConv2d 10 | 11 | 12 | class View(nn.Module): 13 | def __init__(self): 14 | super(View, self).__init__() 15 | 16 | def forward(self, x): 17 | return x.view(-1) 18 | 19 | 20 | class VGGLoss(nn.Module): 21 | def __init__(self, net='caffe', partialconv=False, optimized=False, save_dir='.cache/torch/models'): 22 | super().__init__() 23 | 24 | self.partialconv = partialconv 25 | 26 | if net == 'pytorch': 27 | vgg19 = torchvision.models.vgg19(pretrained=True).features 28 | 29 | self.register_buffer('mean_', torch.FloatTensor([0.485, 0.456, 0.406])[None, :, None, None]) 30 | self.register_buffer('std_', torch.FloatTensor([0.229, 0.224, 0.225])[None, :, None, None]) 31 | 32 | elif net == 'caffe': 33 | if not os.path.exists(join(save_dir, 'vgg_caffe_features.pth')): 34 | vgg_weights = torch.utils.model_zoo.load_url('https://web.eecs.umich.edu/~justincj/models/vgg19-d01eb7cb.pth', model_dir=save_dir) 35 | 36 | map = {'classifier.6.weight':u'classifier.7.weight', 'classifier.6.bias':u'classifier.7.bias'} 37 | vgg_weights = OrderedDict([(map[k] if k in map else k,v) for k,v in vgg_weights.items()]) 38 | 39 | model = torchvision.models.vgg19() 40 | model.classifier = nn.Sequential(View(), *model.classifier._modules.values()) 41 | 42 | model.load_state_dict(vgg_weights) 43 | 44 | vgg19 = model.features 45 | os.makedirs(save_dir, exist_ok=True) 46 | torch.save(vgg19, join(save_dir, 'vgg_caffe_features.pth')) 47 | 48 | self.register_buffer('mean_', torch.FloatTensor([103.939, 116.779, 123.680])[None, :, None, None] / 255.) 49 | self.register_buffer('std_', torch.FloatTensor([1./255, 1./255, 1./255])[None, :, None, None]) 50 | 51 | else: 52 | vgg19 = torch.load(join(save_dir, 'vgg_caffe_features.pth')) 53 | self.register_buffer('mean_', torch.FloatTensor([103.939, 116.779, 123.680])[None, :, None, None] / 255.) 54 | self.register_buffer('std_', torch.FloatTensor([1./255, 1./255, 1./255])[None, :, None, None]) 55 | else: 56 | assert False 57 | 58 | if self.partialconv: 59 | part_conv = PartialConv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 60 | part_conv.weight = vgg19[0].weight 61 | part_conv.bias = vgg19[0].bias 62 | vgg19[0] = part_conv 63 | 64 | vgg19_avg_pooling = [] 65 | 66 | 67 | for weights in vgg19.parameters(): 68 | weights.requires_grad = False 69 | 70 | for module in vgg19.modules(): 71 | if module.__class__.__name__ == 'Sequential': 72 | continue 73 | elif module.__class__.__name__ == 'MaxPool2d': 74 | vgg19_avg_pooling.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0)) 75 | else: 76 | vgg19_avg_pooling.append(module) 77 | 78 | if optimized: 79 | self.layers = [3, 8, 17, 26, 35] 80 | else: 81 | self.layers = [1, 3, 6, 8, 11, 13, 15, 17, 20, 22, 24, 26, 29] 82 | 83 | self.vgg19 = nn.Sequential(*vgg19_avg_pooling) 84 | 85 | # print(self.vgg19) 86 | 87 | def normalize_inputs(self, x): 88 | return (x - self.mean_) / self.std_ 89 | 90 | def forward(self, input, target): 91 | loss = 0 92 | 93 | if self.partialconv: 94 | eps = 1e-9 95 | mask = target.sum(1, True) > eps 96 | mask = mask.float() 97 | 98 | features_input = self.normalize_inputs(input) 99 | features_target = self.normalize_inputs(target) 100 | for i, layer in enumerate(self.vgg19): 101 | if isinstance(layer, PartialConv2d): 102 | features_input = layer(features_input, mask) 103 | features_target = layer(features_target, mask) 104 | else: 105 | features_input = layer(features_input) 106 | features_target = layer(features_target) 107 | 108 | if i in self.layers: 109 | loss = loss + F.l1_loss(features_input, features_target) 110 | 111 | return loss 112 | 113 | 114 | class VGGLossMix(nn.Module): 115 | def __init__(self, weight=0.5): 116 | super(VGGLossMix, self).__init__() 117 | self.l1 = VGGLoss() 118 | self.l2 = VGGLoss(net='caffe') 119 | self.weight = weight 120 | 121 | def forward(self, input, target): 122 | return self.l1(input, target)*self.weight + self.l2(input, target) * (1-self.weight) 123 | -------------------------------------------------------------------------------- /READ/datasets/common.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import yaml 3 | 4 | from functools import lru_cache 5 | 6 | import numpy as np 7 | import cv2 8 | 9 | import torch 10 | 11 | 12 | @lru_cache(maxsize = 1000) 13 | def load_image(path): 14 | img = cv2.imread(path) 15 | assert img is not None, f'could not load {path}' 16 | return img[...,::-1].copy() 17 | 18 | 19 | def any2float(img): 20 | if isinstance(img, np.ndarray): 21 | out = img.astype(np.float32) 22 | if img.dtype == np.uint16: 23 | out /= 65535 24 | elif img.dtype == np.uint8: 25 | out /= 255 26 | elif torch.is_tensor(img): 27 | out = img.float() 28 | if img.dtype == torch.int16: 29 | out /= 65535 30 | elif img.dtype == torch.uint8: 31 | out /= 255 32 | else: 33 | raise TypeError('img must be numpy array or torch tensor') 34 | 35 | return out 36 | 37 | 38 | def rescale_K(K_, sx, sy, keep_fov=True): 39 | K = K_.copy() 40 | K[0, 2] = sx * K[0, 2] 41 | K[1, 2] = sy * K[1, 2] 42 | if keep_fov: 43 | K[0, 0] = sx * K[0, 0] 44 | K[1, 1] = sy * K[1, 1] 45 | return K 46 | 47 | def fit_size(x, d=16): 48 | return x[..., :d*(x.shape[-2]//d), :d*(x.shape[-1]//d)] 49 | 50 | 51 | class ToTensor(object): 52 | def __call__(self, img): 53 | if isinstance(img, np.ndarray): 54 | img = torch.from_numpy(img) 55 | 56 | img = any2float(img) 57 | 58 | return img.permute(2, 0, 1).contiguous() 59 | 60 | def __repr__(self): 61 | return self.__class__.__name__ + '()' 62 | 63 | 64 | def get_dataset_config(yml, dataset): 65 | join = os.path.join 66 | 67 | myhost = os.uname()[1] 68 | if myhost in yml: 69 | data_root = yml[myhost]['data_root'] if 'data_root' in yml[myhost] else '' 70 | else: 71 | data_root = '/' 72 | 73 | ds = yml['datasets'][dataset] 74 | 75 | for k in ds: 76 | if 'path' in k: 77 | ds[k] = join(data_root, ds[k]) 78 | 79 | return ds 80 | 81 | 82 | def split_lists(config, lists): 83 | sz = [len(l) for l in lists] 84 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 85 | 86 | splits = [] 87 | train_inds, val_inds = [], [] 88 | if 'train_ratio' in config: 89 | train_ratio = float(config['train_ratio']) 90 | train_n = int(sz[0] * train_ratio) 91 | train_inds, val_inds = np.split(np.random.permutation(sz[0]), [train_n]) 92 | else: 93 | val_step = int(config['val_step']) 94 | train_drop = int(config['train_drop']) 95 | for i in range(sz[0]): 96 | if i % val_step == 0: 97 | val_inds.append(i) 98 | elif train_drop < i % val_step < val_step - train_drop: 99 | train_inds.append(i) 100 | 101 | # print(train_inds) 102 | # print( val_inds) 103 | 104 | for lst in lists: 105 | lst = np.array(lst) 106 | splits.append([lst[train_inds], lst[val_inds]]) 107 | # print('--') 108 | # [print(x) for x in lst[train_inds]] 109 | # print('-') 110 | # [print(x) for x in lst[val_inds]] 111 | 112 | 113 | return splits -------------------------------------------------------------------------------- /READ/datasets/splitter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | 4 | 5 | 6 | def split_by_ratio(lists, train_ratio): 7 | sz = [len(l) for l in lists] 8 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 9 | 10 | splits = [] 11 | train_inds, val_inds = [], [] 12 | 13 | train_n = int(sz[0] * train_ratio) 14 | train_inds, val_inds = np.split(np.random.permutation(sz[0]), [train_n]) 15 | 16 | print(train_inds) 17 | print( val_inds) 18 | 19 | for lst in lists: 20 | lst = np.array(lst) 21 | splits.append([lst[train_inds], lst[val_inds]]) 22 | 23 | return splits 24 | 25 | 26 | 27 | #100frame test 28 | def split_by_step100(lists, val_step, train_drop): 29 | sz = [len(l) for l in lists] 30 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 31 | 32 | splits = [] 33 | train_inds, val_inds = [], [] 34 | #print('lists',lists) 35 | for i in range(sz[0]): 36 | if 5<(i % 100) < 95 or i<10: 37 | train_inds.append(i) 38 | else: 39 | val_inds.append(i) 40 | 41 | print(train_inds) 42 | print( val_inds) 43 | 44 | for lst in lists: 45 | lst = np.array(lst) 46 | splits.append([lst[train_inds], lst[val_inds]]) 47 | 48 | return splits 49 | 50 | 51 | #正常 52 | def split_by_step(lists, val_step, train_drop): 53 | sz = [len(l) for l in lists] 54 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 55 | 56 | splits = [] 57 | train_inds, val_inds = [], [] 58 | 59 | for i in range(sz[0]): 60 | if i % val_step == 0: 61 | val_inds.append(i) 62 | elif train_drop < i % val_step < val_step - train_drop: 63 | train_inds.append(i) 64 | 65 | print(train_inds) 66 | print( val_inds) 67 | 68 | for lst in lists: 69 | lst = np.array(lst) 70 | splits.append([lst[train_inds], lst[val_inds]]) 71 | 72 | return splits 73 | -------------------------------------------------------------------------------- /READ/gl/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__init__.py -------------------------------------------------------------------------------- /READ/gl/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/camera.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/camera.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/camera.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/camera.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/dataset.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/dataset.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/dataset.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/nn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/nn.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/nn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/nn.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/programs.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/programs.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/programs.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/programs.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/render.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/render.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/render.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/render.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /READ/gl/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/gl/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /READ/gl/dataset.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | from READ.gl.programs import NNScene 4 | 5 | 6 | def parse_input_string_obsolete(what): 7 | config = {} 8 | if '_pr' in what: 9 | config['draw_points'] = True 10 | config['splat_mode'] = True 11 | config['point_size'] = int(what.split('_pr')[-1]) 12 | config['flat_color'] = True 13 | elif '_p' in what: 14 | config['draw_points'] = True 15 | config['splat_mode'] = False 16 | config['point_size'] = int(what.split('_p')[-1]) 17 | config['flat_color'] = True 18 | else: 19 | config['draw_points'] = False 20 | config['splat_mode'] = False 21 | config['point_size'] = 1 22 | config['flat_color'] = False 23 | 24 | if 'colors' in what: 25 | config['mode'] = NNScene.MODE_COLOR, 0 26 | elif 'uv' in what: 27 | m1 = NNScene.UV_TYPE_1D if config['draw_points'] else NNScene.UV_TYPE_2D 28 | config['mode'] = NNScene.MODE_UV, m1 29 | elif 'normals' in what: 30 | nm = what.split('_')[-2] if '_p' in what else what.split('_')[-1] 31 | m1 = ['g', 'r', 'l', 'd'].index(nm) 32 | config['mode'] = NNScene.MODE_NORMALS, m1 33 | elif 'xyz' in what: 34 | config['mode'] = NNScene.MODE_XYZ, 0 35 | 36 | return config 37 | 38 | 39 | def parse_input_string(string): 40 | config = {} 41 | 42 | if re.search('^colors', string): 43 | config['mode'] = NNScene.MODE_COLOR, None 44 | elif re.search('^uv', string): 45 | choices = ['uv_1d', 'uv_2d'] 46 | ch = re.findall('|'.join(choices), string)[-1] 47 | m1 = choices.index(ch) 48 | config['mode'] = NNScene.MODE_UV, m1 49 | elif re.search('^normals', string): 50 | choices = ['normals_m', 'normals_r', 'normals_l', 'normals_d'] 51 | ch = re.findall('|'.join(choices), string)[-1] 52 | m1 = choices.index(ch) 53 | config['mode'] = NNScene.MODE_NORMALS, m1 54 | elif re.search('^xyz', string): 55 | config['mode'] = NNScene.MODE_XYZ, None 56 | elif re.search('^depth', string): 57 | config['mode'] = NNScene.MODE_DEPTH, None 58 | elif re.search('^labels', string): 59 | config['mode'] = NNScene.MODE_LABEL, None 60 | else: 61 | raise ValueError(string) 62 | 63 | res = re.findall('ps[0-9]+|p[0-9]+', string) 64 | if res: 65 | res = res[-1] 66 | config['draw_points'] = True 67 | config['flat_color'] = True 68 | config['point_size'] = int(re.search('[0-9]+', res).group()) 69 | config['splat_mode'] = re.search('^ps', res) is not None 70 | else: 71 | config['draw_points'] = False 72 | config['splat_mode'] = False 73 | config['point_size'] = 1 74 | config['flat_color'] = False 75 | 76 | res = re.findall('ds[0-5]+', string) 77 | if res: 78 | res = res[-1] 79 | config['downscale'] = int(re.search('[0-9]+', res).group()) 80 | 81 | 82 | return config 83 | 84 | 85 | def generate_input_string(config): 86 | s = '' 87 | m0, m1 = config['mode'] 88 | if m0 == NNScene.MODE_COLOR: 89 | s += 'colors' 90 | elif m0 == NNScene.MODE_UV: 91 | s += 'uv' 92 | if m1 == NNScene.UV_TYPE_1D: 93 | s += '_1d' 94 | elif m1 == NNScene.UV_TYPE_2D: 95 | s += '_2d' 96 | else: 97 | raise ValueError 98 | elif m0 == NNScene.MODE_NORMALS: 99 | s += 'normals' 100 | if m1 == NNScene.NORMALS_MODE_MODEL: 101 | s += '_m' 102 | elif m1 == NNScene.NORMALS_MODE_REFLECTION: 103 | s += '_r' 104 | elif m1 == NNScene.NORMALS_MODE_LOCAL: 105 | s += '_l' 106 | elif m1 == NNScene.NORMALS_MODE_DIRECTION: 107 | s += '_d' 108 | elif m0 == NNScene.MODE_XYZ: 109 | s += 'xyz' 110 | elif m0 == NNScene.MODE_DEPTH: 111 | s += 'depth' 112 | 113 | if config['draw_points']: 114 | s += '_p' 115 | if config['splat_mode']: 116 | s += 's' 117 | s += str(config['point_size']) 118 | 119 | if 'downscale' in config: 120 | s += f"_ds{config['downscale']}" 121 | 122 | return s 123 | 124 | 125 | 126 | def test_generate_parse(): 127 | configs = [ 128 | { 129 | 'draw_points': True, 130 | 'splat_mode': True, 131 | 'point_size': 20, 132 | 'flat_color': True, 133 | 134 | 'mode': (NNScene.MODE_UV, NNScene.UV_TYPE_1D) 135 | }, 136 | 137 | { 138 | 'draw_points': True, 139 | 'splat_mode': True, 140 | 'point_size': 20, 141 | 'flat_color': True, 142 | 143 | 'mode': (NNScene.MODE_UV, NNScene.UV_TYPE_2D) 144 | }, 145 | 146 | { 147 | 'draw_points': True, 148 | 'splat_mode': False, 149 | 'point_size': 20, 150 | 'flat_color': True, 151 | 152 | 'mode': (NNScene.MODE_COLOR, None) 153 | }, 154 | 155 | { 156 | 'draw_points': False, 157 | 'splat_mode': False, 158 | 'point_size': 1, 159 | 'flat_color': False, 160 | 161 | 'mode': (NNScene.MODE_NORMALS, NNScene.NORMALS_MODE_REFLECTION) 162 | }, 163 | 164 | { 165 | 'draw_points': False, 166 | 'splat_mode': False, 167 | 'point_size': 1, 168 | 'flat_color': False, 169 | 170 | 'mode': (NNScene.MODE_XYZ, None) 171 | }, 172 | 173 | { 174 | 'draw_points': True, 175 | 'splat_mode': False, 176 | 'point_size': 1, 177 | 'flat_color': True, 178 | 179 | 'mode': (NNScene.MODE_XYZ, None), 180 | 181 | 'downscale': 2 182 | }, 183 | ] 184 | 185 | def check(config): 186 | print('config:\n', config) 187 | s = generate_input_string(config) 188 | print('string:\n', s) 189 | c = parse_input_string(s) 190 | print('generated config:\n', c) 191 | print('-- OK' if config == c else '-- FAIL') 192 | 193 | for c in configs: 194 | check(c) 195 | 196 | 197 | if __name__ == '__main__': 198 | test_generate_parse() -------------------------------------------------------------------------------- /READ/gl/nn.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | try: 3 | import torch 4 | except ImportError: 5 | print('torch is not available') 6 | import numpy as np 7 | 8 | from READ.gl.render import OffscreenRender 9 | # from READ.gl.utils import get_net 10 | 11 | from READ.pipelines import load_pipeline 12 | from READ.datasets.dynamic import MultiscaleRender, default_input_transform 13 | 14 | from scipy.ndimage import gaussian_filter 15 | import torch.nn as nn 16 | 17 | 18 | class GaussianLayer(nn.Module): 19 | _instance = None 20 | 21 | def __init__(self, in_channels, out_channels, kernel_size=21, sigma=3): 22 | super(GaussianLayer, self).__init__() 23 | self.seq = nn.Sequential( 24 | nn.ReflectionPad2d(kernel_size//2), 25 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 26 | ) 27 | 28 | self.weights_init(kernel_size, sigma) 29 | 30 | def forward(self, x): 31 | return self.seq(x) 32 | 33 | def weights_init(self, kernel_size, sigma): 34 | n= np.zeros((kernel_size, kernel_size)) 35 | n[kernel_size//2, kernel_size//2] = 1 36 | k = gaussian_filter(n,sigma=sigma) 37 | for name, f in self.named_parameters(): 38 | f.data.copy_(torch.from_numpy(k)) 39 | 40 | @staticmethod 41 | def get_instance(): 42 | if GaussianLayer._instance is None: 43 | GaussianLayer._instance = GaussianLayer(8, 8, kernel_size=13, sigma=6).cuda() 44 | 45 | return GaussianLayer._instance 46 | 47 | 48 | class BoxFilter(nn.Module): 49 | def __init__(self, in_channels, out_channels, kernel_size=3): 50 | super().__init__() 51 | 52 | self.seq = nn.Sequential( 53 | nn.ReflectionPad2d(kernel_size//2), 54 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 55 | ) 56 | 57 | self.weights_init(kernel_size) 58 | 59 | def forward(self, x): 60 | return self.seq(x) 61 | 62 | def weights_init(self, kernel_size): 63 | kernel = torch.ones((kernel_size, kernel_size)) / kernel_size ** 2 64 | print(self.seq[0].named_parameters()) 65 | 66 | 67 | def to_gpu(data): 68 | if isinstance(data, dict): 69 | for k in data: 70 | data[k] = data[k].cuda() 71 | return data 72 | else: 73 | return data.cuda() 74 | 75 | 76 | class OGL: 77 | def __init__(self, scene, scene_data, viewport_size, net_ckpt, texture_ckpt, out_buffer_location='numpy', supersampling=1, gpu=True, clear_color=None, temporal_average=False): 78 | self.gpu = gpu 79 | 80 | args_upd = { 81 | 'inference': True, 82 | } 83 | 84 | if texture_ckpt: 85 | args_upd['texture_ckpt'] = texture_ckpt 86 | if 'pointcloud' in scene_data: 87 | args_upd['n_points'] = scene_data['pointcloud']['xyz'].shape[0] 88 | 89 | pipeline, args = load_pipeline(net_ckpt, args_to_update=args_upd) 90 | 91 | self.model = pipeline.model 92 | 93 | if args.pipeline == 'READ.pipelines.ogl.TexturePipeline': 94 | self.model.load_textures(0) 95 | 96 | if self.gpu: 97 | self.model.cuda() 98 | self.model.eval() 99 | 100 | if supersampling > 1: 101 | self.model.ss = supersampling 102 | 103 | self.model.temporal_average = temporal_average 104 | 105 | print(f"SUPERSAMPLING: {self.model.ss}") 106 | 107 | factor = 16 108 | assert viewport_size[0] % 16 == 0, f'set width {factor * (viewport_size[0] // factor)}' 109 | assert viewport_size[1] % 16 == 0, f'set height {factor * (viewport_size[1] // factor)}' 110 | 111 | self.renderer = MultiscaleRender(scene, args.input_format, viewport_size, out_buffer_location=out_buffer_location, supersampling=self.model.ss, clear_color=clear_color) 112 | 113 | def infer(self): 114 | input_dict = self.renderer.render() 115 | input_dict = {k: default_input_transform(v)[None] for k, v in input_dict.items()} 116 | if self.gpu: 117 | input_dict = to_gpu(input_dict) 118 | 119 | input_dict['id'] = 0 120 | with torch.set_grad_enabled(False): 121 | out, net_input = self.model(input_dict, return_input=True) 122 | 123 | out = out[0].detach().permute(1, 2, 0) 124 | out = torch.cat( [out, out[:, :, :1] * 0 + 1], 2 ).contiguous() # expected to have 4 channels 125 | 126 | return { 127 | 'output': out, 128 | 'net_input': net_input 129 | } 130 | -------------------------------------------------------------------------------- /READ/gl/render.py: -------------------------------------------------------------------------------- 1 | from glumpy import app, gloo, gl 2 | 3 | from contextlib import contextmanager 4 | import numpy as np 5 | 6 | try: 7 | import pycuda.driver 8 | from pycuda.gl import graphics_map_flags, BufferObject 9 | _PYCUDA = True 10 | except ImportError as err: 11 | print('pycuda import error:', err) 12 | _PYCUDA = False 13 | 14 | import torch 15 | 16 | 17 | class OffscreenRender: 18 | def __init__(self, viewport_size, out_buffer_location='opengl', clear_color=None): 19 | self._init_buffers(viewport_size, out_buffer_location) 20 | 21 | self.clear_color = clear_color if clear_color is not None else (0., 0., 0., 1.) 22 | 23 | def _init_buffers(self, viewport_size, out_buffer_location): 24 | assert out_buffer_location in ['torch', 'opengl', 'numpy'] 25 | 26 | if out_buffer_location == 'torch': 27 | assert _PYCUDA, 'pycuda is not available' 28 | try: 29 | import pycuda.gl.autoinit # this may fails in headless mode 30 | except: 31 | raise RuntimeError('PyCUDA init failed, cannot use torch buffer') 32 | 33 | _ = torch.cuda.FloatTensor(1, 3, 512,512) # needs init here, otherwise does not work 34 | 35 | color_np = np.zeros((viewport_size[1], viewport_size[0], 4), np.float32) 36 | self.color_buf, self.color_buf_cuda = create_shared_texture(color_np) 37 | self.out_buf = torch.zeros((viewport_size[1], viewport_size[0], 4), dtype=torch.float32).cuda() 38 | elif out_buffer_location == 'opengl': 39 | self.color_buf = np.zeros((viewport_size[1], viewport_size[0], 4), dtype=np.float32).view(gloo.TextureFloat2D) 40 | self.out_buf = self.color_buf 41 | elif out_buffer_location == 'numpy': 42 | self.color_buf = np.zeros((viewport_size[1], viewport_size[0], 4), dtype=np.float32).view(gloo.TextureFloat2D) 43 | self.out_buf = np.zeros((viewport_size[1], viewport_size[0], 3), dtype=np.float32) 44 | 45 | self.viewport_size = viewport_size 46 | self.out_buffer_location = out_buffer_location 47 | 48 | self.depth_buf = gloo.DepthBuffer(viewport_size[0], viewport_size[1], gl.GL_DEPTH_COMPONENT32) 49 | 50 | self.fbo = gloo.FrameBuffer(color=self.color_buf, depth=self.depth_buf) 51 | 52 | def render(self, scene, cull_face=True): 53 | self.fbo.activate() 54 | 55 | gl.glEnable(gl.GL_PROGRAM_POINT_SIZE) 56 | gl.glEnable(gl.GL_DEPTH_TEST) 57 | gl.glShadeModel(gl.GL_FLAT) 58 | 59 | if cull_face: 60 | gl.glEnable(gl.GL_CULL_FACE) 61 | gl.glCullFace(gl.GL_BACK) 62 | else: 63 | gl.glDisable(gl.GL_CULL_FACE) 64 | 65 | gl.glClearColor(*self.clear_color) 66 | gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) 67 | gl.glViewport(0, 0, self.viewport_size[0], self.viewport_size[1]) 68 | 69 | if scene.draw_points: 70 | scene.program.draw(gl.GL_POINTS) 71 | else: 72 | assert scene.index_buffer is not None 73 | scene.program.draw(gl.GL_TRIANGLES, scene.index_buffer) 74 | 75 | if self.out_buffer_location == 'torch': 76 | frame = cpy_texture_to_tensor(self.color_buf_cuda, self.out_buf).clone() 77 | elif self.out_buffer_location == 'opengl': 78 | frame = self.out_buf 79 | else: 80 | gl.glReadPixels(0, 0, self.viewport_size[0], self.viewport_size[1], gl.GL_RGB, gl.GL_FLOAT, self.out_buf) 81 | frame = self.out_buf.copy() 82 | 83 | self.fbo.deactivate() 84 | 85 | return frame 86 | 87 | 88 | @contextmanager 89 | def cuda_activate_array(img): 90 | """Context manager simplifying use of pycuda.gl.RegisteredImage""" 91 | mapping = img.map() 92 | yield mapping.array(0,0) 93 | mapping.unmap() 94 | 95 | 96 | @contextmanager 97 | def cuda_activate_buffer(buf): 98 | mapping = buf.map() 99 | yield mapping.device_ptr() 100 | mapping.unmap() 101 | 102 | 103 | def create_shared_texture(arr, map_flags=None): 104 | """Create and return a Texture2D with gloo and pycuda views.""" 105 | 106 | if map_flags is None: 107 | map_flags = graphics_map_flags.WRITE_DISCARD 108 | 109 | gl_view = arr.view(gloo.TextureFloat2D) 110 | gl_view.activate() # force gloo to create on GPU 111 | gl_view.deactivate() 112 | 113 | cuda_view = pycuda.gl.RegisteredImage( 114 | int(gl_view.handle), gl_view.target, map_flags) 115 | 116 | return gl_view, cuda_view 117 | 118 | 119 | def create_shared_buffer(arr): 120 | """Create and return a BufferObject with gloo and pycuda views.""" 121 | gl_view = arr.view(gloo.VertexBuffer) 122 | gl_view.activate() # force gloo to create on GPU 123 | gl_view.deactivate() 124 | cuda_view = BufferObject(np.long(gl_view.handle)) 125 | return gl_view, cuda_view 126 | 127 | 128 | def cpy_texture_to_tensor(texture, tensor): 129 | """Copy GL texture (cuda view) to pytorch tensor""" 130 | with cuda_activate_array(texture) as src: 131 | cpy = pycuda.driver.Memcpy2D() 132 | 133 | cpy.set_src_array(src) 134 | cpy.set_dst_device(tensor.data_ptr()) 135 | cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = tensor.shape[1] * 4 * 4 136 | cpy.height = tensor.shape[0] 137 | cpy(aligned=False) 138 | 139 | torch.cuda.synchronize() 140 | 141 | return tensor 142 | 143 | 144 | def cpy_tensor_to_texture(tensor, texture): 145 | """Copy pytorch tensor to GL texture (cuda view)""" 146 | with cuda_activate_array(texture) as ary: 147 | cpy = pycuda.driver.Memcpy2D() 148 | 149 | cpy.set_src_device(tensor.data_ptr()) 150 | cpy.set_dst_array(ary) 151 | cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = tensor.shape[1] * 4 * 4 152 | cpy.height = tensor.shape[0] 153 | cpy(aligned=False) 154 | 155 | torch.cuda.synchronize() 156 | 157 | return tensor 158 | 159 | 160 | def cpy_buffer_to_tensor(buffer, tensor): 161 | """Copy GL buffer (cuda view) to pytorch tensor""" 162 | n = tensor.numel()*tensor.element_size() 163 | with cuda_activate_buffer(buffer) as buf_ptr: 164 | pycuda.driver.memcpy_dtod(tensor.data_ptr(), buf_ptr, n) 165 | 166 | 167 | def cpy_tensor_to_buffer(tensor, buffer): 168 | """Copy pytorch tensor to GL buffer (cuda view)""" 169 | n = tensor.numel()*tensor.element_size() 170 | with cuda_activate_buffer(buffer) as buf_ptr: 171 | pycuda.driver.memcpy_dtod(buf_ptr, tensor.data_ptr(), n) 172 | 173 | -------------------------------------------------------------------------------- /READ/gl/viewer/camera.py: -------------------------------------------------------------------------------- 1 | from glumpy import app, gl, gloo 2 | import numpy as np 3 | 4 | from READ.gl.viewer.trackball import Trackball 5 | from READ.gl.render import window 6 | 7 | assert window is not None, 'call render.get_window before import' 8 | 9 | 10 | _trackball = None 11 | viewer_flags = { 12 | 'mouse_pressed': False, 13 | 'use_perspective_cam': True, 14 | } 15 | 16 | 17 | @window.event 18 | def on_resize(width, height): 19 | _trackball.resize((width, height)) 20 | 21 | 22 | @window.event 23 | def on_mouse_press(x, y, buttons, modifiers): 24 | """Record an initial mouse press. 25 | """ 26 | 27 | print(buttons, modifiers) 28 | _trackball.set_state(Trackball.STATE_ROTATE) 29 | if (buttons == app.window.mouse.LEFT): 30 | print('left') 31 | ctrl = (modifiers & app.window.key.MOD_CTRL) 32 | shift = (modifiers & app.window.key.MOD_SHIFT) 33 | if (ctrl and shift): 34 | _trackball.set_state(Trackball.STATE_ZOOM) 35 | elif ctrl: 36 | _trackball.set_state(Trackball.STATE_ROLL) 37 | elif shift: 38 | _trackball.set_state(Trackball.STATE_PAN) 39 | elif (buttons == app.window.mouse.MIDDLE): 40 | _trackball.set_state(Trackball.STATE_PAN) 41 | elif (buttons == app.window.mouse.RIGHT): 42 | _trackball.set_state(Trackball.STATE_ZOOM) 43 | 44 | _trackball.down(np.array([x, y])) 45 | 46 | # Stop animating while using the mouse 47 | viewer_flags['mouse_pressed'] = True 48 | 49 | 50 | @window.event 51 | def on_mouse_drag(x, y, dx, dy, buttons): 52 | """Record a mouse drag. 53 | """ 54 | _trackball.drag(np.array([x, y])) 55 | 56 | 57 | @window.event 58 | def on_mouse_release(x, y, button, modifiers): 59 | """Record a mouse release. 60 | """ 61 | viewer_flags['mouse_pressed'] = False 62 | 63 | 64 | @window.event 65 | def on_mouse_scroll(x, y, dx, dy): 66 | """Record a mouse scroll. 67 | """ 68 | if viewer_flags['use_perspective_cam']: 69 | _trackball.scroll(dy) 70 | else: 71 | spfc = 0.95 72 | spbc = 1.0 / 0.95 73 | sf = 1.0 74 | if dy > 0: 75 | sf = spfc * dy 76 | elif dy < 0: 77 | sf = - spbc * dy 78 | 79 | c = self._camera_node.camera 80 | xmag = max(c.xmag * sf, 1e-8) 81 | ymag = max(c.ymag * sf, 1e-8 * c.ymag / c.xmag) 82 | c.xmag = xmag 83 | c.ymag = ymag 84 | 85 | 86 | def get_trackball(init_view, viewport_size, rotation_mode=1): 87 | global _trackball 88 | _trackball = Trackball(init_view, viewport_size, 1, rotation_mode=rotation_mode) 89 | 90 | return _trackball -------------------------------------------------------------------------------- /READ/models/__pycache__/common.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/common.cpython-36.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/common.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/common.cpython-38.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/compose.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/compose.cpython-36.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/compose.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/compose.cpython-38.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/conv.cpython-36.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/conv.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/conv.cpython-38.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/texture.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/texture.cpython-36.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/texture.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/texture.cpython-38.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/unet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/unet.cpython-36.pyc -------------------------------------------------------------------------------- /READ/models/__pycache__/unet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/models/__pycache__/unet.cpython-38.pyc -------------------------------------------------------------------------------- /READ/models/compose.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.ndimage import gaussian_filter 3 | 4 | import torch 5 | import torch.nn as nn 6 | import imageio 7 | import cv2 8 | from PIL import Image 9 | 10 | 11 | 12 | class ModelAndLoss(nn.Module): 13 | def __init__(self, model, loss, use_mask=False): 14 | super().__init__() 15 | self.model = model 16 | self.loss = loss 17 | self.use_mask = use_mask 18 | 19 | def forward(self, *args, **kwargs): 20 | input = args[:-1] 21 | target = args[-1] 22 | if not isinstance(input, (tuple, list)): 23 | input = [input] 24 | 25 | output = self.model(*input, **kwargs) 26 | 27 | if self.use_mask and 'mask' in kwargs and kwargs['mask'] is not None: 28 | loss = self.loss(output * kwargs['mask'], target) 29 | else: 30 | loss = self.loss(output, target) 31 | 32 | return output, loss 33 | 34 | 35 | class BoxFilter(nn.Module): 36 | def __init__(self, in_channels, out_channels, kernel_size=3): 37 | super().__init__() 38 | 39 | self.seq = nn.Sequential( 40 | nn.ReflectionPad2d(kernel_size//2), 41 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None) 42 | ) 43 | 44 | self.weights_init(kernel_size) 45 | 46 | def forward(self, x): 47 | return self.seq(x) 48 | 49 | def weights_init(self, kernel_size): 50 | kernel = torch.ones((kernel_size, kernel_size)) / kernel_size ** 2 51 | self.seq[1].weight.data.copy_(kernel) 52 | 53 | 54 | class GaussianLayer(nn.Module): 55 | _instance = None 56 | 57 | def __init__(self, in_channels, out_channels, kernel_size=1, sigma=3): 58 | super(GaussianLayer, self).__init__() 59 | self.seq = nn.Sequential( 60 | #nn.ReflectionPad2d(kernel_size//2), 61 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 62 | ) 63 | 64 | self.weights_init(kernel_size, sigma) 65 | 66 | def forward(self, x): 67 | return self.seq(x) 68 | 69 | def weights_init(self, kernel_size, sigma): 70 | n= np.zeros((kernel_size, kernel_size)) 71 | n[kernel_size//2, kernel_size//2] = 1 72 | k = gaussian_filter(n,sigma=sigma) 73 | for name, f in self.named_parameters(): 74 | f.data.copy_(torch.from_numpy(k)) 75 | 76 | @staticmethod 77 | def get_instance(): 78 | if GaussianLayer._instance is None: 79 | GaussianLayer._instance = GaussianLayer(8, 8, kernel_size=13, sigma=6).cuda() 80 | 81 | return GaussianLayer._instance 82 | 83 | 84 | class NetAndTexture(nn.Module): 85 | def __init__(self, net, textures, supersampling=1, temporal_average=False): 86 | super().__init__() 87 | 88 | self.net = net 89 | self.ss = supersampling 90 | 91 | try: 92 | textures = dict(textures) 93 | except TypeError: 94 | textures = {0: textures} 95 | 96 | self._textures = {k: v.cpu() for k, v in textures.items()} 97 | self._loaded_textures = [] 98 | 99 | self.last_input = None 100 | self.temporal_average = temporal_average 101 | 102 | def load_textures(self, texture_ids): 103 | if torch.is_tensor(texture_ids): 104 | texture_ids = texture_ids.cpu().tolist() 105 | elif isinstance(texture_ids, int): 106 | texture_ids = [texture_ids] 107 | 108 | for tid in texture_ids: 109 | self._modules[str(tid)] = self._textures[tid] 110 | 111 | self._loaded_textures = texture_ids 112 | 113 | def unload_textures(self): 114 | for tid in self._loaded_textures: 115 | self._modules[str(tid)].cpu() 116 | del self._modules[str(tid)] 117 | 118 | def reg_loss(self): 119 | loss = 0 120 | for tid in self._loaded_textures: 121 | loss += self._modules[str(tid)].reg_loss() 122 | 123 | return loss 124 | 125 | def forward(self, inputs, **kwargs): 126 | out = [] 127 | 128 | texture_ids = inputs['id'] 129 | del inputs['id'] 130 | 131 | if torch.is_tensor(texture_ids): 132 | texture_ids = texture_ids.tolist() 133 | elif isinstance(texture_ids, int): 134 | texture_ids = [texture_ids] 135 | 136 | for i, tid in enumerate(texture_ids): # per item in batch 137 | input = {k: v[i][None] for k, v in inputs.items()} 138 | assert 'uv' in list(input)[0], 'first input must be uv' 139 | 140 | texture = self._modules[str(tid)] 141 | 142 | j = 0 143 | keys = list(input) 144 | input_multiscale = [] 145 | while j < len(keys): # sample texture at multiple scales 146 | tex_sample = None 147 | input_ex = [] 148 | if 'uv' in keys[j]: 149 | tex_sample = texture(input[keys[j]]) 150 | j += 1 151 | while j < len(keys) and 'uv' not in keys[j]: 152 | input_ex.append(input[keys[j]]) 153 | j += 1 154 | assert tex_sample is not None 155 | 156 | input_cat = torch.cat(input_ex + [tex_sample], 1) 157 | 158 | #filter = GaussianLayer(input_cat.shape[1], input_cat.shape[1]).cuda() 159 | 160 | #input_cat = filter(input_cat) 161 | 162 | if self.ss > 1: 163 | input_cat = nn.functional.interpolate(input_cat, scale_factor=1./self.ss, mode='bilinear') 164 | 165 | input_multiscale.append(input_cat) 166 | 167 | if self.temporal_average: 168 | if self.last_input is not None: 169 | for i in range(len(input_multiscale)): 170 | input_multiscale[i] = (input_multiscale[i] + self.last_input[i]) / 2 171 | self.last_input = list(input_multiscale) 172 | 173 | out1 = self.net(*input_multiscale, **kwargs) 174 | out.append(out1) 175 | 176 | out = torch.cat(out, 0) 177 | 178 | if kwargs.get('return_input'): 179 | return out, input_multiscale 180 | else: 181 | return out 182 | 183 | 184 | class MultiscaleNet(nn.Module): 185 | def __init__(self, net, input_modality, supersampling=1): 186 | super().__init__() 187 | 188 | self.net = net 189 | self.input_modality = input_modality 190 | self.ss = supersampling 191 | 192 | def forward(self, inputs, **kwargs): 193 | del inputs['id'] 194 | 195 | modes = len(inputs) 196 | assert modes % self.input_modality == 0 197 | 198 | inputs_ms = [] 199 | input_values = list(inputs.values()) 200 | for i in range(modes // self.input_modality): 201 | i0 = i * self.input_modality 202 | i1 = (i + 1) * self.input_modality 203 | cat = torch.cat(input_values[i0:i1], 1) 204 | if self.ss > 1: 205 | cat = nn.functional.interpolate(cat, scale_factor=1./self.ss, mode='bilinear') 206 | inputs_ms.append(cat) 207 | 208 | out = self.net(*inputs_ms, **kwargs) 209 | 210 | if kwargs.get('return_input'): 211 | return out, inputs_ms 212 | else: 213 | return out 214 | 215 | 216 | class RGBTexture(nn.Module): 217 | def __init__(self, texture, supersampling=1): 218 | super().__init__() 219 | 220 | self.texture = texture 221 | self.ss = supersampling 222 | 223 | def forward(self, inputs, **kwargs): 224 | del inputs['id'] 225 | 226 | assert list(inputs) == ['uv_2d'], 'check input format' 227 | 228 | uv = inputs['uv_2d'] 229 | out = self.texture(uv) 230 | 231 | if kwargs.get('return_input'): 232 | return out, uv 233 | else: 234 | return out 235 | -------------------------------------------------------------------------------- /READ/models/conv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch import nn, cuda 4 | from torch.autograd import Variable 5 | 6 | 7 | ############################################################################### 8 | # BSD 3-Clause License 9 | # 10 | # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. 11 | # 12 | # Author & Contact: Guilin Liu (guilinl@nvidia.com) 13 | ############################################################################### 14 | 15 | class PartialConv2d(nn.Conv2d): 16 | def __init__(self, *args, **kwargs): 17 | 18 | # whether the mask is multi-channel or not 19 | if 'multi_channel' in kwargs: 20 | self.multi_channel = kwargs['multi_channel'] 21 | kwargs.pop('multi_channel') 22 | else: 23 | self.multi_channel = False 24 | 25 | if 'return_mask' in kwargs: 26 | self.return_mask = kwargs['return_mask'] 27 | kwargs.pop('return_mask') 28 | else: 29 | self.return_mask = False 30 | 31 | super(PartialConv2d, self).__init__(*args, **kwargs) 32 | 33 | if self.multi_channel: 34 | self.weight_maskUpdater = torch.ones(self.out_channels, self.in_channels, self.kernel_size[0], self.kernel_size[1]) 35 | else: 36 | self.weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0], self.kernel_size[1]) 37 | 38 | self.slide_winsize = self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2] * self.weight_maskUpdater.shape[3] 39 | 40 | self.last_size = (None, None) 41 | self.update_mask = None 42 | self.mask_ratio = None 43 | 44 | def forward(self, input, mask_in=None): 45 | 46 | if mask_in is not None or self.last_size != (input.data.shape[2], input.data.shape[3]): 47 | self.last_size = (input.data.shape[2], input.data.shape[3]) 48 | 49 | with torch.no_grad(): 50 | if self.weight_maskUpdater.type() != input.type(): 51 | self.weight_maskUpdater = self.weight_maskUpdater.to(input) 52 | 53 | if mask_in is None: 54 | # if mask is not provided, create a mask 55 | if self.multi_channel: 56 | mask = torch.ones(input.data.shape[0], input.data.shape[1], input.data.shape[2], input.data.shape[3]).to(input) 57 | else: 58 | mask = torch.ones(1, 1, input.data.shape[2], input.data.shape[3]).to(input) 59 | else: 60 | mask = mask_in 61 | 62 | self.update_mask = F.conv2d(mask, self.weight_maskUpdater, bias=None, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=1) 63 | 64 | self.mask_ratio = self.slide_winsize/(self.update_mask + 1e-8) 65 | # self.mask_ratio = torch.max(self.update_mask)/(self.update_mask + 1e-8) 66 | self.update_mask = torch.clamp(self.update_mask, 0, 1) 67 | self.mask_ratio = torch.mul(self.mask_ratio, self.update_mask) 68 | 69 | # if self.update_mask.type() != input.type() or self.mask_ratio.type() != input.type(): 70 | # self.update_mask.to(input) 71 | # self.mask_ratio.to(input) 72 | 73 | raw_out = super(PartialConv2d, self).forward(torch.mul(input, mask) if mask_in is not None else input) 74 | 75 | if self.bias is not None: 76 | bias_view = self.bias.view(1, self.out_channels, 1, 1) 77 | output = torch.mul(raw_out - bias_view, self.mask_ratio) + bias_view 78 | output = torch.mul(output, self.update_mask) 79 | else: 80 | output = torch.mul(raw_out, self.mask_ratio) 81 | 82 | 83 | if self.return_mask: 84 | return output, self.update_mask 85 | else: 86 | return output 87 | -------------------------------------------------------------------------------- /READ/models/texture.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | 6 | class Texture(nn.Module): 7 | def null_grad(self): 8 | raise NotImplementedError() 9 | 10 | def reg_loss(self): 11 | return 0. 12 | 13 | 14 | class PointTexture(Texture): 15 | def __init__(self, num_channels, size, activation='none', checkpoint=None, init_method='zeros', reg_weight=0.): 16 | super().__init__() 17 | 18 | assert isinstance(size, int), 'size must be int' 19 | 20 | shape = 1, num_channels, size 21 | 22 | if checkpoint: 23 | self.texture_ = torch.load(checkpoint, map_location='cpu')['texture'].texture_ 24 | else: 25 | if init_method == 'rand': 26 | texture = torch.rand(shape) 27 | elif init_method == 'zeros': 28 | texture = torch.zeros(shape) 29 | else: 30 | raise ValueError(init_method) 31 | self.texture_ = nn.Parameter(texture.float()) 32 | 33 | self.activation = activation 34 | self.reg_weight = reg_weight 35 | 36 | def null_grad(self): 37 | self.texture_.grad = None 38 | 39 | def reg_loss(self): 40 | return self.reg_weight * torch.mean(torch.pow(self.texture_, 2)) 41 | 42 | def forward(self, inputs): 43 | if isinstance(inputs, dict): 44 | ids = None 45 | for f, x in inputs.items(): 46 | if 'uv' in f: 47 | ids = x[:, 0].long() 48 | assert ids is not None, 'Input format does not have uv' 49 | else: 50 | ids = inputs[:, 0] # BxHxW 51 | 52 | sh = ids.shape 53 | n_pts = self.texture_.shape[-1] 54 | 55 | ind = ids.contiguous().view(-1).long() 56 | 57 | texture = self.texture_.permute(1, 0, 2) # Cx1xN 58 | texture = texture.expand(texture.shape[0], sh[0], texture.shape[2]) # CxBxN 59 | texture = texture.contiguous().view(texture.shape[0], -1) # CxB*N 60 | 61 | sample = torch.index_select(texture, 1, ind) # CxB*H*W 62 | sample = sample.contiguous().view(sample.shape[0], sh[0], sh[1], sh[2]) # CxBxHxW 63 | sample = sample.permute(1, 0, 2, 3) # BxCxHxW 64 | 65 | if self.activation == 'sigmoid': 66 | return torch.sigmoid(sample) 67 | elif self.activation == 'tanh': 68 | return torch.tanh(sample) 69 | #print('sample,',sample.shape) 70 | return sample 71 | 72 | 73 | class MeshTexture(Texture): 74 | def __init__(self, num_channels, size, activation='none', init_method='zeros', levels=4, reg_weight=0.): 75 | super().__init__() 76 | 77 | assert isinstance(size, int), f'size must be int not {size}' 78 | 79 | if init_method == 'rand': 80 | init = lambda shape: torch.rand(shape) 81 | elif init_method == 'zeros': 82 | init = lambda shape: torch.zeros(shape) 83 | elif init_method == '0.5': 84 | init = lambda shape: torch.zeros(shape) + 0.5 85 | else: 86 | raise ValueError(init_method) 87 | 88 | assert levels > 0 89 | self.levels = levels 90 | 91 | for i in range(self.levels): 92 | shape = 1, num_channels, size // 2 ** i, size // 2 ** i 93 | tex = nn.Parameter(init(shape)).float() 94 | self.__setattr__(f'texture_{i}', tex) 95 | 96 | self.activation = activation 97 | self.reg_weight = reg_weight 98 | 99 | def null_grad(self): 100 | for i in range(self.levels): 101 | self.__getattr__(f'texture_{i}').grad = None 102 | 103 | def reg_loss(self): 104 | loss = 0. 105 | tex_weight = [8., 2., 1., 0.] 106 | for i in range(self.levels): 107 | tex = self.__getattr__(f'texture_{i}') 108 | loss += self.reg_weight * tex_weight[i] * torch.mean(torch.pow(tex, 2)) 109 | 110 | return loss 111 | 112 | def forward(self, inputs): 113 | uv = (inputs[:, :2] * 2 - 1.0).transpose(1, 3).transpose(1, 2).contiguous() 114 | 115 | samples = [] 116 | for i in range(self.levels): 117 | tex = self.__getattr__(f'texture_{i}') 118 | sample = nn.functional.grid_sample(tex.expand(uv.shape[0], -1, -1, -1), uv) 119 | samples.append(sample) 120 | 121 | out = samples[0] 122 | for i in range(1, self.levels): 123 | out += samples[i] 124 | 125 | if self.activation == 'sigmoid': 126 | return torch.sigmoid(out) 127 | elif self.activation == 'tanh': 128 | return torch.tanh(out) 129 | 130 | return out 131 | -------------------------------------------------------------------------------- /READ/models/unet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | import numpy as np 6 | from functools import partial 7 | 8 | 9 | 10 | 11 | class ResBlock(nn.Module): 12 | def __init__(self, in_channel, out_channel): 13 | super(ResBlock, self).__init__() 14 | self.main = nn.Sequential( 15 | BasicConv(in_channel, out_channel, kernel_size=3, stride=1, relu=True), 16 | BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False) 17 | ) 18 | 19 | def forward(self, x): 20 | return self.main(x) + x 21 | 22 | class BasicConv(nn.Module): 23 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, relu=True, dilation=1, padding_mode='reflect', act_fun=nn.ELU, normalization=nn.BatchNorm2d): 24 | super().__init__() 25 | self.pad_mode = padding_mode 26 | self.filter_size = kernel_size 27 | self.stride = stride 28 | self.dilation = dilation 29 | 30 | n_pad_pxl = int(self.dilation * (self.filter_size - 1) / 2) 31 | self.flag=relu 32 | 33 | # this is for backward campatibility with older model checkpoints 34 | self.block = nn.ModuleDict( 35 | { 36 | 'conv_f': nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=n_pad_pxl), 37 | 'act_f': act_fun(), 38 | 'conv_m': nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=n_pad_pxl), 39 | 'act_m': nn.Sigmoid(), 40 | 'norm': normalization(out_channels) 41 | } 42 | ) 43 | 44 | def forward(self, x, *args, **kwargs): 45 | if self.flag: 46 | features = self.block.act_f(self.block.conv_f(x)) 47 | else: 48 | features = self.block.conv_f(x) 49 | mask = self.block.act_m(self.block.conv_m(x)) 50 | output = features * mask 51 | output = self.block.norm(output) 52 | 53 | return output 54 | 55 | 56 | class EBlock(nn.Module): 57 | def __init__(self, out_channel, num_res=8): 58 | super(EBlock, self).__init__() 59 | 60 | layers = [ResBlock(out_channel, out_channel) for _ in range(num_res)] 61 | 62 | self.layers = nn.Sequential(*layers) 63 | 64 | def forward(self, x): 65 | return self.layers(x) 66 | 67 | 68 | class DBlock(nn.Module): 69 | def __init__(self, channel, num_res=8): 70 | super(DBlock, self).__init__() 71 | 72 | layers = [ResBlock(channel, channel) for _ in range(num_res)] 73 | self.layers = nn.Sequential(*layers) 74 | 75 | def forward(self, x): 76 | return self.layers(x) 77 | 78 | 79 | class AFF(nn.Module): 80 | def __init__(self, in_channel, out_channel): 81 | super(AFF, self).__init__() 82 | self.conv = nn.Sequential( 83 | BasicConv(in_channel, out_channel, kernel_size=1, stride=1, relu=True), 84 | BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False) 85 | ) 86 | 87 | def forward(self, x1, x2, x3, x4): 88 | x = torch.cat([x1, x2, x3, x4], dim=1) 89 | return self.conv(x) 90 | 91 | 92 | class SCM(nn.Module): 93 | def __init__(self, out_plane): 94 | super(SCM, self).__init__() 95 | self.main = nn.Sequential( 96 | BasicConv(8, out_plane//4, kernel_size=3, stride=1, relu=True), 97 | BasicConv(out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True), 98 | BasicConv(out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True), 99 | BasicConv(out_plane // 2, out_plane-8, kernel_size=1, stride=1, relu=True) 100 | ) 101 | 102 | self.conv = BasicConv(out_plane, out_plane, kernel_size=1, stride=1, relu=False) 103 | 104 | def forward(self, x): 105 | x = torch.cat([x, self.main(x)], dim=1) 106 | return self.conv(x) 107 | 108 | 109 | class FAM(nn.Module): 110 | def __init__(self, channel): 111 | super(FAM, self).__init__() 112 | self.merge = BasicConv(channel, channel, kernel_size=3, stride=1, relu=False) 113 | 114 | def forward(self, x1, x2): 115 | x = x1 * x2 116 | out = x1 + self.merge(x) 117 | return out 118 | 119 | 120 | 121 | class UNet(nn.Module): 122 | r""" Rendering network with UNet architecture and multi-scale input. 123 | 124 | Args: 125 | num_input_channels: Number of channels in the input tensor or list of tensors. An integer or a list of integers for each input tensor. 126 | num_output_channels: Number of output channels. 127 | feature_scale: Division factor of number of convolutional channels. The bigger the less parameters in the model. 128 | num_res: Number of block resnet. 129 | """ 130 | def __init__( 131 | self, 132 | num_input_channels=8, 133 | num_output_channels=3, 134 | feature_scale=4, 135 | num_res=4 136 | 137 | ): 138 | super().__init__() 139 | 140 | self.feature_scale = feature_scale 141 | base_channel = 32 142 | 143 | 144 | filters = [64, 128, 256, 512, 1024] 145 | filters = [x // self.feature_scale for x in filters] 146 | 147 | 148 | 149 | self.Encoder = nn.ModuleList([ 150 | EBlock(base_channel, num_res), 151 | EBlock(base_channel*2, num_res), 152 | EBlock(base_channel*4, num_res), 153 | EBlock(base_channel*8, num_res) 154 | ]) 155 | 156 | self.feat_extract = nn.ModuleList([ 157 | BasicConv(8, base_channel, kernel_size=3, relu=True, stride=1), 158 | BasicConv(base_channel, base_channel*2, kernel_size=3, relu=True, stride=2), 159 | BasicConv(base_channel*2, base_channel*4, kernel_size=3, relu=True, stride=2), 160 | BasicConv(base_channel*4, base_channel*2, kernel_size=4, relu=True, stride=2), 161 | BasicConv(base_channel*2, base_channel, kernel_size=4, relu=True, stride=2), 162 | BasicConv(base_channel, 3, kernel_size=3, relu=False, stride=1), 163 | BasicConv(base_channel*4, base_channel*8, kernel_size=3, relu=True, stride=2), 164 | BasicConv(base_channel*8, base_channel*4, kernel_size=4, relu=True, stride=2), 165 | ]) 166 | 167 | self.Decoder = nn.ModuleList([ 168 | DBlock(base_channel * 8, num_res), 169 | DBlock(base_channel * 4, num_res), 170 | DBlock(base_channel * 2, num_res), 171 | DBlock(base_channel, num_res) 172 | ]) 173 | 174 | self.Convs = nn.ModuleList([ 175 | BasicConv(base_channel * 8, base_channel * 4, kernel_size=1, relu=True, stride=1), 176 | BasicConv(base_channel * 4, base_channel * 2, kernel_size=1, relu=True, stride=1), 177 | BasicConv(base_channel * 2, base_channel, kernel_size=1, relu=True, stride=1) 178 | 179 | ]) 180 | 181 | self.ConvsOut = nn.ModuleList( 182 | [ 183 | BasicConv(base_channel * 4, 3, kernel_size=3, relu=False, stride=1), 184 | BasicConv(base_channel * 2, 3, kernel_size=3, relu=False, stride=1), 185 | ] 186 | ) 187 | 188 | self.AFFs = nn.ModuleList([ 189 | AFF(base_channel * 15, base_channel*1), 190 | AFF(base_channel * 15, base_channel*2), 191 | AFF(base_channel * 15, base_channel*4), 192 | ]) 193 | 194 | self.FAM1 = FAM(base_channel * 4) 195 | self.SCM1 = SCM(base_channel * 4) 196 | self.FAM2 = FAM(base_channel * 2) 197 | self.SCM2 = SCM(base_channel * 2) 198 | self.FAM0 = FAM(base_channel * 8) 199 | self.SCM0 = SCM(base_channel * 8) 200 | self.up =nn.Upsample(scale_factor=4, mode='bilinear') 201 | 202 | def forward(self, *inputs, **kwargs): 203 | inputs = list(inputs) 204 | 205 | 206 | 207 | n_input = len(inputs) 208 | 209 | x =inputs[0] 210 | x_2 = inputs[1] 211 | x_4 = inputs[2] 212 | x_8 = inputs[3] 213 | 214 | z2 = self.SCM2(x_2) 215 | z4 = self.SCM1(x_4) 216 | z8 = self.SCM0(x_8) 217 | 218 | 219 | x_ = self.feat_extract[0](x) 220 | res1 = self.Encoder[0](x_) 221 | 222 | 223 | 224 | z = self.feat_extract[1](res1) 225 | z = self.FAM2(z, z2) 226 | res2 = self.Encoder[1](z) 227 | 228 | z = self.feat_extract[2](res2) 229 | z = self.FAM1(z, z4) 230 | res3 = self.Encoder[2](z) 231 | 232 | z = self.feat_extract[6](res3) 233 | 234 | z = self.FAM0(z, z8) 235 | z = self.Encoder[3](z) 236 | 237 | 238 | 239 | z12 = F.interpolate(res1, scale_factor=0.5) 240 | z13 = F.interpolate(res1, scale_factor=0.25) 241 | 242 | z21 = F.interpolate(res2, scale_factor=2) 243 | z23 = F.interpolate(res2, scale_factor=0.5) 244 | 245 | z32 = F.interpolate(res3, scale_factor=2) 246 | z31 = F.interpolate(res3, scale_factor=4) 247 | 248 | z43 = F.interpolate(z, scale_factor=2) 249 | z42 = F.interpolate(z43, scale_factor=2) 250 | z41 = F.interpolate(z42, scale_factor=2) 251 | 252 | res1 = self.AFFs[0](res1, z21, z31, z41)#AFF1 253 | res2 = self.AFFs[1](z12, res2, z32, z42) #AFF2 254 | res3 = self.AFFs[2](z13, z23, res3, z43)#AFF3 255 | 256 | 257 | z = self.Decoder[0](z) 258 | 259 | 260 | z = self.feat_extract[7](z) 261 | z= self.up(z) 262 | 263 | z = torch.cat([z, res3], dim=1) 264 | z = self.Convs[0](z) 265 | z = self.Decoder[1](z) 266 | 267 | 268 | z = self.feat_extract[3](z) 269 | z= self.up(z) 270 | 271 | z = torch.cat([z, res2], dim=1) 272 | z = self.Convs[1](z) 273 | z = self.Decoder[2](z) 274 | 275 | 276 | z = self.feat_extract[4](z) 277 | z= self.up(z) 278 | 279 | z = torch.cat([z, res1], dim=1) 280 | z = self.Convs[2](z) 281 | z = self.Decoder[3](z) 282 | z = self.feat_extract[5](z) 283 | 284 | 285 | return z 286 | -------------------------------------------------------------------------------- /READ/pipelines/__init__.py: -------------------------------------------------------------------------------- 1 | from . pipeline import Pipeline, load_pipeline, save_pipeline -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/ogl.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/ogl.cpython-36.pyc -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/ogl.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/ogl.cpython-38.pyc -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/pipeline.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/pipeline.cpython-36.pyc -------------------------------------------------------------------------------- /READ/pipelines/__pycache__/pipeline.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/pipelines/__pycache__/pipeline.cpython-38.pyc -------------------------------------------------------------------------------- /READ/pipelines/ogl.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | from pathlib import Path 3 | 4 | from torch import autograd, optim 5 | 6 | from READ.pipelines import Pipeline 7 | from READ.datasets.dynamic import get_datasets 8 | from READ.models.texture import PointTexture, MeshTexture 9 | from READ.models.unet import UNet 10 | from READ.models.compose import NetAndTexture, MultiscaleNet, RGBTexture 11 | from READ.criterions.vgg_loss import VGGLoss 12 | from READ.utils.train import to_device, set_requires_grad, save_model, unwrap_model, image_grid, to_numpy, load_model_checkpoint, freeze 13 | from READ.utils.perform import TicToc, AccumDict, Tee 14 | 15 | 16 | TextureOptimizerClass = optim.RMSprop 17 | 18 | 19 | def get_net(input_channels, args): 20 | net = UNet( 21 | num_input_channels=8, 22 | num_output_channels=3, 23 | feature_scale=4, 24 | num_res=4 25 | ) 26 | 27 | return net 28 | 29 | 30 | def get_texture(num_channels, size, args): 31 | if not hasattr(args, 'reg_weight'): 32 | args.reg_weight = 0. 33 | 34 | if args.use_mesh: 35 | texture = MeshTexture(num_channels, size, activation=args.texture_activation, reg_weight=args.reg_weight) 36 | else: 37 | texture = PointTexture(num_channels, size, activation=args.texture_activation, reg_weight=args.reg_weight) 38 | 39 | if args.texture_ckpt: 40 | texture = load_model_checkpoint(args.texture_ckpt, texture) 41 | 42 | return texture 43 | 44 | 45 | def backward_compat(args): 46 | if not hasattr(args, 'input_channels'): 47 | args.input_channels = None 48 | if not hasattr(args, 'conv_block'): 49 | args.conv_block = 'gated' 50 | 51 | if args.pipeline == 'READ.pipelines.ogl.Pix2PixPipeline': 52 | if not hasattr(args, 'input_modality'): 53 | args.input_modality = 1 54 | 55 | return args 56 | 57 | 58 | class TexturePipeline(Pipeline): 59 | def export_args(self, parser): 60 | parser.add_argument('--descriptor_size', type=int, default=8) 61 | parser.add_argument('--texture_size', type=int) 62 | parser.add_argument('--texture_ckpt', type=Path) 63 | parser.add('--texture_lr', type=float, default=1e-1) 64 | parser.add('--texture_activation', type=str, default='none') 65 | parser.add('--n_points', type=int, default=0, help='this is for inference') 66 | 67 | def create(self, args): 68 | args = backward_compat(args) 69 | 70 | if not args.input_channels: 71 | args.input_channels = [args.descriptor_size] * args.num_mipmap 72 | 73 | net = get_net(args.input_channels, args) 74 | 75 | textures = {} 76 | 77 | if args.inference: 78 | if args.use_mesh: 79 | size = args.texture_size 80 | else: 81 | size = args.n_points 82 | textures = { 83 | 0: get_texture(args.descriptor_size, size, args) 84 | } 85 | else: 86 | self.ds_train, self.ds_val = get_datasets(args) 87 | 88 | for ds in self.ds_train: 89 | if args.use_mesh: 90 | assert args.texture_size, 'set texture size' 91 | size = args.texture_size 92 | else: 93 | assert ds.scene_data['pointcloud'] is not None, 'set pointcloud' 94 | size = ds.scene_data['pointcloud']['xyz'].shape[0] 95 | textures[ds.id] = get_texture(args.descriptor_size, size, args) 96 | 97 | self.optimizer = optim.Adam(net.parameters(), lr=args.lr) 98 | 99 | if len(textures) == 1: 100 | self._extra_optimizer = TextureOptimizerClass(textures[0].parameters(), lr=args.texture_lr) 101 | else: 102 | self._extra_optimizer = None 103 | 104 | self.criterion = args.criterion_module(**args.criterion_args).cuda() 105 | 106 | ss = args.supersampling if hasattr(args, 'supersampling') else 1 107 | 108 | self.net = net 109 | self.textures = textures 110 | self.model = NetAndTexture(net, textures, ss) 111 | 112 | self.args = args 113 | 114 | def state_objects(self): 115 | datasets = self.ds_train 116 | 117 | objs = {'net': self.net} 118 | objs.update({ds.name: self.textures[ds.id] for ds in datasets}) 119 | 120 | return objs 121 | 122 | def dataset_load(self, dataset): 123 | self.model.load_textures([ds.id for ds in dataset]) 124 | 125 | for ds in dataset: 126 | ds.load() 127 | 128 | 129 | def extra_optimizer(self, dataset): 130 | # if we have single dataset, don't recreate optimizer 131 | if self._extra_optimizer is not None: 132 | lr_drop = self.optimizer.param_groups[0]['lr'] / self.args.lr 133 | self._extra_optimizer.param_groups[0]['lr'] = self.args.texture_lr * lr_drop 134 | return self._extra_optimizer 135 | 136 | param_group = [] 137 | for ds in dataset: 138 | param_group.append( 139 | {'params': self.textures[ds.id].parameters()} 140 | ) 141 | 142 | lr_drop = self.optimizer.param_groups[0]['lr'] / self.args.lr 143 | 144 | return TextureOptimizerClass(param_group, lr=self.args.texture_lr * lr_drop) 145 | 146 | def dataset_unload(self, dataset): 147 | self.model.unload_textures() 148 | 149 | for ds in dataset: 150 | ds.unload() 151 | self.textures[ds.id].null_grad() 152 | 153 | def get_net(self): 154 | return self.net 155 | 156 | 157 | class Pix2PixPipeline(Pipeline): 158 | def export_args(self, parser): 159 | parser.add('--input_modality', type=int, default=1) 160 | 161 | def create(self, args): 162 | args = backward_compat(args) 163 | 164 | if not args.input_channels: 165 | print('Assume input channels is 3') 166 | args.input_channels = [3] * args.num_mipmap 167 | 168 | net = get_net(args.input_channels, args) 169 | 170 | self.model = MultiscaleNet(net, args.input_modality) 171 | self.net = net 172 | 173 | if not args.inference: 174 | self.ds_train, self.ds_val = get_datasets(args) 175 | 176 | 177 | self.optimizer= optim.Adam(self.model.parameters(), lr=args.lr) 178 | 179 | 180 | self.criterion = args.criterion_module(**args.criterion_args).cuda() 181 | 182 | def state_objects(self): 183 | return {'net': self.net} 184 | 185 | def dataset_load(self, dataset): 186 | for ds in dataset: 187 | ds.load() 188 | 189 | 190 | def dataset_unload(self, dataset): 191 | for ds in dataset: 192 | ds.unload() 193 | 194 | 195 | def get_net(self): 196 | return self.net 197 | 198 | 199 | class RGBTexturePipeline(Pipeline): 200 | def export_args(self, parser): 201 | parser.add('--texture_size', type=int, default=2048) 202 | parser.add('--texture_lr', type=float, default=1e-2) 203 | 204 | def create(self, args): 205 | self.texture = MeshTexture(3, args.texture_size, activation='none', levels=1, reg_weight=0) 206 | self.model = RGBTexture(self.texture) 207 | 208 | if not args.inference: 209 | self.ds_train, self.ds_val = get_datasets(args) 210 | 211 | self.optimizer = TextureOptimizerClass(self.texture.parameters(), lr=args.texture_lr) 212 | 213 | self.criterion = args.criterion_module(**args.criterion_args).cuda() 214 | 215 | def dataset_load(self, dataset): 216 | for ds in dataset: 217 | ds.load() 218 | 219 | def dataset_unload(self, dataset): 220 | for ds in dataset: 221 | ds.unload() 222 | 223 | def state_objects(self): 224 | return {'model': self.model} 225 | 226 | def get_net(self): 227 | return self.model 228 | -------------------------------------------------------------------------------- /READ/pipelines/pipeline.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import munch 3 | 4 | import torch 5 | 6 | from READ.utils.train import get_module, save_model, load_model_checkpoint 7 | from READ.utils.arguments import deval_args 8 | 9 | 10 | class Pipeline: 11 | def export_args(self, parser): 12 | # add arguments of this pipeline to the cmd parser 13 | raise NotImplementedError() 14 | 15 | def create(self, args): 16 | # return dictionary with pipeline components 17 | raise NotImplementedError() 18 | 19 | def dataset_load(self, *args, **kwargs): 20 | # called before running train/val 21 | pass 22 | 23 | def dataset_unload(self, *args, **kwargs): 24 | # called after running train/val 25 | pass 26 | 27 | def get_net(self): 28 | raise NotImplementedError() 29 | 30 | def extra_optimizer(self, *args): 31 | return None 32 | 33 | 34 | def load_pipeline(checkpoint, args_to_update=None): 35 | ckpt = torch.load(checkpoint, map_location='cpu') 36 | 37 | assert 'args' in ckpt 38 | 39 | if args_to_update: 40 | ckpt['args'].update(args_to_update) 41 | ckpt['args']['pipeline']='READ.pipelines.ogl.TexturePipeline' 42 | ckpt['args']['splitter_module']= 'READ.datasets.splitter.split_by_step' 43 | ckpt['args']['criterion_module']= 'READ.criterions.vgg_loss.VGGLoss' 44 | 45 | try: 46 | args = munch.munchify(ckpt['args']) 47 | 48 | pipeline = get_module(args.pipeline)() 49 | pipeline.create(args) 50 | except AttributeError as err: 51 | print('\nERROR: Checkpoint args is incompatible with this version\n', file=sys.stderr) 52 | raise err 53 | 54 | load_model_checkpoint(checkpoint, pipeline.get_net()) 55 | 56 | return pipeline, args 57 | 58 | 59 | def save_pipeline(pipeline, save_dir, epoch, stage, args): 60 | objects = pipeline.state_objects() 61 | 62 | args = deval_args(args) 63 | 64 | for name, obj in objects.items(): 65 | obj_class = obj.__class__.__name__ 66 | filename = f'{obj_class}_stage_{stage}_epoch_{epoch}' 67 | if name: 68 | name = name.replace('/', '_') 69 | filename = f'{filename}_{name}' 70 | save_path = os.path.join(save_dir, filename + '.pth') 71 | save_model(save_path, obj, args=args) 72 | -------------------------------------------------------------------------------- /READ/utils/__pycache__/arguments.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/arguments.cpython-36.pyc -------------------------------------------------------------------------------- /READ/utils/__pycache__/arguments.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/arguments.cpython-38.pyc -------------------------------------------------------------------------------- /READ/utils/__pycache__/perform.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/perform.cpython-36.pyc -------------------------------------------------------------------------------- /READ/utils/__pycache__/perform.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/perform.cpython-38.pyc -------------------------------------------------------------------------------- /READ/utils/__pycache__/train.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/train.cpython-36.pyc -------------------------------------------------------------------------------- /READ/utils/__pycache__/train.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/READ/utils/__pycache__/train.cpython-38.pyc -------------------------------------------------------------------------------- /READ/utils/arguments.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pydoc 3 | from pathlib import Path 4 | import munch 5 | 6 | 7 | class ActionNoYes(argparse.Action): 8 | def __init__(self, 9 | option_strings, 10 | dest, 11 | nargs=0, 12 | const=None, 13 | default=None, 14 | type=None, 15 | choices=None, 16 | required=False, 17 | help="", 18 | metavar=None): 19 | 20 | assert len(option_strings) == 1 21 | assert option_strings[0][:2] == '--' 22 | 23 | name= option_strings[0][2:] 24 | help += f'Use "--{name}" for True, "--no-{name}" for False' 25 | super(ActionNoYes, self).__init__(['--' + name, '--no-' + name], 26 | dest=dest, 27 | nargs=nargs, 28 | const=const, 29 | default=default, 30 | type=type, 31 | choices=choices, 32 | required=required, 33 | help=help, 34 | metavar=metavar) 35 | 36 | def __call__(self, parser, namespace, values, option_string=None): 37 | if option_string.startswith('--no-'): 38 | setattr(namespace, self.dest, False) 39 | else: 40 | setattr(namespace, self.dest, True) 41 | 42 | 43 | class SplitStr(argparse.Action): 44 | 45 | def split(self, x): 46 | if x == '': 47 | return [] 48 | else: 49 | return [self.elem_type(y) for y in x.split(self.delimiter)] 50 | 51 | def __init__(self, 52 | option_strings, 53 | dest, 54 | nargs=None, 55 | const=None, 56 | default=None, 57 | type=None, 58 | choices=None, 59 | required=False, 60 | help="", 61 | metavar=None, 62 | delimiter=',', 63 | elem_type=str): 64 | 65 | self.delimiter = delimiter 66 | self.elem_type = elem_type 67 | 68 | default = self.split(default) 69 | super(SplitStr, self).__init__(option_strings, 70 | dest=dest, 71 | nargs=nargs, 72 | const=const, 73 | default=default, 74 | type=type, 75 | choices=choices, 76 | required=required, 77 | help=help, 78 | metavar=metavar) 79 | 80 | def __call__(self, parser, namespace, values, option_string=None): 81 | print(values) 82 | setattr(namespace, self.dest, self.split(values)) 83 | 84 | 85 | class MyArgumentParser(argparse.ArgumentParser): 86 | def __init__(self, **kwargs): 87 | super(MyArgumentParser, self).__init__(**kwargs) 88 | self.register('action', 'store_bool', ActionNoYes) 89 | self.register('action', 'split_str', SplitStr) 90 | 91 | def add(self, *args, **kwargs): 92 | return self.add_argument(*args, **kwargs) 93 | 94 | 95 | def parse_image_size(string): 96 | error_msg = 'size must have format WxH' 97 | tokens = string.split('x') 98 | if len(tokens) != 2: 99 | raise argparse.ArgumentTypeError(error_msg) 100 | try: 101 | w = int(tokens[0]) 102 | h = int(tokens[1]) 103 | return w, h 104 | except ValueError: 105 | raise argparse.ArgumentTypeError(error_msg) 106 | 107 | 108 | PREFIX = '___' 109 | 110 | 111 | def eval_modules(config): 112 | _upd = {} 113 | for arg in config: 114 | if isinstance(arg, str) and arg.endswith('_module'): 115 | # print(arg, config[arg]) 116 | m = pydoc.locate(config[arg]) 117 | if not m: 118 | raise ValueError(f"module {arg}={config[arg]} not found") 119 | else: 120 | config[arg], _upd[PREFIX + arg] = m, config[arg] 121 | elif isinstance(config[arg], dict): 122 | eval_modules(config[arg]) 123 | config.update(_upd) 124 | 125 | 126 | def eval_paths(config): 127 | _upd = {} 128 | for arg in config: 129 | if isinstance(arg, str) and arg.endswith('_path'): 130 | config[arg], _upd[PREFIX + arg] = Path(config[arg]), config[arg] 131 | elif isinstance(config[arg], dict): 132 | eval_paths(config[arg]) 133 | config.update(_upd) 134 | 135 | 136 | def eval_functions(config): 137 | _upd = {} 138 | for arg in config: 139 | if isinstance(arg, str) and arg.endswith('_func'): 140 | config[arg], _upd[PREFIX + arg] = eval(config[arg]), config[arg] 141 | elif isinstance(config[arg], dict): 142 | eval_functions(config[arg]) 143 | config.update(_upd) 144 | 145 | 146 | def eval_args(args): 147 | args = vars(args) 148 | 149 | eval_modules(args) 150 | eval_paths(args) 151 | eval_functions(args) 152 | 153 | return munch.munchify(args) 154 | 155 | 156 | def deval_args(args): 157 | args = vars(args) 158 | 159 | for key in list(args): 160 | if key.startswith(PREFIX): 161 | args[key.replace(PREFIX, '')] = args[key] 162 | del args[key] 163 | 164 | return args 165 | -------------------------------------------------------------------------------- /READ/utils/perform.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import time 3 | from collections import defaultdict 4 | 5 | import numpy as np 6 | 7 | 8 | class TicToc: 9 | def __init__(self): 10 | self.tic_toc_tic = None 11 | 12 | def tic(self): 13 | self.tic_toc_tic = time.time() 14 | 15 | def toc(self): 16 | assert self.tic_toc_tic, 'You forgot to call tic()' 17 | return (time.time() - self.tic_toc_tic) * 1000 18 | 19 | def tocp(self, str): 20 | print(f"{str} took {toc():.4f}ms") 21 | 22 | @staticmethod 23 | def print_timing(timing, name=''): 24 | print(f'\n=== {name} Timimg ===') 25 | for fn, times in timing.items(): 26 | min, max, mean, p95 = np.min(times), np.max(times), np.mean(times), np.percentile(times, 95) 27 | print(f'{fn}:\tmin: {min:.4f}\tmax: {max:.4f}\tmean: {mean:.4f}ms\tp95: {p95:.4f}ms') 28 | 29 | 30 | class AccumDict: 31 | def __init__(self, num_f=3): 32 | self.d = defaultdict(list) 33 | self.num_f = num_f 34 | 35 | def add(self, k, v): 36 | self.d[k] += [v] 37 | 38 | def __dict__(self): 39 | return self.d 40 | 41 | def __getitem__(self, key): 42 | return self.d[key] 43 | 44 | def __str__(self): 45 | s = '' 46 | for k in self.d: 47 | if not self.d[k]: 48 | continue 49 | cur = self.d[k][-1] 50 | avg = np.mean(self.d[k]) 51 | format_str = '{:.%df}' % self.num_f 52 | cur_str = format_str.format(cur) 53 | avg_str = format_str.format(avg) 54 | s += f'{k} {cur_str} ({avg_str})\t\t' 55 | return s 56 | 57 | def __repr__(self): 58 | return self.__str__() 59 | 60 | 61 | class Tee(object): 62 | def __init__(self, filename): 63 | self.file = open(filename, 'a', buffering=1) 64 | self.terminal = sys.stdout 65 | 66 | def __del__(self): 67 | self.file.close() 68 | 69 | def write(self, data): 70 | self.file.write(data) 71 | self.terminal.write(data) 72 | 73 | def flush(self): 74 | self.file.flush() 75 | -------------------------------------------------------------------------------- /READ/utils/train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import torch 5 | import cv2 6 | import gzip 7 | 8 | import torchvision 9 | 10 | import matplotlib 11 | import matplotlib.cm 12 | 13 | from READ.models.compose import ModelAndLoss 14 | 15 | 16 | def to_device(data, device='cuda:0'): 17 | if isinstance(data, torch.Tensor): 18 | return data.to(device, non_blocking=True) 19 | elif isinstance(data, dict): 20 | for k in data.keys(): 21 | data[k] = to_device(data[k], device) 22 | 23 | return data 24 | elif isinstance(data, (tuple, list)): 25 | for i in range(len(data)): 26 | data[i] = to_device(data[i], device) 27 | 28 | return data 29 | 30 | return data 31 | 32 | 33 | def set_requires_grad(model, value): 34 | for p in model.parameters(): 35 | p.requires_grad = value 36 | 37 | 38 | def freeze(model, b): 39 | set_requires_grad(model, not b) 40 | 41 | 42 | def save_model(save_path, model, args=None, compress=False): 43 | model = unwrap_model(model) 44 | 45 | if not isinstance(args, dict): 46 | args = vars(args) 47 | 48 | dict_to_save = { 49 | 'state_dict': model.state_dict(), 50 | 'args': args 51 | } 52 | 53 | if compress: 54 | with gzip.open(f'{save_path}.gz', 'wb') as f: 55 | torch.save(dict_to_save, f, pickle_protocol=-1) 56 | else: 57 | torch.save(dict_to_save, save_path, pickle_protocol=-1) 58 | 59 | 60 | def load_model_checkpoint(path, model): 61 | ckpt = torch.load(path, map_location='cpu') 62 | 63 | model.load_state_dict(ckpt['state_dict']) 64 | 65 | return model 66 | 67 | 68 | def unwrap_model(model): 69 | model_ = model 70 | while True: 71 | if isinstance(model_, torch.nn.DataParallel): 72 | model_ = model_.module 73 | elif isinstance(model_, ModelAndLoss): 74 | model_ = model_.model 75 | else: 76 | return model_ 77 | 78 | 79 | def colorize(value, vmin=0, vmax=1, cmap='viridis'): 80 | """ 81 | A utility function for Torch/Numpy that maps a grayscale image to a matplotlib 82 | colormap for use with TensorBoard image summaries. 83 | By default it will normalize the input value to the range 0..1 before mapping 84 | to a grayscale colormap. 85 | Arguments: 86 | - value: 2D Tensor of shape [height, width] or 3D Tensor of shape 87 | [height, width, 1]. 88 | - vmin: the minimum value of the range used for normalization. 89 | (Default: value minimum) 90 | - vmax: the maximum value of the range used for normalization. 91 | (Default: value maximum) 92 | - cmap: a valid cmap named for use with matplotlib's `get_cmap`. 93 | (Default: Matplotlib default colormap) 94 | 95 | Returns a 4D uint8 tensor of shape [height, width, 4]. 96 | """ 97 | 98 | # normalize 99 | vmin = value.min() if vmin is None else vmin 100 | vmax = value.max() if vmax is None else vmax 101 | if vmin!=vmax: 102 | value = (value - vmin) / (vmax - vmin) # vmin..vmax 103 | else: 104 | # Avoid 0-division 105 | value = value*0. 106 | # squeeze last dim if it exists 107 | value = value.squeeze() 108 | 109 | cmapper = matplotlib.cm.get_cmap(cmap) 110 | value = cmapper(value,bytes=True) # (nxmx4) 111 | return np.ascontiguousarray(value[:, :, :3].transpose(2, 0, 1)) 112 | 113 | 114 | def resize(imgs, sz=256): 115 | return torch.nn.functional.interpolate(imgs, size=sz) 116 | 117 | 118 | def to_numpy(t, flipy=False, uint8=True, i=0): 119 | out = t[:] 120 | if len(out.shape) == 4: 121 | out = out[i] 122 | out = out.detach().permute(1, 2, 0) 123 | out = out.flip([0]) if flipy else out 124 | out = out.detach().cpu().numpy() 125 | out = (out.clip(0, 1)*255).astype(np.uint8) if uint8 else out 126 | return out 127 | 128 | 129 | def image_grid(*args, sz = 256): 130 | num_img = min( min([len(x) for x in args]), 4) 131 | 132 | grid = [] 133 | for a in args: 134 | b = a[:num_img].detach().cpu().float() 135 | if b.shape[1] == 1: 136 | grid.append(torch.cat( [ torch.from_numpy(colorize(bb)).float()[None, ...]/255 for bb in b ], dim=0 )) 137 | # grid.append(torch.cat( [b, b, b], dim=1 ) ) 138 | else: 139 | grid.append(b[:, :3]) 140 | 141 | # print([x.shape for x in grid ]) 142 | imgs = resize( torch.cat(grid), sz=sz) 143 | x = torchvision.utils.make_grid(imgs, nrow = num_img) 144 | 145 | return x 146 | 147 | 148 | def get_module(path): 149 | import pydoc 150 | 151 | m = pydoc.locate(path) 152 | assert m is not None, f'{path} not found' 153 | 154 | return m 155 | 156 | 157 | class SubsetSampler(torch.utils.data.Sampler): 158 | def __init__(self, dataset): 159 | pass 160 | 161 | def __iter__(self): 162 | pass 163 | 164 | def __len__(self): 165 | pass -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # READ: Large-Scale Neural Scene Rendering for Autonomous Driving 2 | 3 | 4 | 5 | Commercial use is forbidden![](https://img.shields.io/badge/license-GPL-blue) 6 | 7 | Paper: https://arxiv.org/abs/2205.05509 (Old version) 8 | 9 | Video: 10 | [Bilibili](https://www.bilibili.com/video/BV1we411K7ug) [Bilibili2](https://www.bilibili.com/video/BV1C84y1t7pP) [Youtube](https://www.youtube.com/watch?v=2yeH7l4oLYw) 11 | 14 | 15 | 16 | Demo: (Use only one camera view for training) 17 | 18 | All you need is a series of images or videos as input to get the following effect. Scenes are not only limited to driving scenes, but also can be used for tourism scenes, indoor scenes, objects, etc. 19 | 20 |

21 | 22 | 23 | 24 |

25 | 26 |

27 | 28 | 29 |

30 | 31 | 32 | 33 | 34 | 35 | ## Overview: 36 | 37 | This is the code release for our AAAI2023 paper, PyTorch implementation of Large-Scale Neural Scene Rendering for Autonomous Driving(READ), a large-scale neural rendering method is proposed to synthesize the autonomous driving scene~(READ), which makes it possible to synthesize large-scale driving scenarios on a PC. Our model can not only synthesize realistic driving scenes but also stitch and edit driving scenes. 38 | ![contents](./image/main.jpg) 39 | 40 | ## Setup 41 | 42 | The following instructions describe installation of conda environment. Please refer to [requirement](https://github.com/JOP-Lee/READ/blob/main/requirement.sh). 43 | 44 | If you want to set it to headless mode(without X server enabled), see the MyRender in the [src folder](https://github.com/JOP-Lee/READ/tree/main/src). 45 | 46 | Note: This project needs to use the screen to run the script, if not, run the project in [src folder](https://github.com/JOP-Lee/READ/tree/main/src). 47 | 48 | Run this command to install python environment: 49 | ```bash 50 | cd src/MyRender 51 | pip install -v -e . 52 | ``` 53 | 54 | 55 | ## Run 56 | 57 | You can render one of the fitted scenes we provide right away in the real-time viewer or fit your own scene. 58 | 59 | Download fitted scenes and universal rendering network weights from [here](https://zenodo.org/record/7395608#.Y4xv9HZBxPY) and unpack in the Data directory. 60 | 61 | We suppose that you have at least one GeForce GTX 1080 Ti for fitting and inference. 62 | 63 | 64 | 65 | ### Use fitted scene 66 | 67 | Here we show an example how to run fitted scenes in the viewer. 68 | 69 | #### kitti6 70 | ```bash 71 | python viewer.py --config downloads/kitti6.yaml 72 | ``` 73 | 74 | ### Viewer navigation: 75 | 76 | * Move forward: w 77 | * Move backward: s 78 | * Move left: a 79 | * Move right: d 80 | * Rise: q 81 | * Fall: e 82 | * Turn to the left: 1 83 | * Turn to the right: 2 84 | * Turn to the up: 3 85 | * Turn to the down: 4 86 | * Rotation: press left mouse button and drag 87 | * Move: press rigth mouse button and drug / scroll middle mouse botton 88 | * Pan: press middle mouse button and drug 89 | 90 | ## Train 91 | 92 | python train.py --config configs/train_example.yaml --pipeline READ.pipelines.ogl.TexturePipeline --crop_size 256x256 93 | 94 | The size of crop_size depends on your GPU memory, and the parameter train_dataset_args can be adjusted in the configs folder. 95 | 96 | 97 | ## Train with your own data 98 | 99 | It's very simple. You just need a sequence of pictures to do it. 100 | 101 | 1. Use metashape to obtain camera.xml, pointcloud.ply 102 | 2. Place the above files and photos in the Data folder, for example, Data/image/xx.png, Data/camera.xml Data/pointcloud.ply 103 | 3. Change the folder address in configs/paths_example.yaml and load net_ckpt/texture_ckpt model address in configs/train_example.yaml, if any. 104 | 4. Run the code: python train.py --config configs/train_example.yaml --pipeline READ.pipelines.ogl.TexturePipeline --crop_size 256x256 105 | 106 | 107 | 118 | 119 | 120 | ## Scene Editing: 121 | 122 | READ can move and remove the cars in different views. A panorama with larger view can be synthesized by changing the camera parameters. 123 | ![contents](./image/Scene_Editing.jpg) 124 | 125 | 126 | ## Scene Stitching: 127 | 128 | READ is able to synthesize the larger driving scenes and update local areas with obvious changes in road conditions. 129 | ![contents](./image/Scene_Stitching.jpg) 130 | 131 | ## Novel View Synthesis: 132 | 133 | ![contents](./image/NovelView.jpg) 134 | 135 | 136 | ## Acknowledgments 137 | In this code we refer to the following implementations: [npbg](https://github.com/alievk/npbg) and [ 138 | MIMO](https://github.com/chosj95/MIMO-UNet). Great thanks to them! 139 | 140 | 141 | ## Citation 142 | If our work or code helps you, please consider to cite our paper. Thank you! 143 | ```BibTeX 144 | @inproceedings{li2022read, 145 | author = {Li, Zhuopeng and Li, Lu and Zhu, Jianke*}, 146 | title = {READ: Large-Scale Neural Scene Rendering for Autonomous Driving}, 147 | booktitle = {AAAI}, 148 | year = {2023} 149 | } 150 | ``` 151 | 152 | -------------------------------------------------------------------------------- /configs/paths_example.yaml: -------------------------------------------------------------------------------- 1 | datasets: 2 | "scene": 3 | scene_path: scene_path/Data/scene_name/scene.yaml 4 | target_path: scene_path/Data/scene_name/images_undistorted 5 | target_name_func: "lambda i: f'{i}.png'" 6 | 7 | -------------------------------------------------------------------------------- /configs/train_example.yaml: -------------------------------------------------------------------------------- 1 | paths_file: configs/paths_example.yaml 2 | dataset_names: [] 3 | 4 | batch_size: 3 5 | batch_size_val: 3 6 | eval_in_train: True 7 | input_format: uv_1d_p1, uv_1d_p1_ds1, uv_1d_p1_ds2, uv_1d_p1_ds3, uv_1d_p1_ds4 8 | # use_mask: True 9 | 10 | epochs: 100 11 | save_freq: 5 12 | save_dir: data/logs 13 | simple_name: True 14 | 15 | 16 | net_ckpt: scene_path.pth 17 | texture_ckpt: scene_path.pth 18 | 19 | #splitter_module: READ.datasets.splitter.split_by_ratio 20 | #splitter_args: 21 | # train_ratio: 1 22 | 23 | splitter_module: READ.datasets.splitter.split_by_step 24 | splitter_args: 25 | val_step: 10 26 | train_drop: 0.0 27 | 28 | train_dataset_args: 29 | keep_fov: False 30 | random_zoom: [0.7, 2.0] 31 | random_shift: [-1., 1.] 32 | drop_points: 0.0 33 | num_samples: 3000 34 | 35 | val_dataset_args: 36 | keep_fov: False 37 | drop_points: 0.0 38 | 39 | criterion_module: READ.criterions.vgg_loss.VGGLoss 40 | criterion_args: 41 | partialconv: false 42 | -------------------------------------------------------------------------------- /downloads/kitti6.yaml: -------------------------------------------------------------------------------- 1 | viewport_size: [1216, 368] 2 | intrinsic_matrix: Data/kitti6_368_total/camera.xml 3 | view_matrix: Data/kitti6_368_total/camera.xml 4 | pointcloud: Data/kitti6_368_total/pointcloud.ply 5 | 6 | 7 | net_ckpt: data/logs/path 8 | texture_ckpt: data/logs/path 9 | -------------------------------------------------------------------------------- /image/NovelView.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/image/NovelView.jpg -------------------------------------------------------------------------------- /image/Scene_Editing.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/image/Scene_Editing.jpg -------------------------------------------------------------------------------- /image/Scene_Stitching.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/image/Scene_Stitching.jpg -------------------------------------------------------------------------------- /image/main.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/image/main.jpg -------------------------------------------------------------------------------- /requirement.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | conda create -y -n READ python=3.9 4 | conda activate READ 5 | 6 | conda install -y cudnn=8.2.1.32 cudatoolkit-dev=11.2 cudatoolkit=11.2 -c nvidia -c conda-forge 7 | pip install torch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 8 | 9 | #optional 10 | cd MyRender 11 | pip install -v -e . 12 | 13 | pip install \ 14 | numpy \ 15 | pyyaml \ 16 | tensorboardX \ 17 | munch \ 18 | scipy \ 19 | matplotlib \ 20 | Cython \ 21 | PyOpenGL \ 22 | PyOpenGL_accelerate \ 23 | trimesh \ 24 | huepy \ 25 | pillow \ 26 | tqdm \ 27 | scikit-learn 28 | 29 | conda install opencv 30 | 31 | # need to install separately 32 | pip install \ 33 | git+https://github.com/DmitryUlyanov/glumpy \ 34 | numpy-quaternion 35 | 36 | # pycuda 37 | git clone https://github.com/inducer/pycuda 38 | cd pycuda 39 | git submodule update --init 40 | export PATH=$PATH:/usr/local/cuda/bin 41 | ./configure.py --cuda-enable-gl 42 | python setup.py install 43 | cd .. -------------------------------------------------------------------------------- /src/Data/scene.yaml: -------------------------------------------------------------------------------- 1 | viewport_size: [1216, 368] 2 | intrinsic_matrix: camera.xml 3 | view_matrix: camera.xml 4 | pointcloud: pointcloud.ply 5 | 6 | -------------------------------------------------------------------------------- /src/MyRender/CloudProjection/pcpr_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "point_render.cuh" 4 | #include 5 | 6 | 7 | // CUDA forward declarations 8 | 9 | std::vector pcpr_cuda_forward( 10 | torch::Tensor in_points, //(num_points,3) 11 | torch::Tensor , 12 | int tar_width, int tar_height, int block_size 13 | ); 14 | 15 | // C++ interface 16 | 17 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 18 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 19 | #define CHECK_FLOAT(x) AT_ASSERTM(x.type().scalarType()==torch::ScalarType::Float, #x " must be a float tensor") 20 | #define CHECK_Int(x) AT_ASSERTM(x.type().scalarType()==torch::ScalarType::Int, #x " must be a Int tensor") 21 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x); 22 | 23 | std::vector pcpr_cuda_forward( 24 | torch::Tensor in_points, //(num_points,3) 25 | torch::Tensor total_m, 26 | int tar_width, int tar_height, int block_size 27 | ) 28 | { 29 | in_points = in_points.to(torch::kCUDA); 30 | total_m = total_m.to(torch::kCUDA); 31 | 32 | CHECK_INPUT(in_points); CHECK_FLOAT(in_points); 33 | CHECK_INPUT(total_m); CHECK_FLOAT(total_m); 34 | AT_ASSERTM(total_m.sizes().size()==3, "batch_size check"); 35 | 36 | return GPU_PCPR(in_points, total_m, tar_width, tar_height, block_size); 37 | } 38 | 39 | 40 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 41 | m.def("forward", &pcpr_cuda_forward, "PCPR forward (CUDA)"); 42 | } 43 | 44 | -------------------------------------------------------------------------------- /src/MyRender/CloudProjection/point_render.cu: -------------------------------------------------------------------------------- 1 | #include "point_render.cuh" 2 | #include 3 | 4 | #include "helper_math.h" 5 | 6 | #include "cooperative_groups.h" 7 | 8 | struct Matrix4x4 9 | { 10 | public: 11 | float4 col[4]; 12 | __device__ __forceinline__ 13 | Matrix4x4() 14 | { 15 | col[0] = col[1] = col[2] = col[3] = make_float4(0, 0, 0, 0); 16 | } 17 | __device__ __forceinline__ 18 | Matrix4x4(float4 a, float4 b, float4 c, float4 d) 19 | { 20 | col[0].x = a.x; 21 | col[0].y = a.y; 22 | col[0].z = a.z; 23 | col[0].w = a.w; 24 | 25 | col[1].x = b.x; 26 | col[1].y = b.y; 27 | col[1].z = b.z; 28 | col[1].w = b.w; 29 | 30 | col[2].x = c.x; 31 | col[2].y = c.y; 32 | col[2].z = c.z; 33 | col[2].w = c.w; 34 | 35 | col[3].x = d.x; 36 | col[3].y = d.y; 37 | col[3].z = d.z; 38 | col[3].w = d.w; 39 | } 40 | 41 | __device__ __forceinline__ 42 | Matrix4x4 transpose() const 43 | { 44 | Matrix4x4 res; 45 | 46 | res.col[0].x = col[0].x; 47 | res.col[0].y = col[1].x; 48 | res.col[0].z = col[2].x; 49 | res.col[0].w = col[3].x; 50 | 51 | res.col[1].x = col[0].y; 52 | res.col[1].y = col[1].y; 53 | res.col[1].z = col[2].y; 54 | res.col[1].w = col[3].y; 55 | 56 | res.col[2].x = col[0].z; 57 | res.col[2].y = col[1].z; 58 | res.col[2].z = col[2].z; 59 | res.col[2].w = col[3].z; 60 | 61 | res.col[3].x = col[0].w; 62 | res.col[3].y = col[1].w; 63 | res.col[3].z = col[2].w; 64 | res.col[3].w = col[3].w; 65 | return res; 66 | 67 | } 68 | __device__ __forceinline__ 69 | Matrix4x4 inv() const 70 | { 71 | Matrix4x4 res; 72 | res.col[0].x = col[0].x; 73 | res.col[0].y = col[1].x; 74 | res.col[0].z = col[2].x; 75 | res.col[0].w = 0; 76 | 77 | res.col[1].x = col[0].y; 78 | res.col[1].y = col[1].y; 79 | res.col[1].z = col[2].y; 80 | res.col[1].w = 0; 81 | 82 | res.col[2].x = col[0].z; 83 | res.col[2].y = col[1].z; 84 | res.col[2].z = col[2].z; 85 | res.col[2].w = 0; 86 | 87 | res.col[3].x = -dot(col[0], col[3]); 88 | res.col[3].y = -dot(col[1], col[3]); 89 | res.col[3].z = -dot(col[2], col[3]); 90 | res.col[3].w = 1; 91 | return res; 92 | } 93 | }; 94 | 95 | 96 | typedef struct CamMatrix 97 | { 98 | float4 m[4]; 99 | __device__ __forceinline__ 100 | Matrix4x4 getRT() const 101 | { 102 | return Matrix4x4(m[0],m[1],m[2],m[3]); 103 | } 104 | 105 | }; 106 | 107 | namespace math 108 | { 109 | __device__ __forceinline__ 110 | float4 MatrixMul(const Matrix4x4& mat, float4& p) 111 | { 112 | Matrix4x4 res = mat; 113 | float4 ans; 114 | ans.x = dot(res.col[0], p); 115 | ans.y = dot(res.col[1], p); 116 | ans.z = dot(res.col[2], p); 117 | ans.w = dot(res.col[3], p); 118 | 119 | ans = ans / ans.w; 120 | return ans; 121 | } 122 | } 123 | 124 | 125 | __global__ 126 | void DepthProject(float3 * point_clouds, int num_points, int batch_id, 127 | CamMatrix* total_m, int tar_width, int tar_height, 128 | int* mutex_map, float* out_depth, float* out_index) 129 | { 130 | // int ids = blockDim.x * blockIdx.x + threadIdx.x; // index of point 131 | cooperative_groups::grid_group grid = cooperative_groups::this_grid(); 132 | for(int ids = grid.thread_rank(); ids num_points) return; 135 | float4 p = make_float4(point_clouds[ids], 1.0); 136 | 137 | float4 camp = math::MatrixMul(total_m->getRT(), p); 138 | 139 | if(camp.x<-1 || camp.x>1 || camp.y<-1 || camp.y>1 || camp.z<-1 || camp.z>1) return; 140 | 141 | float u = tar_width*(camp.x+1)*0.5; 142 | float v = tar_height*(1-camp.y)*0.5; 143 | float tdepth = ((camp.z+1)*0.5); 144 | 145 | int xx = int(u); 146 | int yy = int(v); 147 | if (xx < 0 || xx >= tar_width || yy < 0 || yy >= tar_height) return; 148 | int ind = batch_id * tar_width * tar_height + yy * tar_width + xx ; 149 | bool isSet = false; 150 | do 151 | { 152 | if ((isSet = atomicCAS(mutex_map + ind, 0, 1)) == false) 153 | { 154 | // critical section goes here 155 | if (out_depth[ind] > tdepth || out_depth[ind]==0) 156 | { 157 | out_depth[ind] = tdepth; 158 | out_index[ind] = (float)ids ; // 0 denote empty 159 | } 160 | } 161 | if (isSet) 162 | { 163 | mutex_map[ind] = 0; 164 | } 165 | } while (!isSet); 166 | } 167 | } 168 | 169 | std::vector GPU_PCPR( 170 | torch::Tensor in_points, //(num_points,3) 171 | torch::Tensor total_m, int tar_width, int tar_height, int block_size) 172 | { 173 | const auto num_points = in_points.size(0); 174 | const auto batch_size = total_m.size(0); 175 | 176 | torch::Tensor out_index = torch::zeros({batch_size, tar_height, tar_width}).to(torch::kCUDA); 177 | torch::Tensor out_depth = torch::zeros({batch_size, tar_height, tar_width}).to(torch::kCUDA); 178 | 179 | int default_block_size = block_size; 180 | int block_num = (num_points+default_block_size-1) / default_block_size; 181 | 182 | int *mutex_map; 183 | cudaMalloc(&mutex_map, sizeof(int) * batch_size * tar_width * tar_height); 184 | cudaMemset(mutex_map, 0, sizeof(int) * batch_size * tar_width * tar_height); 185 | 186 | for(int b=0; b> > ( 188 | (float3*)in_points.data(), num_points, b, 189 | (CamMatrix*)total_m[b].data(), 190 | tar_width, tar_height, mutex_map, 191 | out_depth.data(), out_index.data()); 192 | } 193 | cudaFree(mutex_map); 194 | cudaDeviceSynchronize(); 195 | 196 | out_index = out_index.to(torch::kCPU); 197 | out_depth = out_depth.to(torch::kCPU); 198 | 199 | return {out_index, out_depth}; 200 | } 201 | 202 | -------------------------------------------------------------------------------- /src/MyRender/CloudProjection/point_render.cuh: -------------------------------------------------------------------------------- 1 | #include "cuda_runtime.h" 2 | #include "device_launch_parameters.h" 3 | #include 4 | 5 | std::vector GPU_PCPR( 6 | torch::Tensor in_points, //(num_points,3) 7 | torch::Tensor total_m, 8 | int tar_width, int tar_height, int block_size); // (tar_heigh ,tar_width) 9 | -------------------------------------------------------------------------------- /src/MyRender/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 3 | 4 | setup( 5 | name='pcpr', 6 | version="0.1", 7 | ext_modules=[ 8 | CUDAExtension('pcpr', [ 9 | 'CloudProjection/pcpr_cuda.cpp', 10 | 'CloudProjection/point_render.cu' 11 | ] 12 | ) 13 | ], 14 | cmdclass={ 15 | 'build_ext': BuildExtension 16 | }) -------------------------------------------------------------------------------- /src/READ/criterions/__pycache__/vgg_loss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/criterions/__pycache__/vgg_loss.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/criterions/__pycache__/vgg_loss.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/criterions/__pycache__/vgg_loss.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/criterions/vgg_loss.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | import torch.nn as nn 3 | import torchvision 4 | import torch 5 | from collections import OrderedDict 6 | from os.path import expanduser, join 7 | import os 8 | 9 | from READ.models.conv import PartialConv2d 10 | 11 | 12 | class View(nn.Module): 13 | def __init__(self): 14 | super(View, self).__init__() 15 | 16 | def forward(self, x): 17 | return x.view(-1) 18 | 19 | 20 | class VGGLoss(nn.Module): 21 | def __init__(self, net='caffe', partialconv=False, optimized=False, save_dir='.cache/torch/models'): 22 | super().__init__() 23 | 24 | self.partialconv = partialconv 25 | 26 | if net == 'pytorch': 27 | vgg19 = torchvision.models.vgg19(pretrained=True).features 28 | 29 | self.register_buffer('mean_', torch.FloatTensor([0.485, 0.456, 0.406])[None, :, None, None]) 30 | self.register_buffer('std_', torch.FloatTensor([0.229, 0.224, 0.225])[None, :, None, None]) 31 | 32 | elif net == 'caffe': 33 | if not os.path.exists(join(save_dir, 'vgg_caffe_features.pth')): 34 | vgg_weights = torch.utils.model_zoo.load_url('https://web.eecs.umich.edu/~justincj/models/vgg19-d01eb7cb.pth', model_dir=save_dir) 35 | 36 | map = {'classifier.6.weight':u'classifier.7.weight', 'classifier.6.bias':u'classifier.7.bias'} 37 | vgg_weights = OrderedDict([(map[k] if k in map else k,v) for k,v in vgg_weights.items()]) 38 | 39 | model = torchvision.models.vgg19() 40 | model.classifier = nn.Sequential(View(), *model.classifier._modules.values()) 41 | 42 | model.load_state_dict(vgg_weights) 43 | 44 | vgg19 = model.features 45 | os.makedirs(save_dir, exist_ok=True) 46 | torch.save(vgg19, join(save_dir, 'vgg_caffe_features.pth')) 47 | 48 | self.register_buffer('mean_', torch.FloatTensor([103.939, 116.779, 123.680])[None, :, None, None] / 255.) 49 | self.register_buffer('std_', torch.FloatTensor([1./255, 1./255, 1./255])[None, :, None, None]) 50 | 51 | else: 52 | vgg19 = torch.load(join(save_dir, 'vgg_caffe_features.pth')) 53 | self.register_buffer('mean_', torch.FloatTensor([103.939, 116.779, 123.680])[None, :, None, None] / 255.) 54 | self.register_buffer('std_', torch.FloatTensor([1./255, 1./255, 1./255])[None, :, None, None]) 55 | else: 56 | assert False 57 | 58 | if self.partialconv: 59 | part_conv = PartialConv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 60 | part_conv.weight = vgg19[0].weight 61 | part_conv.bias = vgg19[0].bias 62 | vgg19[0] = part_conv 63 | 64 | vgg19_avg_pooling = [] 65 | 66 | 67 | for weights in vgg19.parameters(): 68 | weights.requires_grad = False 69 | 70 | for module in vgg19.modules(): 71 | if module.__class__.__name__ == 'Sequential': 72 | continue 73 | elif module.__class__.__name__ == 'MaxPool2d': 74 | vgg19_avg_pooling.append(nn.AvgPool2d(kernel_size=2, stride=2, padding=0)) 75 | else: 76 | vgg19_avg_pooling.append(module) 77 | 78 | if optimized: 79 | self.layers = [3, 8, 17, 26, 35] 80 | else: 81 | self.layers = [1, 3, 6, 8, 11, 13, 15, 17, 20, 22, 24, 26, 29] 82 | 83 | self.vgg19 = nn.Sequential(*vgg19_avg_pooling) 84 | 85 | # print(self.vgg19) 86 | 87 | def normalize_inputs(self, x): 88 | return (x - self.mean_) / self.std_ 89 | 90 | def forward(self, input, target): 91 | loss = 0 92 | 93 | if self.partialconv: 94 | eps = 1e-9 95 | mask = target.sum(1, True) > eps 96 | mask = mask.float() 97 | 98 | features_input = self.normalize_inputs(input) 99 | features_target = self.normalize_inputs(target) 100 | for i, layer in enumerate(self.vgg19): 101 | if isinstance(layer, PartialConv2d): 102 | features_input = layer(features_input, mask) 103 | features_target = layer(features_target, mask) 104 | else: 105 | features_input = layer(features_input) 106 | features_target = layer(features_target) 107 | 108 | if i in self.layers: 109 | loss = loss + F.l1_loss(features_input, features_target) 110 | 111 | return loss 112 | 113 | 114 | class VGGLossMix(nn.Module): 115 | def __init__(self, weight=0.5): 116 | super(VGGLossMix, self).__init__() 117 | self.l1 = VGGLoss() 118 | self.l2 = VGGLoss(net='caffe') 119 | self.weight = weight 120 | 121 | def forward(self, input, target): 122 | return self.l1(input, target)*self.weight + self.l2(input, target) * (1-self.weight) 123 | -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/common.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/common.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/common.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/common.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/dynamic.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/dynamic.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/dynamic.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/dynamic.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/splitter.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/splitter.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/datasets/__pycache__/splitter.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/datasets/__pycache__/splitter.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/datasets/common.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import yaml 3 | 4 | from functools import lru_cache 5 | 6 | import numpy as np 7 | import cv2 8 | 9 | import torch 10 | 11 | 12 | @lru_cache(maxsize = 1000) 13 | def load_image(path): 14 | img = cv2.imread(path) 15 | assert img is not None, f'could not load {path}' 16 | return img[...,::-1].copy() 17 | 18 | 19 | def any2float(img): 20 | if isinstance(img, np.ndarray): 21 | out = img.astype(np.float32) 22 | if img.dtype == np.uint16: 23 | out /= 65535 24 | elif img.dtype == np.uint8: 25 | out /= 255 26 | elif torch.is_tensor(img): 27 | out = img.float() 28 | if img.dtype == torch.int16: 29 | out /= 65535 30 | elif img.dtype == torch.uint8: 31 | out /= 255 32 | else: 33 | raise TypeError('img must be numpy array or torch tensor') 34 | 35 | return out 36 | 37 | 38 | def rescale_K(K_, sx, sy, keep_fov=True): 39 | K = K_.copy() 40 | K[0, 2] = sx * K[0, 2] 41 | K[1, 2] = sy * K[1, 2] 42 | if keep_fov: 43 | K[0, 0] = sx * K[0, 0] 44 | K[1, 1] = sy * K[1, 1] 45 | return K 46 | 47 | def fit_size(x, d=16): 48 | return x[..., :d*(x.shape[-2]//d), :d*(x.shape[-1]//d)] 49 | 50 | 51 | class ToTensor(object): 52 | def __call__(self, img): 53 | if isinstance(img, np.ndarray): 54 | img = torch.from_numpy(img) 55 | 56 | img = any2float(img) 57 | 58 | return img.permute(2, 0, 1).contiguous() 59 | 60 | def __repr__(self): 61 | return self.__class__.__name__ + '()' 62 | 63 | 64 | def get_dataset_config(yml, dataset): 65 | join = os.path.join 66 | 67 | myhost = os.uname()[1] 68 | if myhost in yml: 69 | data_root = yml[myhost]['data_root'] if 'data_root' in yml[myhost] else '' 70 | else: 71 | data_root = '/' 72 | 73 | ds = yml['datasets'][dataset] 74 | 75 | return ds 76 | 77 | 78 | def split_lists(config, lists): 79 | sz = [len(l) for l in lists] 80 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 81 | 82 | splits = [] 83 | train_inds, val_inds = [], [] 84 | if 'train_ratio' in config: 85 | train_ratio = float(config['train_ratio']) 86 | train_n = int(sz[0] * train_ratio) 87 | train_inds, val_inds = np.split(np.random.permutation(sz[0]), [train_n]) 88 | else: 89 | val_step = int(config['val_step']) 90 | train_drop = int(config['train_drop']) 91 | for i in range(sz[0]): 92 | if i % val_step == 0: 93 | val_inds.append(i) 94 | elif train_drop < i % val_step < val_step - train_drop: 95 | train_inds.append(i) 96 | 97 | 98 | 99 | for lst in lists: 100 | lst = np.array(lst) 101 | splits.append([lst[train_inds], lst[val_inds]]) 102 | 103 | return splits -------------------------------------------------------------------------------- /src/READ/datasets/splitter.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import pdb 4 | 5 | def split_by_ratio(lists, train_ratio): 6 | sz = [len(l) for l in lists] 7 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 8 | 9 | splits = [] 10 | train_inds, val_inds = [], [] 11 | 12 | train_n = int(sz[0] * train_ratio) 13 | train_inds, val_inds = np.split(np.random.permutation(sz[0]), [train_n]) 14 | 15 | print(train_inds) 16 | print(val_inds) 17 | 18 | for lst in lists: 19 | lst = np.array(lst) 20 | splits.append([lst[train_inds], lst[train_inds]]) 21 | # splits.append([lst[train_inds], lst[val_inds]]) 22 | 23 | return splits 24 | 25 | 26 | def split_by_step(lists, val_step, train_drop): 27 | sz = [len(l) for l in lists] 28 | assert len(set(sz)) == 1, f'list sizes differ {sz}' 29 | 30 | splits = [] 31 | train_inds, val_inds = [], [] 32 | 33 | for i in range(sz[0]): 34 | if i % val_step == 0 and i!=0: 35 | val_inds.append(i) 36 | elif train_drop < i % val_step < val_step - train_drop: 37 | train_inds.append(i) 38 | 39 | val_inds.append(0) 40 | print(train_inds) 41 | print(val_inds) 42 | 43 | for lst in lists: 44 | lst = np.array(lst) 45 | splits.append([lst[train_inds], lst[val_inds]]) 46 | 47 | return splits 48 | 49 | -------------------------------------------------------------------------------- /src/READ/gl/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__init__.py -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/camera.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/camera.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/dataset.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/dataset.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/dataset.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/dataset.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/myrender.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/myrender.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/myrender.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/myrender.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/nn.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/nn.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/programs.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/programs.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/programs.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/programs.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/render.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/render.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/render.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/render.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/gl/__pycache__/utils.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/gl/__pycache__/utils.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/gl/dataset.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | from READ.gl.programs import NNScene 4 | 5 | 6 | def parse_input_string_obsolete(what): 7 | config = {} 8 | if '_pr' in what: 9 | config['draw_points'] = True 10 | config['splat_mode'] = True 11 | config['point_size'] = int(what.split('_pr')[-1]) 12 | config['flat_color'] = True 13 | elif '_p' in what: 14 | config['draw_points'] = True 15 | config['splat_mode'] = False 16 | config['point_size'] = int(what.split('_p')[-1]) 17 | config['flat_color'] = True 18 | else: 19 | config['draw_points'] = False 20 | config['splat_mode'] = False 21 | config['point_size'] = 1 22 | config['flat_color'] = False 23 | 24 | if 'colors' in what: 25 | config['mode'] = NNScene.MODE_COLOR, 0 26 | elif 'uv' in what: 27 | m1 = NNScene.UV_TYPE_1D if config['draw_points'] else NNScene.UV_TYPE_2D 28 | config['mode'] = NNScene.MODE_UV, m1 29 | elif 'normals' in what: 30 | nm = what.split('_')[-2] if '_p' in what else what.split('_')[-1] 31 | m1 = ['g', 'r', 'l', 'd'].index(nm) 32 | config['mode'] = NNScene.MODE_NORMALS, m1 33 | elif 'xyz' in what: 34 | config['mode'] = NNScene.MODE_XYZ, 0 35 | 36 | return config 37 | 38 | 39 | def parse_input_string(string): 40 | config = {} 41 | 42 | if re.search('^colors', string): 43 | config['mode'] = NNScene.MODE_COLOR, None 44 | elif re.search('^uv', string): 45 | choices = ['uv_1d', 'uv_2d'] 46 | ch = re.findall('|'.join(choices), string)[-1] 47 | m1 = choices.index(ch) 48 | config['mode'] = NNScene.MODE_UV, m1 49 | elif re.search('^normals', string): 50 | choices = ['normals_m', 'normals_r', 'normals_l', 'normals_d'] 51 | ch = re.findall('|'.join(choices), string)[-1] 52 | m1 = choices.index(ch) 53 | config['mode'] = NNScene.MODE_NORMALS, m1 54 | elif re.search('^xyz', string): 55 | config['mode'] = NNScene.MODE_XYZ, None 56 | elif re.search('^depth', string): 57 | config['mode'] = NNScene.MODE_DEPTH, None 58 | elif re.search('^labels', string): 59 | config['mode'] = NNScene.MODE_LABEL, None 60 | else: 61 | raise ValueError(string) 62 | 63 | res = re.findall('ps[0-9]+|p[0-9]+', string) 64 | if res: 65 | res = res[-1] 66 | config['draw_points'] = True 67 | config['flat_color'] = True 68 | config['point_size'] = int(re.search('[0-9]+', res).group()) 69 | config['splat_mode'] = re.search('^ps', res) is not None 70 | else: 71 | config['draw_points'] = False 72 | config['splat_mode'] = False 73 | config['point_size'] = 1 74 | config['flat_color'] = False 75 | 76 | res = re.findall('ds[0-5]+', string) 77 | if res: 78 | res = res[-1] 79 | config['downscale'] = int(re.search('[0-9]+', res).group()) 80 | 81 | 82 | return config 83 | 84 | 85 | def generate_input_string(config): 86 | s = '' 87 | m0, m1 = config['mode'] 88 | if m0 == NNScene.MODE_COLOR: 89 | s += 'colors' 90 | elif m0 == NNScene.MODE_UV: 91 | s += 'uv' 92 | if m1 == NNScene.UV_TYPE_1D: 93 | s += '_1d' 94 | elif m1 == NNScene.UV_TYPE_2D: 95 | s += '_2d' 96 | else: 97 | raise ValueError 98 | elif m0 == NNScene.MODE_NORMALS: 99 | s += 'normals' 100 | if m1 == NNScene.NORMALS_MODE_MODEL: 101 | s += '_m' 102 | elif m1 == NNScene.NORMALS_MODE_REFLECTION: 103 | s += '_r' 104 | elif m1 == NNScene.NORMALS_MODE_LOCAL: 105 | s += '_l' 106 | elif m1 == NNScene.NORMALS_MODE_DIRECTION: 107 | s += '_d' 108 | elif m0 == NNScene.MODE_XYZ: 109 | s += 'xyz' 110 | elif m0 == NNScene.MODE_DEPTH: 111 | s += 'depth' 112 | 113 | if config['draw_points']: 114 | s += '_p' 115 | if config['splat_mode']: 116 | s += 's' 117 | s += str(config['point_size']) 118 | 119 | if 'downscale' in config: 120 | s += f"_ds{config['downscale']}" 121 | 122 | return s 123 | 124 | 125 | 126 | def test_generate_parse(): 127 | configs = [ 128 | { 129 | 'draw_points': True, 130 | 'splat_mode': True, 131 | 'point_size': 20, 132 | 'flat_color': True, 133 | 134 | 'mode': (NNScene.MODE_UV, NNScene.UV_TYPE_1D) 135 | }, 136 | 137 | { 138 | 'draw_points': True, 139 | 'splat_mode': True, 140 | 'point_size': 20, 141 | 'flat_color': True, 142 | 143 | 'mode': (NNScene.MODE_UV, NNScene.UV_TYPE_2D) 144 | }, 145 | 146 | { 147 | 'draw_points': True, 148 | 'splat_mode': False, 149 | 'point_size': 20, 150 | 'flat_color': True, 151 | 152 | 'mode': (NNScene.MODE_COLOR, None) 153 | }, 154 | 155 | { 156 | 'draw_points': False, 157 | 'splat_mode': False, 158 | 'point_size': 1, 159 | 'flat_color': False, 160 | 161 | 'mode': (NNScene.MODE_NORMALS, NNScene.NORMALS_MODE_REFLECTION) 162 | }, 163 | 164 | { 165 | 'draw_points': False, 166 | 'splat_mode': False, 167 | 'point_size': 1, 168 | 'flat_color': False, 169 | 170 | 'mode': (NNScene.MODE_XYZ, None) 171 | }, 172 | 173 | { 174 | 'draw_points': True, 175 | 'splat_mode': False, 176 | 'point_size': 1, 177 | 'flat_color': True, 178 | 179 | 'mode': (NNScene.MODE_XYZ, None), 180 | 181 | 'downscale': 2 182 | }, 183 | ] 184 | 185 | def check(config): 186 | print('config:\n', config) 187 | s = generate_input_string(config) 188 | print('string:\n', s) 189 | c = parse_input_string(s) 190 | print('generated config:\n', c) 191 | print('-- OK' if config == c else '-- FAIL') 192 | 193 | for c in configs: 194 | check(c) 195 | 196 | 197 | if __name__ == '__main__': 198 | test_generate_parse() -------------------------------------------------------------------------------- /src/READ/gl/myrender.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | from this import d 4 | import torch 5 | import pcpr 6 | import numpy as np 7 | import pdb 8 | 9 | inv = np.linalg.inv 10 | 11 | 12 | class MyRender: 13 | def __init__(self, ds_list=None): 14 | if ds_list: 15 | self.update_ds(ds_list) 16 | 17 | def update_ds(self, ds_list): 18 | self.ds_list = ds_list 19 | self.ds_ids = [d.id for d in ds_list] 20 | self.tgt_sh = self.ds_list[0].tgt_sh 21 | self.points = {ds.id: torch.from_numpy(np.asarray(ds.scene_data['pointcloud']['xyz']).astype(np.float32)) for ds in ds_list} 22 | 23 | def render(self, data): 24 | input_format = self.ds_list[0].input_format.replace(' ', '').split(',') 25 | out_dict, depth_dict = {}, {} 26 | out_dict['id'] = data['input']['id'] 27 | 28 | proj_matrix = data['proj_matrix'].numpy() 29 | view_matrix = data['view_matrix'].numpy() 30 | total_m = torch.from_numpy( proj_matrix @ inv(view_matrix)) 31 | 32 | for i,k in enumerate(input_format): 33 | w = int(self.tgt_sh[0]*(0.5**i)) 34 | h = int(self.tgt_sh[1]*(0.5**i)) 35 | indexs, depths = torch.zeros(len(data['input']['id']),h,w),torch.zeros(len(data['input']['id']),h,w) 36 | for ds_id in self.ds_ids: 37 | idx = torch.where(data['input']['id']==ds_id)[0] 38 | index_buffer, depth_buffer = pcpr.forward(self.points[ds_id], total_m[idx], w, h, 512) 39 | indexs[idx] = index_buffer 40 | depths[idx] = depth_buffer 41 | out_dict[k] = indexs.unsqueeze(1) 42 | depth_dict[k] = depths.unsqueeze(1) 43 | return out_dict, depth_dict 44 | 45 | -------------------------------------------------------------------------------- /src/READ/gl/nn.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | try: 3 | import torch 4 | except ImportError: 5 | print('torch is not available') 6 | import numpy as np 7 | 8 | from READ.gl.render import OffscreenRender 9 | # from READ.gl.utils import get_net 10 | 11 | from READ.pipelines import load_pipeline 12 | # from READ.datasets.dynamic_tmp import MultiscaleRender, default_input_transform 13 | from READ.datasets.dynamic import MultiscaleRender, default_input_transform 14 | 15 | from scipy.ndimage import gaussian_filter 16 | import torch.nn as nn 17 | 18 | import pdb 19 | 20 | class GaussianLayer(nn.Module): 21 | _instance = None 22 | 23 | def __init__(self, in_channels, out_channels, kernel_size=21, sigma=3): 24 | super(GaussianLayer, self).__init__() 25 | self.seq = nn.Sequential( 26 | nn.ReflectionPad2d(kernel_size//2), 27 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 28 | ) 29 | 30 | self.weights_init(kernel_size, sigma) 31 | 32 | def forward(self, x): 33 | return self.seq(x) 34 | 35 | def weights_init(self, kernel_size, sigma): 36 | n= np.zeros((kernel_size, kernel_size)) 37 | n[kernel_size//2, kernel_size//2] = 1 38 | k = gaussian_filter(n,sigma=sigma) 39 | for name, f in self.named_parameters(): 40 | f.data.copy_(torch.from_numpy(k)) 41 | 42 | @staticmethod 43 | def get_instance(): 44 | if GaussianLayer._instance is None: 45 | GaussianLayer._instance = GaussianLayer(8, 8, kernel_size=13, sigma=6).cuda() 46 | 47 | return GaussianLayer._instance 48 | 49 | 50 | class BoxFilter(nn.Module): 51 | def __init__(self, in_channels, out_channels, kernel_size=3): 52 | super().__init__() 53 | 54 | self.seq = nn.Sequential( 55 | nn.ReflectionPad2d(kernel_size//2), 56 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 57 | ) 58 | 59 | self.weights_init(kernel_size) 60 | 61 | def forward(self, x): 62 | return self.seq(x) 63 | 64 | def weights_init(self, kernel_size): 65 | kernel = torch.ones((kernel_size, kernel_size)) / kernel_size ** 2 66 | print(self.seq[0].named_parameters()) 67 | 68 | 69 | def to_gpu(data): 70 | if isinstance(data, dict): 71 | for k in data: 72 | data[k] = data[k].cuda() 73 | return data 74 | else: 75 | return data.cuda() 76 | 77 | 78 | class OGL: 79 | def __init__(self, scene, scene_data, viewport_size, net_ckpt, texture_ckpt, out_buffer_location='numpy', supersampling=1, gpu=True, clear_color=None, temporal_average=False): 80 | self.gpu = gpu 81 | 82 | args_upd = { 83 | 'inference': True, 84 | } 85 | 86 | if texture_ckpt: 87 | args_upd['texture_ckpt'] = texture_ckpt 88 | if 'pointcloud' in scene_data: 89 | args_upd['n_points'] = scene_data['pointcloud']['xyz'].shape[0] 90 | 91 | pipeline, args = load_pipeline(net_ckpt, args_to_update=args_upd) 92 | 93 | self.model = pipeline.model 94 | 95 | if args.pipeline == 'READ.pipelines.ogl.TexturePipeline': 96 | self.model.load_textures(0) 97 | 98 | if self.gpu: 99 | self.model.cuda() 100 | self.model.eval() 101 | 102 | if supersampling > 1: 103 | self.model.ss = supersampling 104 | 105 | self.model.temporal_average = temporal_average 106 | 107 | print(f"SUPERSAMPLING: {self.model.ss}") 108 | 109 | factor = 16 110 | assert viewport_size[0] % 16 == 0, f'set width {factor * (viewport_size[0] // factor)}' 111 | assert viewport_size[1] % 16 == 0, f'set height {factor * (viewport_size[1] // factor)}' 112 | 113 | self.renderer = MultiscaleRender(scene, args.input_format, viewport_size, out_buffer_location=out_buffer_location, supersampling=self.model.ss, clear_color=clear_color) 114 | 115 | def infer(self, input_dict=None): 116 | if input_dict: 117 | with torch.set_grad_enabled(False): 118 | out, net_input = self.model(input_dict, return_input=True) 119 | else: 120 | input_dict = self.renderer.render() 121 | # pdb.set_trace() 122 | input_dict = {k: default_input_transform(v)[None] for k, v in input_dict.items()} 123 | if self.gpu: 124 | input_dict = to_gpu(input_dict) 125 | 126 | input_dict['id'] = 0 127 | with torch.set_grad_enabled(False): 128 | out, net_input = self.model(input_dict, return_input=True) 129 | #pdb.set_trace() 130 | out = out['im_out'][0].detach().permute(1, 2, 0) 131 | out = torch.cat( [out, out[:, :, :1] * 0 + 1], 2 ).contiguous() # expected to have 4 channels 132 | 133 | return { 134 | 'output': out, 135 | 'net_input': net_input, 136 | 'input': input_dict 137 | } 138 | -------------------------------------------------------------------------------- /src/READ/gl/render.py: -------------------------------------------------------------------------------- 1 | from glumpy import app, gloo, gl 2 | 3 | from contextlib import contextmanager 4 | import numpy as np 5 | 6 | try: 7 | import pycuda.driver 8 | from pycuda.gl import graphics_map_flags, BufferObject 9 | _PYCUDA = True 10 | except ImportError as err: 11 | print('pycuda import error:', err) 12 | _PYCUDA = False 13 | 14 | import torch 15 | import pdb 16 | 17 | class OffscreenRender: 18 | def __init__(self, viewport_size, out_buffer_location='opengl', clear_color=None): 19 | self._init_buffers(viewport_size, out_buffer_location) 20 | 21 | self.clear_color = clear_color if clear_color is not None else (0., 0., 0., 1.) 22 | 23 | def _init_buffers(self, viewport_size, out_buffer_location): 24 | assert out_buffer_location in ['torch', 'opengl', 'numpy'] 25 | 26 | if out_buffer_location == 'torch': 27 | assert _PYCUDA, 'pycuda is not available' 28 | try: 29 | import pycuda.gl.autoinit # this may fails in headless mode 30 | except: 31 | raise RuntimeError('PyCUDA init failed, cannot use torch buffer') 32 | 33 | _ = torch.cuda.FloatTensor(1, 3, 512,512) # needs init here, otherwise does not work 34 | 35 | color_np = np.zeros((viewport_size[1], viewport_size[0], 4), np.float32) 36 | self.color_buf, self.color_buf_cuda = create_shared_texture(color_np) 37 | self.out_buf = torch.zeros((viewport_size[1], viewport_size[0], 4), dtype=torch.float32).cuda() 38 | elif out_buffer_location == 'opengl': 39 | self.color_buf = np.zeros((viewport_size[1], viewport_size[0], 4), dtype=np.float32).view(gloo.TextureFloat2D) 40 | self.out_buf = self.color_buf 41 | elif out_buffer_location == 'numpy': 42 | self.color_buf = np.zeros((viewport_size[1], viewport_size[0], 4), dtype=np.float32).view(gloo.TextureFloat2D) 43 | self.out_buf = np.zeros((viewport_size[1], viewport_size[0], 3), dtype=np.float32) 44 | 45 | self.viewport_size = viewport_size 46 | self.out_buffer_location = out_buffer_location 47 | 48 | self.depth_buf = gloo.DepthBuffer(viewport_size[0], viewport_size[1], gl.GL_DEPTH_COMPONENT32) 49 | 50 | self.fbo = gloo.FrameBuffer(color=self.color_buf, depth=self.depth_buf) 51 | 52 | def render(self, scene, cull_face=True): 53 | self.fbo.activate() 54 | 55 | gl.glEnable(gl.GL_PROGRAM_POINT_SIZE) 56 | gl.glEnable(gl.GL_DEPTH_TEST) 57 | gl.glShadeModel(gl.GL_FLAT) 58 | 59 | if cull_face: 60 | gl.glEnable(gl.GL_CULL_FACE) 61 | gl.glCullFace(gl.GL_BACK) 62 | else: 63 | gl.glDisable(gl.GL_CULL_FACE) 64 | 65 | gl.glClearColor(*self.clear_color) 66 | gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) 67 | gl.glViewport(0, 0, self.viewport_size[0], self.viewport_size[1]) 68 | 69 | if scene.draw_points: 70 | scene.program.draw(gl.GL_POINTS) 71 | else: 72 | assert scene.index_buffer is not None 73 | scene.program.draw(gl.GL_TRIANGLES, scene.index_buffer) 74 | 75 | if self.out_buffer_location == 'torch': 76 | frame = cpy_texture_to_tensor(self.color_buf_cuda, self.out_buf).clone() 77 | elif self.out_buffer_location == 'opengl': 78 | frame = self.out_buf 79 | else: 80 | gl.glReadPixels(0, 0, self.viewport_size[0], self.viewport_size[1], gl.GL_RGB, gl.GL_FLOAT, self.out_buf) 81 | frame = self.out_buf.copy() 82 | # pdb.set_trace() 83 | self.fbo.deactivate() 84 | return frame 85 | 86 | 87 | @contextmanager 88 | def cuda_activate_array(img): 89 | """Context manager simplifying use of pycuda.gl.RegisteredImage""" 90 | mapping = img.map() 91 | yield mapping.array(0,0) 92 | mapping.unmap() 93 | 94 | 95 | @contextmanager 96 | def cuda_activate_buffer(buf): 97 | mapping = buf.map() 98 | yield mapping.device_ptr() 99 | mapping.unmap() 100 | 101 | 102 | def create_shared_texture(arr, map_flags=None): 103 | """Create and return a Texture2D with gloo and pycuda views.""" 104 | 105 | if map_flags is None: 106 | map_flags = graphics_map_flags.WRITE_DISCARD 107 | 108 | gl_view = arr.view(gloo.TextureFloat2D) 109 | gl_view.activate() # force gloo to create on GPU 110 | gl_view.deactivate() 111 | 112 | cuda_view = pycuda.gl.RegisteredImage( 113 | int(gl_view.handle), gl_view.target, map_flags) 114 | 115 | return gl_view, cuda_view 116 | 117 | 118 | def create_shared_buffer(arr): 119 | """Create and return a BufferObject with gloo and pycuda views.""" 120 | gl_view = arr.view(gloo.VertexBuffer) 121 | gl_view.activate() # force gloo to create on GPU 122 | gl_view.deactivate() 123 | cuda_view = BufferObject(np.long(gl_view.handle)) 124 | return gl_view, cuda_view 125 | 126 | 127 | def cpy_texture_to_tensor(texture, tensor): 128 | """Copy GL texture (cuda view) to pytorch tensor""" 129 | with cuda_activate_array(texture) as src: 130 | cpy = pycuda.driver.Memcpy2D() 131 | 132 | cpy.set_src_array(src) 133 | cpy.set_dst_device(tensor.data_ptr()) 134 | cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = tensor.shape[1] * 4 * 4 135 | cpy.height = tensor.shape[0] 136 | cpy(aligned=False) 137 | 138 | torch.cuda.synchronize() 139 | 140 | return tensor 141 | 142 | 143 | def cpy_tensor_to_texture(tensor, texture): 144 | """Copy pytorch tensor to GL texture (cuda view)""" 145 | with cuda_activate_array(texture) as ary: 146 | cpy = pycuda.driver.Memcpy2D() 147 | 148 | cpy.set_src_device(tensor.data_ptr()) 149 | cpy.set_dst_array(ary) 150 | cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = tensor.shape[1] * 4 * 4 151 | cpy.height = tensor.shape[0] 152 | cpy(aligned=False) 153 | 154 | torch.cuda.synchronize() 155 | 156 | return tensor 157 | 158 | 159 | def cpy_buffer_to_tensor(buffer, tensor): 160 | """Copy GL buffer (cuda view) to pytorch tensor""" 161 | n = tensor.numel()*tensor.element_size() 162 | with cuda_activate_buffer(buffer) as buf_ptr: 163 | pycuda.driver.memcpy_dtod(tensor.data_ptr(), buf_ptr, n) 164 | 165 | 166 | def cpy_tensor_to_buffer(tensor, buffer): 167 | """Copy pytorch tensor to GL buffer (cuda view)""" 168 | n = tensor.numel()*tensor.element_size() 169 | with cuda_activate_buffer(buffer) as buf_ptr: 170 | pycuda.driver.memcpy_dtod(buf_ptr, tensor.data_ptr(), n) 171 | 172 | -------------------------------------------------------------------------------- /src/READ/gl/viewer/camera.py: -------------------------------------------------------------------------------- 1 | from glumpy import app, gl, gloo 2 | import numpy as np 3 | 4 | from READ.gl.viewer.trackball import Trackball 5 | from READ.gl.render import window 6 | 7 | assert window is not None, 'call render.get_window before import' 8 | 9 | 10 | _trackball = None 11 | viewer_flags = { 12 | 'mouse_pressed': False, 13 | 'use_perspective_cam': True, 14 | } 15 | 16 | 17 | @window.event 18 | def on_resize(width, height): 19 | _trackball.resize((width, height)) 20 | 21 | 22 | @window.event 23 | def on_mouse_press(x, y, buttons, modifiers): 24 | """Record an initial mouse press. 25 | """ 26 | 27 | print(buttons, modifiers) 28 | _trackball.set_state(Trackball.STATE_ROTATE) 29 | if (buttons == app.window.mouse.LEFT): 30 | print('left') 31 | ctrl = (modifiers & app.window.key.MOD_CTRL) 32 | shift = (modifiers & app.window.key.MOD_SHIFT) 33 | if (ctrl and shift): 34 | _trackball.set_state(Trackball.STATE_ZOOM) 35 | elif ctrl: 36 | _trackball.set_state(Trackball.STATE_ROLL) 37 | elif shift: 38 | _trackball.set_state(Trackball.STATE_PAN) 39 | elif (buttons == app.window.mouse.MIDDLE): 40 | _trackball.set_state(Trackball.STATE_PAN) 41 | elif (buttons == app.window.mouse.RIGHT): 42 | _trackball.set_state(Trackball.STATE_ZOOM) 43 | 44 | _trackball.down(np.array([x, y])) 45 | 46 | # Stop animating while using the mouse 47 | viewer_flags['mouse_pressed'] = True 48 | 49 | 50 | @window.event 51 | def on_mouse_drag(x, y, dx, dy, buttons): 52 | """Record a mouse drag. 53 | """ 54 | _trackball.drag(np.array([x, y])) 55 | 56 | 57 | @window.event 58 | def on_mouse_release(x, y, button, modifiers): 59 | """Record a mouse release. 60 | """ 61 | viewer_flags['mouse_pressed'] = False 62 | 63 | 64 | @window.event 65 | def on_mouse_scroll(x, y, dx, dy): 66 | """Record a mouse scroll. 67 | """ 68 | if viewer_flags['use_perspective_cam']: 69 | _trackball.scroll(dy) 70 | else: 71 | spfc = 0.95 72 | spbc = 1.0 / 0.95 73 | sf = 1.0 74 | if dy > 0: 75 | sf = spfc * dy 76 | elif dy < 0: 77 | sf = - spbc * dy 78 | 79 | c = self._camera_node.camera 80 | xmag = max(c.xmag * sf, 1e-8) 81 | ymag = max(c.ymag * sf, 1e-8 * c.ymag / c.xmag) 82 | c.xmag = xmag 83 | c.ymag = ymag 84 | 85 | 86 | def get_trackball(init_view, viewport_size, rotation_mode=1): 87 | global _trackball 88 | _trackball = Trackball(init_view, viewport_size, 1, rotation_mode=rotation_mode) 89 | 90 | return _trackball -------------------------------------------------------------------------------- /src/READ/models/__pycache__/compose.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/compose.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/compose.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/compose.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/conv.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/conv.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/conv.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/conv.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/texture.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/texture.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/texture.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/texture.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/unet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/unet.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/models/__pycache__/unet.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/models/__pycache__/unet.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/models/app_encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | # E_attr, implicit_mask 5 | class E_attr(nn.Module): 6 | def __init__(self, input_dim_a=3, output_nc=32): 7 | super(E_attr, self).__init__() 8 | dim = 64 9 | self.model = nn.Sequential( 10 | nn.ReflectionPad2d(3), 11 | nn.Conv2d(input_dim_a, dim, 7, 1), 12 | nn.ReLU(inplace=True), ## size 13 | nn.ReflectionPad2d(1), 14 | nn.Conv2d(dim, dim*2, 4, 2), 15 | nn.ReLU(inplace=True), ## size/2 16 | nn.ReflectionPad2d(1), 17 | nn.Conv2d(dim*2, dim*4, 4, 2), 18 | nn.ReLU(inplace=True), ## size/4 19 | nn.ReflectionPad2d(1), 20 | nn.Conv2d(dim*4, dim*4, 4, 2), 21 | nn.ReLU(inplace=True), ## size/8 22 | nn.ReflectionPad2d(1), 23 | nn.Conv2d(dim*4, dim*4, 4, 2), 24 | nn.ReLU(inplace=True), ## size/16 25 | nn.AdaptiveAvgPool2d(1), 26 | nn.Conv2d(dim*4, output_nc, 1, 1, 0)) ## 1*1 27 | 28 | def forward(self, x): 29 | out = self.model(x) 30 | # out = x.view(x.size(0), -1) 31 | return out 32 | 33 | if __name__ == '__main__': 34 | data = torch.randn((1,3,255,255)) 35 | model = E_attr(3) 36 | out = model(data) 37 | import pdb; pdb.set_trace() -------------------------------------------------------------------------------- /src/READ/models/compose.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.ndimage import gaussian_filter 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import imageio 8 | import cv2 9 | from PIL import Image 10 | 11 | import pdb 12 | 13 | 14 | class ModelAndLoss(nn.Module): 15 | def __init__(self, model, loss, use_mask=False): 16 | super().__init__() 17 | self.model = model 18 | self.loss = loss 19 | self.use_mask = use_mask 20 | 21 | def forward(self, *args, **kwargs): 22 | input = args[:-1] 23 | target = args[-1] 24 | # label = 25 | if not isinstance(input, (tuple, list)): 26 | input = [input] 27 | output = self.model(*input, **kwargs) 28 | 29 | im_out = output['im_out'] 30 | 31 | 32 | loss = {} 33 | if self.use_mask and 'mask' in kwargs and kwargs['mask'] is not None: 34 | loss['vgg_loss'] = self.loss(im_out * kwargs['mask'], target) 35 | loss['huber_loss'] = F.huber_loss(im_out * kwargs['mask'], target) 36 | else: 37 | loss['vgg_loss'] = self.loss(im_out, target) 38 | loss['huber_loss'] = F.huber_loss(im_out, target) 39 | if 'seg_out' in output and 'label' in kwargs and kwargs['label'] is not None: 40 | loss['seg_loss'] = F.cross_entropy(output['seg_out'], kwargs['label'], ignore_index=0) 41 | 42 | return output, loss 43 | 44 | 45 | class BoxFilter(nn.Module): 46 | def __init__(self, in_channels, out_channels, kernel_size=3): 47 | super().__init__() 48 | 49 | self.seq = nn.Sequential( 50 | nn.ReflectionPad2d(kernel_size//2), 51 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None) 52 | ) 53 | 54 | self.weights_init(kernel_size) 55 | 56 | def forward(self, x): 57 | return self.seq(x) 58 | 59 | def weights_init(self, kernel_size): 60 | kernel = torch.ones((kernel_size, kernel_size)) / kernel_size ** 2 61 | self.seq[1].weight.data.copy_(kernel) 62 | 63 | 64 | class GaussianLayer(nn.Module): 65 | _instance = None 66 | 67 | def __init__(self, in_channels, out_channels, kernel_size=21, sigma=3): 68 | super(GaussianLayer, self).__init__() 69 | self.seq = nn.Sequential( 70 | nn.ReflectionPad2d(kernel_size//2), 71 | nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, bias=None, groups=8) 72 | ) 73 | 74 | self.weights_init(kernel_size, sigma) 75 | 76 | def forward(self, x): 77 | return self.seq(x) 78 | 79 | def weights_init(self, kernel_size, sigma): 80 | n= np.zeros((kernel_size, kernel_size)) 81 | n[kernel_size//2, kernel_size//2] = 1 82 | k = gaussian_filter(n,sigma=sigma) 83 | for name, f in self.named_parameters(): 84 | f.data.copy_(torch.from_numpy(k)) 85 | 86 | @staticmethod 87 | def get_instance(): 88 | if GaussianLayer._instance is None: 89 | GaussianLayer._instance = GaussianLayer(8, 8, kernel_size=13, sigma=6).cuda() 90 | 91 | return GaussianLayer._instance 92 | 93 | 94 | class NetAndTexture(nn.Module): 95 | def __init__(self, net, textures, supersampling=1, temporal_average=False): 96 | super().__init__() 97 | self.net = net 98 | self.ss = supersampling 99 | 100 | try: 101 | textures = dict(textures) 102 | except TypeError: 103 | textures = {0: textures} 104 | 105 | self._textures = {k: v.cpu() for k, v in textures.items()} 106 | self._loaded_textures = [] 107 | 108 | self.last_input = None 109 | self.temporal_average = temporal_average 110 | 111 | def load_textures(self, texture_ids): 112 | if torch.is_tensor(texture_ids): 113 | texture_ids = texture_ids.cpu().tolist() 114 | elif isinstance(texture_ids, int): 115 | texture_ids = [texture_ids] 116 | 117 | for tid in texture_ids: 118 | self._modules[str(tid)] = self._textures[tid] 119 | self._loaded_textures = texture_ids 120 | 121 | def unload_textures(self): 122 | for tid in self._loaded_textures: 123 | self._modules[str(tid)].cpu() 124 | del self._modules[str(tid)] 125 | 126 | def reg_loss(self): 127 | loss = 0 128 | for tid in self._loaded_textures: 129 | loss += self._modules[str(tid)].reg_loss() 130 | 131 | return loss 132 | 133 | def forward(self, inputs, **kwargs): 134 | outs = {'im_out':[]} 135 | # outs = {'x1':[],'x2':[],'x4':[],} 136 | texture_ids = inputs['id'] 137 | del inputs['id'] 138 | if torch.is_tensor(texture_ids): 139 | texture_ids = texture_ids.tolist() 140 | elif isinstance(texture_ids, int): 141 | texture_ids = [texture_ids] 142 | 143 | for i, tid in enumerate(texture_ids): # per item in batch 144 | input = {k: v[i][None] for k, v in inputs.items()} 145 | assert 'uv' in list(input)[0], 'first input must be uv' 146 | 147 | texture = self._modules[str(tid)] 148 | j = 0 149 | keys = list(input) 150 | input_multiscale = [] 151 | while j < len(keys): # sample texture at multiple scales 152 | tex_sample = None 153 | input_ex = [] 154 | if 'uv' in keys[j]: 155 | tex_sample = texture(input[keys[j]]) 156 | j += 1 157 | while j < len(keys) and 'uv' not in keys[j]: 158 | input_ex.append(input[keys[j]]) 159 | j += 1 160 | assert tex_sample is not None 161 | input_cat = torch.cat(input_ex + [tex_sample], 1) 162 | 163 | 164 | if self.ss > 1: 165 | input_cat = nn.functional.interpolate(input_cat, scale_factor=1./self.ss, mode='bilinear') 166 | 167 | input_multiscale.append(input_cat) 168 | 169 | if self.temporal_average: 170 | if self.last_input is not None: 171 | for i in range(len(input_multiscale)): 172 | input_multiscale[i] = (input_multiscale[i] + self.last_input[i]) / 2 173 | self.last_input = list(input_multiscale) 174 | 175 | out = self.net(*input_multiscale, **kwargs) 176 | outs['im_out'].append(out['im_out']) 177 | if 'seg_out' in out: 178 | if 'seg_out' not in outs: 179 | outs['seg_out'] = [] 180 | outs['seg_out'].append(out['seg_out']) 181 | 182 | 183 | if 'seg_out' in outs and len(outs['seg_out']) == len(outs['im_out']): 184 | outs['seg_out'] = torch.cat(outs['seg_out'], 0) 185 | outs['im_out'] = torch.cat(outs['im_out'], 0) 186 | 187 | 188 | 189 | if kwargs.get('return_input'): 190 | return outs, input_multiscale 191 | else: 192 | return outs 193 | 194 | 195 | class MultiscaleNet(nn.Module): 196 | def __init__(self, net, input_modality, supersampling=1): 197 | super().__init__() 198 | 199 | self.net = net 200 | self.input_modality = input_modality 201 | self.ss = supersampling 202 | 203 | def forward(self, inputs, **kwargs): 204 | del inputs['id'] 205 | 206 | modes = len(inputs) 207 | assert modes % self.input_modality == 0 208 | 209 | inputs_ms = [] 210 | input_values = list(inputs.values()) 211 | for i in range(modes // self.input_modality): 212 | i0 = i * self.input_modality 213 | i1 = (i + 1) * self.input_modality 214 | cat = torch.cat(input_values[i0:i1], 1) 215 | if self.ss > 1: 216 | cat = nn.functional.interpolate(cat, scale_factor=1./self.ss, mode='bilinear') 217 | inputs_ms.append(cat) 218 | 219 | out = self.net(*inputs_ms, **kwargs) 220 | 221 | if kwargs.get('return_input'): 222 | return out, inputs_ms 223 | else: 224 | return out 225 | 226 | 227 | class RGBTexture(nn.Module): 228 | def __init__(self, texture, supersampling=1): 229 | super().__init__() 230 | 231 | self.texture = texture 232 | self.ss = supersampling 233 | 234 | def forward(self, inputs, **kwargs): 235 | del inputs['id'] 236 | 237 | assert list(inputs) == ['uv_2d'], 'check input format' 238 | 239 | uv = inputs['uv_2d'] 240 | out = self.texture(uv) 241 | 242 | if kwargs.get('return_input'): 243 | return out, uv 244 | else: 245 | return out 246 | -------------------------------------------------------------------------------- /src/READ/models/conv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch import nn, cuda 4 | from torch.autograd import Variable 5 | 6 | 7 | ############################################################################### 8 | # BSD 3-Clause License 9 | # 10 | # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. 11 | # 12 | # Author & Contact: Guilin Liu (guilinl@nvidia.com) 13 | ############################################################################### 14 | 15 | class PartialConv2d(nn.Conv2d): 16 | def __init__(self, *args, **kwargs): 17 | 18 | # whether the mask is multi-channel or not 19 | if 'multi_channel' in kwargs: 20 | self.multi_channel = kwargs['multi_channel'] 21 | kwargs.pop('multi_channel') 22 | else: 23 | self.multi_channel = False 24 | 25 | if 'return_mask' in kwargs: 26 | self.return_mask = kwargs['return_mask'] 27 | kwargs.pop('return_mask') 28 | else: 29 | self.return_mask = False 30 | 31 | super(PartialConv2d, self).__init__(*args, **kwargs) 32 | 33 | if self.multi_channel: 34 | self.weight_maskUpdater = torch.ones(self.out_channels, self.in_channels, self.kernel_size[0], self.kernel_size[1]) 35 | else: 36 | self.weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0], self.kernel_size[1]) 37 | 38 | self.slide_winsize = self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2] * self.weight_maskUpdater.shape[3] 39 | 40 | self.last_size = (None, None) 41 | self.update_mask = None 42 | self.mask_ratio = None 43 | 44 | def forward(self, input, mask_in=None): 45 | 46 | if mask_in is not None or self.last_size != (input.data.shape[2], input.data.shape[3]): 47 | self.last_size = (input.data.shape[2], input.data.shape[3]) 48 | 49 | with torch.no_grad(): 50 | if self.weight_maskUpdater.type() != input.type(): 51 | self.weight_maskUpdater = self.weight_maskUpdater.to(input) 52 | 53 | if mask_in is None: 54 | # if mask is not provided, create a mask 55 | if self.multi_channel: 56 | mask = torch.ones(input.data.shape[0], input.data.shape[1], input.data.shape[2], input.data.shape[3]).to(input) 57 | else: 58 | mask = torch.ones(1, 1, input.data.shape[2], input.data.shape[3]).to(input) 59 | else: 60 | mask = mask_in 61 | 62 | self.update_mask = F.conv2d(mask, self.weight_maskUpdater, bias=None, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=1) 63 | 64 | self.mask_ratio = self.slide_winsize/(self.update_mask + 1e-8) 65 | # self.mask_ratio = torch.max(self.update_mask)/(self.update_mask + 1e-8) 66 | self.update_mask = torch.clamp(self.update_mask, 0, 1) 67 | self.mask_ratio = torch.mul(self.mask_ratio, self.update_mask) 68 | 69 | # if self.update_mask.type() != input.type() or self.mask_ratio.type() != input.type(): 70 | # self.update_mask.to(input) 71 | # self.mask_ratio.to(input) 72 | 73 | raw_out = super(PartialConv2d, self).forward(torch.mul(input, mask) if mask_in is not None else input) 74 | 75 | if self.bias is not None: 76 | bias_view = self.bias.view(1, self.out_channels, 1, 1) 77 | output = torch.mul(raw_out - bias_view, self.mask_ratio) + bias_view 78 | output = torch.mul(output, self.update_mask) 79 | else: 80 | output = torch.mul(raw_out, self.mask_ratio) 81 | 82 | 83 | if self.return_mask: 84 | return output, self.update_mask 85 | else: 86 | return output 87 | -------------------------------------------------------------------------------- /src/READ/models/texture.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import pdb 5 | 6 | 7 | class Texture(nn.Module): 8 | def null_grad(self): 9 | raise NotImplementedError() 10 | 11 | def reg_loss(self): 12 | return 0. 13 | 14 | 15 | class PointTexture(Texture): 16 | def __init__(self, num_channels, size, activation='none', checkpoint=None, init_method='zeros', reg_weight=0.): 17 | super().__init__() 18 | 19 | assert isinstance(size, int), 'size must be int' 20 | 21 | shape = 1, num_channels, size 22 | 23 | if checkpoint: 24 | self.texture_ = torch.load(checkpoint, map_location='cpu')['texture'].texture_ 25 | else: 26 | if init_method == 'rand': 27 | texture = torch.rand(shape) 28 | elif init_method == 'zeros': 29 | texture = torch.zeros(shape) 30 | else: 31 | raise ValueError(init_method) 32 | self.texture_ = nn.Parameter(texture.float()) 33 | 34 | self.activation = activation 35 | self.reg_weight = reg_weight 36 | 37 | def null_grad(self): 38 | self.texture_.grad = None 39 | 40 | def reg_loss(self): 41 | return self.reg_weight * torch.mean(torch.pow(self.texture_, 2)) 42 | 43 | def forward(self, inputs): 44 | if isinstance(inputs, dict): 45 | ids = None 46 | for f, x in inputs.items(): 47 | if 'uv' in f: 48 | ids = x[:, 0].long() 49 | assert ids is not None, 'Input format does not have uv' 50 | else: 51 | if len(inputs.shape)==4: # Bx3xHxW 52 | ids = inputs[:, 0] # BxHxW 53 | else: 54 | ids = inputs 55 | 56 | sh = ids.shape 57 | n_pts = self.texture_.shape[-1] 58 | 59 | 60 | ind = ids.contiguous().view(-1).long().cuda() 61 | 62 | texture = self.texture_.permute(1, 0, 2) # Cx1xN 63 | texture = texture.expand(texture.shape[0], sh[0], texture.shape[2]) # CxBxN 64 | texture = texture.contiguous().view(texture.shape[0], -1) # CxB*N 65 | 66 | sample = torch.index_select(texture, 1, ind) # CxB*H*W 67 | sample = sample.contiguous().view(sample.shape[0], sh[0], sh[1], sh[2]) # CxBxHxW 68 | sample = sample.permute(1, 0, 2, 3) # BxCxHxW 69 | 70 | if self.activation == 'sigmoid': 71 | return torch.sigmoid(sample) 72 | elif self.activation == 'tanh': 73 | return torch.tanh(sample) 74 | #print('sample,',sample.shape) 75 | return sample 76 | 77 | 78 | class MeshTexture(Texture): 79 | def __init__(self, num_channels, size, activation='none', init_method='zeros', levels=4, reg_weight=0.): 80 | super().__init__() 81 | 82 | assert isinstance(size, int), f'size must be int not {size}' 83 | 84 | if init_method == 'rand': 85 | init = lambda shape: torch.rand(shape) 86 | elif init_method == 'zeros': 87 | init = lambda shape: torch.zeros(shape) 88 | elif init_method == '0.5': 89 | init = lambda shape: torch.zeros(shape) + 0.5 90 | else: 91 | raise ValueError(init_method) 92 | 93 | assert levels > 0 94 | self.levels = levels 95 | 96 | for i in range(self.levels): 97 | shape = 1, num_channels, size // 2 ** i, size // 2 ** i 98 | tex = nn.Parameter(init(shape)).float() 99 | self.__setattr__(f'texture_{i}', tex) 100 | 101 | self.activation = activation 102 | self.reg_weight = reg_weight 103 | 104 | def null_grad(self): 105 | for i in range(self.levels): 106 | self.__getattr__(f'texture_{i}').grad = None 107 | 108 | def reg_loss(self): 109 | loss = 0. 110 | tex_weight = [8., 2., 1., 0.] 111 | for i in range(self.levels): 112 | tex = self.__getattr__(f'texture_{i}') 113 | loss += self.reg_weight * tex_weight[i] * torch.mean(torch.pow(tex, 2)) 114 | 115 | return loss 116 | 117 | def forward(self, inputs): 118 | uv = (inputs[:, :2] * 2 - 1.0).transpose(1, 3).transpose(1, 2).contiguous() 119 | 120 | samples = [] 121 | for i in range(self.levels): 122 | tex = self.__getattr__(f'texture_{i}') 123 | sample = nn.functional.grid_sample(tex.expand(uv.shape[0], -1, -1, -1), uv) 124 | samples.append(sample) 125 | 126 | out = samples[0] 127 | for i in range(1, self.levels): 128 | out += samples[i] 129 | 130 | if self.activation == 'sigmoid': 131 | return torch.sigmoid(out) 132 | elif self.activation == 'tanh': 133 | return torch.tanh(out) 134 | 135 | return out 136 | -------------------------------------------------------------------------------- /src/READ/models/unet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | import numpy as np 6 | from functools import partial 7 | 8 | 9 | 10 | class BasicConv(nn.Module): 11 | # Gated_conv 12 | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, relu=True, dilation=1, padding_mode='reflect', act_fun=nn.ELU, normalization=nn.BatchNorm2d): 13 | super().__init__() 14 | self.pad_mode = padding_mode 15 | self.filter_size = kernel_size 16 | self.stride = stride 17 | self.dilation = dilation 18 | 19 | n_pad_pxl = int(self.dilation * (self.filter_size - 1) / 2) 20 | self.flag=relu 21 | 22 | # this is for backward campatibility with older model checkpoints 23 | self.block = nn.ModuleDict( 24 | { 25 | 'conv_f': nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=n_pad_pxl), 26 | 'act_f': act_fun(), 27 | 'conv_m': nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=n_pad_pxl), 28 | 'act_m': nn.Sigmoid(), 29 | 'norm': normalization(out_channels) 30 | } 31 | ) 32 | 33 | def forward(self, x, *args, **kwargs): 34 | if self.flag: 35 | features = self.block.act_f(self.block.conv_f(x)) 36 | else: 37 | features = self.block.conv_f(x) 38 | mask = self.block.act_m(self.block.conv_m(x)) 39 | output = features * mask 40 | output = self.block.norm(output) 41 | 42 | return output 43 | 44 | class ResBlock(nn.Module): 45 | def __init__(self, in_channel, out_channel): 46 | super(ResBlock, self).__init__() 47 | self.main = nn.Sequential( 48 | BasicConv(in_channel, out_channel, kernel_size=3, stride=1, relu=True), 49 | BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False) 50 | ) 51 | 52 | def forward(self, x): 53 | return self.main(x) + x 54 | 55 | 56 | class SCM(nn.Module): 57 | def __init__(self, out_plane): 58 | super(SCM, self).__init__() 59 | self.main = nn.Sequential( 60 | BasicConv(8, out_plane//4, kernel_size=3, stride=1, relu=True), 61 | BasicConv(out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True), 62 | BasicConv(out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True), 63 | BasicConv(out_plane // 2, out_plane-8, kernel_size=1, stride=1, relu=True) 64 | ) 65 | 66 | self.conv = BasicConv(out_plane, out_plane, kernel_size=1, stride=1, relu=False) 67 | 68 | def forward(self, x): 69 | x = torch.cat([x, self.main(x)], dim=1) 70 | return self.conv(x) 71 | 72 | 73 | class EBlock(nn.Module): 74 | def __init__(self, out_channel, num_res=8): 75 | super(EBlock, self).__init__() 76 | 77 | layers = [ResBlock(out_channel, out_channel) for _ in range(num_res)] 78 | 79 | self.layers = nn.Sequential(*layers) 80 | 81 | def forward(self, x): 82 | return self.layers(x) 83 | 84 | 85 | class DBlock(nn.Module): 86 | def __init__(self, channel, num_res=8): 87 | super(DBlock, self).__init__() 88 | 89 | layers = [ResBlock(channel, channel) for _ in range(num_res)] 90 | self.layers = nn.Sequential(*layers) 91 | 92 | def forward(self, x): 93 | return self.layers(x) 94 | 95 | 96 | 97 | class FAM(nn.Module): 98 | def __init__(self, channel): 99 | super(FAM, self).__init__() 100 | self.merge = BasicConv(channel, channel, kernel_size=3, stride=1, relu=False) 101 | 102 | def forward(self, x1, x2): 103 | x = x1 * x2 104 | out = x1 + self.merge(x) 105 | return out 106 | 107 | class AFF(nn.Module): 108 | def __init__(self, in_channel, out_channel): 109 | super(AFF, self).__init__() 110 | self.conv = nn.Sequential( 111 | BasicConv(in_channel, out_channel, kernel_size=1, stride=1, relu=True), 112 | BasicConv(out_channel, out_channel, kernel_size=3, stride=1, relu=False) 113 | ) 114 | 115 | def forward(self, x1, x2, x3, x4): 116 | x = torch.cat([x1, x2, x3, x4], dim=1) 117 | return self.conv(x) 118 | 119 | 120 | 121 | 122 | class UNet(nn.Module): 123 | r""" Rendering network with UNet architecture and multi-scale input. 124 | 125 | Args: 126 | num_input_channels: Number of channels in the input tensor or list of tensors. An integer or a list of integers for each input tensor. 127 | num_output_channels: Number of output channels. 128 | feature_scale: Division factor of number of convolutional channels. The bigger the less parameters in the model. 129 | num_res: Number of block resnet. 130 | """ 131 | def __init__( 132 | self, 133 | num_input_channels=8, 134 | num_output_channels=3, 135 | feature_scale=4, 136 | num_res=4 137 | 138 | ): 139 | super().__init__() 140 | 141 | self.feature_scale = feature_scale 142 | base_channel = 32 143 | 144 | 145 | filters = [64, 128, 256, 512, 1024] 146 | filters = [x // self.feature_scale for x in filters] 147 | 148 | base_channel = 32 149 | 150 | 151 | self.feat_extract = nn.ModuleList([ 152 | BasicConv(8, base_channel, kernel_size=3, relu=True, stride=1), 153 | BasicConv(base_channel, base_channel*2, kernel_size=3, relu=True, stride=2), 154 | BasicConv(base_channel*2, base_channel*4, kernel_size=3, relu=True, stride=2), 155 | BasicConv(base_channel*4, base_channel*2, kernel_size=4, relu=True, stride=2), 156 | BasicConv(base_channel*2, base_channel, kernel_size=4, relu=True, stride=2), 157 | BasicConv(base_channel, 3, kernel_size=3, relu=False, stride=1), 158 | BasicConv(base_channel*4, base_channel*8, kernel_size=3, relu=True, stride=2), 159 | BasicConv(base_channel*8, base_channel*4, kernel_size=4, relu=True, stride=2), 160 | ]) 161 | 162 | self.SCM0 = SCM(base_channel * 8) 163 | self.SCM1 = SCM(base_channel * 4) 164 | self.SCM2 = SCM(base_channel * 2) 165 | 166 | self.FAM0 = FAM(base_channel * 8) 167 | self.FAM1 = FAM(base_channel * 4) 168 | self.FAM2 = FAM(base_channel * 2) 169 | 170 | self.AFFs = nn.ModuleList([ 171 | AFF(base_channel * 15, base_channel*1), 172 | AFF(base_channel * 15, base_channel*2), 173 | AFF(base_channel * 15, base_channel*4), 174 | ]) 175 | 176 | self.Encoder = nn.ModuleList([ 177 | EBlock(base_channel, num_res), 178 | EBlock(base_channel*2, num_res), 179 | EBlock(base_channel*4, num_res), 180 | EBlock(base_channel*8, num_res) 181 | ]) 182 | 183 | self.Decoder = nn.ModuleList([ 184 | DBlock(base_channel * 8, num_res), 185 | DBlock(base_channel * 4, num_res), 186 | DBlock(base_channel * 2, num_res), 187 | DBlock(base_channel, num_res) 188 | ]) 189 | 190 | self.Convs = nn.ModuleList([ 191 | BasicConv(base_channel * 8, base_channel * 4, kernel_size=1, relu=True, stride=1), 192 | BasicConv(base_channel * 4, base_channel * 2, kernel_size=1, relu=True, stride=1), 193 | BasicConv(base_channel * 2, base_channel, kernel_size=1, relu=True, stride=1) 194 | 195 | ]) 196 | 197 | self.ConvsOut = nn.ModuleList( 198 | [ 199 | BasicConv(base_channel * 4, 3, kernel_size=3, relu=False, stride=1), 200 | BasicConv(base_channel * 2, 3, kernel_size=3, relu=False, stride=1), 201 | ] 202 | ) 203 | 204 | self.up =nn.Upsample(scale_factor=4, mode='bilinear') 205 | 206 | 207 | def forward(self, *inputs, **kwargs): 208 | inputs = list(inputs) 209 | 210 | n_input = len(inputs) 211 | 212 | 213 | x =inputs[0] 214 | x_2 = inputs[1] 215 | x_4 = inputs[2] 216 | x_8 = inputs[3] 217 | 218 | z2 = self.SCM2(x_2) 219 | z4 = self.SCM1(x_4) 220 | z8 = self.SCM0(x_8) 221 | 222 | x_ = self.feat_extract[0](x) 223 | res1 = self.Encoder[0](x_) 224 | 225 | z = self.feat_extract[1](res1) 226 | z = self.FAM2(z, z2) 227 | res2 = self.Encoder[1](z) 228 | 229 | z = self.feat_extract[2](res2) 230 | z = self.FAM1(z, z4) 231 | res3 = self.Encoder[2](z) 232 | 233 | z = self.feat_extract[6](res3) 234 | z = self.FAM0(z, z8) 235 | z = self.Encoder[3](z) 236 | 237 | z12 = F.interpolate(res1, scale_factor=0.5) 238 | z13 = F.interpolate(res1, scale_factor=0.25) 239 | 240 | z21 = F.interpolate(res2, scale_factor=2) 241 | z23 = F.interpolate(res2, scale_factor=0.5) 242 | 243 | z32 = F.interpolate(res3, scale_factor=2) 244 | z31 = F.interpolate(res3, scale_factor=4) 245 | 246 | z43 = F.interpolate(z, scale_factor=2) 247 | z42 = F.interpolate(z43, scale_factor=2) 248 | z41 = F.interpolate(z42, scale_factor=2) 249 | 250 | res1 = self.AFFs[0](res1, z21, z31, z41) 251 | res2 = self.AFFs[1](z12, res2, z32, z42) 252 | res3 = self.AFFs[2](z13, z23, res3, z43) 253 | z = self.Decoder[0](z) 254 | 255 | 256 | z = self.feat_extract[7](z) 257 | 258 | z= self.up(z) 259 | z = torch.cat([z, res3], dim=1) 260 | z = self.Convs[0](z) 261 | z = self.Decoder[1](z) 262 | 263 | 264 | z = self.feat_extract[3](z) 265 | z= self.up(z) 266 | 267 | z = torch.cat([z, res2], dim=1) 268 | z = self.Convs[1](z) 269 | z = self.Decoder[2](z) 270 | 271 | 272 | z = self.feat_extract[4](z) 273 | z= self.up(z) 274 | 275 | z = torch.cat([z, res1], dim=1) 276 | z = self.Convs[2](z) 277 | z = self.Decoder[3](z) 278 | z = self.feat_extract[5](z) 279 | 280 | return {'im_out':z} 281 | 282 | if __name__ == '__main__': 283 | import pdb 284 | import time 285 | import numpy as np 286 | 287 | model = UNet().to('cuda') 288 | input = [] 289 | img_sh = [256,256] 290 | sh_unit = 32 291 | img_sh = list(map(lambda a: a-a%sh_unit+sh_unit if a%sh_unit!=0 else a, img_sh)) 292 | down = lambda a,b : a//2**b 293 | input.append(torch.zeros((1,8,down(img_sh[0],0), down(img_sh[1],0)), requires_grad=True).cuda()) 294 | input.append(F.interpolate(input[0],scale_factor=0.5)) 295 | input.append(F.interpolate(input[1],scale_factor=0.5)) 296 | input.append(F.interpolate(input[2],scale_factor=0.5)) 297 | 298 | 299 | model.eval() 300 | st = time.time() 301 | print(input[0].max(),input[0].min()) 302 | with torch.set_grad_enabled(False): 303 | out = model(*input) 304 | pdb.set_trace() 305 | print('model',time.time()-st) 306 | model.to('cpu') 307 | -------------------------------------------------------------------------------- /src/READ/pipelines/__init__.py: -------------------------------------------------------------------------------- 1 | from . pipeline import Pipeline, load_pipeline, save_pipeline -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/ogl.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/ogl.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/ogl.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/ogl.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/pipeline.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/pipeline.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/__pycache__/pipeline.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/pipelines/__pycache__/pipeline.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/pipelines/ogl.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | from pathlib import Path 3 | 4 | import torch 5 | import torch.nn.functional as F 6 | from torch import autograd, optim 7 | 8 | from READ.pipelines import Pipeline 9 | # from READ.datasets.dynamic_tmp import get_datasets 10 | from READ.datasets.dynamic import get_datasets 11 | from READ.models.texture import PointTexture, MeshTexture 12 | from READ.models.compose import NetAndTexture, MultiscaleNet, RGBTexture 13 | from READ.criterions.vgg_loss import VGGLoss 14 | from READ.utils.train import to_device, set_requires_grad, save_model, unwrap_model, image_grid, to_numpy, load_model_checkpoint, freeze 15 | from READ.utils.perform import TicToc, AccumDict, Tee 16 | 17 | import pdb 18 | 19 | TextureOptimizerClass = optim.RMSprop 20 | 21 | 22 | def get_net(input_channels, args): 23 | import importlib 24 | 25 | UNet = getattr(importlib.import_module(f'READ.models.{args.model}'), "UNet") 26 | 27 | net = UNet( 28 | num_input_channels=8, 29 | num_output_channels=3, 30 | feature_scale=4, 31 | num_res=4 32 | ) 33 | return net 34 | 35 | 36 | 37 | 38 | def get_texture(num_channels, size, args, texture_ckpt=None): 39 | if not hasattr(args, 'reg_weight'): 40 | args.reg_weight = 0. 41 | if args.use_mesh: 42 | texture = MeshTexture(num_channels, size, activation=args.texture_activation, reg_weight=args.reg_weight) 43 | else: 44 | texture = PointTexture(num_channels, size, activation=args.texture_activation, reg_weight=args.reg_weight) 45 | 46 | if texture_ckpt is not None: 47 | texture = load_model_checkpoint(texture_ckpt, texture) 48 | return texture 49 | 50 | 51 | def backward_compat(args): 52 | if not hasattr(args, 'input_channels'): 53 | args.input_channels = None 54 | if not hasattr(args, 'conv_block'): 55 | args.conv_block = 'gated' 56 | 57 | if args.pipeline == 'READ.pipelines.ogl.Pix2PixPipeline': 58 | if not hasattr(args, 'input_modality'): 59 | args.input_modality = 1 60 | 61 | return args 62 | 63 | 64 | class TexturePipeline(Pipeline): 65 | def export_args(self, parser): 66 | parser.add_argument('--descriptor_size', type=int, default=8) 67 | parser.add_argument('--texture_size', type=int) 68 | parser.add_argument('--texture_ckpt', type=Path) 69 | parser.add('--texture_lr', type=float, default=1e-1) 70 | parser.add('--texture_activation', type=str, default='none') 71 | parser.add('--n_points', type=int, default=0, help='this is for inference') 72 | 73 | def create(self, args): 74 | args = backward_compat(args) 75 | 76 | if not args.input_channels: 77 | args.input_channels = args.descriptor_size 78 | # [args.descriptor_size] * args.num_mipmap 79 | 80 | net = get_net(args.input_channels, args) 81 | 82 | textures = {} 83 | 84 | if args.inference: 85 | if args.use_mesh: 86 | size = args.texture_size 87 | else: 88 | size = args.n_points 89 | textures = { 90 | 0: get_texture(args.descriptor_size, size, args) 91 | } 92 | else: 93 | # self.ds_train, self.ds_val = get_datasets(args) 94 | self.ds_train, self.ds_val, self.texture_ckpts = get_datasets(args) 95 | 96 | for ds in self.ds_train: 97 | if args.use_mesh: 98 | assert args.texture_size, 'set texture size' 99 | size = args.texture_size 100 | else: 101 | assert ds.scene_data['pointcloud'] is not None, 'set pointcloud' 102 | size = ds.scene_data['pointcloud']['xyz'].shape[0] 103 | # textures[ds.id] = get_texture(args.descriptor_size, size, args) 104 | textures[ds.id] = get_texture(args.descriptor_size, size, args, self.texture_ckpts[ds.id]) 105 | 106 | self.optimizer = optim.Adam(net.parameters(), lr=args.lr) 107 | if len(textures) == 1: 108 | self._extra_optimizer = TextureOptimizerClass(textures[0].parameters(), lr=args.texture_lr) 109 | else: 110 | self._extra_optimizer = None 111 | 112 | self.criterion = args.criterion_module(**args.criterion_args).cuda() 113 | 114 | ss = args.supersampling if hasattr(args, 'supersampling') else 1 115 | 116 | self.net = net 117 | self.textures = textures 118 | self.model = NetAndTexture(net, textures, ss) 119 | 120 | self.args = args 121 | 122 | def state_objects(self): 123 | datasets = self.ds_train 124 | 125 | objs = {'net': self.net} 126 | objs.update({ds.name: self.textures[ds.id] for ds in datasets}) 127 | 128 | return objs 129 | 130 | def dataset_load(self, dataset): 131 | self.model.load_textures([ds.id for ds in dataset]) 132 | 133 | for ds in dataset: 134 | ds.load() 135 | 136 | 137 | def extra_optimizer(self, dataset): 138 | # if we have single dataset, don't recreate optimizer 139 | if self._extra_optimizer is not None: 140 | lr_drop = self.optimizer.param_groups[0]['lr'] / self.args.lr 141 | self._extra_optimizer.param_groups[0]['lr'] = self.args.texture_lr * lr_drop 142 | return self._extra_optimizer 143 | 144 | param_group = [] 145 | for ds in dataset: 146 | param_group.append( 147 | {'params': self.textures[ds.id].parameters()} 148 | ) 149 | 150 | lr_drop = self.optimizer.param_groups[0]['lr'] / self.args.lr 151 | return TextureOptimizerClass(param_group, lr=self.args.texture_lr * lr_drop) 152 | 153 | def dataset_unload(self, dataset): 154 | self.model.unload_textures() 155 | 156 | for ds in dataset: 157 | ds.unload() 158 | self.textures[ds.id].null_grad() 159 | 160 | def get_net(self): 161 | return self.net 162 | 163 | -------------------------------------------------------------------------------- /src/READ/pipelines/pipeline.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import munch 3 | 4 | import torch 5 | 6 | from READ.utils.train import get_module, save_model, load_model_checkpoint 7 | from READ.utils.arguments import deval_args 8 | 9 | 10 | class Pipeline: 11 | def export_args(self, parser): 12 | # add arguments of this pipeline to the cmd parser 13 | raise NotImplementedError() 14 | 15 | def create(self, args): 16 | # return dictionary with pipeline components 17 | raise NotImplementedError() 18 | 19 | def dataset_load(self, *args, **kwargs): 20 | # called before running train/val 21 | pass 22 | 23 | def dataset_unload(self, *args, **kwargs): 24 | # called after running train/val 25 | pass 26 | 27 | def get_net(self): 28 | raise NotImplementedError() 29 | 30 | def extra_optimizer(self, *args): 31 | return None 32 | 33 | 34 | def load_pipeline(checkpoint, args_to_update=None): 35 | if os.path.exists(checkpoint): 36 | ckpt = torch.load(checkpoint, map_location='cpu') 37 | 38 | assert 'args' in ckpt 39 | 40 | if args_to_update: 41 | ckpt['args'].update(args_to_update) 42 | 43 | try: 44 | args = munch.munchify(ckpt['args']) 45 | 46 | pipeline = get_module(args.pipeline)() 47 | pipeline.create(args) 48 | except AttributeError as err: 49 | print('\nERROR: Checkpoint args is incompatible with this version\n', file=sys.stderr) 50 | raise err 51 | 52 | if checkpoint is not None: 53 | load_model_checkpoint(checkpoint, pipeline.get_net()) 54 | 55 | return pipeline, args 56 | 57 | 58 | def save_pipeline(pipeline, save_dir, epoch, stage, args): 59 | objects = pipeline.state_objects() 60 | args_ = deval_args(args) 61 | 62 | for name, obj in objects.items(): 63 | if name=='net' and args.freeze_net: 64 | continue 65 | obj_class = obj.__class__.__name__ 66 | filename = f'{obj_class}' 67 | if epoch is not None: 68 | filename += f'_latest_{epoch}' 69 | # '_stage_{stage} 70 | if name: 71 | name = name.replace('/', '_') 72 | filename = f'{filename}_{name}' 73 | save_path = os.path.join(save_dir, filename + '.pth') 74 | save_model(save_path, obj, args=args_) 75 | -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/arguments.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/arguments.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/arguments.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/arguments.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/perform.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/perform.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/perform.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/perform.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/train.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/train.cpython-36.pyc -------------------------------------------------------------------------------- /src/READ/utils/__pycache__/train.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/JOP-Lee/READ/77bbfa369124b988e98745ca735cd43b990fb1f9/src/READ/utils/__pycache__/train.cpython-39.pyc -------------------------------------------------------------------------------- /src/READ/utils/arguments.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pydoc 3 | from pathlib import Path 4 | import munch 5 | 6 | 7 | class ActionNoYes(argparse.Action): 8 | def __init__(self, 9 | option_strings, 10 | dest, 11 | nargs=0, 12 | const=None, 13 | default=None, 14 | type=None, 15 | choices=None, 16 | required=False, 17 | help="", 18 | metavar=None): 19 | 20 | assert len(option_strings) == 1 21 | assert option_strings[0][:2] == '--' 22 | 23 | name= option_strings[0][2:] 24 | help += f'Use "--{name}" for True, "--no-{name}" for False' 25 | super(ActionNoYes, self).__init__(['--' + name, '--no-' + name], 26 | dest=dest, 27 | nargs=nargs, 28 | const=const, 29 | default=default, 30 | type=type, 31 | choices=choices, 32 | required=required, 33 | help=help, 34 | metavar=metavar) 35 | 36 | def __call__(self, parser, namespace, values, option_string=None): 37 | if option_string.startswith('--no-'): 38 | setattr(namespace, self.dest, False) 39 | else: 40 | setattr(namespace, self.dest, True) 41 | 42 | 43 | class SplitStr(argparse.Action): 44 | 45 | def split(self, x): 46 | if x == '': 47 | return [] 48 | else: 49 | return [self.elem_type(y) for y in x.split(self.delimiter)] 50 | 51 | def __init__(self, 52 | option_strings, 53 | dest, 54 | nargs=None, 55 | const=None, 56 | default=None, 57 | type=None, 58 | choices=None, 59 | required=False, 60 | help="", 61 | metavar=None, 62 | delimiter=',', 63 | elem_type=str): 64 | 65 | self.delimiter = delimiter 66 | self.elem_type = elem_type 67 | 68 | default = self.split(default) 69 | super(SplitStr, self).__init__(option_strings, 70 | dest=dest, 71 | nargs=nargs, 72 | const=const, 73 | default=default, 74 | type=type, 75 | choices=choices, 76 | required=required, 77 | help=help, 78 | metavar=metavar) 79 | 80 | def __call__(self, parser, namespace, values, option_string=None): 81 | print(values) 82 | setattr(namespace, self.dest, self.split(values)) 83 | 84 | 85 | class MyArgumentParser(argparse.ArgumentParser): 86 | def __init__(self, **kwargs): 87 | super(MyArgumentParser, self).__init__(**kwargs) 88 | self.register('action', 'store_bool', ActionNoYes) 89 | self.register('action', 'split_str', SplitStr) 90 | 91 | def add(self, *args, **kwargs): 92 | return self.add_argument(*args, **kwargs) 93 | 94 | 95 | def parse_image_size(string): 96 | error_msg = 'size must have format WxH' 97 | tokens = string.split('x') 98 | if len(tokens) != 2: 99 | raise argparse.ArgumentTypeError(error_msg) 100 | try: 101 | w = int(tokens[0]) 102 | h = int(tokens[1]) 103 | return w, h 104 | except ValueError: 105 | raise argparse.ArgumentTypeError(error_msg) 106 | 107 | 108 | PREFIX = '___' 109 | 110 | 111 | def eval_modules(config): 112 | _upd = {} 113 | for arg in config: 114 | if isinstance(arg, str) and arg.endswith('_module'): 115 | # print(arg, config[arg]) 116 | m = pydoc.locate(config[arg]) 117 | if not m: 118 | raise ValueError(f"module {arg}={config[arg]} not found") 119 | else: 120 | config[arg], _upd[PREFIX + arg] = m, config[arg] 121 | elif isinstance(config[arg], dict): 122 | eval_modules(config[arg]) 123 | config.update(_upd) 124 | 125 | 126 | def eval_paths(config): 127 | _upd = {} 128 | for arg in config: 129 | if isinstance(arg, str) and arg.endswith('_path'): 130 | config[arg], _upd[PREFIX + arg] = Path(config[arg]), config[arg] 131 | elif isinstance(config[arg], dict): 132 | eval_paths(config[arg]) 133 | config.update(_upd) 134 | 135 | 136 | def eval_functions(config): 137 | _upd = {} 138 | for arg in config: 139 | if isinstance(arg, str) and arg.endswith('_func'): 140 | config[arg], _upd[PREFIX + arg] = eval(config[arg]), config[arg] 141 | elif isinstance(config[arg], dict): 142 | eval_functions(config[arg]) 143 | config.update(_upd) 144 | 145 | 146 | def eval_args(args): 147 | args = vars(args) 148 | 149 | eval_modules(args) 150 | eval_paths(args) 151 | eval_functions(args) 152 | 153 | return munch.munchify(args) 154 | 155 | 156 | def deval_args(args): 157 | args = vars(args) 158 | 159 | for key in list(args): 160 | if key.startswith(PREFIX): 161 | args[key.replace(PREFIX, '')] = args[key] 162 | del args[key] 163 | 164 | return args 165 | -------------------------------------------------------------------------------- /src/READ/utils/perform.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import time 3 | from collections import defaultdict 4 | 5 | import numpy as np 6 | 7 | 8 | class TicToc: 9 | def __init__(self): 10 | self.tic_toc_tic = None 11 | 12 | def tic(self): 13 | self.tic_toc_tic = time.time() 14 | 15 | def toc(self): 16 | assert self.tic_toc_tic, 'You forgot to call tic()' 17 | return (time.time() - self.tic_toc_tic) * 1000 18 | 19 | def tocp(self, str): 20 | print(f"{str} took {self.toc():.4f}ms") 21 | 22 | @staticmethod 23 | def print_timing(timing, name=''): 24 | print(f'\n=== {name} Timimg ===') 25 | for fn, times in timing.items(): 26 | min, max, mean, p95 = np.min(times), np.max(times), np.mean(times), np.percentile(times, 95) 27 | print(f'{fn}:\tmin: {min:.4f}\tmax: {max:.4f}\tmean: {mean:.4f}ms\tp95: {p95:.4f}ms') 28 | 29 | 30 | class AccumDict: 31 | def __init__(self, num_f=3): 32 | self.d = defaultdict(list) 33 | self.num_f = num_f 34 | 35 | def add(self, k, v): 36 | self.d[k] += [v] 37 | 38 | def __dict__(self): 39 | return self.d 40 | 41 | def __getitem__(self, key): 42 | return self.d[key] 43 | 44 | def __str__(self): 45 | s = '' 46 | for k in self.d: 47 | if not self.d[k]: 48 | continue 49 | cur = self.d[k][-1] 50 | avg = np.mean(self.d[k]) 51 | format_str = '{:.%df}' % self.num_f 52 | cur_str = format_str.format(cur) 53 | avg_str = format_str.format(avg) 54 | s += f'{k} {cur_str} ({avg_str})\t\t' 55 | return s 56 | 57 | def __repr__(self): 58 | return self.__str__() 59 | 60 | 61 | class Tee(object): 62 | def __init__(self, filename): 63 | self.file = open(filename, 'a', buffering=1) 64 | self.terminal = sys.stdout 65 | 66 | def __del__(self): 67 | self.file.close() 68 | 69 | def write(self, data): 70 | self.file.write(data) 71 | self.terminal.write(data) 72 | 73 | def flush(self): 74 | self.file.flush() 75 | -------------------------------------------------------------------------------- /src/READ/utils/train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import random 4 | import torch 5 | import cv2 6 | import gzip 7 | 8 | import torchvision 9 | 10 | import matplotlib 11 | import matplotlib.cm 12 | 13 | from READ.models.compose import ModelAndLoss 14 | 15 | def get_color_map(num_classes): 16 | color_map = {} 17 | for i in range(0, num_classes): 18 | j = 0 19 | lab = i 20 | tmp = [0,0,0] 21 | while lab: 22 | tmp[0] |= (((lab >> 0) & 1) << (7 - j)) 23 | tmp[1] |= (((lab >> 1) & 1) << (7 - j)) 24 | tmp[2] |= (((lab >> 2) & 1) << (7 - j)) 25 | j += 1 26 | lab >>= 3 27 | color_map[i] = tmp 28 | return color_map 29 | 30 | def label_to_color(label, color_map): 31 | b,h,w = label.shape 32 | trans = True 33 | if isinstance(label, torch.Tensor): 34 | label = label.detach().cpu().numpy() 35 | data = np.zeros((b,h,w,3)) 36 | for i,c in color_map.items(): 37 | data[label==i] = c 38 | return data.transpose(0,-1,1,2) # bchw 39 | 40 | def to_device(data, device='cuda:0'): 41 | if isinstance(data, torch.Tensor): 42 | return data.to(device, non_blocking=True) 43 | elif isinstance(data, dict): 44 | for k in data.keys(): 45 | data[k] = to_device(data[k], device) 46 | 47 | return data 48 | elif isinstance(data, (tuple, list)): 49 | for i in range(len(data)): 50 | data[i] = to_device(data[i], device) 51 | 52 | return data 53 | 54 | return data 55 | 56 | 57 | def set_requires_grad(model, value): 58 | for p in model.parameters(): 59 | p.requires_grad = value 60 | 61 | 62 | def freeze(model, b): 63 | set_requires_grad(model, not b) 64 | 65 | 66 | def save_model(save_path, model, args=None, compress=False): 67 | model = unwrap_model(model) 68 | 69 | if not isinstance(args, dict): 70 | args = vars(args) 71 | 72 | dict_to_save = { 73 | 'state_dict': model.state_dict(), 74 | 'args': args 75 | } 76 | 77 | if compress: 78 | with gzip.open(f'{save_path}.gz', 'wb') as f: 79 | torch.save(dict_to_save, f, pickle_protocol=-1) 80 | else: 81 | torch.save(dict_to_save, save_path, pickle_protocol=-1) 82 | 83 | 84 | def load_model_checkpoint(path, model): 85 | if os.path.exists(path): 86 | ckpt = torch.load(path, map_location='cpu') 87 | model.load_state_dict(ckpt['state_dict']) 88 | return model 89 | 90 | 91 | def unwrap_model(model): 92 | model_ = model 93 | while True: 94 | if isinstance(model_, torch.nn.DataParallel): 95 | model_ = model_.module 96 | elif isinstance(model_, ModelAndLoss): 97 | model_ = model_.model 98 | else: 99 | return model_ 100 | 101 | 102 | def colorize(value, vmin=0, vmax=1, cmap='viridis'): 103 | """ 104 | A utility function for Torch/Numpy that maps a grayscale image to a matplotlib 105 | colormap for use with TensorBoard image summaries. 106 | By default it will normalize the input value to the range 0..1 before mapping 107 | to a grayscale colormap. 108 | Arguments: 109 | - value: 2D Tensor of shape [height, width] or 3D Tensor of shape 110 | [height, width, 1]. 111 | - vmin: the minimum value of the range used for normalization. 112 | (Default: value minimum) 113 | - vmax: the maximum value of the range used for normalization. 114 | (Default: value maximum) 115 | - cmap: a valid cmap named for use with matplotlib's `get_cmap`. 116 | (Default: Matplotlib default colormap) 117 | 118 | Returns a 4D uint8 tensor of shape [height, width, 4]. 119 | """ 120 | 121 | # normalize 122 | vmin = value.min() if vmin is None else vmin 123 | vmax = value.max() if vmax is None else vmax 124 | if vmin!=vmax: 125 | value = (value - vmin) / (vmax - vmin) # vmin..vmax 126 | else: 127 | # Avoid 0-division 128 | value = value*0. 129 | # squeeze last dim if it exists 130 | value = value.squeeze() 131 | 132 | cmapper = matplotlib.cm.get_cmap(cmap) 133 | value = cmapper(value,bytes=True) # (nxmx4) 134 | return np.ascontiguousarray(value[:, :, :3].transpose(2, 0, 1)) 135 | 136 | 137 | def resize(imgs, sz=256): 138 | return torch.nn.functional.interpolate(imgs, size=sz) 139 | 140 | 141 | def to_numpy(t, flipy=False, uint8=True, i=0): 142 | out = t[:] 143 | if len(out.shape) == 4: 144 | out = out[i] 145 | out = out.detach().permute(1, 2, 0) # HWC 146 | out = out.flip([0]) if flipy else out 147 | out = out.detach().cpu().numpy() 148 | out = (out.clip(0, 1)*255).astype(np.uint8) if uint8 else out 149 | return out 150 | 151 | 152 | def image_grid(*args, sz = 256): 153 | num_img = min( min([len(x) for x in args]), 4) 154 | 155 | grid = [] 156 | for a in args: 157 | b = a[:num_img].detach().cpu().float() 158 | if b.shape[1] == 1: 159 | grid.append(torch.cat( [ torch.from_numpy(colorize(bb)).float()[None, ...]/255 for bb in b ], dim=0 )) 160 | # grid.append(torch.cat( [b, b, b], dim=1 ) ) 161 | else: 162 | grid.append(b[:, :3]) 163 | 164 | # print([x.shape for x in grid ]) 165 | imgs = resize( torch.cat(grid), sz=sz) 166 | x = torchvision.utils.make_grid(imgs, nrow = num_img) 167 | 168 | return x 169 | 170 | 171 | def get_module(path): 172 | import pydoc 173 | 174 | m = pydoc.locate(path) 175 | assert m is not None, f'{path} not found' 176 | 177 | return m 178 | 179 | 180 | class SubsetSampler(torch.utils.data.Sampler): 181 | def __init__(self, dataset): 182 | pass 183 | 184 | def __iter__(self): 185 | pass 186 | 187 | def __len__(self): 188 | pass -------------------------------------------------------------------------------- /src/configs/paths_example.yaml: -------------------------------------------------------------------------------- 1 | datasets: 2 | 3 | 4 | "kitti6_368_total": 5 | scene_path: downloads/kitti6_368_total_sample.yaml 6 | target_path: path/READ/src/Data/kitti6_368_total/images_undistorted 7 | target_name_func: "lambda i: f'{i}.png'" 8 | 9 | -------------------------------------------------------------------------------- /src/configs/train_example.yaml: -------------------------------------------------------------------------------- 1 | paths_file: configs/paths_example.yaml 2 | dataset_names: [kitti6_368_total] 3 | 4 | batch_size: 2 5 | batch_size_val: 2 6 | eval_in_train: True 7 | input_format: uv_1d_p1, uv_1d_p1_ds1, uv_1d_p1_ds2, uv_1d_p1_ds3, uv_1d_p1_ds4 8 | # use_mask: True 9 | 10 | epochs: 100 11 | save_freq: 10 12 | save_dir: data/logs 13 | simple_name: True 14 | 15 | # net_ckpt: data/logs/path.pth 16 | 17 | #splitter_module: READ.datasets.splitter.split_by_ratio 18 | #splitter_args: 19 | # train_ratio: 0.9 20 | 21 | splitter_module: READ.datasets.splitter.split_by_step 22 | splitter_args: 23 | val_step: 10 24 | train_drop: 0.0 25 | 26 | train_dataset_args: 27 | keep_fov: False 28 | random_zoom: [0.2,2] 29 | # [0.7, 2] 30 | random_shift: [-1., 1.] 31 | drop_points: 0.0 32 | num_samples: 4 33 | inner_batch: 4 34 | 35 | val_dataset_args: 36 | keep_fov: False 37 | drop_points: 0.0 38 | 39 | criterion_module: READ.criterions.vgg_loss.VGGLoss 40 | criterion_args: 41 | partialconv: false 42 | -------------------------------------------------------------------------------- /src/docker/Dockerfile: -------------------------------------------------------------------------------- 1 | # Base packages for ubuntu 2 | # FROM nvidia/vulkan:1.1.121-cuda-10.1--ubuntu18.04 3 | # FROM nvidia/cudagl:10.1-devel-ubuntu18.04 4 | 5 | FROM nvidia/cudagl:11.3.1-devel-ubuntu18.04 6 | 7 | # clean the libs list 8 | RUN apt-get clean 9 | 10 | # RUN apt-key adv --recv-keys --keyserver keyserver.ubuntu.com A4B469963BF863CC 11 | # RUN packages='libsdl2-2.0 xserver-xorg libvulkan1' \ 12 | # && apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y $packages --no-install-recommends \ 13 | # && VULKAN_API_VERSION=`dpkg -s libvulkan1 | grep -oP 'Version: [0-9|\.]+' | grep -oP '[0-9|\.]+'` && \ 14 | # mkdir -p /etc/vulkan/icd.d/ && \ 15 | # echo \ 16 | # "{\ 17 | # \"file_format_version\" : \"1.0.0\",\ 18 | # \"ICD\": {\ 19 | # \"library_path\": \"libGLX_nvidia.so.0\",\ 20 | # \"api_version\" : \"${VULKAN_API_VERSION}\"\ 21 | # }\ 22 | # }" > /etc/vulkan/icd.d/nvidia_icd.json \ 23 | # && rm -rf /var/lib/apt/lists/* 24 | 25 | RUN apt-get update -q 26 | RUN apt-get install -y \ 27 | git \ 28 | wget \ 29 | bzip2 \ 30 | htop \ 31 | vim \ 32 | nano \ 33 | g++ \ 34 | make \ 35 | build-essential \ 36 | software-properties-common \ 37 | apt-transport-https \ 38 | sudo \ 39 | gosu \ 40 | libgl1-mesa-glx \ 41 | graphviz \ 42 | curl \ 43 | libglew-dev \ 44 | libglfw3-dev \ 45 | libjpeg-dev \ 46 | libjsoncpp-dev \ 47 | libpng-dev \ 48 | mesa-utils \ 49 | xorg-dev \ 50 | xvfb \ 51 | ne 52 | # lxde \ 53 | # x11vnc \ 54 | 55 | ## Download and install miniconda 56 | RUN wget https://repo.continuum.io/miniconda/Miniconda3-py37_4.8.3-Linux-x86_64.sh -O /tmp/miniconda.sh 57 | 58 | RUN /bin/bash /tmp/miniconda.sh -b -p /opt/conda && \ 59 | rm /tmp/miniconda.sh 60 | 61 | ENV PATH=/opt/conda/bin:$PATH 62 | ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH 63 | 64 | RUN conda list 65 | 66 | ## Additional packages 67 | RUN pip install torch==1.10.2+cu113 \ 68 | torchvision==0.11.3+cu113 \ 69 | torchaudio==0.10.2+cu113 \ 70 | -f https://download.pytorch.org/whl/torch_stable.html 71 | 72 | RUN pip install tensorboardX \ 73 | numpy \ 74 | pyyaml \ 75 | pillow \ 76 | munch \ 77 | scipy \ 78 | scikit-image \ 79 | scikit-learn \ 80 | matplotlib \ 81 | trimesh \ 82 | tqdm \ 83 | h5py \ 84 | Cython \ 85 | opencv-python \ 86 | PyOpenGL \ 87 | PyOpenGL_accelerate \ 88 | huepy \ 89 | pyglet 90 | # -i https://pypi.tuna.tsinghua.edu.cn/simple 91 | 92 | # RUN pip install glumpy numpy-quaternion 93 | RUN pip install git+https://github.com/DmitryUlyanov/glumpy \ 94 | numpy-quaternion 95 | 96 | 97 | RUN echo "carla ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers 98 | RUN useradd -m carla 99 | USER carla 100 | ENV SDL_VIDEODRIVER=x11 101 | WORKDIR /home/carla 102 | -------------------------------------------------------------------------------- /src/downloads/kitti6.yaml: -------------------------------------------------------------------------------- 1 | viewport_size: [1216, 368] 2 | intrinsic_matrix: Data/kitti6_368_total/camera.xml 3 | view_matrix: Data/kitti6_368_total/camera.xml 4 | pointcloud: Data/kitti6_368_total/pointcloud.ply 5 | 6 | 7 | data_ratio: 0.01 8 | -------------------------------------------------------------------------------- /src/scripts/train_data.sh: -------------------------------------------------------------------------------- 1 | pip install -v -e ./MyRender 2 | 3 | python train.py --config configs/train_example.yaml --pipeline READ.pipelines.ogl.TexturePipeline --name kitti6_test --headless --------------------------------------------------------------------------------