├── LICENSE ├── README.md ├── calibration ├── data │ ├── capture_1_0_inch.mat │ ├── capture_1_5_inch.mat │ ├── capture_2_0_inch.mat │ ├── capture_2_5_inch.mat │ ├── capture_3_0_inch.mat │ ├── capture_3_5_inch.mat │ ├── capture_4_0_inch.mat │ ├── capture_4_5_inch.mat │ ├── capture_5_0_inch.mat │ ├── capture_5_5_inch.mat │ ├── capture_6_0_inch.mat │ ├── capture_6_5_inch.mat │ ├── capture_7_0_inch.mat │ ├── capture_7_5_inch.mat │ ├── capture_8_0_inch.mat │ └── capture_direct.mat ├── diffusion_model.py ├── optimize.py └── sweep_n.py ├── cdt_reconstruction.py ├── data ├── cones.mat ├── letter_s.mat ├── letter_t.mat ├── letter_u_50.mat ├── letter_u_52.mat ├── letter_u_54.mat ├── letter_u_56.mat ├── letter_u_58.mat ├── letter_u_60.mat ├── letter_u_62.mat ├── letter_u_64.mat ├── letter_u_66.mat ├── letter_u_68.mat ├── letter_u_70.mat ├── letter_u_72.mat ├── letter_u_74.mat ├── letter_u_76.mat ├── letter_u_78.mat ├── letter_u_80.mat ├── letters_ut.mat ├── mannequin.mat ├── resolution_50.mat └── resolution_70.mat ├── main.py ├── requirements.txt ├── setup.bash └── utils.py /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | Copyright (c) 2020, Stanford University 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 5 | 6 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 7 | 8 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 9 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # README 2 | 3 | This repository contains the code to reproduce the results presented in "Three-dimensional imaging through scattering media based on confocal diffuse tomography" by David B. Lindell and Gordon Wetzstein. 4 | 5 | [Project Webpage](http://www.computationalimaging.org/publications/confocal-diffuse-tomography/) 6 | 7 | Code prerequisites: Linux-based operating system, installation of Anaconda or Miniconda. Developed and tested with Python 3.6 on Arch Linux kernel 5.6.8-arch1-1 with an Intel Core i7-9750H CPU. 8 | 9 | The code is developed in Pytorch and supports both CPU and GPU execution. Most NVIDIA GPUs should be automatically detected and used. 10 | 11 | # List of contents 12 | * `main.py` - wrapper program called to generate results 13 | * `cdt_reconstruction.py` - contains processing code to run the reconstructions and reproduce results of main 14 | paper. 15 | * `data/cones.mat` - data file for Fig. 3 16 | * `data/letter_s.mat` - data file for Fig. 2 17 | * `data/letters_ut.mat` - data file for Fig. 3 18 | * `data/letter_t.mat` - data file for Fig. 3 19 | * `data/letter_u*.mat` - data files shown in Supplementary Movie 2 for an object at varying distances behind the media 20 | * `data/mannequin.mat` - data file for Fig. 3 21 | * `data/resolution*.mat` - data files for Supplementary Fig. 17 22 | * `README.txt` - this file 23 | * `requirements.txt` - list of requisite Python packages 24 | * `setup.bash` - example script to install a conda environment and required packages 25 | and run the code 26 | * `utils.py` - contains helper functions for the processing 27 | 28 | # Instructions 29 | To run the demo code, follow the instructions in the setup.bash script. This script assumes that you have Anaconda or Miniconda installed. If not, follow the provided link for instructions on how to do this. The script sets up a new Python 3.6 environment, installs the required Python packages (listed in `requirements.txt`), and runs the reconstruction code. 30 | 31 | The install time of the setup.bash script is less than 5 minutes on the tested configuration. Runtime of the demo program `main.py` is less than one minute on the tested configuration and outputs. 32 | 33 | The expected output of the `main.py` program is a figure showing maximum intensity projections of the selected 3D measurement volume and reconstruction. 34 | 35 | Execute `python main.py --help` for a list of commandline options for running the demo. 36 | 37 | [Anaconda installation instructions](https://docs.anaconda.com/anaconda/install/) 38 | 39 | # Calibration 40 | 41 | Code to estimate the parameters of the scattering medium is included in the `calibration` folder. The data capture and optimization procedure are described in the Methods and Supplementary Information of the paper. We jointly optimize the parameters over captured time-resolved transmittance profiles for varying thicknesses of the scattering material (from 1 to 8 inches). Since the optimization is non-linear, we vary the initialized values of the reduced scattering coefficient and refractive index and select the values that best fit the data. Analyzing the distribution of optimized values is also useful for understanding the variance of the resulting estimate. 42 | 43 | * `sweep_n.py` - script to run the optimizations over all initializations 44 | * `optimize.py` - main optimization script 45 | * `diffusion_model.py` - wrapper class that contains the diffusion model and loads and processes captured measurements 46 | 47 | Please direct questions to lindell@stanford.edu and 48 | gordon.wetzstein@stanford.edu. 49 | -------------------------------------------------------------------------------- /calibration/data/capture_1_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_1_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_1_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_1_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_2_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_2_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_2_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_2_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_3_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_3_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_3_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_3_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_4_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_4_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_4_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_4_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_5_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_5_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_5_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_5_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_6_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_6_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_6_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_6_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_7_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_7_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_7_5_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_7_5_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_8_0_inch.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_8_0_inch.mat -------------------------------------------------------------------------------- /calibration/data/capture_direct.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/calibration/data/capture_direct.mat -------------------------------------------------------------------------------- /calibration/diffusion_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import h5py 3 | import numpy as np 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | device = torch.device('cpu') 9 | torch.set_default_dtype(torch.float64) 10 | 11 | 12 | class DiffusionModel(nn.Module): 13 | def __init__(self, us_init=150., ua_init=0., n_init=1.0, t0_init=0., datadir='./data'): 14 | super().__init__() 15 | self.bin_size = 8e-12 16 | self.t = torch.arange(1, 65536)*self.bin_size 17 | self.t = self.t[self.t < 35e-9] 18 | self.c0 = 3e8 19 | self.c = self.c0 / n_init 20 | self.thicknesses = torch.arange(1, 8.5, 0.5) 21 | self.datadir = datadir 22 | 23 | self.t = self.t.to(device) 24 | self.thicknesses = self.thicknesses.to(device) 25 | 26 | # get diffuse reflectance 27 | self.R = self.calculate_reflection_coeff(n_init) 28 | 29 | # all values are scaled for numerical accuracy so optimized value should be ~0 to 10 30 | self.us_scale = 100. 31 | self.t0_scale = 1e9 32 | 33 | self.us_init = us_init / self.us_scale # originally m^-1 34 | self.ua_init = ua_init # m^-1 35 | self.n_init = n_init # mm 36 | self.t0_init = t0_init # ns 37 | 38 | self.laser, _ = self.load_impulse_response() 39 | self.data = self.preprocess_data() 40 | 41 | # optimizable parameters 42 | self.ua = nn.Parameter(torch.Tensor([self.ua_init])).to(device) 43 | self.us = nn.Parameter(torch.Tensor([self.us_init])).to(device) 44 | self.n = torch.Tensor([self.n_init]).to(device) 45 | self.t0 = nn.Parameter(self.t0_init * torch.ones_like(self.thicknesses)).to(device) 46 | 47 | def fresnel(self, n, theta_i): 48 | n0 = 1. 49 | with np.errstate(invalid='ignore'): 50 | theta_t = np.arcsin(n*np.sin(theta_i) / n0) 51 | R = 0.5 * (np.sin(theta_i - theta_t)**2 / np.sin(theta_i + theta_t)**2 + 52 | np.tan(theta_i - theta_t)**2 / np.tan(theta_i + theta_t)**2) 53 | R[theta_i == 0] = (n - n0)**2 / (n + n0)**2 54 | R[np.arcsin(n0/n) < theta_i] = 1. 55 | return R 56 | 57 | def calculate_reflection_coeff(self, n1): 58 | ''' 59 | calculate reflection coefficient given the refractive indices of the 60 | two materials. This is derived in 61 | 62 | Zhu, J. X., D. J. Pine, and D. A. Weitz. 63 | "Internal reflection of diffusive light in random media." 64 | Physical Review A 44.6 (1991): 3948. 65 | ''' 66 | # integrate to calculate c1 and c2 67 | theta = np.linspace(0., np.pi/2, 501) 68 | 69 | c1 = abs(np.trapz(self.fresnel(n1, theta)*np.sin(theta)*np.cos(theta), theta)) 70 | theta = -np.linspace(-np.pi/2, 0., 501) 71 | c2 = abs(np.trapz(self.fresnel(n1, theta)*np.sin(theta)*np.cos(theta)**2, theta)) 72 | 73 | R = (3*c2 + 2*c1) / (3*c2 - 2*c1 + 2) 74 | return R 75 | 76 | def load_impulse_response(self): 77 | N = 129 78 | 79 | with h5py.File(self.datadir + '/capture_direct.mat', 'r') as f: 80 | data = np.array(f['out']).astype(np.float64) 81 | data = torch.from_numpy(data).to(device) 82 | 83 | # subtract dc component 84 | data = data[data > 0] 85 | data = data - torch.median(data) 86 | data = torch.clamp(data, 0.) 87 | 88 | shift = torch.where(data > 100)[0][0].int() 89 | data = data[shift:shift+N].squeeze() 90 | data = data / torch.sum(data) 91 | 92 | # we need to flip manually since pytorch 93 | # convolution is implemented as correlation 94 | data = torch.flip(data, dims=(0,)) 95 | 96 | return data[None, None, :], shift 97 | 98 | def preprocess_data(self): 99 | data = [] 100 | for idx, thickness in enumerate(self.thicknesses): 101 | switcher = { 102 | 1.0: self.datadir + '/capture_1_0_inch.mat', 103 | 1.5: self.datadir + '/capture_1_5_inch.mat', 104 | 2.0: self.datadir + '/capture_2_0_inch.mat', 105 | 2.5: self.datadir + '/capture_2_5_inch.mat', 106 | 3.0: self.datadir + '/capture_3_0_inch.mat', 107 | 3.5: self.datadir + '/capture_3_5_inch.mat', 108 | 4.0: self.datadir + '/capture_4_0_inch.mat', 109 | 4.5: self.datadir + '/capture_4_5_inch.mat', 110 | 5.0: self.datadir + '/capture_5_0_inch.mat', 111 | 5.5: self.datadir + '/capture_5_5_inch.mat', 112 | 6.0: self.datadir + '/capture_6_0_inch.mat', 113 | 6.5: self.datadir + '/capture_6_5_inch.mat', 114 | 7.0: self.datadir + '/capture_7_0_inch.mat', 115 | 7.5: self.datadir + '/capture_7_5_inch.mat', 116 | 8.0: self.datadir + '/capture_8_0_inch.mat'} 117 | 118 | d = 0.0254 * thickness 119 | with h5py.File(switcher[thickness.item()], 'r') as f: 120 | foam = np.array(f['out']).astype(np.float64) 121 | foam = torch.from_numpy(foam).to(device) 122 | 123 | # subtract dc component 124 | foam = foam[foam > 0] 125 | foam = foam - torch.median(foam) 126 | foam = torch.clamp(foam, 0.) 127 | 128 | # find calibrated shift to start of scattering layer 129 | laser, shift = self.load_impulse_response() 130 | 131 | # time bins for data 132 | propagation_distance = d 133 | 134 | # subtract the distance measured to the SPAD using the hardware setup 135 | shift = -shift + propagation_distance / 3e8 / self.bin_size 136 | 137 | foam = torch.roll(foam, shift.int().item()) 138 | foam = foam[:len(self.t)] 139 | foam = foam / torch.max(foam) 140 | 141 | # align peaks to initial guess of model coefficients 142 | model = self.forward(d, self.ua_init, self.us_init) 143 | model = F.conv1d(model, laser, padding=(laser.shape[-1])) 144 | max_vals, _ = torch.max(model, dim=2, keepdim=True) 145 | model = model / max_vals 146 | model_idx = torch.argmax(model) 147 | foam_idx = torch.argmax(foam) 148 | shift = model_idx - foam_idx 149 | foam = torch.roll(foam, shift.int().item()) 150 | foam = F.pad(foam, (0, laser.shape[-1]+1)) 151 | 152 | data.append(foam[None, None, :]) 153 | 154 | data = torch.cat(data, dim=0) 155 | return data 156 | 157 | def forward(self, d, ua, us, t0=0.): 158 | ''' 159 | Returns the diffusion model for a slab with finite thickness given by 160 | Michael S. Patterson, B. Chance, and B. C. Wilson, 161 | "Time resolved reflectance and transmittance for the noninvasive 162 | measurement of tissue optical properties," 163 | Appl. Opt. 28, 2331-2336 (1989) 164 | 165 | parameters: 166 | d - thickness 167 | ua - absorption coefficient 168 | us - reduced scattering coefficient 169 | ze - extrapolation distance 170 | ''' 171 | 172 | t = self.t[None, :, None] 173 | c = self.c 174 | us = us * self.us_scale 175 | t0 = t0 / self.t0_scale 176 | ze = 2/3 * 1/us * (1 + self.R) / (1 - self.R) 177 | tshift = torch.clamp(t-t0, 8e-12) 178 | 179 | z0 = 1 / us 180 | D = 1 / (3 * (ua + us)) 181 | 182 | # Photon migration through a turbid slab described by a model 183 | # based on diffusion approximation. 184 | # https://www.osapublishing.org/ao/abstract.cfm?uri=ao-36-19-4587 185 | n_dipoles = 20 186 | ii = torch.arange(-n_dipoles, n_dipoles+1)[None, None, :].to(device) 187 | z1 = d * (1 - 2 * ii) - 4*ii*ze - z0 188 | z2 = d * (1 - 2 * ii) - (4*ii - 2)*ze + z0 189 | 190 | dipole_term = z1 * torch.exp(-(z1**2) / (4*D*c*(tshift))) - \ 191 | z2 * torch.exp(-(z2**2) / (4*D*c*(tshift))) 192 | 193 | dipole_term = torch.sum(dipole_term, dim=-1, keepdim=True) # sum over dipoles 194 | 195 | model = (4*np.pi*D*c)**(-3/2) * torch.clamp(t-t0, 1e-14)**(-5/2) \ 196 | * torch.exp(-ua * c * (t-t0)) \ 197 | * dipole_term 198 | model = model.squeeze() 199 | 200 | return model[None, None, :] 201 | -------------------------------------------------------------------------------- /calibration/optimize.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from diffusion_model import DiffusionModel 4 | import torch.nn.functional as F 5 | import torch.optim 6 | from torch.utils.tensorboard import SummaryWriter 7 | import os 8 | import matplotlib.pyplot as plt 9 | import configargparse 10 | 11 | torch.set_num_threads(2) 12 | device = torch.device('cpu') 13 | torch.set_default_dtype(torch.float64) 14 | 15 | p = configargparse.ArgumentParser() 16 | 17 | p.add_argument('--experiment_name', type=str, required=True, 18 | help='path to directory where checkpoints & tensorboard events will be saved.') 19 | p.add_argument('--lr', type=float, default=1e-2, help='learning rate. default=1e-2') 20 | p.add_argument('--num_iters', type=int, default=5000, 21 | help='Number of iterations to run.') 22 | p.add_argument('--convergence', type=float, default=1e-6, 23 | help='Number of iterations to run.') 24 | p.add_argument('--steps_til_summary', type=int, default=50, 25 | help='Iterations until tensorboard summary is saved.') 26 | p.add_argument('--us_init', type=float, default=150, 27 | help='Initial reduced scattering coefficient value. (default 150 1/m).') 28 | p.add_argument('--ua_init', type=float, default=0, 29 | help='Initial absorption coefficient value. (default 0 1/m).') 30 | p.add_argument('--n_init', type=float, default=1.0, 31 | help='Value for refractive index of scattering media. (default 1.5).') 32 | 33 | 34 | p.add_argument('--logging_root', type=str, default='./log', help='root for logging') 35 | opt = p.parse_args() 36 | 37 | for arg in vars(opt): 38 | print(arg, getattr(opt, arg)) 39 | 40 | 41 | def cond_mkdir(path): 42 | if not os.path.exists(path): 43 | os.makedirs(path) 44 | 45 | 46 | def write_summary(writer, m, mse, model, data, total_steps): 47 | plt.switch_backend('agg') 48 | fig = plt.figure() 49 | 50 | t = np.arange(1, model.shape[-1]+1) * m.bin_size 51 | for i in range(model.shape[0]): 52 | plt.plot(t, model[i, :].detach().cpu().numpy().squeeze()) 53 | plt.xlim([0, 10e-9]) 54 | plt.ylim([0, 1.1]) 55 | 56 | plt.gca().set_prop_cycle(None) 57 | for i in range(data.shape[0]): 58 | plt.plot(t, data[i, :].detach().cpu().numpy().squeeze(), '.', markersize=1) 59 | plt.xlim([0, 10e-9]) 60 | plt.ylim([0, 1.1]) 61 | 62 | writer.add_figure('plots', fig, global_step=total_steps) 63 | writer.add_scalar('us', m.us * m.us_scale, total_steps) 64 | writer.add_scalar('ua', m.ua, total_steps) 65 | 66 | 67 | def optimize(): 68 | log_dir = opt.logging_root 69 | cond_mkdir(log_dir) 70 | 71 | summaries_dir = os.path.join(log_dir, opt.experiment_name) 72 | cond_mkdir(summaries_dir) 73 | 74 | m = DiffusionModel(us_init=opt.us_init, ua_init=opt.ua_init, 75 | n_init=opt.n_init) 76 | m.to(device) 77 | 78 | optim = torch.optim.Adam(m.parameters(), lr=opt.lr, amsgrad=True) 79 | 80 | writer = SummaryWriter(summaries_dir) 81 | 82 | converged = False 83 | converged_eps = opt.convergence 84 | prev_loss = 1e6 85 | for ii in range(opt.num_iters): 86 | 87 | loss, model, data = objective(m) 88 | optim.zero_grad() 89 | 90 | # write summary 91 | writer.add_scalar('mse', loss, ii) 92 | if not ii % opt.steps_til_summary: 93 | write_summary(writer, m, loss, model, data, ii) 94 | print(f'{ii}: {loss.detach().cpu().numpy():03f}') 95 | 96 | loss.backward() 97 | optim.step() 98 | 99 | # values should be non-negative 100 | def clamp_nonnegative(m): 101 | m.ua.data = torch.clamp(m.ua.data, min=0) 102 | m.us.data = torch.clamp(m.us.data, min=0) 103 | 104 | m.apply(clamp_nonnegative) 105 | 106 | if torch.abs(loss - prev_loss) < converged_eps: 107 | converged = True 108 | break 109 | 110 | prev_loss = loss.clone() 111 | 112 | out = {'us': m.us.detach().cpu().numpy().squeeze().item() * m.us_scale, 113 | 'ua': m.ua.detach().cpu().numpy().squeeze().item(), 114 | 't0': m.t0.detach().cpu().numpy().squeeze(), 115 | 'us_init': opt.us_init, 116 | 'ua_init': opt.ua_init, 117 | 'n_init': opt.n_init, 118 | 'mse': loss.detach().cpu().numpy().squeeze().item(), 119 | 'iters': ii, 120 | 'converged': converged, 121 | 'converged_eps': converged_eps} 122 | np.save(os.path.join(summaries_dir, 'out.npy'), out) 123 | 124 | 125 | def objective(m: DiffusionModel): 126 | 127 | # get predicted diffusion model 128 | model = [] 129 | for idx, thickness in enumerate(m.thicknesses): 130 | d = 0.0254 * thickness 131 | model.append(m.forward(d, m.ua, m.us, m.t0[idx])) 132 | model = torch.cat(model, dim=0) 133 | 134 | # convolve model with measured impulse response 135 | model = F.conv1d(model, m.laser, padding=(m.laser.shape[-1])) 136 | max_vals, _ = torch.max(model, dim=2, keepdim=True) 137 | model = model / max_vals 138 | 139 | # get data 140 | data = m.data.clone() 141 | 142 | # calculate error 143 | mse = torch.sum((model - data)**2) 144 | return mse, model, data 145 | 146 | 147 | optimize() 148 | -------------------------------------------------------------------------------- /calibration/sweep_n.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import numpy as np 3 | import os 4 | 5 | processes = set() 6 | max_processes = 10 7 | vals = [(us_init, n) for us_init in np.linspace(200, 300, 11) for n in np.linspace(1.01, 1.23, 23)] 8 | 9 | for val in vals: 10 | us_init = val[0] 11 | n = val[1] 12 | processes.add(subprocess.Popen(["python", "optimize.py", "--experiment_name", f"us_{us_init}_n_{n:.03f}", 13 | "--n_init", str(n), "--us_init", str(us_init)])) 14 | 15 | if len(processes) >= max_processes: 16 | os.wait() 17 | processes.difference_update( 18 | [p for p in processes if p.poll() is not None]) 19 | 20 | # Check if all the child processes were closed 21 | for p in processes: 22 | if p.poll() is None: 23 | p.wait() 24 | -------------------------------------------------------------------------------- /cdt_reconstruction.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Code for "Three-dimensional imaging through scattering media based on confocal diffuse tomography" 3 | David B. Lindell and Gordon Wetzstein 4 | 5 | See README file in this directory for instructions on how to setup and run the code 6 | ''' 7 | 8 | import h5py 9 | import time 10 | import numpy as np 11 | from numpy.fft import ifftn, fftn 12 | import matplotlib.pyplot as plt 13 | import torch 14 | import torch.nn.functional as TorchF 15 | from utils import fk, compl_mul, conj 16 | 17 | 18 | class CDTReconstruction(): 19 | # class to define scattering parameters and perform 20 | # reconstruction using confocal diffuse tomography 21 | 22 | def __init__(self, scene, mu_s=None, zd=None, pause=5, device=torch.device('cuda:0')): 23 | 24 | self.device = device 25 | self.scene = scene 26 | self.pause = pause 27 | 28 | # set hyper parameters 29 | if scene == 'letter_s': 30 | self.snr = 1e4 # SNR parameter for Wiener deconvolution 31 | self.scan_size = 0.6 # size of scanned area 32 | self.size_calibration = 1.06 # calibrated scanned area scaling for reconstruction 33 | self.exposure_time = 60 / 32**2 # per pixel exposure time, seconds 34 | elif scene == 'mannequin': 35 | self.snr = 2e4 36 | self.scan_size = 0.7 37 | self.size_calibration = 0.87 38 | self.exposure_time = 720 / 32**2 39 | elif scene == 'letters_ut': 40 | self.snr = 1e4 41 | self.scan_size = 0.7 42 | self.size_calibration = 1.0 43 | self.exposure_time = 600 / 32**2 44 | elif scene == 'letter_t': 45 | self.snr = 2.5e3 46 | self.scan_size = 0.7 47 | self.size_calibration = 1.0 48 | self.exposure_time = 3600 / 32**2 49 | elif scene == 'cones': 50 | self.snr = 2e4 51 | self.scan_size = 0.7 52 | self.size_calibration = 1.0 53 | self.exposure_time = 400 / 32**2 54 | elif scene == 'resolution_50': 55 | self.snr = 1.5e4 56 | self.scan_size = 0.7 57 | self.size_calibration = 1.02 58 | self.exposure_time = 80 / 32**2 59 | elif scene == 'resolution_70': 60 | self.snr = 1.5e4 61 | self.scan_size = 0.7 62 | self.size_calibration = 1.04 63 | self.exposure_time = 80 / 32**2 64 | elif 'letter_u' in scene: 65 | self.snr = 5e3 66 | self.scan_size = 0.7 67 | self.size_calibration = 1.0 68 | self.exposure_time = 60 / 32**2 69 | else: 70 | raise ValueError('Unexpected input to scene parameter.') 71 | 72 | # physical parameters 73 | # found by minimizing model fit error to calibration data 74 | self.c0 = 3e8 75 | self.n = 1.12 76 | self.c = self.c0/self.n 77 | self.mu_a = 0.53 78 | self.mu_s = 262 79 | self.ze = 0.0036 80 | 81 | # volume dimensions 82 | self.Nx = 32 83 | self.Ny = 32 84 | self.Nz = 128 85 | self.xmin = -self.size_calibration * self.scan_size / 2 86 | self.xmax = self.size_calibration * self.scan_size / 2 87 | self.ymin = -self.size_calibration * self.scan_size / 2 88 | self.ymax = self.size_calibration * self.scan_size / 2 89 | self.zmin = 0 90 | self.zmax = 2 # maximum path length in hidden volume (meters) 91 | 92 | self.x = np.linspace(self.xmin, self.xmax, self.Nx) 93 | self.y = np.linspace(self.ymin, self.ymax, self.Ny) 94 | self.z = np.linspace(self.zmin, self.zmax, self.Nz) 95 | self.X, self.Z, self.Y = np.meshgrid(self.x, self.z, self.y) 96 | 97 | # laser position 98 | self.xl = 0 99 | self.yl = 0 100 | self.zl = 0 101 | 102 | # diffuser positioning 103 | self.xd = np.linspace(2*self.xmin, 2*self.xmax, 2*self.Nx)[None, :, None] 104 | self.yd = np.linspace(2*self.ymin, 2*self.ymax, 2*self.Ny)[None, None, :] 105 | self.t = np.linspace(0, 2*self.zmax, 2*self.Nz) / self.c 106 | self.t = self.t[:, None, None] 107 | self.zd = 0.0254 # thickness of diffuser 108 | 109 | # allow optional override of these parameters 110 | if zd: 111 | self.zd = zd 112 | if mu_s: 113 | self.mu_s = mu_s 114 | 115 | # set diffusion kernel 116 | self.diffusion_fpsf = [] 117 | self.setDiffusionKernel(self.c, self.t, self.xl, self.yl, self.zl, 118 | self.xd, self.yd, self.zd, self.ze, 119 | self.mu_s, self.mu_a) 120 | 121 | def setDiffusionKernel(self, v, t, xl, yl, zl, xd, yd, zd, ze, mu_s, mu_a): 122 | 123 | ''' 124 | Returns the diffusion model for a slab with finite thickness given by 125 | Michael S. Patterson, B. Chance, and B. C. Wilson, 126 | "Time resolved reflectance and transmittance for the noninvasive 127 | measurement of tissue optical properties," 128 | Appl. Opt. 28, 2331-2336 (1989) 129 | ''' 130 | 131 | t[0, :] = 1 132 | d = zd - zl 133 | z0 = 1 / mu_s 134 | D = 1 / (3 * (mu_a + mu_s)) 135 | rho = np.sqrt((xd-xl)**2 + (yd - yl)**2) 136 | 137 | # Photon migration through a turbid slab described by a model 138 | # based on diffusion approximation. 139 | # https://www.osapublishing.org/ao/abstract.cfm?uri=ao-36-19-4587 140 | n_dipoles = 20 141 | ii = np.arange(-n_dipoles, n_dipoles+1)[None, None, :] 142 | z1 = d * (1 - 2 * ii) - 4*ii*ze - z0 143 | z2 = d * (1 - 2 * ii) - (4*ii - 2)*ze + z0 144 | 145 | dipole_term = z1 * np.exp(-(z1**2) / (4*D*v*t)) - \ 146 | z2 * np.exp(-(z2**2) / (4*D*v*t)) 147 | 148 | dipole_term = np.sum(dipole_term, axis=-1)[..., None] # sum over dipoles 149 | 150 | diff_kernel = (4*np.pi*D*v)**(-3/2) * t**(-5/2) \ 151 | * np.exp(-mu_a * v * t - rho**2 / (4*D*v*t)) \ 152 | * dipole_term 153 | 154 | psf = diff_kernel 155 | 156 | diffusion_psf = psf / np.sum(psf) 157 | diffusion_psf = np.roll(diffusion_psf, -xd.shape[1]//2, axis=1) 158 | diffusion_psf = np.roll(diffusion_psf, -yd.shape[2]//2, axis=2) 159 | diffusion_psf = fftn(diffusion_psf) * fftn(diffusion_psf) 160 | diffusion_psf = abs(ifftn(diffusion_psf)) 161 | 162 | # convert to pytorch and take fft 163 | self.diffusion_fpsf = torch.from_numpy(diffusion_psf.astype(np.float32)).to(self.device)[None, None, :, :, :] 164 | self.diffusion_fpsf = self.diffusion_fpsf.rfft(3, onesided=False) 165 | return 166 | 167 | def AT(self, x): 168 | # wrapper function for f--k migration 169 | 170 | return fk(x, 2*self.xmax, 2*self.zmax) 171 | 172 | def M(self, x): 173 | # trimming function 174 | 175 | return x[:, :, :self.Nz, :self.Nx, :self.Ny] 176 | 177 | def MT(self, x): 178 | # padding function 179 | 180 | return TorchF.pad(x, (0, self.Ny, 0, self.Nx, 0, self.Nz)) 181 | 182 | def run(self): 183 | # run confocal diffuse tomography reconstruction 184 | 185 | with h5py.File('./data/' + self.scene + '.mat', 'r') as f: 186 | meas = np.array(f['meas']).transpose(2, 1, 0) 187 | f.close() 188 | 189 | # trim scene to 1 meter along the z-dimension 190 | # and downsample to ~50 ps time binning from 16 ps 191 | b = meas[:417, :, :] 192 | downsampled = np.zeros((self.Nz, 32, 32)) 193 | for i in range(meas.shape[1]): 194 | for j in range(meas.shape[2]): 195 | x = np.linspace(0, 1, self.Nz) 196 | xp = np.linspace(0, 1, 417) 197 | yp = b[:, i, j].squeeze() 198 | downsampled[:, i, j] = np.interp(x, xp, yp) 199 | b = downsampled 200 | b /= np.max(b) # normalize to 0 to 1 201 | 202 | # initialize pytorch arrays 203 | b = torch.from_numpy(b).to(self.device)[None, None, :, :, :].float() 204 | x = torch.zeros(b.size()[0], 1, 2*self.Nz, 2*self.Nx, 2*self.Ny).to(self.device) 205 | 206 | # construct inverse psf for Wiener filtering 207 | tmp = compl_mul(self.diffusion_fpsf, conj(self.diffusion_fpsf)) 208 | tmp = tmp + 1/self.snr 209 | invpsf = compl_mul(conj(self.diffusion_fpsf), 1/tmp) 210 | 211 | # measure inversion runtime 212 | if self.device.type == 'cpu': 213 | start = time.time() 214 | else: 215 | start = torch.cuda.Event(enable_timing=True) 216 | end = torch.cuda.Event(enable_timing=True) 217 | start.record() 218 | 219 | # pad measurements 220 | x = self.MT(b) 221 | 222 | # perform f-k migration on measurements 223 | x_fk = self.AT(x) 224 | 225 | # perform deconvolution 226 | x_deconv = compl_mul(x.rfft(3, onesided=False), invpsf).ifft(3)[:, :, :, :, :, 0] 227 | 228 | # confocal inverse filter 229 | x = self.AT(x_deconv) 230 | 231 | # measure elapsed time 232 | if self.device.type == 'cpu': 233 | stop = time.time() 234 | print('Elapsed time: %.02f ms' % (1000 * (stop - start))) 235 | else: 236 | end.record() 237 | torch.cuda.synchronize() 238 | print('Elapsed time: %.02f ms' % (start.elapsed_time(end))) 239 | 240 | # plot results 241 | x_npy = x.cpu().data.numpy().squeeze()[:self.Nz, :self.Nx, :self.Ny] 242 | b_npy = b.cpu().data.numpy().squeeze() 243 | x_deconv_npy = x_deconv.cpu().data.numpy().squeeze()[:self.Nz, :self.Nx, :self.Ny] 244 | x_fk_npy = x_fk.cpu().data.numpy().squeeze()[:self.Nz, :self.Nx, :self.Ny] 245 | 246 | # trim any amplified noise at the very end of the volume 247 | x_npy[-15:, :, :] = 0 248 | 249 | if self.pause > 0: 250 | plt.suptitle('Measurements and reconstruction') 251 | plt.subplot(231) 252 | plt.imshow(np.max(b_npy, axis=0), cmap='gray', extent=[self.xmin, self.xmax, self.ymin, self.ymax]) 253 | plt.xlabel('x (m)') 254 | plt.ylabel('y (m)') 255 | plt.subplot(232) 256 | plt.imshow(np.max(b_npy, axis=1), aspect=(self.xmax-self.xmin)/(self.zmax/3e8*1e9), cmap='gray', 257 | extent=[self.xmin, self.xmax, self.zmax/3e8*1e9, self.zmin]) 258 | plt.xlabel('x (m)') 259 | plt.ylabel('t (ns)') 260 | plt.subplot(233) 261 | plt.imshow(np.max(b_npy, axis=2), aspect=(self.ymax-self.ymin)/(self.zmax/3e8*1e9), cmap='gray', 262 | extent=[self.ymin, self.ymax, self.zmax/3e8*1e9, self.zmin]) 263 | plt.xlabel('y (m)') 264 | plt.ylabel('t (ns)') 265 | 266 | plt.subplot(234) 267 | plt.imshow(np.max(x_npy, axis=0), cmap='gray', extent=[self.xmin, self.xmax, self.ymin, self.ymax]) 268 | plt.xlabel('x (m)') 269 | plt.ylabel('y (m)') 270 | plt.subplot(235) 271 | plt.imshow(np.max(x_npy, axis=1), aspect=(self.xmax-self.xmin)/(self.zmax/2), cmap='gray', 272 | extent=[self.xmin, self.xmax, self.zmax/2, self.zmin]) 273 | plt.xlabel('x (m)') 274 | plt.ylabel('z (m)') 275 | plt.subplot(236) 276 | plt.imshow(np.max(x_npy, axis=2), aspect=(self.ymax-self.ymin)/(self.zmax/2), cmap='gray', 277 | extent=[self.ymin, self.ymax, self.zmax/2, self.zmin]) 278 | plt.xlabel('y (m)') 279 | plt.ylabel('z (m)') 280 | plt.tight_layout() 281 | 282 | plt.pause(self.pause) 283 | 284 | # return measurements, deconvolved meas, reconstruction 285 | return b_npy, x_fk_npy, x_deconv_npy, x_npy 286 | -------------------------------------------------------------------------------- /data/cones.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/cones.mat -------------------------------------------------------------------------------- /data/letter_s.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_s.mat -------------------------------------------------------------------------------- /data/letter_t.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_t.mat -------------------------------------------------------------------------------- /data/letter_u_50.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_50.mat -------------------------------------------------------------------------------- /data/letter_u_52.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_52.mat -------------------------------------------------------------------------------- /data/letter_u_54.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_54.mat -------------------------------------------------------------------------------- /data/letter_u_56.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_56.mat -------------------------------------------------------------------------------- /data/letter_u_58.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_58.mat -------------------------------------------------------------------------------- /data/letter_u_60.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_60.mat -------------------------------------------------------------------------------- /data/letter_u_62.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_62.mat -------------------------------------------------------------------------------- /data/letter_u_64.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_64.mat -------------------------------------------------------------------------------- /data/letter_u_66.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_66.mat -------------------------------------------------------------------------------- /data/letter_u_68.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_68.mat -------------------------------------------------------------------------------- /data/letter_u_70.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_70.mat -------------------------------------------------------------------------------- /data/letter_u_72.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_72.mat -------------------------------------------------------------------------------- /data/letter_u_74.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_74.mat -------------------------------------------------------------------------------- /data/letter_u_76.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_76.mat -------------------------------------------------------------------------------- /data/letter_u_78.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_78.mat -------------------------------------------------------------------------------- /data/letter_u_80.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letter_u_80.mat -------------------------------------------------------------------------------- /data/letters_ut.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/letters_ut.mat -------------------------------------------------------------------------------- /data/mannequin.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/mannequin.mat -------------------------------------------------------------------------------- /data/resolution_50.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/resolution_50.mat -------------------------------------------------------------------------------- /data/resolution_70.mat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computational-imaging/confocal-diffuse-tomography/0992bf5a58f972533800e944e4a606f64d7283e0/data/resolution_70.mat -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import argparse 4 | import torch.backends.cudnn as cudnn 5 | import numpy as np 6 | from cdt_reconstruction import CDTReconstruction 7 | cudnn.benchmark = True 8 | 9 | scene_choices = ['cones', 'letter_s', 'letters_ut', 'mannequin', 10 | 'letter_t', 'resolution_50', 'resolution_70'] +\ 11 | [f'letter_u_{x}' for x in np.arange(50, 82, 2)] 12 | 13 | # Argument parsing 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--scene', type=str, default='letter_s', 16 | choices=scene_choices, 17 | help='name of scene to reconstruct.') 18 | parser.add_argument('--cpu', action='store_true', default=False, help='Force run on CPU, default=False') 19 | parser.add_argument('--gpu_id', type=int, default=0, help='index of which GPU to run on, default=0') 20 | parser.add_argument('--pause', type=int, default=5, help='how long to display figure, default=5 (seconds)') 21 | opt = parser.parse_args() 22 | print('Confocal diffuse tomography reconstruction') 23 | print('\n'.join(["\t%s: %s" % (key, value) for key, value in vars(opt).items()])) 24 | 25 | # check to see if GPU is available 26 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 27 | os.environ["CUDA_VISIBLE_DEVICES"] = "{}".format(opt.gpu_id) 28 | if opt.cpu: 29 | device = torch.device('cpu') 30 | else: 31 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 32 | 33 | 34 | def main(): 35 | cdt = CDTReconstruction(opt.scene, pause=opt.pause, device=device) 36 | cdt.run() 37 | return 38 | 39 | 40 | if __name__ == '__main__': 41 | main() 42 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==0.11.0 2 | cachetools==4.1.1 3 | certifi==2020.12.5 4 | chardet==3.0.4 5 | ConfigArgParse==1.2.3 6 | cycler==0.10.0 7 | future==0.18.2 8 | google-auth==1.23.0 9 | google-auth-oauthlib==0.4.2 10 | grpcio==1.34.0 11 | h5py==2.10.0 12 | idna==2.10 13 | importlib-metadata==3.1.1 14 | kiwisolver==1.3.1 15 | Markdown==3.3.3 16 | matplotlib==3.1.1 17 | numpy==1.16.4 18 | oauthlib==3.1.0 19 | Pillow==8.0.1 20 | protobuf==3.14.0 21 | pyasn1==0.4.8 22 | pyasn1-modules==0.2.8 23 | pyparsing==2.4.7 24 | python-dateutil==2.8.1 25 | requests==2.25.0 26 | requests-oauthlib==1.3.0 27 | rsa==4.6 28 | six==1.15.0 29 | tensorboard==2.4.0 30 | tensorboard-plugin-wit==1.7.0 31 | torch==1.5.0 32 | urllib3==1.26.2 33 | Werkzeug==1.0.1 34 | zipp==3.4.0 35 | -------------------------------------------------------------------------------- /setup.bash: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Script to set up python environment to run the confocal diffuse tomography code 4 | # note that this requires a system installaton of anaconda or miniconda 5 | # if you don't have this, follow the instructions here for miniconda 6 | # https://docs.conda.io/en/latest/miniconda.html 7 | 8 | # create a new anaconda environment for this project 9 | conda create -n cdt python=3.6 10 | eval "$(conda shell.bash hook)" # this is just required to activate an anaconda env in a bash script 11 | conda activate cdt 12 | pip install -r requirements.txt 13 | 14 | # run the reconstructions on captured data and display each for 2 sec. 15 | python main.py --scene letter_s --pause 2 # reconstruct scene from Fig. 2 16 | python main.py --scene mannequin --pause 2 # reconstruct scene from Fig. 3 17 | python main.py --scene letters_ut --pause 2 # reconstruct scene from Fig. 3 18 | python main.py --scene letter_t --pause 2 # reconstruct scene from Fig. 3 19 | python main.py --scene cones --pause 2 # reconstruct scene from Fig. 3 20 | 21 | # optionally remove environment when finished with code 22 | # conda deactivate 23 | # conda env remove -n cdt 24 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def interpolate(grid, lin_ind_frustrum, voxel_coords, device_id): 5 | """ linear interpolation for frequency-wavenumber migration 6 | adapted from https://github.com/vsitzmann/deepvoxels/blob/49369e243001658ccc8ba3be97d87c85273c9f15/projection.py 7 | """ 8 | 9 | depth, width, height = grid.shape 10 | 11 | lin_ind_frustrum = lin_ind_frustrum.long() 12 | 13 | x_indices = voxel_coords[1, :] 14 | y_indices = voxel_coords[2, :] 15 | z_indices = voxel_coords[0, :] 16 | 17 | mask = ((x_indices < 0) | (y_indices < 0) | (z_indices < 0) | 18 | (x_indices > width-1) | (y_indices > height-1) | (z_indices > depth-1)).to(device_id) 19 | 20 | x0 = x_indices.floor().long() 21 | y0 = y_indices.floor().long() 22 | z0 = z_indices.floor().long() 23 | 24 | x0 = torch.clamp(x0, 0, width - 1) 25 | y0 = torch.clamp(y0, 0, height - 1) 26 | z0 = torch.clamp(z0, 0, depth - 1) 27 | z1 = (z0 + 1).long() 28 | z1 = torch.clamp(z1, 0, depth - 1) 29 | 30 | x_indices = torch.clamp(x_indices, 0, width - 1) 31 | y_indices = torch.clamp(y_indices, 0, height - 1) 32 | z_indices = torch.clamp(z_indices, 0, depth - 1) 33 | 34 | x = x_indices - x0.float() 35 | y = y_indices - y0.float() 36 | z = z_indices - z0.float() 37 | 38 | output = torch.zeros(height * width * depth).to(device_id) 39 | tmp1 = grid[z0, x0, y0] * (1 - z) * (1 - x) * (1 - y) 40 | tmp2 = grid[z1, x0, y0] * z * (1 - x) * (1 - y) 41 | output[lin_ind_frustrum] = tmp1 + tmp2 42 | 43 | output = output * (1 - mask.float()) 44 | output = output.contiguous().view(depth, width, height) 45 | 46 | return output 47 | 48 | 49 | def roll_n(X, axis, n): 50 | """ circular shift function """ 51 | 52 | f_idx = tuple(slice(None, None, None) if i != axis else slice(0, n, None) for i in range(X.dim())) 53 | b_idx = tuple(slice(None, None, None) if i != axis else slice(n, None, None) for i in range(X.dim())) 54 | front = X[f_idx] 55 | back = X[b_idx] 56 | return torch.cat([back, front], axis) 57 | 58 | 59 | def fftshift(x): 60 | real, imag = torch.unbind(x, -1) 61 | 62 | if real.ndim > 3: 63 | dim_start = 2 64 | else: 65 | dim_start = 0 66 | 67 | for dim in range(dim_start, len(real.size())): 68 | n_shift = real.size(dim)//2 69 | if real.size(dim) % 2 != 0: 70 | n_shift += 1 # for odd-sized images 71 | real = roll_n(real, axis=dim, n=n_shift) 72 | imag = roll_n(imag, axis=dim, n=n_shift) 73 | return torch.stack((real, imag), -1) # last dim=2 (real&imag) 74 | 75 | 76 | def ifftshift(x): 77 | real, imag = torch.unbind(x, -1) 78 | 79 | if real.ndim > 3: 80 | dim_stop = 1 81 | else: 82 | dim_stop = -1 83 | 84 | for dim in range(len(real.size()) - 1, dim_stop, -1): 85 | real = roll_n(real, axis=dim, n=real.size(dim)//2) 86 | imag = roll_n(imag, axis=dim, n=imag.size(dim)//2) 87 | return torch.stack((real, imag), -1) # last dim=2 (real&imag) 88 | 89 | 90 | def compl_mul(X, Y): 91 | """ complex multiplication for pytorch; real and imaginary parts are 92 | stored in the last channel of the arrays 93 | see https://discuss.pytorch.org/t/aten-cuda-implementation-of-complex-multiply/17215/2 94 | """ 95 | 96 | assert X.shape[-1] == 2 and Y.shape[-1] == 2, 'Last dimension must be 2' 97 | return torch.stack( 98 | (X[..., 0] * Y[..., 0] - X[..., 1] * Y[..., 1], 99 | X[..., 0] * Y[..., 1] + X[..., 1] * Y[..., 0]), 100 | dim=-1) 101 | 102 | 103 | def conj(x): 104 | # complex conjugation for pytorch 105 | 106 | tmp = x.clone() 107 | tmp[:, :, :, :, :, 1] = tmp[:, :, :, :, :, 1] * -1 108 | return tmp 109 | 110 | 111 | def fk(meas, width, mrange): 112 | """ perform f--k migration """ 113 | 114 | device = meas.device 115 | meas = meas.squeeze() 116 | width = torch.FloatTensor([width]).to(device) 117 | mrange = torch.FloatTensor([mrange]).to(device) 118 | 119 | N = meas.size()[1]//2 # spatial resolution 120 | M = meas.size()[0]//2 # temporal resolution 121 | data = torch.sqrt(torch.clamp(meas, 0)) 122 | 123 | M_grid = torch.arange(-M, M).to(device) 124 | N_grid = torch.arange(-N, N).to(device) 125 | [z, x, y] = torch.meshgrid(M_grid, N_grid, N_grid) 126 | z = (z.type(torch.FloatTensor) / M).to(device) 127 | x = (x.type(torch.FloatTensor) / N).to(device) 128 | y = (y.type(torch.FloatTensor) / N).to(device) 129 | 130 | # pad data 131 | tdata = data 132 | 133 | # fourier transform 134 | if tdata.ndim > 3: 135 | tdata = fftshift(tdata.fft(3)) 136 | else: 137 | tdata = fftshift(tdata.rfft(3, onesided=False)) 138 | 139 | tdata_real, tdata_imag = torch.unbind(tdata, -1) 140 | 141 | # interpolation coordinates 142 | z_interp = torch.sqrt(abs((((N * mrange) / (M * width * 4))**2) * 143 | (x**2 + y**2) + z**2)) 144 | coords = torch.stack((z_interp.flatten(), x.flatten(), y.flatten()), 0) 145 | lin_ind = torch.arange(z.numel()).to(device) 146 | coords[0, :] = (coords[0, :] + 1) * M 147 | coords[1, :] = (coords[1, :] + 1) * N 148 | coords[2, :] = (coords[2, :] + 1) * N 149 | 150 | # run interpolation 151 | tvol_real = interpolate(tdata_real, lin_ind, coords, device) 152 | tvol_imag = interpolate(tdata_imag, lin_ind, coords, device) 153 | tvol = torch.stack((tvol_real, tvol_imag), -1) 154 | 155 | # zero out redundant spectrum 156 | x = x[:, :, :, None] 157 | y = y[:, :, :, None] 158 | z = z[:, :, :, None] 159 | tvol = tvol * abs(z) / torch.clamp(torch.sqrt(abs((((N * mrange) / (M * width * 4))**2) * 160 | (x**2 + y**2)+z**2)), 1e-8) 161 | tvol = tvol * (z > 0).type(torch.FloatTensor).to(device) 162 | 163 | # inverse fourier transform and crop 164 | tvol = ifftshift(tvol).ifft(3).squeeze() 165 | geom = tvol[:, :, :, 0]**2 + tvol[:, :, :, 1]**2 166 | geom = geom[None, None, :, :, :] 167 | 168 | return geom 169 | --------------------------------------------------------------------------------