Add 'link' feature to Python (for 2D and 3D data)

6 files changed, 107 insertions, 9 deletions

diff --git a/python/astra/CFloat32CustomPython.h b/python/astra/CFloat32CustomPython.h
new file mode 100644
index 0000000..d8593fc
--- /dev/null
+++ b/python/astra/CFloat32CustomPython.h
@@ -0,0 +1,17 @@
+class CFloat32CustomPython : public astra::CFloat32CustomMemory {
+public:
+    CFloat32CustomPython(PyObject * arrIn)
+    {
+        arr = arrIn;
+        // Set pointer to numpy data pointer
+        m_fPtr = (float *)PyArray_DATA(arr);
+        // Increase reference count since ASTRA has a reference
+        Py_INCREF(arr);
+    }
+    virtual ~CFloat32CustomPython() {
+        // Decrease reference count since ASTRA object is destroyed
+        Py_DECREF(arr);
+    }
+private:
+    PyObject* arr;
+};
\ No newline at end of file
diff --git a/python/astra/PyIncludes.pxd b/python/astra/PyIncludes.pxd
index 7df02c5..13329d1 100644
--- a/python/astra/PyIncludes.pxd
+++ b/python/astra/PyIncludes.pxd
@@ -63,10 +63,14 @@ cdef extern from "astra/VolumeGeometry2D.h" namespace "astra":
 		float32 getWindowMaxY()
 		Config* getConfiguration()
 
+cdef extern from "astra/Float32Data2D.h" namespace "astra":
+	cdef cppclass CFloat32CustomMemory:
+		pass
 
 cdef extern from "astra/Float32VolumeData2D.h" namespace "astra":
 	cdef cppclass CFloat32VolumeData2D:
 		CFloat32VolumeData2D(CVolumeGeometry2D*)
+		CFloat32VolumeData2D(CVolumeGeometry2D*, CFloat32CustomMemory*)
 		CVolumeGeometry2D * getGeometry()
 		int getWidth()
 		int getHeight()
@@ -130,6 +134,7 @@ cdef extern from "astra/ParallelProjectionGeometry2D.h" namespace "astra":
 cdef extern from "astra/Float32ProjectionData2D.h" namespace "astra":
 	cdef cppclass CFloat32ProjectionData2D:
 		CFloat32ProjectionData2D(CProjectionGeometry2D*)
+		CFloat32ProjectionData2D(CProjectionGeometry2D*, CFloat32CustomMemory*)
 		CProjectionGeometry2D * getGeometry()
 		void changeGeometry(CProjectionGeometry2D*)
 		int getDetectorCount()
@@ -207,6 +212,7 @@ cdef extern from "astra/ProjectionGeometry3D.h" namespace "astra":
 cdef extern from "astra/Float32VolumeData3DMemory.h" namespace "astra":
 	cdef cppclass CFloat32VolumeData3DMemory:
 		CFloat32VolumeData3DMemory(CVolumeGeometry3D*)
+		CFloat32VolumeData3DMemory(CVolumeGeometry3D*, CFloat32CustomMemory*)
 		CVolumeGeometry3D* getGeometry()
 
 
@@ -231,6 +237,8 @@ cdef extern from "astra/Float32ProjectionData3DMemory.h" namespace "astra":
 	cdef cppclass CFloat32ProjectionData3DMemory:
 		CFloat32ProjectionData3DMemory(CProjectionGeometry3D*)
 		CFloat32ProjectionData3DMemory(CConeProjectionGeometry3D*)
+		CFloat32ProjectionData3DMemory(CProjectionGeometry3D*, CFloat32CustomMemory*)
+		CFloat32ProjectionData3DMemory(CConeProjectionGeometry3D*, CFloat32CustomMemory*)
 		CProjectionGeometry3D* getGeometry()
 
 cdef extern from "astra/Float32Data3D.h" namespace "astra":
diff --git a/python/astra/data2d.py b/python/astra/data2d.py
index 8c4be03..f119f05 100644
--- a/python/astra/data2d.py
+++ b/python/astra/data2d.py
@@ -24,6 +24,7 @@
 #
 #-----------------------------------------------------------------------
 from . import data2d_c as d
