├── reconstructions └── placeholder.txt ├── teaser.jpg ├── captured_data ├── E │ ├── direct_after.mat │ ├── direct_before.mat │ ├── scan_10-14-19_17-37.mat │ ├── LongExpNoObject_after.mat │ └── LongExpNoObject_before.mat ├── K │ ├── direct_after.mat │ ├── direct_before.mat │ ├── scan_10-15-19_21-12.mat │ ├── LongExpNoObject_after.mat │ └── LongExpNoObject_before.mat ├── Y │ ├── direct_after.mat │ ├── direct_before.mat │ ├── scan_10-15-19_21-57.mat │ ├── LongExpNoObject_after.mat │ └── LongExpNoObject_before.mat ├── Mannequin │ ├── direct_after.mat │ ├── direct_before.mat │ ├── scan_10-30-19_12-55.mat │ ├── LongExpNoObject_after.mat │ └── LongExpNoObject_before.mat └── Mannequin_Assymetric │ ├── direct_after.mat │ ├── direct_before.mat │ ├── scan_10-30-19_16-28.mat │ ├── LongExpNoObject_after.mat │ └── LongExpNoObject_before.mat ├── KeyholeEnvironment.yml ├── LICENSE ├── README.md ├── utils.py └── Demo.py /reconstructions/placeholder.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/teaser.jpg -------------------------------------------------------------------------------- /captured_data/E/direct_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/E/direct_after.mat -------------------------------------------------------------------------------- /captured_data/E/direct_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/E/direct_before.mat -------------------------------------------------------------------------------- /captured_data/K/direct_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/K/direct_after.mat -------------------------------------------------------------------------------- /captured_data/K/direct_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/K/direct_before.mat -------------------------------------------------------------------------------- /captured_data/Y/direct_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Y/direct_after.mat -------------------------------------------------------------------------------- /captured_data/Y/direct_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Y/direct_before.mat -------------------------------------------------------------------------------- /captured_data/E/scan_10-14-19_17-37.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/E/scan_10-14-19_17-37.mat -------------------------------------------------------------------------------- /captured_data/K/scan_10-15-19_21-12.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/K/scan_10-15-19_21-12.mat -------------------------------------------------------------------------------- /captured_data/Mannequin/direct_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin/direct_after.mat -------------------------------------------------------------------------------- /captured_data/Y/scan_10-15-19_21-57.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Y/scan_10-15-19_21-57.mat -------------------------------------------------------------------------------- /captured_data/E/LongExpNoObject_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/E/LongExpNoObject_after.mat -------------------------------------------------------------------------------- /captured_data/E/LongExpNoObject_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/E/LongExpNoObject_before.mat -------------------------------------------------------------------------------- /captured_data/K/LongExpNoObject_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/K/LongExpNoObject_after.mat -------------------------------------------------------------------------------- /captured_data/K/LongExpNoObject_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/K/LongExpNoObject_before.mat -------------------------------------------------------------------------------- /captured_data/Mannequin/direct_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin/direct_before.mat -------------------------------------------------------------------------------- /captured_data/Y/LongExpNoObject_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Y/LongExpNoObject_after.mat -------------------------------------------------------------------------------- /captured_data/Y/LongExpNoObject_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Y/LongExpNoObject_before.mat -------------------------------------------------------------------------------- /captured_data/Mannequin/scan_10-30-19_12-55.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin/scan_10-30-19_12-55.mat -------------------------------------------------------------------------------- /captured_data/Mannequin/LongExpNoObject_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin/LongExpNoObject_after.mat -------------------------------------------------------------------------------- /captured_data/Mannequin/LongExpNoObject_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin/LongExpNoObject_before.mat -------------------------------------------------------------------------------- /captured_data/Mannequin_Assymetric/direct_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin_Assymetric/direct_after.mat -------------------------------------------------------------------------------- /captured_data/Mannequin_Assymetric/direct_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin_Assymetric/direct_before.mat -------------------------------------------------------------------------------- /captured_data/Mannequin_Assymetric/scan_10-30-19_16-28.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin_Assymetric/scan_10-30-19_16-28.mat -------------------------------------------------------------------------------- /captured_data/Mannequin_Assymetric/LongExpNoObject_after.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin_Assymetric/LongExpNoObject_after.mat -------------------------------------------------------------------------------- /captured_data/Mannequin_Assymetric/LongExpNoObject_before.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/KeyholeImaging/HEAD/captured_data/Mannequin_Assymetric/LongExpNoObject_before.mat -------------------------------------------------------------------------------- /KeyholeEnvironment.yml: -------------------------------------------------------------------------------- 1 | name: Keyhole 2 | channels: 3 | - pytorch 4 | - defaults 5 | dependencies: 6 | - blas=1.0=mkl 7 | - ca-certificates=2020.12.8 8 | - certifi=2020.12.5 9 | - cudatoolkit=11.0.221 10 | - cycler=0.10.0 11 | - freetype=2.10.4 12 | - h5py=2.10.0 13 | - hdf5=1.10.4 14 | - icu=58.2 15 | - intel-openmp=2020.2 16 | - jpeg=9b 17 | - kiwisolver=1.3.0 18 | - libpng=1.6.37 19 | - libtiff=4.1.0 20 | - libuv=1.40.0 21 | - lz4-c=1.9.2 22 | - matplotlib=3.3.2 23 | - matplotlib-base=3.3.2 24 | - mkl=2020.2 25 | - mkl-service=2.3.0 26 | - mkl_fft=1.2.0 27 | - mkl_random=1.1.1 28 | - ninja=1.10.2 29 | - numpy=1.19.2 30 | - numpy-base=1.19.2 31 | - olefile=0.46 32 | - openssl=1.1.1i 33 | - pillow=8.0.1 34 | - pip=20.3.1 35 | - pyparsing=2.4.7 36 | - pyqt=5.9.2 37 | - python=3.8.5 38 | - python-dateutil=2.8.1 39 | - pytorch=1.7.1 40 | - qt=5.9.7 41 | - setuptools=51.0.0 42 | - sip=4.19.13 43 | - six=1.15.0 44 | - sqlite=3.33.0 45 | - tk=8.6.10 46 | - torchaudio=0.7.2 47 | - torchvision=0.8.2 48 | - tornado=6.1 49 | - typing_extensions=3.7.4.3 50 | - wheel=0.36.2 51 | - xz=5.2.5 52 | - zlib=1.2.11 53 | - zstd=1.4.5 54 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020, Christopher Metzler. 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without modification, 5 | are permitted provided that the following conditions are met: 6 | 7 | 1. Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 10 | 2. Redistributions in binary form must reproduce the above copyright notice, 11 | this list of conditions and the following disclaimer in the documentation 12 | and/or other materials provided with the distribution. 13 | 14 | 3. Neither the name of the copyright holder nor the names of its contributors 15 | may be used to endorse or promote products derived from this software without 16 | specific prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 19 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 20 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 21 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 22 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 23 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 24 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 25 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 26 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 27 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Keyhole Imaging Code & Dataset 2 | 3 | Code associated with the paper "Keyhole Imaging: Non-Line-of-Sight Imaging and Tracking of Moving Objects Along a Single Optical Path" by Chris Metzler, David Lindell, and Gordon Wetzstein. 4 | 5 | 6 | ![Teaser](./teaser.jpg) 7 | 8 | ## Abstract 9 | Non-line-of-sight (NLOS) imaging and tracking is an emerging technology that allows the shape or position of objects around corners or behind diffusers to be recovered from transient, time-of-flight, measurements. However, existing NLOS approaches require the imaging system to scan a large area on a visible surface, where the indirect light paths of hidden objects are sampled. In many applications, such as robotic vision or autonomous driving, optical access to a large scanning area may not be available, which severely limits the practicality of existing NLOS techniques. Here, we propose a new approach, dubbed keyhole imaging, that captures a sequence of transient measurements along a single optical path, for example, through a keyhole. Assuming that the hidden object of interest moves during the acquisition time, we effectively capture a series of time-resolved projections of the object's shape from unknown viewpoints. We derive inverse methods based on expectation-maximization to recover the object's shape and location using these measurements. Then, with the help of long exposure times and retroreflective tape, we demonstrate successful experimental results with a prototype keyhole imaging system. 10 | 11 | 12 | ## Dependencies 13 | Dependencies are best handled using Anaconda. All dependencies for the testing code can be installed by running "conda env create -f KeyholeEnvironment.yml". 14 | 15 | ## Running the Code 16 | Demo.py will reconstruct the objects and their trajectories. Modify the "reconstruction" argument to set which object/trajectory is reconstructed. 17 | 18 | Cuda is disabled by default. Enabling it will dramatically speed up the reconstructions, but requires a GPU with around 10GB of RAM. 19 | 20 | ## Comments and Questions 21 | Contact metzler@umd.edu 22 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | #Chris Metzler 2020 2 | import torch 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | import time 6 | import torch 7 | 8 | c = 3e8 9 | bin_resolution_t = 16e-12 10 | 11 | # cuda = True if torch.cuda.is_available() else False 12 | cuda=False 13 | Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor 14 | 15 | def TotalVariation(img): 16 | tv_h = ((img[1:,:] - img[:-1,:]).abs()).sum() 17 | tv_w = ((img[:,1:] - img[:,:-1]).abs()).sum() 18 | return (tv_h + tv_w) 19 | 20 | def Laplacian_and_Hessian(img): 21 | if img.is_cuda: 22 | laplacian_filter = torch.tensor([[1., 1., 1.], [1., -8., 1.], [1., 1., 1.]]).cuda() 23 | else: 24 | laplacian_filter = torch.tensor([[1., 1., 1.], [1., -8., 1.], [1., 1., 1.]]) 25 | L = torch.conv2d(img.reshape(1,1,img.shape[0],img.shape[1]),laplacian_filter.reshape(1,1,laplacian_filter.shape[0],laplacian_filter.shape[1])) 26 | H = torch.conv2d(L,laplacian_filter.reshape(1,1,laplacian_filter.shape[0],laplacian_filter.shape[1])) 27 | return L,H 28 | 29 | def Sample_Lambertian(voxel_values=np.ones((32**2,1)),sampling_coordinates=[[0,0,-2],[1,0,-2],[0,2,0]],voxel_coordinates=[],n_t=512,bin_resolution_t=bin_resolution_t,cuda=cuda,jitter=False,filter=False): 30 | Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor 31 | bin_resolution_d = bin_resolution_t*c 32 | 33 | if not voxel_coordinates: 34 | #Use default grid 35 | for j in range(voxel_values.shape[0]): 36 | for i in range(voxel_values.shape[1]): 37 | x_loc=i/(voxel_values.shape[0]-1+.00000000001) 38 | y_loc=1.-j/(voxel_values.shape[1]-1+.00000000001)#Largest y should go first, images are indexed from the top down 39 | z_loc=0 40 | voxel_coordinates+=[[x_loc,y_loc,z_loc]] 41 | if cuda: 42 | voxel_values = voxel_values.reshape(-1,).cuda() 43 | else: 44 | voxel_values=voxel_values.reshape(-1,).cpu() 45 | 46 | n_s=len(sampling_coordinates) 47 | if cuda: 48 | measurement=torch.zeros((n_s,n_t)).type(torch.cuda.FloatTensor) 49 | else: 50 | measurement = torch.zeros((n_s, n_t)) 51 | 52 | sampling_coordinates=Tensor(np.array(sampling_coordinates)) 53 | voxel_coordinates=Tensor(np.array(voxel_coordinates)) 54 | 55 | A=torch.abs(sampling_coordinates[:,2][:,None]-(voxel_coordinates[:,2][None,:]))#[:,None] and [None,:] perform broadcasting 56 | C=torch.cdist(sampling_coordinates,voxel_coordinates) 57 | 58 | cos_theta=A/C 59 | LambDropoff = cos_theta**4 60 | 61 | drop_off_r=1./C**4 62 | 63 | Y = C*2#Need to account for path there and back 64 | 65 | if cuda: 66 | time_coords = (torch.ceil(Y / bin_resolution_d)).type(torch.cuda.LongTensor) 67 | else: 68 | time_coords=(torch.ceil(Y/bin_resolution_d)).type(torch.LongTensor) 69 | 70 | 71 | if jitter: 72 | if cuda: 73 | time_coords=time_coords+(2*torch.rand(size=time_coords.shape)-1).type(torch.LongTensor).cuda() 74 | else: 75 | time_coords = time_coords + (2 * torch.rand(size=time_coords.shape) - 1).type(torch.LongTensor) 76 | 77 | time_coords[time_coords>=measurement.shape[1]]=0#Write all distances that are too far to the first index 78 | 79 | 80 | # start = time.time() 81 | for ind_samp in range(len(sampling_coordinates)): 82 | measurement[ind_samp].put_(time_coords[ind_samp], voxel_values*LambDropoff[ind_samp,:]*drop_off_r[ind_samp,:], accumulate=True) 83 | 84 | 85 | if filter: 86 | kernel_size=3 87 | myfilter = torch.nn.Conv1d(in_channels=n_s, out_channels=n_s, kernel_size=3, stride=1, padding=0, dilation=0, groups=n_s, bias=False, padding_mode='zeros') 88 | if cuda: 89 | myfilter.weight.data = torch.zeros(size=(n_s,1,kernel_size)).cuda()#Should be out_channel x in_channels/groups x kernel_size 90 | else: 91 | myfilter.weight.data = torch.zeros(size=(n_s, 1, kernel_size)) 92 | myfilter.weight.data[:,0,0]=1/3 93 | myfilter.weight.data[:, 0, 1] = 1 / 3 94 | myfilter.weight.data[:, 0, 2] = 1 / 3 95 | myfilter.weight.requires_grad = False 96 | measurement=myfilter(measurement.unsqueeze(0)).squeeze() 97 | 98 | measurement[:,0]=0 99 | return measurement 100 | 101 | def Sample_Retroreflective(voxel_values=np.ones((32**2,1)),sampling_coordinates=[[0,0,-2],[1,0,-2],[0,2,0]],voxel_coordinates=[],n_t=512,bin_resolution_t=bin_resolution_t,cuda=cuda,jitter=False,filter=False): 102 | Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor 103 | bin_resolution_d = bin_resolution_t*c 104 | 105 | if not voxel_coordinates: 106 | #Use default grid 107 | for j in range(voxel_values.shape[0]): 108 | for i in range(voxel_values.shape[1]): 109 | x_loc=i/(voxel_values.shape[0]-1+.00000000001) 110 | y_loc=1.-j/(voxel_values.shape[1]-1+.00000000001)#Largest y should go first, images are indexed from the top down 111 | z_loc=0 112 | voxel_coordinates+=[[x_loc,y_loc,z_loc]] 113 | if cuda: 114 | voxel_values = voxel_values.reshape(-1,).cuda() 115 | else: 116 | voxel_values=voxel_values.reshape(-1,).cpu() 117 | 118 | n_s=len(sampling_coordinates) 119 | if cuda: 120 | measurement=torch.zeros((n_s,n_t)).type(torch.cuda.FloatTensor) 121 | else: 122 | measurement = torch.zeros((n_s, n_t)) 123 | 124 | sampling_coordinates=Tensor(np.array(sampling_coordinates)) 125 | voxel_coordinates=Tensor(np.array(voxel_coordinates)) 126 | 127 | A=torch.abs(sampling_coordinates[:,2][:,None]-(voxel_coordinates[:,2][None,:]))#[:,None] and [None,:] perform broadcasting 128 | C=torch.cdist(sampling_coordinates,voxel_coordinates) 129 | 130 | cos_theta=A/C 131 | LambDropoff = cos_theta**2#4 # Lets model > lambertian drop-off 132 | 133 | drop_off_r=1./C**2 134 | 135 | 136 | Y = C*2#Need to account for path there and back 137 | 138 | if cuda: 139 | time_coords = (torch.ceil(Y / bin_resolution_d)).type(torch.cuda.LongTensor) 140 | else: 141 | time_coords=(torch.ceil(Y/bin_resolution_d)).type(torch.LongTensor) 142 | 143 | 144 | if jitter: 145 | if cuda: 146 | time_coords=time_coords+(2*torch.rand(size=time_coords.shape)-1).type(torch.LongTensor).cuda() 147 | else: 148 | time_coords = time_coords + (2 * torch.rand(size=time_coords.shape) - 1).type(torch.LongTensor) 149 | 150 | time_coords[time_coords>=measurement.shape[1]]=0#Write all distances that are too far to the first index 151 | 152 | for ind_samp in range(len(sampling_coordinates)): 153 | measurement[ind_samp].put_(time_coords[ind_samp], voxel_values*LambDropoff[ind_samp,:]*drop_off_r[ind_samp,:], accumulate=True) 154 | 155 | 156 | if filter: 157 | kernel_size=3 158 | myfilter = torch.nn.Conv1d(in_channels=n_s, out_channels=n_s, kernel_size=3, stride=1, padding=0, dilation=0, groups=n_s, bias=False, padding_mode='zeros') 159 | if cuda: 160 | myfilter.weight.data = torch.zeros(size=(n_s,1,kernel_size)).cuda()#Should be out_channel x in_channels/groups x kernel_size 161 | else: 162 | myfilter.weight.data = torch.zeros(size=(n_s, 1, kernel_size)) 163 | myfilter.weight.data[:,0,0]=1/3 164 | myfilter.weight.data[:, 0, 1] = 1 / 3 165 | myfilter.weight.data[:, 0, 2] = 1 / 3 166 | myfilter.weight.requires_grad = False 167 | measurement=myfilter(measurement.unsqueeze(0)).squeeze() 168 | 169 | 170 | measurement[:,0]=0#Discard the first index 171 | return measurement 172 | 173 | 174 | if __name__ == "__main__": 175 | #Target is by default places along a 32x32, 1mx1m grid at the origin 176 | voxel_values=np.zeros((32,32)) 177 | voxel_values[0,:]=1 178 | voxel_values[10,:]=1 179 | voxel_values[:,0]=1 180 | voxel_values=np.ones((32,32)) 181 | 182 | voxel_coordinates=[] 183 | for j in range(32): 184 | for i in range(32): 185 | x_loc = i / (32-1) 186 | y_loc = 1 - j / (32-1)#Largest y should go first, images are indexed from the top down 187 | z_loc = 0 188 | voxel_coordinates += [[x_loc, y_loc, z_loc]] 189 | 190 | #sampling_coordinates should trace out a path along a 2m x 2m wall that is 2m from the target 191 | sampling_coordinates=[] 192 | for i in range(32): 193 | for j in range(32): 194 | x_samp=-1+2*i/(32.-1) 195 | y_samp=-1+2*j/(32.-1) 196 | z_samp=-2. 197 | sampling_coordinates+=[[x_samp,y_samp,z_samp]] 198 | 199 | 200 | voxel_values=Tensor(voxel_values) 201 | start = time.time() 202 | measurement_torch=Sample_Lambertian(voxel_values,sampling_coordinates,voxel_coordinates=voxel_coordinates) 203 | end = time.time() 204 | print("torch overall: ", end-start) 205 | 206 | measurement_torch=measurement_torch.cpu().data.numpy() 207 | 208 | plt.plot(np.squeeze(measurement_torch[0,:])) 209 | plt.show() 210 | 1 -------------------------------------------------------------------------------- /Demo.py: -------------------------------------------------------------------------------- 1 | #Demo associated with Keyhole Imaging 2 | #Chris Metzler 2020 3 | import argparse 4 | import os 5 | import numpy as np 6 | import utils as utils 7 | import time 8 | import matplotlib.pyplot as plt 9 | from matplotlib.pyplot import imsave 10 | from torch.autograd import Variable 11 | import torch.nn as nn 12 | import torch 13 | import h5py 14 | 15 | os.makedirs("reconstructions", exist_ok=True) 16 | 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument("--lr", type=float, default=0.1, help="adam: learning rate") 19 | parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient") 20 | parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient") 21 | parser.add_argument("--img_size", type=int, default=256, help="size of each image dimension") 22 | parser.add_argument("--padded_size_recon", type=int, default=256, help="size of each image dimension used for forward model") 23 | parser.add_argument("--reconstruction", type=str, default='Mannequin_Assymetric', help="Dataset to reconstruct") 24 | parser.add_argument("--extra_suffix", type=str, default='', help="Suffix to be appended to the end of the filename") 25 | parser.add_argument("--closeup", type=bool, default=True, help="Reconstruct only 60cm x 60cm") 26 | parser.add_argument("--init_var_scalar", type=float, default=1, help="How much to scale the initial variance") 27 | opt = parser.parse_args() 28 | print(opt) 29 | 30 | # cuda = True if torch.cuda.is_available() else False 31 | cuda=False#Requires nvidia GPU. Works with 10GB of ram. Does not work with 6GB of ram. 32 | 33 | n_t=768#Number of time bins used 34 | bin_resolution_t=16e-12#Size of each time bin, seconds 35 | 36 | EM_iters = 30 37 | M_iters = (2+np.arange(0,1*EM_iters,1)).tolist()#Don't optimize for very long at the early iterations where you have a poor estimate of the object's position. 38 | optimizer_reset_rate=1000#Resets at minimum once every M_iters 39 | n_restarts = 1 40 | Beta_rate=1.3 41 | 42 | reconstruction=opt.reconstruction 43 | 44 | #Reconstruction parameters/priors 45 | TV_weight = 0 # 1e3 46 | LaplacianL1_weight = 2e3 47 | L1_weight = 2e3 48 | HessianL1_weight = 0 49 | sigma = 200 50 | object_y_offset = .1 #During reconstruction, how far to place the object grid above the floor 51 | object_scale = 1 # How many meters tall and wide is the object. Overwritten by closeup 52 | background_sub=True 53 | 54 | x_range = 1 #How many meters can the object move along the x-axis (horizontal) 55 | x_offset = -x_range / 2 56 | z_offset = .64 # At it's closest point, how far away is the object 57 | y_offset = 1.13 # Laser's height above the floor 58 | clip_until = 260#Ignore any counts that occur before this time bin 59 | if reconstruction=='E': 60 | with h5py.File('captured_data/E/scan_10-14-19_17-37.mat', 'r') as f: 61 | measurements_np = np.array(f['data']) 62 | xpos = np.array(f['xpos']) 63 | zpos = np.array(f['zpos']) 64 | with h5py.File('captured_data/E/direct_after.mat', 'r') as f: 65 | direct_np = np.array(f['direct_after'])#Direct measurement of wall, used to estimate peak_center 66 | with h5py.File('captured_data/E/LongExpNoObject_after.mat', 'r') as f: 67 | NoObjectMeasurement = np.array(f['LongExpNoObject_after']) 68 | peak_center = 630 69 | elif reconstruction == 'K': 70 | with h5py.File('captured_data/K/scan_10-15-19_21-12.mat', 'r') as f: 71 | measurements_np = np.array(f['data']) 72 | xpos = np.array(f['xpos']) 73 | zpos = np.array(f['zpos']) 74 | with h5py.File('captured_data/K/direct_after.mat', 'r') as f: 75 | direct_np = np.array(f['direct_after'])#Direct measurement of wall, used to estimate peak_center 76 | with h5py.File('captured_data/K/LongExpNoObject_after.mat', 'r') as f: 77 | NoObjectMeasurement = np.array(f['LongExpNoObject_after']) 78 | peak_center = 635 79 | elif reconstruction == 'Y': 80 | with h5py.File('captured_data/Y/scan_10-15-19_21-57.mat', 'r') as f: 81 | measurements_np = np.array(f['data']) 82 | xpos = np.array(f['xpos']) 83 | zpos = np.array(f['zpos']) 84 | with h5py.File('captured_data/Y/direct_after.mat', 'r') as f: 85 | direct_np = np.array(f['direct_after'])#Direct measurement of wall, used to estimate peak_center 86 | with h5py.File('captured_data/Y/LongExpNoObject_after.mat', 'r') as f: 87 | NoObjectMeasurement = np.array(f['LongExpNoObject_after']) 88 | peak_center = 636 89 | elif reconstruction=='Mannequin': 90 | with h5py.File('captured_data/Mannequin/scan_10-30-19_12-55.mat', 'r') as f: 91 | measurements_np = np.array(f['data']) 92 | xpos = np.array(f['xpos']) 93 | zpos = np.array(f['zpos']) 94 | with h5py.File('captured_data/Mannequin/direct_after.mat', 'r') as f: 95 | direct_np = np.array(f['direct_after'])#Direct measurement of wall, used to estimate peak_center 96 | with h5py.File('captured_data/Mannequin/LongExpNoObject_after.mat', 'r') as f: 97 | NoObjectMeasurement = np.array(f['LongExpNoObject_after']) 98 | peak_center = 631 99 | elif reconstruction=='Mannequin_Assymetric': 100 | with h5py.File('captured_data/Mannequin_Assymetric/scan_10-30-19_16-28.mat', 'r') as f: 101 | measurements_np = np.array(f['data']) 102 | xpos = np.array(f['xpos']) 103 | zpos = np.array(f['zpos']) 104 | with h5py.File('captured_data/Mannequin_Assymetric/direct_after.mat', 'r') as f: 105 | direct_np = np.array(f['direct_after'])#Direct measurement of wall, used to estimate peak_center 106 | with h5py.File('captured_data/Mannequin_Assymetric/LongExpNoObject_after.mat', 'r') as f: 107 | NoObjectMeasurement = np.array(f['LongExpNoObject_after']) 108 | peak_center = 630 109 | else: 110 | raise Exception("unrecognized reconstruction") 111 | 112 | measurements_np = measurements_np[:, 0:4096] 113 | if background_sub: 114 | NoObjectMeasurement = NoObjectMeasurement[0:4096] 115 | for i in range(measurements_np.shape[0]): 116 | measurements_np[i, :] = measurements_np[i, :] - NoObjectMeasurement.squeeze() 117 | measurements_np=np.expand_dims(measurements_np, axis=1) 118 | 119 | if opt.closeup: 120 | opt.extra_suffix=opt.extra_suffix + 'Closeup' 121 | object_y_offset =.5 122 | object_scale = .6 123 | 124 | display_some=True#Display only important details 125 | display=False#Display most details 126 | display_many=False#Display all details 127 | plot_trajectory=True 128 | 129 | measurements_np=measurements_np[:,:,0:2048] 130 | 131 | direct_np=direct_np[0:2048] 132 | if display_some: 133 | plt.plot(direct_np[peak_center-4:peak_center+5]) 134 | plt.show() 135 | 136 | measurements_np=measurements_np[:,:,peak_center:peak_center+n_t] 137 | 138 | measurements_np[:,:,0:clip_until]=0 139 | measurements_np=np.reshape(measurements_np,(-1,measurements_np.shape[-1])) 140 | 141 | if display: 142 | plt.plot(np.transpose(measurements_np[:,clip_until:])) 143 | plt.show() 144 | plt.plot(np.transpose(measurements_np[:,:])) 145 | plt.show() 146 | 147 | Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor 148 | 149 | 150 | j_fixed=True 151 | 152 | sampling_coordinates = [] 153 | i_max=33 154 | j_max=20. 155 | k_max=33 156 | 157 | y_range=0#How many meters can the object move along the y-axis (vertical) 158 | z_range=.15#How many meters can the object move along the z-axis (depth) 159 | for k in range(int(k_max)): 160 | if j_fixed: 161 | js=[int(np.floor(j_max/2))] 162 | else: 163 | js=range(int(j_max)) 164 | for j in js: 165 | for i in range(int(i_max)): 166 | x_samp = x_offset+x_range * i / (i_max-1) 167 | y_samp = y_offset+y_range * j / (j_max-1) 168 | z_samp = z_range * k / (k_max-1) 169 | sampling_coordinates += [[x_samp, y_samp, z_samp]] 170 | 171 | 172 | voxel_coordinates=[] 173 | for j in range(opt.padded_size_recon):#Outer indexes must be y (vertical) because that is how images are stored 174 | for i in range(opt.padded_size_recon): 175 | x_loc=-object_scale/2 + i/(opt.padded_size_recon-1)*object_scale#put the center of the object at 0 176 | y_loc=object_y_offset + object_scale - j/(opt.padded_size_recon-1)*object_scale#Largest y should go first, images are indexed from the top down 177 | z_loc=z_offset+z_range 178 | voxel_coordinates+=[[x_loc,y_loc,z_loc]] 179 | 180 | sampling_coordinates_np=np.array(sampling_coordinates) 181 | voxel_coordinates_np=np.array(voxel_coordinates) 182 | 183 | if display: 184 | fig = plt.figure() 185 | ax = fig.add_subplot(111, projection='3d') 186 | ax.scatter(sampling_coordinates_np[:, 0], sampling_coordinates_np[:, 2], sampling_coordinates_np[:, 1], c='r', marker='o') 187 | ax.scatter(voxel_coordinates_np[:, 0], voxel_coordinates_np[:, 2], voxel_coordinates_np[:, 1], c='b', marker='o') 188 | ax.set_xlabel('Width') 189 | ax.set_ylabel('Depth') 190 | ax.set_zlabel('Height') 191 | plt.show() 192 | 193 | 194 | measurements = Tensor(measurements_np) 195 | measurements[:,0]=0#My forward model discards the first index 196 | 197 | 198 | measurements_np=measurements.cpu().data.numpy() 199 | n_samples = measurements.shape[0] 200 | 201 | if display_some: 202 | plt.plot(np.transpose(measurements_np[:,:])) 203 | plt.show() 204 | 205 | 206 | scale_factor_recon=int(opt.padded_size_recon/opt.img_size) 207 | upsampler=torch.nn.Upsample(scale_factor=scale_factor_recon,mode='nearest') 208 | 209 | ## Reconstruct using EM algorithm 210 | def Q_loss(new_recon, sigma, old_recon=None, W_in=None, Beta=1): 211 | if W_in is not None: 212 | W = W_in.clone() 213 | else: 214 | assert old_recon is not None, "Either old_recon or W_in must be provided" 215 | W = torch.zeros(n_samples, len(sampling_coordinates)) 216 | if cuda: 217 | W = W.cuda() 218 | Q = 0 219 | if W_in is None: 220 | with torch.no_grad(): 221 | f_x_old_allthetas = utils.Sample_Lambertian(voxel_values=old_recon, 222 | sampling_coordinates=sampling_coordinates, 223 | voxel_coordinates=voxel_coordinates, n_t=n_t, 224 | bin_resolution_t=bin_resolution_t, cuda=cuda) 225 | if cuda: 226 | f_x_old_allthetas = f_x_old_allthetas.cuda() 227 | f_x_new_allthetas = utils.Sample_Lambertian(voxel_values=new_recon, sampling_coordinates=sampling_coordinates, 228 | voxel_coordinates=voxel_coordinates, n_t=n_t, 229 | bin_resolution_t=bin_resolution_t, cuda=cuda) 230 | if cuda: 231 | f_x_new_allthetas = f_x_new_allthetas.cuda() 232 | for l in range(n_samples): # Sum over all the measurements 233 | measurement_l = measurements[l, :] 234 | if W_in is None: 235 | with torch.no_grad(): 236 | tmp = -torch.sum((measurement_l - f_x_old_allthetas) ** 2, dim=1) 237 | Soft = nn.Softmax(dim=0) 238 | w = Soft(tmp / (2 * sigma ** 2) * Beta) 239 | W[l, :] = w.clone() 240 | Q = Q + torch.sum(W[l, :] * (torch.sum((measurement_l - f_x_new_allthetas) ** 2, dim=1))) # Note the sign is flipped compared withthe minorizing loss. This is the slowest part of my code. 241 | if W_in is None and display == True: 242 | plt.imshow(W.cpu().data.numpy()) 243 | plt.show() 244 | if plot_trajectory: 245 | W_np = W.cpu().data.numpy() 246 | Trajectory = np.zeros((33, 33)) 247 | for i in range(W_np.shape[0]): 248 | Trajectory += W_np[i].reshape((33, 33)) 249 | plt.imshow(Trajectory) 250 | plt.show() 251 | return Q, W 252 | 253 | 254 | Q_best = np.inf 255 | for attempt in range(n_restarts): 256 | if cuda: 257 | recons_var = Variable( np.sqrt(opt.init_var_scalar)*torch.randn([opt.img_size, opt.img_size]).cuda(),requires_grad=True) # Init has implicitly had its square root taken 258 | else: 259 | recons_var = Variable( np.sqrt(opt.init_var_scalar)*torch.randn([opt.img_size, opt.img_size]),requires_grad=True) 260 | 261 | Beta = 1 / (Beta_rate ** EM_iters) 262 | t000 = time.time() 263 | for E_iter in range(EM_iters): 264 | recons = torch.squeeze( 265 | upsampler(recons_var.view(1, 1, opt.img_size, opt.img_size) ** 2)) # Enforces non-negativity on recons 266 | # E Step 267 | # Create a function that majorizes the negative log likelihood 268 | old_recon = recons.clone() 269 | sigma_this_iter = sigma 270 | print("Sigma this iter is " + str(sigma_this_iter)) 271 | warm_start=False#Use oracle object positions as initial trajecotry 272 | # M Step 273 | if E_iter==0 and warm_start: 274 | W_init = torch.zeros(n_samples, len(sampling_coordinates)) 275 | if cuda: 276 | W_init = W_init.cuda() 277 | for i in range(len(xpos)): 278 | x_samp = xpos[i][0] - .5 279 | y_samp = y_offset 280 | z_samp = zpos[i][0] 281 | true_samp = [x_samp, y_samp, z_samp] 282 | for j in range(len(sampling_coordinates)): 283 | samp_coord=sampling_coordinates[j] 284 | if np.linalg.norm(np.array(true_samp)-np.array(samp_coord))<.01: 285 | W_init[i,j]=1 286 | [Q, W] = Q_loss(recons, old_recon=old_recon, sigma=sigma_this_iter, Beta=Beta,W_in=W_init) 287 | else: 288 | [Q, W] = Q_loss(recons, old_recon=old_recon, sigma=sigma_this_iter, Beta=Beta) # Changing sigma over time is equivalent to deterministic annealing https://papers.nips.cc/paper/941-deterministic-annealing-variant-of-the-em-algorithm.pdf 289 | Beta = Beta * Beta_rate 290 | if display: 291 | tmp1 = np.squeeze(recons[:, :].cpu().data.numpy()) 292 | plt.imshow(tmp1) 293 | plt.title("Current Recon") 294 | plt.show() 295 | f_pred = utils.Sample_Lambertian(voxel_values=recons, sampling_coordinates=sampling_coordinates, 296 | voxel_coordinates=voxel_coordinates, n_t=n_t, 297 | bin_resolution_t=bin_resolution_t, cuda=cuda) 298 | t00 = time.time() 299 | for M_iter in range(M_iters[E_iter]): 300 | if M_iter % optimizer_reset_rate == 0: # Every optimizer_reset_the optimizer: resets momentum terms. 301 | optimizer = torch.optim.Adam([recons_var], lr=opt.lr, betas=(opt.b1, opt.b2)) # Decay learning rate over the EM iterations 302 | optimizer.zero_grad() 303 | recons = torch.squeeze( 304 | upsampler(recons_var.view(1, 1, opt.img_size, opt.img_size) ** 2)) # Enforces non-negativity on recons 305 | [Q, _] = Q_loss(recons, sigma=sigma_this_iter, W_in=W) # Need to increase sigma to avoid all probabilities being 0 306 | loss = Q 307 | TV_loss = TV_weight * utils.TotalVariation(recons) 308 | L1_loss = L1_weight * recons.abs().sum() 309 | [L, H] = utils.Laplacian_and_Hessian(recons) 310 | LaplacianL1_loss = LaplacianL1_weight * L.abs().sum() 311 | HessianL1_loss = HessianL1_weight * H.abs().sum() 312 | loss = loss + TV_loss + L1_loss + LaplacianL1_loss + HessianL1_loss 313 | loss.backward() 314 | optimizer.step() 315 | if M_iter % 10 == 0: 316 | print("E iteration %d of %d, M iteration %d of %d" % (E_iter + 1, EM_iters, M_iter, M_iters[E_iter])) 317 | print("Q loss is ", Q.cpu().data.numpy()) 318 | if display_many: 319 | tmp1 = np.squeeze(recons[:, :].cpu().data.numpy()) 320 | plt.imshow(tmp1) 321 | plt.show() 322 | print("%d M Iterations took %f seconds" % (M_iters[E_iter], time.time() - t00)) 323 | if Q.cpu().data.numpy() < Q_best: 324 | Q_best = Q.cpu().data.numpy() 325 | recon_best = recons.cpu().data.numpy() 326 | complete_time = time.time() - t000 327 | print("Total time: ", str(complete_time)) 328 | ReconstructedObject = np.squeeze(recon_best) 329 | if display_some: 330 | plt.imshow(np.fliplr(ReconstructedObject),cmap='inferno') 331 | plt.title("Best Recon") 332 | plt.show() 333 | 334 | 335 | #Based on W at the last attempt, visualize the best guess trajectory of the object 336 | if plot_trajectory: 337 | W_np=W.cpu().data.numpy() 338 | # Display trajectory as sum of pdfs 339 | Trajectory = np.zeros((33, 33)) 340 | for i in range(W_np.shape[0]): 341 | Trajectory += W_np[i].reshape((33, 33)) 342 | Trajectory[Trajectory > 1] = 1 343 | if display_some: 344 | plt.imshow((Trajectory)) 345 | plt.axis('off') 346 | cbar = plt.colorbar() 347 | cbar.ax.tick_params(labelsize=20) 348 | plt.show() 349 | # Display trajectory as sum of one-hot vectors 350 | Trajectory=np.zeros((33,33)) 351 | for i in range(W_np.shape[0]): 352 | tmp=np.zeros((33*33,)) 353 | tmp[np.argmax(W_np[i])]=1 354 | tmp=np.reshape(tmp,(33,33)) 355 | Trajectory=Trajectory+tmp 356 | Trajectory[Trajectory>1]=1 357 | if display_some: 358 | plt.imshow((Trajectory)) 359 | plt.axis('off') 360 | cbar = plt.colorbar() 361 | cbar.ax.tick_params(labelsize=20) 362 | plt.show() 363 | 364 | 365 | #Save numpy arrays with the reconstruction and the trajectory 366 | Save_dir = 'reconstructions' 367 | suffix = '' 368 | if TV_weight is not 0: 369 | suffix = suffix + '_TV' + str(TV_weight) 370 | if LaplacianL1_weight is not 0: 371 | suffix = suffix + '_LapL1' + str(LaplacianL1_weight) 372 | if L1_weight is not 0: 373 | suffix = suffix + '_L1' + str(L1_weight) 374 | if HessianL1_weight is not 0: 375 | suffix = suffix + '_MyHess' + str(HessianL1_weight) 376 | 377 | suffix = suffix + '_n' + str(opt.padded_size_recon) + '_sigma' + str(sigma) 378 | savename = './' + Save_dir + '/'+reconstruction + '_UnknownLocation' + suffix + opt.extra_suffix +'.npz' 379 | if plot_trajectory: 380 | np.savez(savename, ReconstructedObject=ReconstructedObject, Trajectory=Trajectory,W_np=W_np,xpos=xpos,zpos=zpos)#W_np records the trajectory in order 381 | else: 382 | np.savez(savename, ReconstructedObject=ReconstructedObject,W_np=W_np,xpos=xpos,zpos=zpos) 383 | 384 | #Plot reconstructions and trajectories 385 | def LoadPlotandSave(filename,unknown=False,flip=False,legend=False): 386 | a=np.load(filename) 387 | newfilename=filename 388 | recon=a['ReconstructedObject'] 389 | recon=recon[4:251,4:251] 390 | recon_sorted=recon.copy().flatten() 391 | recon_sorted.sort() 392 | recon[recon>recon_sorted[-100]]=recon_sorted[-100] 393 | recon=np.fliplr(recon) 394 | cmaps=['inferno'] 395 | i=0 396 | plt.imshow(recon**.5,cmap=cmaps[i]) 397 | plt.show() 398 | imsave(newfilename[:-4]+'_Object.png', recon,cmap=cmaps[i]) 399 | 400 | if unknown: 401 | Trajectory=a['Trajectory'] 402 | print('Test') 403 | W_np=a['W_np'] 404 | xpos=a['xpos'] 405 | zpos=a['zpos'] 406 | W_np_onehot = np.zeros((W_np.shape[0],33*33)) 407 | for i in range(W_np.shape[0]): 408 | W_np_onehot[i,np.argmax(W_np[i])]=1 409 | W_np_onehot=np.reshape(W_np_onehot,(W_np.shape[0],33,33)) 410 | 411 | W_np_zs=np.array([np.nonzero(W_np_onehot[i])[0] for i in range(W_np.shape[0])]) 412 | W_np_xs = np.array([np.nonzero(W_np_onehot[i])[1] for i in range(W_np.shape[0])]) 413 | if flip: 414 | W_np_xs = 32-W_np_xs 415 | import matplotlib.collections as collections 416 | 417 | fig = plt.figure(num=None, figsize=(4, 4), dpi=200, facecolor='w', edgecolor='k') 418 | ax = fig.add_subplot(1, 1, 1) 419 | gt_call = plt.scatter(-zpos,xpos,s=10,c=(0,0,0)) 420 | for i in range(W_np.shape[0]): 421 | min_alpha=.2 422 | max_alpha=.8 423 | alpha=min_alpha + (max_alpha-min_alpha)*i/W_np.shape[0] 424 | c = ((1 - i / W_np.shape[0]), 0, (i / W_np.shape[0]),alpha) 425 | if i==0: 426 | init_call = ax.scatter(-W_np_zs[i] / 32. * .15, W_np_xs[i] / 32., s=100, c=c)#, alpha=alpha) 427 | else: 428 | final_call = ax.scatter(-W_np_zs[i] / 32. * .15, W_np_xs[i] / 32., s=100, c=c)#, alpha=alpha) 429 | alpha_start=min_alpha 430 | alpha_mid=min_alpha + (max_alpha-min_alpha)*np.floor(W_np.shape[0]/2)/W_np.shape[0] 431 | alpha_end=min_alpha + (max_alpha-min_alpha)*(W_np.shape[0]-1)/W_np.shape[0] 432 | c_start=((1-0/W_np.shape[0]),0,(0/W_np.shape[0]),alpha_start) 433 | c_mid=((1-np.floor(W_np.shape[0]/2)/W_np.shape[0]),0,np.floor(W_np.shape[0]/2)/W_np.shape[0],alpha_mid) 434 | c_end=((1-(W_np.shape[0]-1)/W_np.shape[0]),0,((W_np.shape[0]-1)/W_np.shape[0]),alpha_end) 435 | circles = collections.CircleCollection(sizes=[100, 100, 100], facecolors=[c_start, c_mid, c_end]) 436 | circles.set_alpha(.5) 437 | frame1 = plt.gca() 438 | frame1.axes.get_xaxis().set_visible(False) 439 | frame1.axes.get_yaxis().set_visible(False) 440 | plt.savefig(newfilename[:-4] + '_Trajectory.png', bbox_inches="tight") 441 | if legend: 442 | leg = ax.legend((gt_call, circles), ('Truth', 'Estimate'), fontsize='xx-large',scatterpoints=3, scatteryoffsets=[.5], handlelength=2) 443 | plt.savefig(newfilename[:-4] + '_Trajectory_wLegend.png', bbox_inches="tight") 444 | plt.show() 445 | 446 | flip=False#Use for display only 447 | LoadPlotandSave(savename,unknown=True,flip=flip,legend=False) 448 | 449 | print("FINISHED RECONSTRUCTION") 450 | 451 | --------------------------------------------------------------------------------