├── .gitignore ├── LICENSE ├── README.md ├── model.py ├── run_slam_tum.py └── visualizer.py /.gitignore: -------------------------------------------------------------------------------- 1 | render* 2 | __pycache__* -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 ueda0319 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # iMAP-pytorch 2 | This is reproduction implementation for iMAP(https://edgarsucar.github.io/iMAP/) 3 | 4 | ![tum_test](https://github.com/ueda0319/iMAP_pytorch/wiki/imgs/teddy.gif) 5 | ## dependency 6 | - PyTorch 1.8 7 | - opencv-python 8 | - numpy 9 | 10 | For visualizer(WIP, optional) 11 | - pymcubes 12 | - pyopengl 13 | - pangolin (requires offline installation, check it out here: https://github.com/uoip/pangolin) 14 | 15 | 16 | ## Quick start 17 | Download RGBD dataset from TUM(https://vision.in.tum.de/data/datasets/rgbd-dataset/download) 18 | 19 | run with following commands 20 | ``` 21 | python run_slam_tum.py path/to/rgbd_dataset_freiburg1_teddy/ 22 | ``` 23 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | import torch.nn.functional as F 5 | import numpy as np 6 | import cv2 7 | import copy 8 | import time 9 | #from visualizer import Visualizer 10 | #import mcubes 11 | @torch.jit.script 12 | def mish(x): 13 | return x * torch.tanh(F.softplus(x)) 14 | class IMAP(nn.Module): 15 | def __init__(self): 16 | super(IMAP, self).__init__() 17 | self.B = nn.Linear(3,93, bias=False).cuda() 18 | nn.init.normal_(self.B.weight, 0.0, 25.0) 19 | self.fc1 = nn.Linear(93,256).cuda() 20 | self.fc2 = nn.Linear(256,256).cuda() 21 | self.fc3 = nn.Linear(256+93,256).cuda() 22 | self.fc4 = nn.Linear(256,256).cuda() 23 | self.fc5 = nn.Linear(256,4, bias=False).cuda() 24 | self.fc5.weight.data[3,:]*=0.1 25 | def forward(self, pos): 26 | # Position embedding 27 | gamma =torch.sin(self.B(pos)) 28 | # NeRF model 29 | h1 = F.relu(self.fc1(gamma)) 30 | h2 = torch.cat([F.relu(self.fc2(h1)), gamma],dim=1) 31 | h3 = F.relu(self.fc3(h2)) 32 | h4 = F.relu(self.fc4(h3)) 33 | out = self.fc5(h4) 34 | return out 35 | 36 | class Camera(): 37 | def __init__(self, rgb, depth, px,py,pz,rx,ry,rz,a=0.0, b=0.0,fx=525.0, fy=525.0, cx=319.5, cy=239.5): 38 | self.params = torch.tensor([rx,ry,rz,px,py,pz,a,b]).detach().cuda().requires_grad_(True) 39 | # Camera Calibrations 40 | self.fx = fx 41 | self.fy = fy 42 | self.cx = cx 43 | self.cy = cy 44 | self.K = torch.from_numpy(np.array([ 45 | [fx, 0.0, cx], 46 | [0.0, fy, cy], 47 | [0.0, 0.0, 1.0], 48 | ]).astype(np.float32)).cuda().requires_grad_(False) 49 | self.Ki = torch.from_numpy(np.array([ 50 | [1.0/fx, 0.0, -cx/fx], 51 | [0.0, 1.0/fy, -cy/fy], 52 | [0.0, 0.0, 1.0], 53 | ]).astype(np.float32)).cuda().requires_grad_(False) 54 | # For convert depth from 16bit color 55 | self.factor = 1 / 50000.0 56 | # Images cache 57 | self.rgb = torch.from_numpy((rgb).astype(np.float32)).cuda()*(1.0/256) 58 | self.depth = torch.from_numpy(depth.astype(np.float32)).cuda()*self.factor 59 | # Lighting parameter 60 | self.exp_a = torch.cuda.FloatTensor(1) 61 | self.R = torch.cuda.FloatTensor(3,3).fill_(0) 62 | self.T = torch.cuda.FloatTensor(3,3).fill_(0) 63 | self.Li = torch.cuda.FloatTensor(64).fill_(1.0/64) 64 | self.size = depth.shape 65 | self.update_transform() 66 | self.optimizer = optim.Adam([self.params], lr=0.005) 67 | def setImages(self, rgb, depth): 68 | self.rgb = torch.from_numpy((rgb).astype(np.float32)).cuda()*(1.0/256) 69 | self.depth = torch.from_numpy(depth.astype(np.float32)).cuda()*self.factor 70 | # Calc Transform from camera parameters 71 | def update_transform(self): 72 | i = torch.cuda.FloatTensor(3,3).fill_(0) 73 | i[0,0] = 1 74 | i[1,1] = 1 75 | i[2,2] = 1 76 | w1 = torch.cuda.FloatTensor(3,3).fill_(0) 77 | w1[1, 2] = -1 78 | w1[2, 1] = 1 79 | w2 = torch.cuda.FloatTensor(3,3).fill_(0) 80 | w2[2, 0] = -1 81 | w2[0, 2] = 1 82 | w3 = torch.cuda.FloatTensor(3,3).fill_(0) 83 | w3[0, 1] = -1 84 | w3[1, 0] = 1 85 | 86 | th = torch.norm(self.params[0:3]) 87 | thi = 1.0/(th+1e-12) 88 | n = thi * self.params[0:3] 89 | c = torch.cos(th) 90 | s = torch.sin(th) 91 | w = n[0]*w1 + n[1]*w2 + n[2]*w3 92 | ww = torch.matmul(w,w) 93 | R1 = i + s * w 94 | self.R = R1 + (1.0-c)*ww 95 | self.T = self.params[3:6] 96 | self.exp_a = torch.exp(self.params[6]) 97 | 98 | def rays_batch(self, u, v): 99 | batch_size = u.shape[0] 100 | p = torch.cuda.FloatTensor(batch_size, 3,1).fill_(1) 101 | p[:, 0, 0] = u 102 | p[:, 1, 0] = v 103 | ray = torch.matmul(self.R, torch.matmul(self.Ki, p))[:,:,0] 104 | # Normalize Ray Vector 105 | with torch.no_grad(): 106 | ray_li = (1.0/torch.norm(ray, dim=1).reshape(batch_size,1).expand(batch_size,3)) 107 | rayn = ray * ray_li 108 | return rayn, self.T.reshape(1,3).expand(batch_size,3) 109 | 110 | # Hierarchical sampling 111 | def sample_pdf(bins, weights, N_samples): 112 | # Get pdf 113 | weights = weights + 1e-5 114 | pdf = weights * torch.reciprocal(torch.sum(weights, -1, keepdim=True)) 115 | cdf = torch.cumsum(pdf, -1) 116 | cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) 117 | 118 | u = torch.rand(list(cdf.shape[:-1]) + [N_samples]).cuda() 119 | 120 | # Invert CDF 121 | u = u.contiguous() 122 | inds = torch.searchsorted(cdf, u, right=True) 123 | below = torch.max(torch.zeros_like(inds-1), inds-1) 124 | above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds) 125 | inds_g = torch.stack([below, above], -1) 126 | 127 | matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]] 128 | cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g) 129 | bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g) 130 | 131 | denom = (cdf_g[...,1]-cdf_g[...,0]) 132 | denom = torch.where(denom<1e-5, torch.ones_like(denom), denom) 133 | t = (u-cdf_g[...,0])/denom 134 | samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0]) 135 | samples_cat,_ = torch.sort(torch.cat([samples, bins], -1), dim=-1) 136 | return samples_cat 137 | 138 | 139 | class Mapper(): 140 | def __init__(self): 141 | self.model = IMAP().cuda() 142 | self.model_tracking = IMAP().cuda() 143 | self.cameras = [] 144 | self.optimizer = optim.Adam(self.model.parameters(), lr=0.005) 145 | self.render_id=0 146 | def updateModelForTracking(self): 147 | self.model_tracking.load_state_dict(copy.deepcopy(self.model.state_dict())) 148 | def addCamera(self, rgb_filename, depth_filename, px,py,pz,rx,ry,rz, a, b): 149 | rgb = cv2.imread(rgb_filename, cv2.IMREAD_COLOR) 150 | depth = cv2.imread(depth_filename, cv2.IMREAD_ANYDEPTH) 151 | camera = Camera(rgb, depth, px,py,pz,rx,ry,rz,a,b) 152 | self.cameras.append(camera) 153 | ''' 154 | def render_marching_cube(self, voxel_size=64, threshold=30.0): 155 | with torch.no_grad(): 156 | vs = 2.4/voxel_size 157 | t = np.linspace(-1.2, 1.2, voxel_size+1) 158 | query_pts = np.stack(np.meshgrid(t, t, t), -1).astype(np.float32) 159 | sampling = torch.from_numpy(query_pts.reshape([-1,3])).cuda() 160 | print(sampling.cpu()) 161 | out = self.model(sampling) 162 | sigma=out[:,3].detach().cpu().numpy().reshape(voxel_size+1,voxel_size+1,voxel_size+1) 163 | print(np.min(sigma), np.max(sigma)) 164 | vertices, triangles = mcubes.marching_cubes(sigma, threshold) 165 | if vertices.shape[0]==0: 166 | return np.zeros((3,3), np.float32), np.zeros((3,3), np.float32) 167 | print(vertices.shape) 168 | color_sampling = torch.from_numpy(np.stack(vertices, -1).astype(np.float32).reshape([-1,3]))*vs-1.2 169 | 170 | out = self.model(color_sampling.cuda()) 171 | colors = out[:,:3].detach().cpu().numpy() 172 | 173 | vt=vertices[triangles.flatten()]*vs 174 | cl = colors[triangles.flatten()] 175 | return vt,cl 176 | ''' 177 | 178 | def volume_render(self, dists, sigmas): 179 | max_dist = 1.5 180 | batch_size = dists.shape[0] 181 | step = dists.shape[1] 182 | deltas = dists[:, 1:] - dists[:, :-1] 183 | o = 1-torch.exp(-sigmas[:, :-1,3]*deltas) 184 | wo = torch.cumprod((1-o), 1) 185 | w = torch.cuda.FloatTensor(batch_size, step-1).fill_(0) 186 | w[:, 1:] = o[:, 1:] * wo[:, :-1] 187 | #print(w[0,:]) 188 | dh = w * dists[:, :-1] 189 | ih = w.reshape(batch_size, -1,1).expand(batch_size, -1,3) * sigmas[:, :-1, :3] 190 | d = torch.sum(dh,dim=1) 191 | i = torch.sum(ih, dim=1) 192 | dv = torch.sum(w*torch.square(dists[:, :-1]-d.reshape(batch_size,1).expand(batch_size, step-1)), dim=1) 193 | 194 | # BlackBack 195 | d += wo[:,-1] *max_dist 196 | #i += wo[:,-1].reshape(batch_size,1).expand(batch_size,3) 197 | # dv *= torch.reciprocal(torch.sum(w,axis=1)) 198 | return d,i,dv 199 | 200 | def render_rays(self, u, v, camera, n_coarse=32, n_fine=12, model_freeze=False): 201 | if model_freeze: 202 | model=self.model_tracking 203 | else: 204 | model=self.model 205 | batch_size = u.shape[0] 206 | ray, orig= camera.rays_batch(u,v) 207 | 208 | with torch.no_grad(): 209 | ds = torch.linspace(0.0001,1.2,n_coarse).cuda().reshape(1, n_coarse).expand(batch_size, n_coarse) 210 | rays = orig.reshape(batch_size, 1,3).expand(batch_size, n_coarse,3) + ray.reshape(batch_size, 1,3).expand(batch_size, n_coarse,3)* ds.reshape(batch_size, n_coarse,1).expand(batch_size, n_coarse,3) 211 | sigmas = model(rays.reshape(-1,3)).reshape(batch_size,n_coarse, 4) 212 | 213 | delta = ds[0,1]-ds[0,0] 214 | o = 1-torch.exp(-sigmas[:, :,3]*delta)[:,1:] 215 | t = 1-torch.exp(-torch.cumsum(sigmas[:, :,3]*delta, 1)[:,:-1]) 216 | w = o*t 217 | ds_fine = sample_pdf(ds, w, n_fine) 218 | rays_fine = orig.reshape(batch_size, 1,3).expand(batch_size, n_coarse+n_fine,3) + ray.reshape(batch_size, 1,3).expand(batch_size, n_coarse+n_fine,3)* ds_fine.reshape(batch_size, n_coarse+n_fine,1).expand(batch_size, n_coarse+n_fine,3) 219 | 220 | sigmas_fine = model(rays_fine.reshape(-1,3)).reshape(batch_size,n_coarse+n_fine, 4) 221 | d_f,i_f,dv_f = self.volume_render(ds_fine, sigmas_fine) 222 | return d_f,camera.exp_a*i_f+camera.params[7],dv_f 223 | # render Full image 224 | def render(self, camera): 225 | with torch.no_grad(): 226 | h = camera.size[0] 227 | w = camera.size[1] 228 | depth = torch.cuda.FloatTensor(h,w).fill_(0) 229 | rgb = torch.cuda.FloatTensor(h,w,3).fill_(0) 230 | vstep=40 231 | for v in range(0,h,vstep): 232 | vs = factor * torch.arange(v,v+vstep).reshape(vstep,1).expand(vstep,w).reshape(-1).cuda() 233 | us = factor * torch.arange(w).reshape(1,w).expand(vstep,w).reshape(-1).cuda() 234 | d,r, dv, d_f,i_f,dv_f = self.render_rays(us, vs, camera) 235 | depth[v:v+vstep,:] = d_f.reshape(-1,w) 236 | rgb[v:v+vstep,:,:] = i_f.reshape(-1,w,3) 237 | 238 | rgb_cv = torch.clamp(rgb*255, 0, 255).detach().cpu().numpy().astype(np.uint8) 239 | depth_cv = torch.clamp(depth*50000/256, 0, 255).detach().cpu().numpy().astype(np.uint8) 240 | return rgb_cv, depth_cv 241 | # render preview image 242 | def render_small(self, camera, label): 243 | with torch.no_grad(): 244 | camera.update_transform() 245 | factor=5 246 | h = int(camera.size[0]/factor) 247 | w = int(camera.size[1]/factor) 248 | depth = torch.cuda.FloatTensor(h,w).fill_(0) 249 | rgb = torch.cuda.FloatTensor(h,w,3).fill_(0) 250 | vs = factor * torch.arange(h).reshape(h,1).expand(h,w).reshape(-1).cuda() 251 | us = factor * torch.arange(w).reshape(1,w).expand(h,w).reshape(-1).cuda() 252 | d_f,i_f,dv_f = self.render_rays(us, vs, camera) 253 | depth = d_f.reshape(-1,w) 254 | rgb = i_f.reshape(-1,w,3) 255 | rgb_cv = torch.clamp(rgb*255, 0, 255).detach().cpu().numpy().astype(np.uint8) 256 | depth_cv = torch.clamp(depth*50000/256, 0, 255).detach().cpu().numpy().astype(np.uint8) 257 | 258 | rgb_gt = torch.clamp(camera.rgb*255, 0, 255).detach().cpu().numpy().astype(np.uint8) 259 | depth_gt = torch.clamp(camera.depth*50000/256, 0, 255).detach().cpu().numpy().astype(np.uint8) 260 | prev_rgb = cv2.hconcat([cv2.resize(rgb_cv, (camera.size[1], camera.size[0])), rgb_gt]) 261 | prev_depth = cv2.cvtColor(cv2.hconcat([cv2.resize(depth_cv, (camera.size[1], camera.size[0])), depth_gt]), cv2.COLOR_GRAY2RGB) 262 | prev = cv2.vconcat([prev_rgb, prev_depth]) 263 | cv2.imwrite("render/{}_{:04}.png".format(label,self.render_id), prev) 264 | self.render_id+=1 265 | cv2.imshow("{}_rgb".format(label), prev) 266 | cv2.waitKey(1) 267 | # Mapping step 268 | def mapping(self, batch_size = 200, activeSampling=True): 269 | if len(self.cameras)<5: 270 | camera_ids = np.arange(len(self.cameras)) 271 | else: 272 | camera_ids = np.random.randint(0,len(self.cameras)-2,5) 273 | camera_ids[3] = len(self.cameras)-1 274 | camera_ids[4] = len(self.cameras)-2 275 | 276 | for camera_id in camera_ids: 277 | self.optimizer.zero_grad() 278 | self.cameras[camera_id].optimizer.zero_grad() 279 | self.cameras[camera_id].update_transform() 280 | 281 | 282 | h = self.cameras[camera_id].size[0] 283 | w = self.cameras[camera_id].size[1] 284 | # ActiveSampling 285 | if activeSampling: 286 | with torch.no_grad(): 287 | sh = int(h/8) 288 | sw = int(w/8) 289 | ul=[] 290 | vl=[] 291 | ri = torch.cuda.IntTensor(64).fill_(0) 292 | for i in range(64): 293 | ni = int(batch_size*self.cameras[camera_id].Li[i]) 294 | if ni<1: 295 | ni=1 296 | ri[i] = ni 297 | ul.append((torch.rand(ni)*(sw-1)).to(torch.int16).cuda() + (i%8)*sw) 298 | vl.append((torch.rand(ni)*(sh-1)).to(torch.int16).cuda() + int(i/8)*sh) 299 | us = torch.cat(ul)#(torch.rand(batch_size)*(w-1)).to(torch.int16).cuda() 300 | vs = torch.cat(vl)#(torch.rand(batch_size)*(h-1)).to(torch.int16).cuda() 301 | else: 302 | us = (torch.rand(batch_size)*(w-1)).to(torch.int16).cuda() 303 | vs = (torch.rand(batch_size)*(h-1)).to(torch.int16).cuda() 304 | depth, rgb, depth_var = self.render_rays(us, vs, self.cameras[camera_id]) 305 | rgb_gt = torch.cat([self.cameras[camera_id].rgb[v, u, :].unsqueeze(0) for u, v in zip(us, vs)]) 306 | depth_gt = torch.cat([self.cameras[camera_id].depth[v, u].unsqueeze(0) for u, v in zip(us, vs)]) 307 | depth[depth_gt==0]=0 308 | with torch.no_grad(): 309 | ivar = torch.reciprocal(torch.sqrt(depth_var)) 310 | ivar[ivar.isinf()]=1 311 | ivar[ivar.isnan()]=1 312 | lg = torch.mean(torch.abs(depth-depth_gt)*ivar) 313 | lp = 5*torch.mean(torch.abs(rgb-rgb_gt)) 314 | 315 | 316 | loss = lg+lp 317 | loss.backward()#retain_graph=True) 318 | self.optimizer.step() 319 | #print(lg.detach().cpu(), lp.detach().cpu()) 320 | if camera_id > 0: 321 | self.cameras[camera_id].optimizer.step() 322 | 323 | # Update ActiveSampling 324 | if activeSampling: 325 | with torch.no_grad(): 326 | e = torch.abs(depth-depth_gt) + torch.sum(torch.abs(rgb -rgb_gt), 1) 327 | ris = torch.cumsum(ri, 0) 328 | Li = torch.cuda.FloatTensor(64).fill_(0) 329 | Li[0] = torch.mean(e[:ris[0]]) 330 | for i in range(1,64): 331 | Li[i] = torch.mean(e[ris[i-1]:ris[i]]) 332 | LiS = 1.0 / torch.sum(Li) 333 | self.cameras[camera_id].Li = LiS * Li 334 | # Tracking step 335 | def track(self, camera, batch_size = 200): 336 | self.updateModelForTracking() 337 | for iter in range(20): 338 | camera.optimizer.zero_grad() 339 | camera.update_transform() 340 | h = camera.size[0] 341 | w = camera.size[1] 342 | us = (torch.rand(batch_size)*(w-1)).to(torch.int16).cuda() 343 | vs = (torch.rand(batch_size)*(h-1)).to(torch.int16).cuda() 344 | depth, rgb, depth_var = self.render_rays(us, vs, camera, model_freeze=True) 345 | rgb_gt = torch.cat([camera.rgb[v, u, :].unsqueeze(0) for u, v in zip(us, vs)]) 346 | depth_gt = torch.cat([camera.depth[v, u].unsqueeze(0) for u, v in zip(us, vs)]) 347 | depth[depth_gt==0]=0 348 | with torch.no_grad(): 349 | ivar = torch.reciprocal(torch.sqrt(depth_var)) 350 | ivar[ivar.isinf()]=1 351 | ivar[ivar.isnan()]=1 352 | lg = torch.mean(torch.abs(depth-depth_gt)*ivar) 353 | lp = 5*torch.mean(torch.abs(rgb-rgb_gt)) 354 | loss = lg+lp 355 | loss.backward() 356 | camera.optimizer.step() 357 | p = float(torch.sum(((torch.abs(depth - depth_gt) * torch.reciprocal(depth_gt+1e-12)) < 0.1).int()).cpu().item()) /batch_size 358 | if p>0.8: 359 | break 360 | print("Tracking: P=", p) 361 | #del camera.optimizer 362 | return p -------------------------------------------------------------------------------- /run_slam_tum.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | import glob 4 | import os 5 | import csv 6 | import cv2 7 | import threading 8 | from model import Camera, Mapper 9 | 10 | def read_files(folder_path=r"/home/itsuki/RGBD/rgbd_dataset_freiburg1_teddy/"): 11 | csv_file = open(folder_path + "rgb.txt", "r") 12 | f = csv.reader(csv_file, delimiter=" ") 13 | next(f) 14 | next(f) 15 | next(f) 16 | rgb_filenames = [] 17 | for row in f: 18 | rgb_filenames.append("{}{}".format(folder_path, row[1])) 19 | csv_file = open(folder_path + "depth.txt", "r") 20 | f = csv.reader(csv_file, delimiter=" ") 21 | next(f) 22 | next(f) 23 | next(f) 24 | depth_filenames = [] 25 | for row in f: 26 | depth_filenames.append("{}{}".format(folder_path, row[1])) 27 | return rgb_filenames, depth_filenames 28 | 29 | def mappingThread(mapper): 30 | while True: 31 | mapper.mapping() 32 | time.sleep(0.1) 33 | print("update map") 34 | def main(): 35 | mapper = Mapper() 36 | if 2<= len(sys.argv): 37 | rgb_filenames, depth_filenames = read_files(sys.argv[1]) 38 | else: 39 | rgb_filenames, depth_filenames = read_files() 40 | frame_length = min(len(rgb_filenames), len(depth_filenames)) 41 | mapper.addCamera(rgb_filenames[0], 42 | depth_filenames[0], 43 | 0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0) 44 | fixed_camera = Camera(cv2.imread(rgb_filenames[0], cv2.IMREAD_COLOR), 45 | cv2.imread(depth_filenames[0], cv2.IMREAD_ANYDEPTH), 46 | 0.0,0.0,0.0,1e-8,1e-8,1e-8,0.0,0.0) 47 | tracking_camera = Camera(cv2.imread(rgb_filenames[0], cv2.IMREAD_COLOR), 48 | cv2.imread(depth_filenames[0], cv2.IMREAD_ANYDEPTH), 49 | 0.0,0.0,0.0,1e-8,1e-8,1e-8,0.0,0.0) 50 | # Initialize Map 51 | for i in range(200): 52 | mapper.mapping(batch_size=200, activeSampling=False) 53 | 54 | # For calc kinematics for camera motion 55 | last_pose = tracking_camera.params 56 | camera_vel = torch.tensor([0.0,0.0,0.0,0.0,0.0,0.0, 0.0, 0.0]).detach().cuda().requires_grad_(True) 57 | 58 | last_kf=0 59 | 60 | mapping_thread = threading.Thread(target=mappingThread, args=(mapper,)) 61 | mapping_thread.start() 62 | for frame in range(1,frame_length): 63 | tracking_camera.params.data += camera_vel 64 | tracking_camera.setImages(cv2.imread(rgb_filenames[frame], cv2.IMREAD_COLOR), 65 | cv2.imread(depth_filenames[frame], cv2.IMREAD_ANYDEPTH)) 66 | pe = mapper.track(tracking_camera) 67 | camera_vel = 0.2 * camera_vel + 0.8*(tracking_camera.params-last_pose) 68 | last_pose = tracking_camera.params 69 | if pe < 0.65 and frame-last_kf>5: 70 | p = tracking_camera.params 71 | mapper.addCamera(rgb_filenames[frame], 72 | depth_filenames[frame], 73 | last_pose[3],last_pose[4],last_pose[5], 74 | last_pose[0],last_pose[1],last_pose[2], 75 | last_pose[6],last_pose[7]) 76 | print("Add keyframe") 77 | print(last_pose.cpu()) 78 | last_kf=frame 79 | # Render from tracking camera 80 | mapper.render_small(tracking_camera, "view") 81 | # Render from fixed camera 82 | #mapper.render_small(fixed_camera, "fixed_camera") 83 | mapping_thread.join() 84 | 85 | main() -------------------------------------------------------------------------------- /visualizer.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import mcubes 3 | import OpenGL.GL as gl 4 | import pangolin 5 | import threading 6 | import time 7 | class Visualizer(threading.Thread): 8 | def __init__(self): 9 | super(Visualizer, self).__init__() 10 | self.points = np.zeros((3,3), np.float32) 11 | self.colors = np.zeros((3,3), np.float32) 12 | self.is_run=True 13 | def run(self): 14 | 15 | pangolin.CreateWindowAndBind('Main', 640, 480) 16 | gl.glEnable(gl.GL_DEPTH_TEST) 17 | 18 | # Define Projection and initial ModelView matrix 19 | self.scam = pangolin.OpenGlRenderState( 20 | pangolin.ProjectionMatrix(640, 480, 420, 420, 320, 240, 0.2, 100), 21 | pangolin.ModelViewLookAt(-2, 2, -2, 0, 0, 0, pangolin.AxisDirection.AxisY)) 22 | handler = pangolin.Handler3D(self.scam) 23 | 24 | # Create Interactive View in window 25 | self.dcam = pangolin.CreateDisplay() 26 | self.dcam.SetBounds(0.0, 1.0, 0.0, 1.0, -640.0/480.0) 27 | self.dcam.SetHandler(handler) 28 | while not pangolin.ShouldQuit() and self.is_run: 29 | 30 | gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT) 31 | gl.glClearColor(1.0, 1.0, 1.0, 1.0) 32 | self.dcam.Activate(self.scam) 33 | 34 | # Render OpenGL Cube 35 | #pangolin.glDrawColouredCube() 36 | 37 | # Draw Point Cloud 38 | #points = np.random.random((100000, 3)) * 10 39 | gl.glPointSize(2) 40 | gl.glColor3f(1.0, 0.0, 0.0) 41 | 42 | #colors = map(float, colors.flat) 43 | 44 | gl.glEnableClientState(gl.GL_VERTEX_ARRAY) 45 | gl.glVertexPointerf(self.points) 46 | gl.glEnableClientState(gl.GL_COLOR_ARRAY) 47 | gl.glColorPointer(3, gl.GL_FLOAT, 0, self.colors) 48 | gl.glDrawArrays(gl.GL_TRIANGLES, 0, self.points.shape[0]) 49 | #pangolin.glDrawVertices(points.shape[0], points, gl.GL_TRIANGLES) 50 | #pangolin.DrawPoints(points) 51 | 52 | pangolin.FinishFrame() 53 | time.sleep(0.05) 54 | self.is_run=False 55 | def test(): 56 | X, Y, Z = np.mgrid[:30, :30, :30] 57 | u = ((X-15)**2 + (Y-15)**2 + (Z-15)**2 - 8**2) 58 | print(u.shape) 59 | vertices, triangles = mcubes.marching_cubes(u, 0) 60 | colors = (vertices-np.min(vertices)) / (np.max(vertices)-np.min(vertices)) 61 | colors[:, 0] = 1.0-colors[:,0] 62 | vt = vertices[triangles.flatten()].astype(np.float32)*0.1 63 | cl = colors[triangles.flatten()].astype(np.float32) 64 | vis = Visualizer() 65 | vis.points = vt 66 | vis.colors = cl 67 | vis.start() 68 | #vis.start() 69 | t = 0.0 70 | while vis.is_run:#not pangolin.ShouldQuit(): 71 | # vis.draw() 72 | # print("b") 73 | vis.points = vt +np.sin(t) 74 | vis.colors = cl 75 | t+=0.01 76 | time.sleep(0.01) 77 | vis.join() 78 | if __name__ == "__main__": 79 | test() --------------------------------------------------------------------------------