├── .gitignore
├── README.md
├── examples
    └── rendering.py
├── img
    └── Example.png
├── pydrr
    ├── Detector.py
    ├── GeometryContext.py
    ├── GpuVolumeContext.py
    ├── Kernel.py
    ├── KernelManager.py
    ├── KernelModule.py
    ├── Projector.py
    ├── VolumeContext.py
    ├── __init__.py
    ├── autoinit.py
    ├── ext
    │   ├── CupyContextBridge.py
    │   └── ExampleContextBridge.py
    ├── kernels
    │   └── __init__.py
    ├── simple
    │   └── __init__.py
    └── utils.py
└── requirements.txt


/.gitignore:
--------------------------------------------------------------------------------
1 | .vscode/*
2 | __pycache__


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # DRR generation by pycuda
2 | PyDRR lets you generate Digitaly Reconstructed Radiography image using CUDA.
3 | 
4 | ![Example](./img/Example.png)
5 | 
6 | # Requirements
7 | - numpy  1.13.1 or later
8 | - PyCUDA 2017.1.1 or later
9 | 


--------------------------------------------------------------------------------
/examples/rendering.py:
--------------------------------------------------------------------------------
  1 | import pydrr
  2 | import pydrr.autoinit
  3 | import SimpleITK as sitk
  4 | import matplotlib.pyplot as plt
  5 | import mpl_toolkits.axes_grid1
  6 | import numpy as np
  7 | import sys
  8 | from pydrr import utils
  9 | 
 10 | def main():
 11 |     if len(sys.argv) < 2:
 12 |         print('rendering.py <volume path>')
 13 |         return
 14 | 
 15 |     mhd_filename = sys.argv[0]
 16 |     # Load materials
 17 |     itkimage = sitk.ReadImage(mhd_filename)
 18 |     volume  = sitk.GetArrayFromImage(itkimage)
 19 |     spacing = itkimage.GetSpacing()
 20 |     spacing = spacing[::-1]
 21 | 
 22 |     volume = pydrr.utils.HU2Myu(volume - 700, 0.2683)
 23 | 
 24 |     pm_Nx3x4, image_size, image_spacing = load_test_projection_matrix()
 25 |     T_Nx4x4 = load_test_transform_matrix()
 26 | 
 27 |     # Construct objects
 28 |     volume_context = pydrr.VolumeContext(volume.astype(np.float32), spacing)
 29 |     geometry_context = pydrr.GeometryContext()
 30 |     geometry_context.projection_matrix = pm_Nx3x4
 31 | 
 32 |     n_channels = T_Nx4x4.shape[0] * pm_Nx3x4.shape[0]
 33 |     detector = pydrr.Detector(pydrr.Detector.make_detector_size(image_size, n_channels), image_spacing)
 34 |     # detector = pydrr.Detector.from_geometry(geometry_context, T_Nx4x4) # You can use from_geometry if you set pixel_size and image_size.
 35 |     projector = pydrr.Projector(detector, 1.0).to_gpu()
 36 | 
 37 |     # Host memory -> (Device memory) -> Texture memory
 38 |     t_volume_context = volume_context.to_texture()
 39 |     
 40 |     d_image = projector.project(t_volume_context, geometry_context, T_Nx4x4)
 41 | 
 42 |     # Device memory -> Host memory
 43 |     image = d_image.get()
 44 |     print('Result image shape:', image.shape)
 45 |     plt.figure(figsize=(16,9))
 46 |     n_show_channels = 3
 47 |     for i in range(min(image.shape[2], n_show_channels)):
 48 |         ax = plt.subplot(1, min(image.shape[2], n_show_channels), i+1)
 49 |         divider = mpl_toolkits.axes_grid1.make_axes_locatable(ax)
 50 |         cax = divider.append_axes('right', '5%', pad='3%')
 51 |         im = ax.imshow(image[:, :, i], interpolation='none', cmap='gray')
 52 |         fig.colorbar(im, cax=cax)
 53 |     plt.show()
 54 | 
 55 | def load_test_projection_matrix(SDD=2000, SOD=1800, image_size=[1280, 1280], spacing=[0.287, 0.287] ):
 56 | 
 57 |     if isinstance(image_size, list):
 58 |         image_size = np.array(image_size)
 59 | 
 60 |     if isinstance(spacing, list):
 61 |         spacing = np.array(spacing)
 62 | 
 63 |     extrinsic_R = utils.convertTransRotTo4x4([[0,0,0,90,0,0],
 64 |                                               [0,0,0,0,90,0],
 65 |                                               [0,0,0,0,0,90]])
 66 | 
 67 |     print('extrinsic_R:', extrinsic_R)
 68 |     print('extrinsic_R.shape:', extrinsic_R.shape)
 69 | 
 70 |     extrinsic_T = utils.convertTransRotTo4x4([0,0,-SOD,0,0,0])
 71 | 
 72 |     print('extrinsic_T:', extrinsic_T)
 73 |     print('extrinsic_T.shape:', extrinsic_T.shape)
 74 | 
 75 | 
 76 |     extrinsic = utils.concatenate4x4(extrinsic_T, extrinsic_R)
 77 | 
 78 |     print('extrinsic:', extrinsic)
 79 |     print('extrinsic.shape:', extrinsic.shape)
 80 | 
 81 | 
 82 |     intrinsic = np.array([[-SDD/spacing[0], 0, image_size[0]/2.0], # unit: [pixel]
 83 |                           [0, -SDD/spacing[1], image_size[1]/2.0],
 84 |                           [0,                0,               1]])
 85 | 
 86 |     print('intrinsic:', intrinsic)
 87 |     print('intrinsic.shape:', intrinsic.shape)
 88 | 
 89 | 
 90 |     pm_Nx3x4 = utils.constructProjectionMatrix(intrinsic, extrinsic)
 91 |     #pm_Nx3x4 = np.repeat(pm_Nx3x4, 400, axis=0)
 92 |     
 93 |     print('pm_Nx3x4:', pm_Nx3x4)
 94 |     print('pm_Nx3x4.shape:', pm_Nx3x4.shape)
 95 | 
 96 |     return pm_Nx3x4, image_size, spacing
 97 | 
 98 | def load_test_transform_matrix(n_channels=1):
 99 |     T_Nx6 = np.array([0,0,0,90,0,0])
100 |     T_Nx6 = np.expand_dims(T_Nx6, axis=0)
101 |     T_Nx6 = np.repeat(T_Nx6, n_channels, axis=0)
102 |     T_Nx4x4 = utils.convertTransRotTo4x4(T_Nx6)
103 | 
104 |     print('T_Nx4x4:', T_Nx4x4)
105 |     print('T_Nx4x4.shape:', T_Nx4x4.shape)
106 | 
107 |     return T_Nx4x4
108 | 
109 | if __name__ == '__main__':
110 |     main()
111 | 


--------------------------------------------------------------------------------
/img/Example.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/elda27/pydrr/4bf56a760f760066a051fa2e5306e09e2b4b3f1a/img/Example.png


--------------------------------------------------------------------------------
/pydrr/Detector.py:
--------------------------------------------------------------------------------
 1 | from pycuda import gpuarray
 2 | from functools import wraps
 3 | from . import KernelManager
 4 | from . import utils
 5 | import numpy as np
 6 | 
 7 | class Detector:
 8 |     def __init__(self, image_size, pixel_spacing, image=None, *, cpu=None):
 9 |         self.cpu = cpu
10 |         if self.is_cpu() and len(image_size) == 2:
11 |             image_size = Detector.make_detector_size(image_size, 1)
12 |         #self.image = self.alloc(image_size) if image is None else image
13 |         self.image = np.ascontiguousarray(np.zeros(image_size, dtype=np.float32)) if image is None else image
14 |         self.image_size = image_size
15 |         self.pixel_spacing = pixel_spacing
16 | 
17 |     def to_cpu(self):
18 |         #assert self.cpu is not None
19 |         if self.is_cpu():
20 |             return self
21 |         return self.cpu
22 | 
23 |     def to_gpu(self):
24 |         #assert self.cpu is None
25 |         if self.is_gpu():
26 |             return self
27 |         image_size = KernelManager.Module.get_global(
28 |             'd_image_size', 
29 |             np.array(self.image_size, dtype=np.float32)
30 |             )
31 |         return Detector(image_size, self.pixel_spacing, gpuarray.to_gpu(self.image), cpu=self)
32 | 
33 |     def is_cpu(self):
34 |         return self.cpu is None
35 | 
36 |     def is_gpu(self):
37 |         return self.cpu is not None
38 | 
39 |     def alloc(self, image_size):
40 |         return np.transpose(np.zeros((image_size[2],image_size[1],image_size[0]), dtype=np.float32))
41 | 
42 | 
43 |     @staticmethod
44 |     def from_geometry(geometry_context, T_Nx4x4):
45 |         n_proj = geometry_context.projection_matrix.shape[0] if geometry_context.projection_matrix.ndim == 3 else 1
46 |         n_T = T_Nx4x4.shape[0] if T_Nx4x4.ndim == 3 else 1
47 | 
48 |         image_size = Detector.make_detector_size(geometry_context.image_size, n_proj * n_T)
49 |         return Detector(image_size, geometry_context.pixel_spacing)
50 | 
51 |     @staticmethod
52 |     def make_detector_size(image_size, n_channels):
53 |         return np.array((image_size[0], image_size[1], n_channels), dtype=np.int32)
54 | 


--------------------------------------------------------------------------------
/pydrr/GeometryContext.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | from . import utils
  3 | 
  4 | class GeometryContext:
  5 |     def __init__(self):
  6 |         self.SOD_ = 0.0
  7 |         self.SDD_ = 0.0
  8 |         self.pixel_spacing_ = (1.0, 1.0)
  9 |         self.image_size_ = (1024, 1024)
 10 |         self.view_matrix_ = np.eye(4,4, dtype=np.float32)
 11 |         
 12 |         self.intrinsic_ = None
 13 |         self.extrinsic_ = None
 14 | 
 15 |         self.projection_matrix_ = None
 16 | 
 17 |     @property
 18 |     def SOD(self):
 19 |         return self.SOD_
 20 | 
 21 |     @SOD.setter
 22 |     def SOD(self, value):
 23 |         self.intrinsic = None
 24 |         self.SOD_ = value
 25 | 
 26 |     @property
 27 |     def SDD(self):
 28 |         return self.SDD_
 29 | 
 30 |     @SDD.setter
 31 |     def SDD(self, value):
 32 |         self.intrinsic = None
 33 |         self.extrinsic = None
 34 |         self.SDD_ = value
 35 | 
 36 |     @property
 37 |     def pixel_spacing(self):
 38 |         return self.pixel_spacing_
 39 | 
 40 |     @pixel_spacing.setter
 41 |     def pixel_spacing(self, value):
 42 |         self.intrinsic = None
 43 |         self.pixel_spacing_ = value
 44 | 
 45 |     @property
 46 |     def image_size(self):
 47 |         return self.image_size_
 48 | 
 49 |     @image_size.setter
 50 |     def image_size(self, value):
 51 |         self.intrinsic = None
 52 |         self.image_size_ = value
 53 |         
 54 |     @property
 55 |     def view_matrix(self):
 56 |         return self.view_matrix_
 57 | 
 58 |     @view_matrix.setter
 59 |     def view_matrix(self, value):
 60 |         self.extrinsic = None
 61 |         self.view_matrix_ = value
 62 | 
 63 |     @property
 64 |     def intrinsic(self):
 65 |         if self.intrinsic_ is None:
 66 |             self.intrinsic_ = np.array([
 67 |                 [ self.SOD / self.pixel_spacing[0], 0, self.image_size[0] / 2 ], 
 68 |                 [ 0, self.SOD / self.pixel_spacing[1], self.image_size[1] / 2 ],
 69 |                 [0, 0, 1]
 70 |             ])
 71 |         return self.intrinsic_
 72 |     
 73 |     @intrinsic.setter
 74 |     def intrinsic(self, new_intrinsic):
 75 |         self.projection_matrix = None
 76 |         self.intrinsic_ = new_intrinsic
 77 |     
 78 |     @property
 79 |     def extrinsic(self):
 80 |         if self.extrinsic_ is None:
 81 |             extrinsic_T = utils.convertTransRotTo4x4([0, 0, -self.SOD, 0, 0, 0])
 82 |             self.extrinsic_ = utils.concatenate4x4(extrinsic_T, self.view_matrix)
 83 |         return self.extrinsic_
 84 | 
 85 |     @extrinsic.setter
 86 |     def extrinsic(self, new_extrinsic):
 87 |         self.projection_matrix = None
 88 |         self.extrinsic_ = new_extrinsic
 89 |     
 90 |     @property
 91 |     def projection_matrix(self):
 92 |         if self.projection_matrix_ is None:
 93 |             self.projection_matrix_ = utils.constructProjectionMatrix(self.intrinsic, self.extrinsic)
 94 |         return self.projection_matrix_
 95 | 
 96 |     @projection_matrix.setter
 97 |     def projection_matrix(self, value):
 98 |         self.projection_matrix_ = value
 99 | 
100 |     


--------------------------------------------------------------------------------
/pydrr/GpuVolumeContext.py:
--------------------------------------------------------------------------------
 1 | from .VolumeContext import VolumeContext
 2 | import numpy as np
 3 | 
 4 | class GpuVolumeContext(VolumeContext):
 5 |     def __init__(self, gpu_object, spacing, bridge):
 6 |         self.volume_object = gpu_object
 7 |         self.bridge = bridge
 8 | 
 9 |         self.volume = self.bridge.get_pointer(gpu_object)
10 |         self.spacing = spacing
11 |         self.volume_shape = self.bridge.get_shape(gpu_object)
12 | 
13 |         volume_size = np.asarray(self.volume_shape, dtype=np.uint32)
14 |         self.spacing = np.asarray(spacing, dtype=np.float32)
15 |         self.volume_corner_mm = np.array(volume_size * self.spacing / 2.0, dtype=np.float32)
16 |         self.cpu = None
17 |         self.texture = None
18 | 
19 |     def to_cpu(self):
20 |         if self.cpu is None:
21 |             self.cpu = VolumeContext(self.bridge.get_cpu_array(self.volume_object), self.spacing)
22 |         return self.cpu
23 | 
24 |     def to_gpu(self):
25 |         assert self.volume is not None
26 |         return self
27 | 
28 |     def is_cpu(self):
29 |         cpu = self.to_cpu()
30 |         return cpu is not None
31 | 
32 |     def is_gpu(self):
33 |         return True # Always true
34 | 
35 |     def is_texture(self):
36 |         return False
37 | 


--------------------------------------------------------------------------------
/pydrr/Kernel.py:
--------------------------------------------------------------------------------
 1 | from pycuda import driver, compiler, gpuarray, tools 
 2 | from . import KernelManager
 3 | 
 4 | class Kernel:
 5 |     def __init__(self, module, func_name, attrs):
 6 |         self.parent_module = module
 7 |         self.kernel = module.get_function(func_name)
 8 |         self.attributes = attrs
 9 |         self.setCurrent()
10 | 
11 |     def invoke(self, *args, **kwargs):
12 |         flag, not_founds = self.parent_module.verify_attributes(self.attributes['global'])
13 |         assert flag, 'Following global attributes are not initialized: {}'.format(not_founds)
14 |         flag, not_founds = self.parent_module.verify_texture_attributes(self.attributes['texture'])
15 |         assert flag, 'Following texture attributes are not initialized: {}'.format(not_founds)
16 | 
17 |         if 'grid' not in kwargs:
18 |             kwargs['grid'] = (1, 1, 1)
19 |         if 'block' not in kwargs:
20 |             kwargs['block'] = (1, 1, 1)
21 | 
22 |         return self.kernel(*args, **kwargs)
23 | 
24 |     def setCurrent(self):
25 |         self.parent_module.setCurrentModule()
26 |         KernelManager.Kernel = self


--------------------------------------------------------------------------------
/pydrr/KernelManager.py:
--------------------------------------------------------------------------------
 1 | from pycuda import driver, compiler, gpuarray, tools 
 2 | 
 3 | from pydrr import kernels
 4 | from .KernelModule import KernelModule
 5 | 
 6 | class KernelManager:
 7 |     Kernel = None
 8 |     Module = None
 9 |     Modules = []
10 | 
11 |     def __init__(self, default_kernel, *kernel_codes):
12 |         for kernel_code, kernel_info in kernel_codes:
13 |             KernelManager.Modules.append(KernelModule(kernel_code, kernel_info)) 
14 | 
15 |         KernelManager.Module = KernelManager.Modules[0]
16 |         KernelManager.Kernel = KernelManager.Module.get_kernel(default_kernel)
17 | 
18 | _manager = None
19 | 
20 | def initialize():
21 |     global _manager
22 |     _manager = KernelManager('render_with_linear_interp',
23 |         (
24 |             kernels.render_kernel,
25 |             {
26 |                 'render_with_linear_interp':
27 |                 {
28 |                     'global': [
29 |                         'd_step_size_mm',
30 |                         'd_image_size',
31 |                         'd_volume_spacing',
32 |                         'd_volume_corner_mm',
33 |                     ],
34 |                     'texture' : [
35 |                         't_volume',
36 |                         't_proj_param_Nx12',
37 |                     ],
38 |                 },
39 |                 'print_device_params': { 'global':[], 'texture':['t_proj_param_Nx12'] }
40 |             }
41 |         ),
42 |     )
43 | 
44 | 


--------------------------------------------------------------------------------
/pydrr/KernelModule.py:
--------------------------------------------------------------------------------
 1 | from pycuda import driver, compiler, gpuarray, tools 
 2 | from .Kernel import Kernel
 3 | from . import KernelManager
 4 | 
 5 | class KernelModule:
 6 |     Interpolations = {
 7 |         'linear': driver.filter_mode.LINEAR
 8 |     }
 9 |     def __init__(self, source, info):
10 |         self.module = compiler.SourceModule(source, options=['--ptxas-options=-v'], cache_dir=None, keep=False)
11 |         self.kernels = dict()
12 |         self.attributes = dict()
13 |         self.texture_attributes = dict()
14 |         self.setCurrentModule()
15 |         for func_name, attrs in info.items():
16 |             self.attributes.update(dict(zip(attrs['global'], [ False for _ in attrs['global'] ])))
17 |             self.texture_attributes.update(dict(zip(attrs['texture'], [ False for _ in attrs['texture'] ])))
18 |             self.kernels[func_name] = Kernel(self, func_name, attrs)
19 | 
20 |     def get_function(self, name):
21 |         kernel_func = self.module.get_function(name)
22 |         return kernel_func
23 | 
24 |     def get_kernel(self, name):
25 |         return self.kernels[name]
26 | 
27 |     def verify_attributes(self, attrs):
28 |         if not attrs: # Attributes are empty
29 |             return True, []
30 | 
31 |         founds = [ (self.attributes[name], name) for name in attrs ]
32 |         founds, names = map(list, zip(*founds))
33 |         return all(founds), [ name for name, found in zip(names, founds) if not found ]
34 | 
35 |     def verify_texture_attributes(self, attrs):
36 |         if not attrs: # Attributes are empty
37 |             return True, []
38 | 
39 |         founds = [ (self.texture_attributes[name], name) for name in attrs ]
40 |         founds, names = map(list, zip(*founds))
41 |         return all(founds), [ name for name, found in zip(names, founds) if not found ]
42 | 
43 |     def get_global(self, name, host_obj):
44 |         assert name in self.attributes, 'Unknown global atrribute: {}'.format(name)
45 |         self.attributes[name] = True
46 | 
47 |         device_obj = self.module.get_global(name)[0]
48 |         driver.memcpy_htod(device_obj, host_obj)
49 |         return device_obj
50 | 
51 |     def get_texture(self, name, device_obj, interpolation=None):
52 |         assert name in self.texture_attributes, 'Unknown texture atrribute: {}'.format(name)
53 |         self.texture_attributes[name] = True
54 | 
55 |         texture_obj = self.module.get_texref(name)
56 |         if interpolation is not None and interpolation in KernelModule.Interpolations:
57 |             texture_obj.set_filter_mode(KernelModule.Interpolations[interpolation])
58 |         texture_obj.set_array(device_obj)
59 |         return texture_obj
60 | 
61 |     def setCurrentModule(self):
62 |         KernelManager.Module = self
63 | 


--------------------------------------------------------------------------------
/pydrr/Projector.py:
--------------------------------------------------------------------------------
 1 | from functools import wraps
 2 | import numpy as np
 3 | from . import KernelManager
 4 | from . import utils
 5 | from pycuda import driver
 6 | 
 7 | class Projector:
 8 |     
 9 |     # set block and grid size
10 |     block = (32, 32, 1)
11 |     grid  = None
12 | 
13 |     def __init__(self, target_detector, step_size_mm = 1, *, cpu=None):
14 |         self.target_detector = target_detector
15 |         self.step_size_mm = step_size_mm
16 |         self.cpu = cpu
17 | 
18 |     def project(self, volume_context, geometry_context, T_Nx4x4):
19 |         assert self.cpu is not None, 'CPU ray casting is not supported.'
20 | 
21 |         image_size = self.target_detector.to_cpu().image_size
22 |         pm_Nx3x4 = geometry_context.projection_matrix
23 | 
24 |         p_Nx12 = utils.constructProjectionParameter(pm_Nx3x4, np.array(image_size[:2]), T_Nx4x4)
25 |         
26 |         assert self.target_detector.cpu.image_size[2] == p_Nx12.shape[0], 'Unmatched detector channel and pose parameter channel.(Actual: {} != {})'.format(self.target_detector.cpu.image_size[2], p_Nx12.shape[0])
27 | 
28 |         h_p_Nx12 = p_Nx12.astype(np.float32)
29 |         d_p_Nx12 = driver.np_to_array(h_p_Nx12, order='C')
30 |         t_p_Nx12 = KernelManager.Module.get_texture('t_proj_param_Nx12', d_p_Nx12)
31 | 
32 |         grid = (16, 16, 1)
33 |         if Projector.grid is None:
34 |             grid = tuple(np.uint32(np.ceil(image_size / Projector.block)).tolist())
35 |         
36 |         KernelManager.Kernel.invoke(
37 |             self.target_detector.image.gpudata,
38 |             texrefs=[volume_context.volume, t_p_Nx12],
39 |             block=Projector.block, grid=grid
40 |         )
41 |         # Display debug info
42 |         # print_kernel = KernelManager.Module.get_kernel('print_device_params')
43 |         # print_kernel.invoke(texrefs=[t_p_Nx12])
44 | 
45 |         return self.target_detector.image
46 | 
47 |     def to_gpu(self):
48 |         assert self.cpu is None
49 | 
50 |         step_size_mm = KernelManager.Module.get_global('d_step_size_mm', np.float32(self.step_size_mm))
51 |         return Projector(self.target_detector.to_gpu(), step_size_mm, cpu=self)
52 | 


--------------------------------------------------------------------------------
/pydrr/VolumeContext.py:
--------------------------------------------------------------------------------
 1 | import numpy as np
 2 | from pycuda import driver
 3 | from . import utils
 4 | from . import KernelManager
 5 | 
 6 | class VolumeContext:
 7 |     def __init__(self, volume, spacing = (1,1,1), *, cpu=None, gpu=None):
 8 |         self.cpu = cpu
 9 |         self.gpu = gpu
10 |         if volume is None:
11 |             return
12 |         
13 |         if not volume.flags['C_CONTIGUOUS']: # The array must be contiguous array for gpu copy.
14 |             volume = np.ascontiguousarray(volume, dtype=np.float32)
15 |         
16 |         self.volume = volume
17 |         volume_size = np.asarray(self.volume.shape, dtype=np.uint32)
18 |         self.spacing = np.asarray(spacing, dtype=np.float32)
19 |         self.volume_corner_mm = np.array(volume_size * self.spacing / 2.0, dtype=np.float32)
20 | 
21 |     def to_cpu(self):
22 |         assert self.cpu is not None
23 |         return self.cpu
24 | 
25 |     def to_gpu(self):
26 |         assert self.volume is not None
27 | 
28 |         if self.is_gpu():
29 |             return self
30 |         elif self.is_texture():
31 |             return self.gpu
32 | 
33 |         obj = VolumeContext(None, cpu=self)
34 |         obj.volume = driver.np_to_array(self.volume, order='C')
35 | 
36 |         obj.spacing = KernelManager.Module.get_global('d_volume_spacing', self.spacing)
37 |         obj.volume_corner_mm = KernelManager.Module.get_global('d_volume_corner_mm', self.volume_corner_mm)
38 |         return obj
39 | 
40 |     def to_texture(self, interpolation = 'linear'):
41 |         cpu = None
42 |         gpu = None
43 |         if self.is_cpu():
44 |             cpu = self
45 |             gpu = self.to_gpu()
46 |         elif self.is_gpu():
47 |             cpu = self.cpu
48 |             gpu = self
49 |         else:
50 |             return self
51 |         obj = VolumeContext(None, cpu=cpu, gpu=gpu)
52 | 
53 |         obj.volume = KernelManager.Module.get_texture('t_volume', gpu.volume, interpolation)
54 | 
55 |         obj.spacing = gpu.spacing
56 |         obj.volume_corner_mm = gpu.volume_corner_mm
57 |         return obj
58 | 
59 |     def is_cpu(self):
60 |         return self.cpu is None and self.gpu is None
61 | 
62 |     def is_gpu(self):
63 |         return self.cpu is not None and self.gpu is None
64 | 
65 |     def is_texture(self):
66 |         return self.cpu is not None and self.gpu is not None
67 | 


--------------------------------------------------------------------------------
/pydrr/__init__.py:
--------------------------------------------------------------------------------
1 | import pydrr.utils
2 | from pydrr.KernelManager import KernelManager, initialize
3 | from pydrr.VolumeContext import VolumeContext
4 | from pydrr.GeometryContext import GeometryContext
5 | from pydrr.Detector import Detector
6 | from pydrr.Projector import Projector


--------------------------------------------------------------------------------
/pydrr/autoinit.py:
--------------------------------------------------------------------------------
1 | import pycuda.autoinit
2 | from pydrr.KernelManager import initialize
3 | 
4 | initialize()
5 | 


--------------------------------------------------------------------------------
/pydrr/ext/CupyContextBridge.py:
--------------------------------------------------------------------------------
 1 | import cupy
 2 | 
 3 | class CupyContextBridge:
 4 |     @classmethod
 5 |     def get_pointer(cls, array):
 6 |         return array.data.ptr
 7 | 
 8 |     @classmethod
 9 |     def get_shape(cls, array):
10 |         return array.shape
11 | 
12 |     @classmethod
13 |     def get_cpu_array(cls, array):
14 |         return cupy.asnumpy(array)


--------------------------------------------------------------------------------
/pydrr/ext/ExampleContextBridge.py:
--------------------------------------------------------------------------------
 1 | 
 2 | class ExampleContextBridge:
 3 |     """Example bridge object for GpuVolumeContext
 4 |     
 5 |     A bridge object works get attribute from user type to construct VolumeContext.
 6 |     """
 7 | 
 8 |     @classmethod
 9 |     def get_pointer(cls, array):
10 |         """Get device pointer from user type.
11 |         """
12 | 
13 |         raise NotImplementedError
14 | 
15 |     @classmethod
16 |     def get_shape(cls, array):
17 |         """Get array shape from user type.
18 |         """
19 |         raise NotImplementedError
20 | 
21 |     @classmethod
22 |     def get_cpu_array(cls, array):
23 |         """Construct numpy array from user type.
24 |         """
25 |         raise NotImplementedError
26 | 
27 | 


--------------------------------------------------------------------------------
/pydrr/kernels/__init__.py:
--------------------------------------------------------------------------------
  1 | 
  2 | render_kernel=r'''
  3 | texture<float, cudaTextureType3D, cudaReadModeElementType> t_volume;
  4 | texture<float, cudaTextureType2D, cudaReadModeElementType> t_proj_param_Nx12;
  5 | 
  6 | __device__ __constant__ float  d_step_size_mm;
  7 | __device__ __constant__ float3 d_image_size;
  8 | __device__ __constant__ float3 d_volume_spacing;
  9 | __device__ __constant__ float3 d_volume_corner_mm;
 10 | 
 11 | // negate
 12 | __device__ float3 operator-(float3 &a)
 13 | {
 14 |     return make_float3(-a.x, -a.y, -a.z);
 15 | }
 16 | 
 17 | // min
 18 | static __device__ float3 fminf(float3 a, float3 b)
 19 | {
 20 | 	return make_float3(fminf(a.x,b.x), fminf(a.y,b.y), fminf(a.z,b.z));
 21 | }
 22 | 
 23 | // max
 24 | static __device__ float3 fmaxf(float3 a, float3 b)
 25 | {
 26 | 	return make_float3(fmaxf(a.x,b.x), fmaxf(a.y,b.y), fmaxf(a.z,b.z));
 27 | }
 28 | 
 29 | // addition
 30 | static __device__ float3 operator+(float3 a, float3 b)
 31 | {
 32 |     return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
 33 | }
 34 | 
 35 | __device__ float3 operator+(float3 a, float b)
 36 | {
 37 |     return make_float3(a.x + b, a.y + b, a.z + b);
 38 | }
 39 | __device__ void operator+=(float3 &a, float3 b)
 40 | {
 41 |     a.x += b.x; a.y += b.y; a.z += b.z;
 42 | }
 43 | 
 44 | // subtract
 45 | static __device__ float3 operator-(float3 a, float3 b)
 46 | {
 47 |     return make_float3(a.x - b.x, a.y - b.y, a.z - b.z);
 48 | }
 49 | static __device__ float3 operator-(float3 a, float b)
 50 | {
 51 |     return make_float3(a.x - b, a.y - b, a.z - b);
 52 | }
 53 | __device__ void operator-=(float3 &a, float3 b)
 54 | {
 55 |     a.x -= b.x; a.y -= b.y; a.z -= b.z;
 56 | }
 57 | 
 58 | // multiply
 59 | __device__ float3 operator*(float3 a, float3 b)
 60 | {
 61 |     return make_float3(a.x * b.x, a.y * b.y, a.z * b.z);
 62 | }
 63 | 
 64 | static __device__ float3 operator*(float3 a, float s)
 65 | {
 66 |     return make_float3(a.x * s, a.y * s, a.z * s);
 67 | }
 68 | 
 69 | static __device__ float3 operator*(float s, float3 a)
 70 | {
 71 |     return make_float3(a.x * s, a.y * s, a.z * s);
 72 | }
 73 | __device__ void operator*=(float3 &a, float s)
 74 | {
 75 |     a.x *= s; a.y *= s; a.z *= s;
 76 | }
 77 | 
 78 | // divide
 79 | static __device__ float3 operator/(float3 a, float3 b)
 80 | {
 81 |     return make_float3(a.x / b.x, a.y / b.y, a.z / b.z);
 82 | }
 83 | static __device__ float3 operator/(float3 a, float s)
 84 | {
 85 |     float inv = 1.0f / s;
 86 |     return a * inv;
 87 | }
 88 | __device__ float3 operator/(float s, float3 a)
 89 | {
 90 |     float inv = 1.0f / s;
 91 |     return a * inv;
 92 | }
 93 | __device__ void operator/=(float3 &a, float s)
 94 | {
 95 |     float inv = 1.0f / s;
 96 |     a *= inv;
 97 | }
 98 | 
 99 | // dot product
100 | __device__ float dot(float3 a, float3 b)
101 | { 
102 |     return a.x * b.x + a.y * b.y + a.z * b.z;
103 | }
104 | 
105 | // normalize
106 | __device__ float3 normalize(float3 v)
107 | {
108 |     float invLen = rsqrtf(dot(v, v));
109 |     return v * invLen;
110 | }
111 | 
112 | struct Ray {
113 | 	float3 o;	// origin
114 | 	float3 d;	// direction
115 | };
116 | 
117 | __device__ Ray computeNormalizedRay(const float x, const float y, const int ch)
118 | {
119 |     // compute a unit vector connecting the source and the normalized pixel (x, y) on the imaging plane using pre-computed corner points.
120 |     Ray ray;
121 |     ray.d = normalize( make_float3( tex2D(t_proj_param_Nx12, 0, ch)+tex2D(t_proj_param_Nx12, 3, ch)*x+tex2D(t_proj_param_Nx12, 6, ch)*y, 
122 |                                     tex2D(t_proj_param_Nx12, 1, ch)+tex2D(t_proj_param_Nx12, 4, ch)*x+tex2D(t_proj_param_Nx12, 7, ch)*y, 
123 |                                     tex2D(t_proj_param_Nx12, 2, ch)+tex2D(t_proj_param_Nx12, 5, ch)*x+tex2D(t_proj_param_Nx12, 8, ch)*y ) );
124 |     ray.o = make_float3( tex2D(t_proj_param_Nx12, 9, ch), tex2D(t_proj_param_Nx12, 10, ch), tex2D(t_proj_param_Nx12, 11, ch) );
125 |     return ray;
126 | }
127 | 
128 | // intersect ray with a box
129 | // http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm
130 | // This code is based on volumeRender demo in CUDA SDK
131 | __device__ bool intersectBoxRay(float3 box, const Ray ray, float &tnear, float &tfar)
132 | {
133 |     // compute intersection of ray with all six planes
134 |     float3 tbot = (-box - ray.o) / ray.d;
135 |     float3 ttop = ( box - ray.o) / ray.d;
136 | 
137 |     // re-order intersections to find smallest and largest on each axis
138 |     float3 tmin = fminf(ttop, tbot);
139 |     float3 tmax = fmaxf(ttop, tbot);
140 | 
141 |     // find the largest tmin and the smallest tmax
142 |     tnear = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));
143 |     tfar  = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));
144 |     return tfar > tnear;
145 | }
146 | 
147 | 
148 | __global__ void render_with_linear_interp(float *d_image_N)
149 | {
150 | 
151 |     int x = blockDim.x * blockIdx.x + threadIdx.x;
152 |     int y = blockDim.y * blockIdx.y + threadIdx.y;
153 |     int z = blockDim.z * blockIdx.z + threadIdx.z;
154 | 
155 |     int index = x*int(d_image_size.z)*int(d_image_size.y)+y*int(d_image_size.z)+z;
156 |     //if (index >= (d_image_size.x * d_image_size.y * d_image_size.z)) return;
157 |     if (index >= (d_image_size.x * d_image_size.y * d_image_size.z) ||
158 |         (x >= d_image_size.x || y >= d_image_size.y || z >= d_image_size.z)
159 |         )
160 |         return;
161 |         
162 |     Ray ray = computeNormalizedRay(
163 |         ((float)x+0.5f)/d_image_size.x, 
164 |         ((float)y+0.5f)/d_image_size.y,
165 |         z
166 |         );
167 |     
168 |     float tnear = 0.0f, tfar = 0.0f, RPL = 0.0f; 
169 |     if(!intersectBoxRay(d_volume_corner_mm, ray, tnear, tfar)) return;
170 | 
171 |     // compute Radiological Path Length (RPL) by trilinear interpolation (texture fetching)
172 |     float3 cur = (ray.o + tnear * ray.d + d_volume_corner_mm) / d_volume_spacing;  // object coordinate (mm) -> texture (voxel) coordinate
173 |     float3 delta_dir = d_step_size_mm * ray.d / d_volume_spacing;                  // object coordinate (mm) -> texture (voxel) coordinate
174 | 
175 |     for(float travelled_length = 0; travelled_length < (tfar-tnear); travelled_length += d_step_size_mm, cur += delta_dir){
176 |         // pick the density value at the current point and accumulate it (Note: currently consider only single input volume case)
177 |         // access to register memory and texture memory (filterMode of texture should be 'linear')
178 |         RPL += tex3D(t_volume, cur.z, cur.y, cur.x) * d_step_size_mm;
179 |     }
180 | 
181 |     d_image_N[index] += RPL;
182 | }
183 | 
184 | // Debug print
185 | __global__ void print_device_params()
186 | {
187 |     int x = blockDim.x * blockIdx.x + threadIdx.x;
188 |     int y = blockDim.y * blockIdx.y + threadIdx.y;
189 |     int z = blockDim.z * blockIdx.z + threadIdx.z;
190 | 
191 |     if (x != 0 || y != 0 || z != 0)
192 |     {
193 |         return;
194 |     }
195 | 
196 |     printf("step_size_mm: %f\n", d_step_size_mm);
197 |     printf("image_size: %f, %f, %f\n", d_image_size.x, d_image_size.y, d_image_size.z);
198 |     printf("volume_spacing: %f, %f, %f\n", d_volume_spacing.x, d_volume_spacing.y, d_volume_spacing.z);
199 |     printf("volume_corner_mm: %f, %f, %f\n", d_volume_corner_mm.x, d_volume_corner_mm.y, d_volume_corner_mm.z);
200 | 
201 |     for (int i = 0; i < 3; ++i)
202 |     {
203 |         printf("t_proj_parameter[%d]: [%f, %f, %f %f; %f, %f, %f, %f; %f, %f, %f, %f;]\n", i, 
204 |             tex2D(t_proj_param_Nx12, 0, i), tex2D(t_proj_param_Nx12, 1, i), tex2D(t_proj_param_Nx12, 2, i),
205 |             tex2D(t_proj_param_Nx12, 3, i), tex2D(t_proj_param_Nx12, 4, i), tex2D(t_proj_param_Nx12, 5, i),
206 |             tex2D(t_proj_param_Nx12, 6, i), tex2D(t_proj_param_Nx12, 7, i), tex2D(t_proj_param_Nx12, 8, i),
207 |             tex2D(t_proj_param_Nx12, 9, i), tex2D(t_proj_param_Nx12, 10, i), tex2D(t_proj_param_Nx12, 11, i)
208 |             );
209 |     }
210 | }
211 | 
212 | '''


--------------------------------------------------------------------------------
/pydrr/simple/__init__.py:
--------------------------------------------------------------------------------
 1 | import pydrr
 2 | import numpy as np
 3 | import pycuda.autoinit
 4 | import pydrr.autoinit
 5 | 
 6 | def generate_drr(volume, spacing, **kwargs):
 7 |     """Generate DRR(Digitaly Reconstruct Radiography) image from volume.
 8 |     
 9 |     Args:
10 |         volume (numpy.array): three-dimensional array. This volume is aligned by z, y, x
11 |         spacing (tuple[float]): volume spacing. This argument is aligned by z, y, x. 
12 |     
13 |     Keyword Args:
14 |         is_hu_volume: Input volume is HU volume so it will be converted myu volume
15 |         SOD (float): Source to object distance (Default: 1800).
16 |         SDD (float): Source to detector distance (Default: 2000).
17 |         view (tuple[float]): View direction by x, y, z angles. (Default: [90, 0, 0]).
18 |         pixel_size (tuple[float]): Pixel size of the detector.
19 |         image_size (tuple[float]): Image size of the detector equal to result image size.
20 |         extrinsic (numpy.array): Extrinsic matrix Nx4x4. If None, this value will compute automatically from view, SOD and SDD.
21 |         intrinsic (numpy.array): Intrinsic matrix Nx3x3. If None, this vlaue will compute automatically from pixel_size and image_size.
22 |         pose (array-like):  Translation and rotation of the volume. (Default:)
23 |         projector (pydrr.Projector): Projector object. If None, Projector object is automatically construct.
24 |                                      When you invoke this function many times, you might set this argument by performance reason.
25 |     Returns:
26 |         numpy.array: DRR image
27 |     """
28 | 
29 |     args = {
30 |         'SOD': 1800,
31 |         'SDD': 2000,
32 |         'view': [90, 0, 0],
33 |         'extrinsic': None,
34 |         'intrinsic': None,
35 |         'projection': None,
36 |         'pixel_size': [0.417, 0.417],
37 |         'image_size': [1024, 1024],
38 |         'pose': [0, 0, 0, 0, 0, 0],
39 |         'projector': None,
40 |         'is_hu_volume': False
41 |     }
42 | 
43 |     args.update(kwargs)
44 |     
45 |     if args['is_hu_volume']:
46 |         volume = pydrr.utils.HU2Myu(volume, 0.02)
47 |     
48 |     if args['extrinsic'] is None:
49 |         args['extrinsic'] = pydrr.utils.convertTransRotTo4x4(
50 |             np.asarray((0, 0, -args['SOD']) + tuple(args['view']))
51 |             )
52 |     if args['intrinsic'] is None:
53 |         args['intrinsic'] = np.array(
54 |             [[-args['SDD']/args['pixel_size'][0], 0, args['image_size'][0]/2.0], # unit: [pixel]
55 |              [0, -args['SDD']/args['pixel_size'][1], args['image_size'][1]/2.0],
56 |              [0,                0,               1]]
57 |              )
58 |     
59 |     pm_Nx3x4 = pydrr.utils.constructProjectionMatrix(args['intrinsic'], args['extrinsic'])
60 |     if pm_Nx3x4.ndim == 2:
61 |         pm_Nx3x4 = pm_Nx3x4[np.newaxis, :, :]
62 | 
63 |     T_Nx4x4 = pydrr.utils.convertTransRotTo4x4(np.asarray(args['pose']))
64 |     if T_Nx4x4.ndim == 2:
65 |         T_Nx4x4 = T_Nx4x4[np.newaxis, :, :]
66 | 
67 |     # Define contexts.
68 |     volume_context = pydrr.VolumeContext(volume.astype(np.float32), spacing)
69 |     geometry_context = pydrr.GeometryContext()
70 |     geometry_context.projection_matrix = pm_Nx3x4
71 | 
72 |     n_channels = T_Nx4x4.shape[0] * pm_Nx3x4.shape[0]
73 | 
74 |     if args['projector'] is None:
75 |         detector = pydrr.Detector(
76 |             pydrr.Detector.make_detector_size(args['image_size'], n_channels), 
77 |             args['pixel_size']
78 |             )
79 | 
80 |         projector = pydrr.Projector(detector, 1.0).to_gpu()
81 | 
82 |     # Host memory -> (Device memory) -> Texture memory
83 |     t_volume_context = volume_context.to_texture()
84 |     
85 |     d_image = projector.project(t_volume_context, geometry_context, T_Nx4x4)
86 | 
87 |     # Device memory -> Host memory
88 |     return d_image.get()
89 | 


--------------------------------------------------------------------------------
/pydrr/utils.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import math
  3 | import time
  4 | from pycuda import driver
  5 | 
  6 | class Timer(object):
  7 | 
  8 |     def __init__(self, name=None):
  9 |         self.name = name
 10 | 
 11 |     def __enter__(self):
 12 |         self.tstart = time.time()
 13 | 
 14 |     def __exit__(self, type, value, traceback):
 15 |         if self.name:
 16 |             print('[%s]' % self.name, end=' ')
 17 |         print('Elapsed: %s [sec]' % (time.time() - self.tstart))
 18 | 
 19 | def get_global(module, name, host_obj):
 20 |     device_obj = module.get_global(name)[0]
 21 |     driver.memcpy_htod(device_obj, host_obj)
 22 |     return device_obj
 23 | 
 24 | 
 25 | def load_volume(filename, axis_order='xyz'):
 26 |     import SimpleITK as sitk
 27 |     itkimage = sitk.ReadImage(filename)
 28 |     volume = sitk.GetArrayFromImage(itkimage)
 29 |     spacing = np.array(list(itkimage.GetSpacing()))
 30 |     origin = np.array(list(itkimage.GetOrigin()))
 31 | 
 32 |     if axis_order == 'xyz':
 33 |         volume = np.transpose(volume, (2, 1, 0))
 34 |         volume = np.ascontiguousarray(volume)
 35 |     elif axis_order == 'zyx':
 36 |         spacing = spacing[::-1]
 37 |         origin = origin[::-1]
 38 |     else:
 39 |         raise ValueError('unexpected axis order')
 40 | 
 41 |     return volume, spacing, origin
 42 | 
 43 | 
 44 | def HU2Myu(HU_images, myu_water):
 45 |     """
 46 |     convert CT images represented in HU to linear attenuation coefficient
 47 |         :param HU_images: images in HU
 48 |         :param myu_water: linear attenuation coefficient of water at the effective
 49 |                           energy in unit of 'mm2/g'. The effective energy
 50 |                           is generally close to 30% or 40% of peak energy
 51 |                           (see http://www.sprawls.org/ppmi2/RADPEN/)
 52 |                           see http: // physics.nist.gov/PhysRefData/XrayMassCoef/ComTab/water.html
 53 |                           for list of attenuation coefficient of water.
 54 |                           (note that the unit of(myu/rho) listed here is
 55 |                           cm2/g which is cm-1*(cm3/g), and we convert
 56 |                           the unit for myu and rho into mm-1 and mm3/g
 57 |                           respectively, so the value should be
 58 |                           devided by 10)
 59 |                           e.g.
 60 |                             120kVP(peak energy) -> about 40keV(4e-2MeV)(effective energy)
 61 |                             -> 0.2683cm2/g -> 0.02683mm2/g
 62 | 
 63 |                           definition of Hounsfield Unit(air: -1000, water: 0)
 64 |                             HU = 1000 * (myu - myu_water) / myu_water
 65 |                             ->
 66 |                             myu = HU * myu_water / 1000 + myu_water
 67 |     """
 68 |     myu_images = np.fmax((1000.0 + np.float32(HU_images)) * myu_water / 1000.0, 0.0)  # we clamp negative value
 69 |     return myu_images
 70 | 
 71 | 
 72 | def concatenate4x4(*matrix_Nx4x4):
 73 | 
 74 |     matrix = np.identity(4)
 75 |     for m in matrix_Nx4x4:
 76 |         matrix = np.matmul(matrix, m)
 77 |     return matrix
 78 | 
 79 | 
 80 | def constructProjectionMatrix(intrinsic_Nx4x4, extrinsic_Nx4x4):
 81 | 
 82 |     ndim = intrinsic_Nx4x4.ndim
 83 | 
 84 |     if ndim == 3 and intrinsic_Nx4x4.shape[1:] == (3, 3):
 85 |         N = intrinsic_Nx4x4.shape[0]
 86 |         matrix = np.zeros((N, 3, 4))
 87 |         matrix[:, :3, :3] = intrinsic_Nx4x4
 88 |     elif ndim == 2 and intrinsic_Nx4x4.shape == (3, 3):
 89 |         matrix = np.zeros((3, 4))
 90 |         matrix[:3, :3] = intrinsic_Nx4x4
 91 |     else:
 92 |         raise ValueError('unexpected shape')
 93 | 
 94 |     matrix = np.matmul(matrix, extrinsic_Nx4x4)
 95 |     return matrix
 96 | 
 97 | 
 98 | def constructProjectionParameter(pm_Nx3x4, image_size, T_Nx4x4=np.eye(4)):
 99 | 
100 |     # check pm_Nx3x4
101 |     ndim = pm_Nx3x4.ndim
102 |     if ndim == 3 and pm_Nx3x4.shape[1:] == (3, 4):
103 |         pass
104 |     elif ndim == 2 and pm_Nx3x4.shape == (3, 4):
105 |         pm_Nx3x4 = np.expand_dims(pm_Nx3x4, axis=0)
106 |     else:
107 |         raise ValueError('unexpected shape')
108 | 
109 |     N = pm_Nx3x4.shape[0]
110 | 
111 |     # check T_Nx4x4
112 |     ndim = T_Nx4x4.ndim
113 |     if ndim == 3 and T_Nx4x4.shape[1:] == (4, 4):
114 |         if T_Nx4x4.shape[0] == N:
115 |             pass
116 |         elif T_Nx4x4.shape[0] != 1 and N == 1:
117 |             N = T_Nx4x4.shape[0]
118 |             pm_Nx3x4 = np.repeat(pm_Nx3x4, N, axis=0)
119 |         elif T_Nx4x4.shape[0] == 1 and N != 1:
120 |             T_Nx4x4 = np.repeat(T_Nx4x4, N, axis=0)
121 |         else:
122 |             raise ValueError('unexpected shape')
123 |     elif ndim == 2 and T_Nx4x4.shape == (4, 4):
124 |         T_Nx4x4 = np.expand_dims(T_Nx4x4, axis=0)
125 |         T_Nx4x4 = np.repeat(T_Nx4x4, N, axis=0)
126 |     else:
127 |         raise ValueError('unexpected shape')
128 | 
129 |     # check image_size
130 |     if isinstance(image_size, list):
131 |         image_size = np.array(image_size)
132 | 
133 |     ndim = image_size.ndim
134 |     if ndim == 2 and image_size.shape[1] == (2):
135 |         pass
136 |     elif ndim == 1 and image_size.shape[0] == (2):
137 |         image_size = np.expand_dims(image_size, axis=0)
138 |     else:
139 |         raise ValueError('unexpected shape')
140 | 
141 |     if N != image_size.shape[0] and 1 != image_size.shape[0]:
142 |         raise ValueError('unexpected shape')
143 | 
144 |     # run construction
145 |     p_Nx12 = np.zeros((N, 12))
146 | 
147 |     pm_Nx3x4 = np.matmul(pm_Nx3x4, T_Nx4x4)
148 |     rot_Nx3x3 = pm_Nx3x4[:, :3, :3]
149 |     inv_rot_Nx3x3 = np.zeros_like(rot_Nx3x3)
150 |     for i in range(N):
151 |         inv_rot_Nx3x3[i, :, :] = np.linalg.inv(rot_Nx3x3[i, :, :])
152 |     
153 |     ray_s = pm_Nx3x4[:, :, 3]
154 |     ray_s = np.expand_dims(ray_s, axis=2)
155 |     ray_s = np.matmul(inv_rot_Nx3x3, ray_s)
156 |     ray_s *= -1
157 | 
158 |     p_Nx12[:, :3] = inv_rot_Nx3x3[:, :, 2]
159 |     p_Nx12[:, 3:6] = inv_rot_Nx3x3[:, :, 0] * np.float32(image_size[:, 0])
160 |     p_Nx12[:, 6:9] = inv_rot_Nx3x3[:, :, 1] * np.float32(image_size[:, 1])
161 |     p_Nx12[:, 9:] = ray_s[:, :, 0]
162 | 
163 |     return p_Nx12
164 | 
165 | 
166 | def convertTransRotTo4x4(transrot_Nx6, is_radians=False):
167 | 
168 |     if isinstance(transrot_Nx6, list):
169 |         transrot_Nx6 = np.array(transrot_Nx6)
170 | 
171 |     ndim = transrot_Nx6.ndim
172 | 
173 |     if ndim == 2 and transrot_Nx6.shape[1] == 6:
174 |         pass
175 |     elif ndim == 1 and transrot_Nx6.shape[0] == 6:
176 |         transrot_Nx6 = np.expand_dims(transrot_Nx6, axis=0)
177 |     else:
178 |         raise ValueError('unexpected shape: ndim is 1 or 2 and first or second shape is 6. Actual:{}'.format(transrot_Nx6))
179 | 
180 |     N = transrot_Nx6.shape[0]
181 | 
182 |     angle_rad = transrot_Nx6[:, 3:]
183 |     if not is_radians:
184 |         angle_rad = (np.pi/180.0) * angle_rad
185 | 
186 |     cos_Nx3 = np.cos(angle_rad)
187 |     sin_Nx3 = np.sin(angle_rad)
188 | 
189 |     matrix_Nx4x4 = np.zeros((N, 4, 4))
190 |     matrix_Nx4x4[:, 0, 0] = cos_Nx3[:, 1] * cos_Nx3[:, 2]
191 |     matrix_Nx4x4[:, 0, 1] = -cos_Nx3[:, 0] * sin_Nx3[:, 2] + \
192 |         sin_Nx3[:, 0] * sin_Nx3[:, 1] * cos_Nx3[:, 2]
193 |     matrix_Nx4x4[:, 0, 2] = sin_Nx3[:, 0] * sin_Nx3[:, 2] + \
194 |         cos_Nx3[:, 0] * sin_Nx3[:, 1] * cos_Nx3[:, 2]
195 |     matrix_Nx4x4[:, 0, 3] = transrot_Nx6[:, 0]
196 |     matrix_Nx4x4[:, 1, 0] = cos_Nx3[:, 1] * sin_Nx3[:, 2]
197 |     matrix_Nx4x4[:, 1, 1] = cos_Nx3[:, 0] * cos_Nx3[:, 2] + \
198 |         sin_Nx3[:, 0] * sin_Nx3[:, 1] * sin_Nx3[:, 2]
199 |     matrix_Nx4x4[:, 1, 2] = -sin_Nx3[:, 0] * cos_Nx3[:, 2] + \
200 |         cos_Nx3[:, 0] * sin_Nx3[:, 1] * sin_Nx3[:, 2]
201 |     matrix_Nx4x4[:, 1, 3] = transrot_Nx6[:, 1]
202 |     matrix_Nx4x4[:, 2, 0] = -sin_Nx3[:, 1]
203 |     matrix_Nx4x4[:, 2, 1] = sin_Nx3[:, 0] * cos_Nx3[:, 1]
204 |     matrix_Nx4x4[:, 2, 2] = cos_Nx3[:, 0] * cos_Nx3[:, 1]
205 |     matrix_Nx4x4[:, 2, 3] = transrot_Nx6[:, 2]
206 |     matrix_Nx4x4[:, 3, 3] = np.ones((N,))
207 | 
208 |     if ndim == 1:
209 |         matrix_Nx4x4 = matrix_Nx4x4[0]
210 | 
211 |     return matrix_Nx4x4
212 | 
213 | 
214 | def convert4x4ToTransRot(matrix_Nx4x4, is_radians=False, eps=np.finfo(float).eps):
215 | 
216 |     ndim = matrix_Nx4x4.ndim
217 | 
218 |     if ndim == 3 and matrix_Nx4x4.shape[1:] == (4, 4):
219 |         pass
220 |     elif ndim == 2 and matrix_Nx4x4.shape == (4, 4):
221 |         matrix_Nx4x4 = np.expand_dims(matrix_Nx4x4, axis=0)
222 |     else:
223 |         raise ValueError('unexpected shape')
224 | 
225 |     R_Nx3x3 = matrix_Nx4x4[:, :3, :3]
226 |     rot_Nx3 = convertRot3x3ToRPY(R_Nx3x3, is_radians=is_radians, eps=eps)
227 |     trans_Nx3 = matrix_Nx4x4[:, :3, 3]
228 | 
229 |     transrot_Nx6 = np.concatenate((trans_Nx3, rot_Nx3), axis=1)
230 | 
231 |     if ndim == 2:
232 |         transrot_Nx6 = transrot_Nx6[0]
233 | 
234 |     return transrot_Nx6
235 | 
236 | 
237 | def convertRot3x3ToRPY(R_Nx3x3, is_radians=False, eps=np.finfo(float).eps):
238 | 
239 |     ndim = R_Nx3x3.ndim
240 | 
241 |     if ndim == 3 and R_Nx3x3.shape[1:] == (3, 3):
242 |         pass
243 |     elif ndim == 2 and R_Nx3x3.shape == (3, 3):
244 |         R_Nx3x3 = np.expand_dims(R_Nx3x3, axis=0)
245 |     else:
246 |         raise ValueError('unexpected shape')
247 | 
248 |     N = R_Nx3x3.shape[0]
249 |     R_Nx9 = np.reshape(R_Nx3x3, (N, 9), order='C')
250 | 
251 |     pitch = np.arctan2(-R_Nx9[:, 6],
252 |                        np.sqrt(np.power(R_Nx9[:, 0], 2*np.ones(N,)) + np.power(R_Nx9[:, 3], 2*np.ones(N,))))
253 | 
254 |     yaw = np.zeros((N,))
255 |     roll = np.zeros((N,))
256 | 
257 |     index = np.abs(pitch - np.pi/2.0) < eps
258 |     roll[index] = np.arctan2(R_Nx9[index, 1], R_Nx9[index, 4])
259 | 
260 |     index = np.abs(pitch + np.pi/2.0) < eps
261 |     roll[index] = -np.arctan2(R_Nx9[index, 1], R_Nx9[index, 4])
262 | 
263 |     index = np.logical_and(np.abs(pitch - np.pi/2.0) >= eps,
264 |                            np.abs(pitch + np.pi/2.0) >= eps)
265 |     yaw[index] = np.arctan2(R_Nx9[index, 3]/np.cos(pitch[index]),
266 |                             R_Nx9[index, 0]/np.cos(pitch[index]))
267 |     roll[index] = np.arctan2(R_Nx9[index, 7]/np.cos(pitch[index]),
268 |                              R_Nx9[index, 8]/np.cos(pitch[index]))
269 | 
270 |     roll = np.expand_dims(roll,  axis=1)
271 |     pitch = np.expand_dims(pitch, axis=1)
272 |     yaw = np.expand_dims(yaw,   axis=1)
273 | 
274 |     rot_Nx3 = np.concatenate((roll, pitch, yaw), axis=1)
275 |     if not is_radians:
276 |         rot_Nx3 = 180.0/np.pi * rot_Nx3
277 | 
278 |     if ndim == 2:
279 |         rot_Nx3 = rot_Nx3[0]
280 | 
281 |     return rot_Nx3
282 | 
283 | def matTranslation(vector):
284 |     assert len(vector) == 3
285 |     e = np.eye(4)
286 |     e[3, 0:3] = vector
287 |     return e
288 | 
289 | def test_convertTransRotTo4x4():
290 | 
291 |     transrot_Nx6 = np.array([10, 20, 30, 1, 10, -90])
292 |     transrot_Nx6 = np.expand_dims(transrot_Nx6, axis=0)
293 |     transrot_Nx6 = np.repeat(transrot_Nx6, 10000, axis=0)
294 | 
295 |     print(transrot_Nx6.ndim)
296 | 
297 |     print('transrot_Nx6:', transrot_Nx6)
298 |     print('transrot_Nx6.shape:', transrot_Nx6.shape)
299 | 
300 |     with Timer('convertTransRotTo4x4'):
301 |         matrix_Nx4x4 = convertTransRotTo4x4(transrot_Nx6)
302 |     print('matrix_Nx4x4:', matrix_Nx4x4)
303 | 
304 |     with Timer('convert4x4ToTransRot'):
305 |         transrot_Nx6 = convert4x4ToTransRot(matrix_Nx4x4)
306 |     print('transrot_Nx6:', transrot_Nx6)
307 | 
308 |     np.testing.assert_array_almost_equal(
309 |         transrot_Nx6, convert4x4ToTransRot(matrix_Nx4x4), verbose=True)
310 | 
311 | if __name__ == "__main__":
312 |     test_convertTransRotTo4x4()
313 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | numpy
2 | pycuda


--------------------------------------------------------------------------------