Merge pull request #30 from stuwilkins/add_image_functions

Add image functions

.travis.yml

@@ -26,7 +26,7 @@ install:
 - conda install conda-build
 - conda create -n testenv pip python=$TRAVIS_PYTHON_VERSION six "numpy<1.10" pandas
   h5py coverage jsonschema jinja2 sphinx pyyaml ipywidgets tzlocal prettytable boltons
-  pymongo=2.9 pims tifffile databroker sphinx-bootstrap-theme pandoc
+  pymongo=2.9 pims=0.3.3 tifffile databroker sphinx-bootstrap-theme pandoc
 - source activate testenv
 - ln -s /home/travis/mc/envs/testenv/lib/libgmp.so.10 /home/travis/mc/envs/testenv/lib/libgmp.so.3
 - pip install coveralls

csxtools/image/__init__.py

@@ -3,8 +3,9 @@
 __version__ = get_versions()['version']
 del get_versions
 
-__all__ = ['rotate90']
+__all__ = ['rotate90', 'stackmean', 'stacksum']
 
 # Now import useful functions
 
 from .transform import rotate90
+from .stack import stackmean, stacksum

csxtools/image/stack.py

@@ -0,0 +1,45 @@
+from ..ext import image as extimage
+
+
+def stackmean(array):
+    """Cacluate the mean of a stack
+
+    This function calculates the mean of a stack of images (or any array).
+    It ignores values that are np.NAN and does not include them in the mean
+    calculation. It assumes an array of shape (.. i, j, x, y) where x and y
+    are the size of the returned array (x, y).
+
+    Parameters
+    ----------
+    array : array_like
+        Input array of at least 3 dimensions.
+
+    Returns
+    -------
+    array
+        2D Array of mean of stack.
+    """
+    X, Y = extimage.stackmean(array, 1)
+    return X
+
+
+def stacksum(array):
+    """Cacluate the sum of a stack
+
+    This function calculates the sum of a stack of images (or any array).
+    It ignores values that are np.NAN and does not include them in the sum
+    calculation. It assumes an array of shape (.. i, j, x, y) where x and y
+    are the size of the returned array (x, y).
+
+    Parameters
+    ----------
+    array : array_like
+        Input array of at least 3 dimensions.
+
+    Returns
+    -------
+    tuple
+        tuple of 2 arrays of the sum and number of points in the sum
+    """
+    X, Y = extimage.stackmean(array, 0)
+    return X, Y

csxtools/utils.py

@@ -4,7 +4,7 @@
 from pims import pipeline
 
 from .fastccd import correct_images
-from .image import rotate90
+from .image import rotate90, stackmean
 from .settings import detectors
 
 import logging
@@ -84,10 +84,8 @@ def get_fastccd_images(light_header, dark_headers=None,
                             ttime.time() - tt)
 
                 b = correct_images(b, gain=(1, 1, 1))
-                b = b.reshape((-1, b.shape[-2], b.shape[-1]))
-
                 tt = ttime.time()
-                b = np.nanmean(b, axis=0)
+                b = stackmean(b)
                 logger.info("Mean of image stack took %.3f seconds",
                             ttime.time() - tt)
 

src/image.c

@@ -87,3 +87,111 @@ void rotate90(data_t *in, data_t *out, int ndims, index_t *dims, int sense){
   free(map);
 }
 
+
+int stackmean(data_t *in, data_t *mout, long int *nout, int ndims, index_t *dims, int norm){
+  index_t nimages = dims[0];
+  index_t M = dims[ndims-1];
+  index_t N = dims[ndims-2];
+  index_t imsize = N*M;
+
+  int retval = 0;
+
+  long int **nvalues;
+  data_t **mean;
+
+  int num_threads;
+
+  int x;
+  for(x=1;x<(ndims-2);x++){
+    nimages = nimages * dims[x];
+  }
+
+  // Get the maximum threads 
+
+  int max_threads = omp_get_max_threads(); 
+  if(!(nvalues = malloc(sizeof(long int *) * max_threads))){
+    return 1;
+  }
+  if(!(mean = malloc(sizeof(data_t *) * max_threads))){
+    free(nvalues);
+    return 1;
+  }
+
+#pragma omp parallel shared(nvalues, mean, num_threads, imsize, in, retval)
+  {
+    // Allocate both a result array and an array for the number of values
+
+#pragma omp single
+    num_threads = omp_get_num_threads();
+
+    int thread_num = omp_get_thread_num();
+
+    data_t *_mean = calloc(imsize, sizeof(data_t));
+    long int *_nvalues = calloc(imsize, sizeof(long int));
+    mean[thread_num] = _mean;
+    nvalues[thread_num] = _nvalues;
+
+    // Test if we have the memory allocated
+    if((_mean != NULL) && (_nvalues != NULL)){
+
+      // Now do the actual mean calculation
+      int i;
+#pragma omp for private(i) 
+      for(i=0;i<nimages;i++){
+        int j;
+        for(j=0;j<imsize;j++){
+          data_t ival = in[(i * imsize + j)];
+          if(!isnan(ival)){
+            _mean[j] = _mean[j] + ival;
+            _nvalues[j]++;
+          }
+        }
+      }
+    }
+
+  } // pragma omp paralell
+
+  // Now calculate the mean 
+
+  int n,i;
+  for(n=1;n<num_threads;n++){
+    for(i=0;i<imsize;i++){
+      mean[0][i] += mean[n][i];
+      nvalues[0][i] += nvalues[n][i];
+    }
+  }
+
+  for(i=0;i<imsize;i++){
+    nout[i] = nvalues[0][i];
+    if(norm){
+      if(nvalues[0][i]){
+        mout[i] = mean[0][i] / nvalues[0][i];
+      } else {
+        mout[i] = 0.0;
+        nout[i] = 0;
+      }
+    } else {
+      mout[i] = mean[0][i];
+    }
+  }
+
+  // free up all memory
+
+  for(n=0;n<num_threads;n++){
+    if(mean[n]) {
+      free(mean[n]);
+    } else {
+      retval = 2; 
+    }
+    if(nvalues[n]) {
+      free(nvalues[n]);
+    } else {
+      retval = 2;
+    }
+  }
+
+  free(mean);
+  free(nvalues);
+
+  return retval;
+}

