Add python data3d.change_geometry

3 files changed, 81 insertions, 0 deletions

diff --git a/python/astra/PyIncludes.pxd b/python/astra/PyIncludes.pxd
index 7df02c5..a581f88 100644
--- a/python/astra/PyIncludes.pxd
+++ b/python/astra/PyIncludes.pxd
@@ -196,18 +196,29 @@ cdef extern from "astra/VolumeGeometry3D.h" namespace "astra":
 		CVolumeGeometry3D()
 		bool initialize(Config)
 		Config * getConfiguration()
+		int getGridColCount()
+		int getGridRowCount()
+		int getGridSliceCount()
 
 cdef extern from "astra/ProjectionGeometry3D.h" namespace "astra":
 	cdef cppclass CProjectionGeometry3D:
 		CProjectionGeometry3D()
 		bool initialize(Config)
 		Config * getConfiguration()
+		int getProjectionCount()
+		int getDetectorColCount()
+		int getDetectorRowCount()
 
 
 cdef extern from "astra/Float32VolumeData3DMemory.h" namespace "astra":
 	cdef cppclass CFloat32VolumeData3DMemory:
 		CFloat32VolumeData3DMemory(CVolumeGeometry3D*)
 		CVolumeGeometry3D* getGeometry()
+		void changeGeometry(CVolumeGeometry3D*)
+		int getRowCount()
+		int getColCount()
+		int getSliceCount()
+
 
 
 cdef extern from "astra/ParallelProjectionGeometry3D.h" namespace "astra":
@@ -232,6 +243,10 @@ cdef extern from "astra/Float32ProjectionData3DMemory.h" namespace "astra":
 		CFloat32ProjectionData3DMemory(CProjectionGeometry3D*)
 		CFloat32ProjectionData3DMemory(CConeProjectionGeometry3D*)
 		CProjectionGeometry3D* getGeometry()
+		void changeGeometry(CProjectionGeometry3D*)
+		int getDetectorColCount()
+		int getDetectorRowCount()
+		int getAngleCount()
 
 cdef extern from "astra/Float32Data3D.h" namespace "astra":
 	cdef cppclass CFloat32Data3D:
diff --git a/python/astra/data3d.py b/python/astra/data3d.py
index a2e9201..4679489 100644
--- a/python/astra/data3d.py
+++ b/python/astra/data3d.py
@@ -90,6 +90,17 @@ def get_geometry(i):
     """
     return d.get_geometry(i)
 
+def change_geometry(i, geometry):
+    """Change the geometry of a 3D object.
+
+    :param i: ID of object.
+    :type i: :class:`int`
+    :param geometry: Volume or projection geometry.
+    :type geometry: :class:`dict`
+
+    """
+    return d.change_geometry(i, geometry)
+
 def dimensions(i):
     """Get dimensions of a 3D object.
 
diff --git a/python/astra/data3d_c.pyx b/python/astra/data3d_c.pyx
index 4b069f7..48af032 100644
--- a/python/astra/data3d_c.pyx
+++ b/python/astra/data3d_c.pyx
@@ -122,6 +122,61 @@ def get_geometry(i):
         raise Exception("Not a known data object")
     return geom
 
+def change_geometry(i, geom):
+    cdef CFloat32Data3DMemory * pDataObject = dynamic_cast_mem(getObject(i))
+    cdef CFloat32ProjectionData3DMemory * pDataObject2
+    cdef CFloat32VolumeData3DMemory * pDataObject3
+    if pDataObject.getType() == THREEPROJECTION:
+        pDataObject2 = <CFloat32ProjectionData3DMemory * >pDataObject
+        # TODO: Reduce code duplication here
+        cfg = utils.dictToConfig(six.b('ProjectionGeometry'), geom)
+        tpe = wrap_from_bytes(cfg.self.getAttribute(six.b('type')))
+        if (tpe == "parallel3d"):
+            ppGeometry = <CProjectionGeometry3D*> new CParallelProjectionGeometry3D();
+        elif (tpe == "parallel3d_vec"):
+            ppGeometry = <CProjectionGeometry3D*> new CParallelVecProjectionGeometry3D();
+        elif (tpe == "cone"):
+            ppGeometry = <CProjectionGeometry3D*> new CConeProjectionGeometry3D();
+        elif (tpe == "cone_vec"):
+            ppGeometry = <CProjectionGeometry3D*> new CConeVecProjectionGeometry3D();
+        else:
+            raise Exception("Invalid geometry type.")
+        if not ppGeometry.initialize(cfg[0]):
+            del cfg
+            del ppGeometry
+            raise Exception('Geometry class not initialized.')
+        del cfg
+        if (ppGeometry.getDetectorColCount() != pDataObject2.getDetectorColCount() or \
+            ppGeometry.getProjectionCount() != pDataObject2.getAngleCount() or \
+            ppGeometry.getDetectorRowCount() != pDataObject2.getDetectorRowCount()):
+            del ppGeometry
+            raise Exception(
+                "The dimensions of the data do not match those specified in the geometry.")
+        pDataObject2.changeGeometry(ppGeometry)
+        del ppGeometry
+
+    elif pDataObject.getType() == THREEVOLUME:
+        pDataObject3 = <CFloat32VolumeData3DMemory * >pDataObject
+        cfg = utils.dictToConfig(six.b('VolumeGeometry'), geom)
+        pGeometry = new CVolumeGeometry3D()
+        if not pGeometry.initialize(cfg[0]):
+            del cfg
+            del pGeometry
+            raise Exception('Geometry class not initialized.')
+        del cfg
+        if (pGeometry.getGridColCount() != pDataObject3.getColCount() or \
+            pGeometry.getGridRowCount() != pDataObject3.getRowCount() or \
+            pGeometry.getGridSliceCount() != pDataObject3.getSliceCount()):
+            del pGeometry
+            raise Exception(
+                "The dimensions of the data do not match those specified in the geometry.")
+        pDataObject3.changeGeometry(pGeometry)
+        del pGeometry
+
+    else:
+        raise Exception("Not a known data object")
+
+
 cdef fillDataObject(CFloat32Data3DMemory * obj, data):
     if data is None:
         fillDataObjectScalar(obj, 0)