+import numpy as np
 
 def clear():
     """Clear all 2D data objects."""
@@ -52,6 +53,26 @@ def create(datatype, geometry, data=None):
     """
     return d.create(datatype,geometry,data)
 
+def link(datatype, geometry, data):
+    """Link a 2D numpy array with the toolbox.
+        
+    :param datatype: Data object type, '-vol' or '-sino'.
+    :type datatype: :class:`string`
+    :param geometry: Volume or projection geometry.
+    :type geometry: :class:`dict`
+    :param data: Numpy array to link
+    :type data: :class:`numpy.ndarray`
+    :returns: :class:`int` -- the ID of the constructed object.
+    
+    """
+    if not isinstance(data,np.ndarray):
+        raise ValueError("Input should be a numpy array")
+    if not data.dtype==np.float32:
+        raise ValueError("Numpy array should be float32")
+    if not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
+        raise ValueError("Numpy array should be C_CONTIGUOUS and ALIGNED")
+    return d.create(datatype,geometry,data,True)
+
 def store(i, data):
     """Fill existing 2D object with data.
     
diff --git a/python/astra/data2d_c.pyx b/python/astra/data2d_c.pyx
index b9c105e..ac54898 100644
--- a/python/astra/data2d_c.pyx
+++ b/python/astra/data2d_c.pyx
@@ -49,6 +49,10 @@ from .utils import wrap_from_bytes
 
 cdef CData2DManager * man2d = <CData2DManager * >PyData2DManager.getSingletonPtr()
 
+cdef extern from "CFloat32CustomPython.h":
+    cdef cppclass CFloat32CustomPython:
+        CFloat32CustomPython(arrIn)
+
 def clear():
     man2d.clear()
 
@@ -61,11 +65,12 @@ def delete(ids):
         man2d.remove(ids)
 
 
-def create(datatype, geometry, data=None):
+def create(datatype, geometry, data=None, link=False):
     cdef Config *cfg
     cdef CVolumeGeometry2D * pGeometry
     cdef CProjectionGeometry2D * ppGeometry
     cdef CFloat32Data2D * pDataObject2D
+    cdef CFloat32CustomMemory * pCustom
     if datatype == '-vol':
         cfg = utils.dictToConfig(six.b('VolumeGeometry'), geometry)
         pGeometry = new CVolumeGeometry2D()
@@ -73,7 +78,11 @@ def create(datatype, geometry, data=None):
             del cfg
             del pGeometry
             raise Exception('Geometry class not initialized.')
-        pDataObject2D = <CFloat32Data2D * > new CFloat32VolumeData2D(pGeometry)
+        if link:
+            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+            pDataObject2D = <CFloat32Data2D * > new CFloat32VolumeData2D(pGeometry, pCustom)
+        else:
+            pDataObject2D = <CFloat32Data2D * > new CFloat32VolumeData2D(pGeometry)
         del cfg
         del pGeometry
     elif datatype == '-sino':
@@ -91,7 +100,11 @@ def create(datatype, geometry, data=None):
             del cfg
             del ppGeometry
             raise Exception('Geometry class not initialized.')
-        pDataObject2D = <CFloat32Data2D * > new CFloat32ProjectionData2D(ppGeometry)
+        if link:
+            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+            pDataObject2D = <CFloat32Data2D * > new CFloat32ProjectionData2D(ppGeometry, pCustom)
+        else:
+            pDataObject2D = <CFloat32Data2D * > new CFloat32ProjectionData2D(ppGeometry)
         del ppGeometry
         del cfg
     else:
@@ -101,7 +114,7 @@ def create(datatype, geometry, data=None):
         del pDataObject2D
         raise Exception("Couldn't initialize data object.")
 
-    fillDataObject(pDataObject2D, data)
+    if not link: fillDataObject(pDataObject2D, data)
 
     return man2d.store(pDataObject2D)
 
