Add ReconstructionAlgorithm2D/3D classes for plugins (matching C++ classes)

3 files changed, 55 insertions, 1 deletions

diff --git a/python/astra/PyIncludes.pxd b/python/astra/PyIncludes.pxd
index 909f58f..a099c31 100644
--- a/python/astra/PyIncludes.pxd
+++ b/python/astra/PyIncludes.pxd
@@ -62,6 +62,7 @@ cdef extern from "astra/VolumeGeometry2D.h" namespace "astra":
 		float32 getWindowMaxX()
 		float32 getWindowMaxY()
 		Config* getConfiguration()
+		bool isEqual(CVolumeGeometry2D*)
 
 cdef extern from "astra/Float32Data2D.h" namespace "astra":
 	cdef cppclass CFloat32CustomMemory:
@@ -89,6 +90,7 @@ cdef extern from "astra/ProjectionGeometry2D.h" namespace "astra":
 		float32 getProjectionAngle(int)
 		float32 getDetectorWidth()
 		Config* getConfiguration()
+		bool isEqual(CProjectionGeometry2D*)
 
 cdef extern from "astra/Float32Data2D.h" namespace "astra::CFloat32Data2D":
 	cdef enum TWOEDataType "astra::CFloat32Data2D::EDataType":
diff --git a/python/astra/data2d_c.pyx b/python/astra/data2d_c.pyx
index 4919bf2..801fd8e 100644
--- a/python/astra/data2d_c.pyx
+++ b/python/astra/data2d_c.pyx
@@ -34,6 +34,9 @@ from cython cimport view
 cimport PyData2DManager
 from .PyData2DManager cimport CData2DManager
 
+cimport PyProjector2DManager
+from .PyProjector2DManager cimport CProjector2DManager
+
 cimport PyXMLDocument
 from .PyXMLDocument cimport XMLDocument
 
@@ -54,6 +57,8 @@ import operator
 from six.moves import reduce
 
 cdef CData2DManager * man2d = <CData2DManager * >PyData2DManager.getSingletonPtr()
+cdef CProjector2DManager * manProj = <CProjector2DManager * >PyProjector2DManager.getSingletonPtr()
+
 
 cdef extern from "CFloat32CustomPython.h":
     cdef cppclass CFloat32CustomPython:
@@ -164,7 +169,6 @@ def store(i, data):
     cdef CFloat32Data2D * pDataObject = getObject(i)
     fillDataObject(pDataObject, data)
 
-
 def get_geometry(i):
     cdef CFloat32Data2D * pDataObject = getObject(i)
     cdef CFloat32ProjectionData2D * pDataObject2
@@ -179,6 +183,27 @@ def get_geometry(i):
         raise Exception("Not a known data object")
     return geom
 
+cdef CProjector2D * getProjector(i) except NULL:
+    cdef CProjector2D * proj = manProj.get(i)
+    if proj == NULL:
+        raise Exception("Projector not initialized.")
+    if not proj.isInitialized():
+        raise Exception("Projector not initialized.")
+    return proj
+
+def check_compatible(i, proj_id):
+    cdef CProjector2D * proj = getProjector(proj_id)
+    cdef CFloat32Data2D * pDataObject = getObject(i)
+    cdef CFloat32ProjectionData2D * pDataObject2
+    cdef CFloat32VolumeData2D * pDataObject3
+    if pDataObject.getType() == TWOPROJECTION:
+        pDataObject2 = <CFloat32ProjectionData2D * >pDataObject
+        return pDataObject2.getGeometry().isEqual(proj.getProjectionGeometry())
+    elif pDataObject.getType() == TWOVOLUME:
+        pDataObject3 = <CFloat32VolumeData2D * >pDataObject
+        return pDataObject3.getGeometry().isEqual(proj.getVolumeGeometry())
+    else:
+        raise Exception("Not a known data object")
 
 def change_geometry(i, geom):
     cdef Config *cfg
diff --git a/python/astra/plugin.py b/python/astra/plugin.py
index 4b32e6e..11cc5cc 100644
--- a/python/astra/plugin.py
+++ b/python/astra/plugin.py
@@ -26,6 +26,10 @@
 
 from . import plugin_c as p
 from . import log
+from . import data2d
+from . import data2d_c
+from . import data3d
+from . import projector
 import inspect
 import traceback
 
@@ -65,6 +69,29 @@ class base(object):
         kwargs = dict((k,optDict[k]) for k in opt if k in optDict)
         self.initialize(cfg, *args, **kwargs)
 
+class ReconstructionAlgorithm2D(base):
+
+    def astra_init(self, cfg):
+        self.pid = cfg['ProjectorId']
+        self.s = data2d.get_shared(cfg['ProjectionDataId'])
+        self.v = data2d.get_shared(cfg['ReconstructionDataId'])
+        self.vg = projector.volume_geometry(self.pid)
+        self.pg = projector.projection_geometry(self.pid)
+        if not data2d_c.check_compatible(cfg['ProjectionDataId'], self.pid):
+            raise ValueError("Projection data and projector not compatible")
+        if not data2d_c.check_compatible(cfg['ReconstructionDataId'], self.pid):
+            raise ValueError("Reconstruction data and projector not compatible")
+        super(ReconstructionAlgorithm2D,self).astra_init(cfg)
+
+class ReconstructionAlgorithm3D(base):
+
+    def astra_init(self, cfg):
+        self.s = data3d.get_shared(cfg['ProjectionDataId'])
+        self.v = data3d.get_shared(cfg['ReconstructionDataId'])
+        self.vg = data3d.get_geometry(cfg['ReconstructionDataId'])
+        self.pg = data3d.get_geometry(cfg['ProjectionDataId'])
+        super(ReconstructionAlgorithm3D,self).astra_init(cfg)
+
 def register(className):
     """Register plugin with ASTRA.