infer dimension from geometry (#42)

* infer dimension from geometry

closes #41

* use default pixel_num_v

1 files changed, 20 insertions, 4 deletions

diff --git a/Wrappers/Python/ccpi/astra/utils/convert_geometry_to_astra.py b/Wrappers/Python/ccpi/astra/utils/convert_geometry_to_astra.py
index 2b19305..0c4312b 100644
--- a/Wrappers/Python/ccpi/astra/utils/convert_geometry_to_astra.py
+++ b/Wrappers/Python/ccpi/astra/utils/convert_geometry_to_astra.py
@@ -2,8 +2,24 @@ import astra
 import numpy as np
 
 def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
-    # Set up ASTRA Volume and projection geometry, not stored
-    if sinogram_geometry.dimension == '2D':
+    '''Set up ASTRA Volume and projection geometry, not stored
+
+       :param volume_geometry: ccpi.framework.ImageGeometry
+       :param sinogram_geometry: ccpi.framework.AcquisitionGeometry
+       
+       :returns ASTRA volume and sinogram geometry'''
+
+    # determine if the geometry is 2D or 3D
+    horiz = sinogram_geometry.pixel_num_h
+    vert  = sinogram_geometry.pixel_num_v
+    if vert >= 1:
+        dimension = '3D'
+    elif vert == 0:
+        dimension = '2D'
+    else:
+        raise ValueError('Number of pixels at detector on the vertical axis must be >= 0. Got {}'.format(vert))
+    
+    if dimension == '2D':
         vol_geom = astra.create_vol_geom(volume_geometry.voxel_num_x, 
                                          volume_geometry.voxel_num_y, 
                                          volume_geometry.get_min_x(), 
@@ -26,7 +42,7 @@ def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
         else:
             NotImplemented
             
-    elif sinogram_geometry.dimension == '3D':
+    elif dimension == '3D':
         vol_geom = astra.create_vol_geom(volume_geometry.voxel_num_x, 
                                          volume_geometry.voxel_num_y, 
                                          volume_geometry.voxel_num_z, 
@@ -59,4 +75,4 @@ def convert_geometry_to_astra(volume_geometry, sinogram_geometry):
     else:
         NotImplemented
     
-    return vol_geom, proj_geom
\ No newline at end of file
+    return vol_geom, proj_geom