Compile c++ shared object as pyronn_torch_cpp

src/pyronn_torch/__init__.py

 # -*- coding: utf-8 -*-
 
-from os.path import dirname, join
+import sys
+from os.path import dirname
 
-import torch
 from pkg_resources import DistributionNotFound, get_distribution
 
+from pyronn_torch.conebeam import ConeBeamProjector
+
 try:
     # Change here if project is renamed and does not equal the package name
     dist_name = 'pyronn-torch'
@@ -16,8 +18,12 @@ finally:
 
 
 try:
-    cpp_extension = torch.ops.load_library(join(dirname(__file__), 'pyronn_torch.so'))
-except Exception:
+    sys.path.append(dirname(__file__))
+    cpp_extension = __import__('pyronn_torch_cpp')
+except Exception as e:
+    import warnings
+    warnings.warn(str(e))
     import pyronn_torch.codegen
-    cpp_extension = pyronn_torch.codegen.compile_shared_object()
+    cpp_extension = pyronn_torch.codegen.generate_shared_object()
 
+__all__ = ['ConeBeamProjector', 'cpp_extension']

src/pyronn_torch/codegen.py

@@ -151,7 +151,7 @@ def generate_shared_object(output_folder=None, source_files=None, show_code=Fals
 
     object_cache = get_cache_config()['object_cache']
 
-    module_name = 'PYRO_NN'
+    module_name = 'pyronn_torch_cpp'
 
     if not output_folder:
         output_folder = dirname(__file__)
@@ -176,10 +176,13 @@ def generate_shared_object(output_folder=None, source_files=None, show_code=Fals
     if show_code:
         pystencils.show_code(module, custom_backend=FrameworkIntegrationPrinter())
 
-    extension = module.compile(extra_source_files=cuda_sources, extra_cuda_flags=['-arch=sm_35'], with_cuda=True)
+    extension = module.compile(extra_source_files=cuda_sources,
+                               extra_cuda_flags=['-arch=sm_35'],
+                               with_cuda=True,
+                               compile_module_name=module_name)
 
-    shared_object_file = module.compiled_file.replace('.cpp', '.so')
-    copyfile(shared_object_file, join(output_folder, 'pyronn_torch.so'))
+    shared_object_file = module.compiled_file
+    copyfile(shared_object_file, join(output_folder, module_name + '.so'))
     copyfile(module.compiled_file, join(output_folder, 'pyronn_torch.cpp'))
 
     return extension

src/pyronn_torch/conebeam.py

+#
+# Copyright © 2020 Stephan Seitz <stephan.seitz@fau.de>
+#
+# Distributed under terms of the GPLv3 license.
+
+"""
+
+"""
+import numpy as np
+import sympy as sp
+import torch
+
+import pyronn_torch
+
+
+class ConeBeamProjector:
+    class ForwardProjection(torch.autograd.Function):
+
+        def __init__(self, projection_shape,
+                     source_points,
+                     inverse_matrices,
+                     projection_matrices,
+                     volume_origin,
+                     volume_spacing,
+                     projection_multiplier,
+                     step_size=1.,
+                     with_texture=True):
+            self.projection_shape = projection_shape
+            self.source_points = source_points
+            self.inverse_matrices = inverse_matrices
+            self.projection_matrices = projection_matrices
+            self.volume_origin = volume_origin
+            self.volume_spacing = volume_spacing
+            self.projection_multiplier = projection_multiplier
+            self.with_texture = with_texture
+            self.step_size = step_size
+
+        def forward(self, volume):
+            volume = volume.cuda().contiguous()
+            projection = torch.zeros(self.projection_shape, device='cuda')
+
+            assert pyronn_torch.cpp_extension
+            if self.with_texture:
+                pyronn_torch.cpp_extension.call_Cone_Projection_Kernel_Tex_Interp_Launcher(
+                    self.inverse_matrices,
+                    projection,
+                    self.source_points,
+                    self.step_size,
+                    volume,
+                    *reversed(self.volume_spacing))
+            else:
+                pyronn_torch.cpp_extension.call_Cone_Projection_Kernel_Launcher(
+                    self.inverse_matrices,
+                    projection,
+                    self.source_points,
+                    self.step_size,
+                    volume,
+                    *reversed(self.volume_spacing))
+
+            return projection,
+
+        def backward(self, *projection_grad):
+            projection_grad = projection_grad[0]
+            self.projection_matrices
+            volume_grad = torch.zeros(self.volume_shape, device='cuda')
+
+            assert pyronn_torch.cpp_extension
+            pyronn_torch.cpp_extension.call_Cone_Backprojection3D_Kernel_Launcher(
+                self.projection_matrices,
+                projection_grad,
+                self.projection_multiplier,
+                volume_grad,
+                *reversed(self.volume_origin),
+                *reversed(self.volume_spacing))
+
+            return volume_grad,
+
+    class BackwardProjection(torch.autograd.Function):
+        pass
+
+    # def __init__(self,
+        # volume_shape,
+        # volume_spacing,
+        # volume_origin,
+        # projection_shape,
+        # projection_spacing,
+        # projection_origin,
+        # projection_matrices):
+        # self._volume_shape = volume_shape
+        # self._volume_origin = volume_origin
+        # self._volume_spacing = volume_spacing
+        # self._projection_shape = projection_shape
+        # self._projection_matrices = projection_matrices
+        # self._projection_spacing = projection_spacing
+        # self._projection_origin = projection_origin
+        # self._calc_inverse_matrices()
+
+    def __init__(self):
+        import pyconrad.autoinit
+        import pyconrad.config
+        self._volume_shape = pyconrad.config.get_reco_shape()
+        self._volume_spacing = pyconrad.config.get_reco_spacing()
+        self._volume_origin = pyconrad.config.get_reco_origin()
+        self._projection_shape = pyconrad.config.get_sino_shape()
+        self._projection_spacing = [pyconrad.config.get_geometry().getPixelDimensionY(),
+                                    pyconrad.config.get_geometry().getPixelDimensionX()]
+        self._projection_origin = [pyconrad.config.get_geometry().getDetectorOffsetV(),
+                                   pyconrad.config.get_geometry().getDetectorOffsetU()]
+        self._projection_matrices_numpy = pyconrad.config.get_projection_matrices()
+
+        self._calc_inverse_matrices()
+
+    def new_volume_tensor(self, requires_grad=False):
+        return torch.zeros(self._volume_shape, requires_grad=requires_grad).cuda()
+
+    def new_projection_tensor(self, requires_grad=False):
+        return torch.zeros(self._projection_shape, requires_grad=requires_grad).cuda()
+
+    def project_forward(self, volume, step_size=1., use_texture=True):
+        return self.ForwardProjection(self._projection_shape,
+                                      self._source_points,
+                                      self._inverse_matrices,
+                                      self._projection_matrices,
+                                      self._volume_origin,
+                                      self._volume_shape,
+                                      self._projection_multiplier,
+                                      step_size,
+                                      use_texture).forward(volume)
+
+    def project_backward(self, projection_stack):
+        return self.BackwardProjection(projection_stack)
+
+    def _calc_inverse_matrices(self):
+        self._projection_matrices = torch.stack(tuple(
+            map(torch.from_numpy, self._projection_matrices_numpy))).cuda().contiguous()
+
+        inv_spacing = np.array([1/s for s in reversed(self._volume_spacing)], np.float32)
+
+        camera_centers = map(lambda x: np.array(sp.Matrix(x).nullspace(), np.float32), self._projection_matrices_numpy)
+
+        source_points = map(lambda x: (x[0, :3] / x[0, 3] * inv_spacing - np.array(list(reversed(self._volume_origin)))
+                                       * inv_spacing).astype(np.float32), camera_centers)
+
+        inv_matrices = map(lambda x: (np.linalg.inv(x[:3, :3]) *
+                                      inv_spacing).astype(np.float32), self._projection_matrices_numpy)
+
+        self._inverse_matrices = torch.stack(tuple(map(torch.from_numpy, inv_matrices))).cuda().contiguous()
+        self._source_points = torch.stack(tuple(map(torch.from_numpy, source_points))).cuda().contiguous()
+        self._projection_multiplier = 1.

tests/test_projection.py

+#
+# Copyright © 2020 Stephan Seitz <stephan.seitz@fau.de>
+#
+# Distributed under terms of the GPLv3 license.
+
+"""
+
+"""
+import pyronn_torch
+
+
+def init():
+    assert pyronn_torch.cpp_extension
+
+
+def test_projection():
+    breakpoint()
+    projector = pyronn_torch.ConeBeamProjector()
+
+    volume = projector.new_volume_tensor()
+
+    volume += 1.
+    result = projector.project_forward(volume, use_texture=False)
+
+    print(result)