diff --git a/python/astra/data3d.py b/python/astra/data3d.py
index a2e9201..4fdf9d7 100644
--- a/python/astra/data3d.py
+++ b/python/astra/data3d.py
@@ -24,6 +24,7 @@
 #
 #-----------------------------------------------------------------------
 from . import data3d_c as d
+import numpy as np
 
 def create(datatype,geometry,data=None):
     """Create a 3D object.
@@ -39,6 +40,27 @@ def create(datatype,geometry,data=None):
     """
     return d.create(datatype,geometry,data)
 
+def link(datatype, geometry, data):
+    """Link a 3D numpy array with the toolbox.
+        
+    :param datatype: Data object type, '-vol' or '-sino'.
+    :type datatype: :class:`string`
+    :param geometry: Volume or projection geometry.
+    :type geometry: :class:`dict`
+    :param data: Numpy array to link
+    :type data: :class:`numpy.ndarray`
+    :returns: :class:`int` -- the ID of the constructed object.
+    
+    """
+    if not isinstance(data,np.ndarray):
+        raise ValueError("Input should be a numpy array")
+    if not data.dtype==np.float32:
+        raise ValueError("Numpy array should be float32")
+    if not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
+        raise ValueError("Numpy array should be C_CONTIGUOUS and ALIGNED")
+    return d.create(datatype,geometry,data,True)
+
+
 def get(i):
     """Get a 3D object.
 
diff --git a/python/astra/data3d_c.pyx b/python/astra/data3d_c.pyx
index 4b069f7..f2c6e26 100644
--- a/python/astra/data3d_c.pyx
+++ b/python/astra/data3d_c.pyx
@@ -50,12 +50,17 @@ cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr
 cdef extern from *:
     CFloat32Data3DMemory * dynamic_cast_mem "dynamic_cast<astra::CFloat32Data3DMemory*>" (CFloat32Data3D * ) except NULL
 
-def create(datatype,geometry,data=None):
+cdef extern from "CFloat32CustomPython.h":
+    cdef cppclass CFloat32CustomPython:
+        CFloat32CustomPython(arrIn)
+
+def create(datatype,geometry,data=None, link=False):
     cdef Config *cfg
     cdef CVolumeGeometry3D * pGeometry
     cdef CProjectionGeometry3D * ppGeometry
     cdef CFloat32Data3DMemory * pDataObject3D
     cdef CConeProjectionGeometry3D* pppGeometry
+    cdef CFloat32CustomMemory * pCustom
     if datatype == '-vol':
         cfg = utils.dictToConfig(six.b('VolumeGeometry'), geometry)
         pGeometry = new CVolumeGeometry3D()
@@ -63,7 +68,11 @@ def create(datatype,geometry,data=None):
             del cfg
             del pGeometry
             raise Exception('Geometry class not initialized.')
-        pDataObject3D = <CFloat32Data3DMemory * > new CFloat32VolumeData3DMemory(pGeometry)
+        if link:
+            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32VolumeData3DMemory(pGeometry, pCustom)
+        else:
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32VolumeData3DMemory(pGeometry)
         del cfg
         del pGeometry
     elif datatype == '-sino' or datatype == '-proj3d':
@@ -84,7 +93,11 @@ def create(datatype,geometry,data=None):
             del cfg
             del ppGeometry
             raise Exception('Geometry class not initialized.')
-        pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry)
+        if link:
+            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry, pCustom)
+        else:
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry)
         del ppGeometry
         del cfg
     elif datatype == "-sinocone":
@@ -94,7 +107,11 @@ def create(datatype,geometry,data=None):
             del cfg
             del pppGeometry
             raise Exception('Geometry class not initialized.')
-        pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(pppGeometry)
+        if link:
+            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(pppGeometry, pCustom)
+        else:
+            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(pppGeometry)
     else:
         raise Exception("Invalid datatype.  Please specify '-vol' or '-proj3d'.")
 
@@ -102,7 +119,7 @@ def create(datatype,geometry,data=None):
         del pDataObject3D
         raise Exception("Couldn't initialize data object.")
 
-    fillDataObject(pDataObject3D, data)
+    if not link: fillDataObject(pDataObject3D, data)
 
     pDataObject3D.updateStatistics()