src/image.h

@@ -43,5 +43,6 @@ typedef float data_t;
 
 // Function prototypes
 void rotate90(data_t *in, data_t *out, int ndims, index_t *dims, int sense);
+int stackmean(data_t *in, data_t *mout, long int *nout, int ndims, index_t *dims, int norm);
 
 #endif

src/imagemodule.c

@@ -92,9 +92,62 @@ static PyObject* image_rotate90(PyObject *self, PyObject *args){
   return NULL;
 }
 
+static PyObject* image_stackmean(PyObject *self, PyObject *args){
+  PyObject *_input = NULL;
+  PyArrayObject *input = NULL;
+  PyArrayObject *nout = NULL;
+  PyArrayObject *mout = NULL;
+  npy_intp *dims;
+  npy_intp newdims[2];
+  int ndims;
+  int norm;
+
+  if(!PyArg_ParseTuple(args, "Oi", &_input, &norm)){
+    return NULL;
+  }
+
+  input = (PyArrayObject*)PyArray_FROMANY(_input, NPY_FLOAT, 3, 0,NPY_ARRAY_IN_ARRAY);
+  if(!input){
+    goto error;
+  }
+
+  ndims = PyArray_NDIM(input);
+  dims = PyArray_DIMS(input);
+
+  // Just make a new 2D array
+  newdims[0] = dims[ndims-2];
+  newdims[1] = dims[ndims-1];
+
+  mout = (PyArrayObject*)PyArray_SimpleNew(2, newdims, NPY_FLOAT);
+  if(!mout){
+    goto error;
+  }
+  nout = (PyArrayObject*)PyArray_SimpleNew(2, newdims, NPY_LONG);
+  if(!nout){
+    goto error;
+  }
+
+  if(stackmean((data_t*)PyArray_DATA(input), (data_t*)PyArray_DATA(mout),
+               (long int*)PyArray_DATA(nout), ndims, dims, norm)){
+    PyErr_SetString(PyExc_MemoryError, "Could not allocate memory");
+    goto error;
+  }
+
+  Py_XDECREF(input);
+  return Py_BuildValue("(NN)", mout, nout);
+
+error:
+  Py_XDECREF(input);
+  Py_XDECREF(nout);
+  Py_XDECREF(mout);
+  return NULL;
+}
+
 static PyMethodDef imageMethods[] = {
   { "rotate90", image_rotate90, METH_VARARGS,
     "Rotate stack of images 90 degrees (with sense)"},
+  { "stackmean", image_stackmean, METH_VARARGS,
+    "Calculate mean of an image stack"},
   {NULL, NULL, 0, NULL}
 };
 

tests/test_image.py

@@ -1,4 +1,4 @@
-from csxtools.image import rotate90
+from csxtools.image import rotate90, stackmean, stacksum
 import numpy as np
 from numpy.testing import assert_array_equal
 
@@ -12,3 +12,45 @@ def test_rotate90():
     y = rotate90(np.array([x, x, x, x]), 'cw')
     for i in y:
         assert_array_equal(i, np.rot90(x, -1))
+
+
+def test_stackmean():
+    x = np.ones((1, 100, 100), dtype=np.float32) * np.nan
+    m = stackmean(x)
+    assert_array_equal(m, np.zeros((100, 100), dtype=np.float32))
+
+    x = np.ones((1000, 100, 100), dtype=np.float32) * 52.0
+    m = stackmean(x)
+    assert_array_equal(m, np.ones((100, 100), dtype=np.float32) * 52.0)
+
+    # Now test with nans
+
+    x = np.ones((1000, 100, 100), dtype=np.float32) * 52.0
+    x[10] = np.nan
+    x[23] = np.nan
+    x[40] = np.nan
+    m = stackmean(x)
+    assert_array_equal(m, np.ones((100, 100), dtype=np.float32) * 52.0)
+
+
+def test_stacksum():
+    x = np.ones((1, 100, 100), dtype=np.float32) * np.nan
+    m, n = stacksum(x)
+    assert_array_equal(m, np.zeros((100, 100), dtype=np.float32))
+    assert_array_equal(n, np.zeros((100, 100), dtype=np.float32))
+
+    x = np.ones((1000, 100, 100), dtype=np.float32) * 52.0
+    m, n = stacksum(x)
+    assert_array_equal(m, np.ones((100, 100), dtype=np.float32) * 52.0 * 1000)
+    assert_array_equal(n, np.ones((100, 100), dtype=np.float32) * 1000.0)
+
+    # Now test with nans
+
+    x = np.ones((1000, 100, 100), dtype=np.float32) * 2
+    x[10] = np.nan
+    x[23] = np.nan
+    x[40] = np.nan
+    m, n = stacksum(x)
+    assert_array_equal(m, np.ones((100, 100), dtype=np.float32) *
+                       2 * (1000 - 3))
+    assert_array_equal(n, np.ones((100, 100), dtype=np.float32) * (1000 - 3))