Protect from crash as in issue #29

pydensecrf/densecrf.pxd

@@ -86,6 +86,8 @@ cdef extern from "densecrf/include/densecrf.h":
 
 cdef class DenseCRF:
     cdef c_DenseCRF *_this
+    cdef int _nlabel
+    cdef int _nvar
 
 
 cdef class DenseCRF2D(DenseCRF):

pydensecrf/densecrf.pyx

@@ -59,6 +59,9 @@ cdef class DenseCRF:
         else:
             self._this = NULL
 
+        self._nvar = nvar
+        self._nlabel = nlabels
+
     def __dealloc__(self):
         # Because destructors are virtual, this is enough to delete any object
         # of child classes too.
@@ -66,12 +69,20 @@ cdef class DenseCRF:
             del self._this
 
     def addPairwiseEnergy(self, float[:,::1] features not None, compat, KernelType kernel=DIAG_KERNEL, NormalizationType normalization=NORMALIZE_SYMMETRIC):
+        if features.shape[0] != self._nlabel or features.shape[1] != self._nvar:
+            raise ValueError("Bad shape for pairwise energy (Need {}, got {})".format((self._nlabel, self._nvar), (features.shape[0], features.shape[1])))
+
         self._this.addPairwiseEnergy(eigen.c_matrixXf(features), _labelcomp(compat), kernel, normalization)
 
     def setUnary(self, Unary u):
         self._this.setUnaryEnergy(u.move())
 
     def setUnaryEnergy(self, float[:,::1] u not None, float[:,::1] f = None):
+        if u.shape[0] != self._nlabel or u.shape[1] != self._nvar:
+            raise ValueError("Bad shape for unary energy (Need {}, got {})".format((self._nlabel, self._nvar), (u.shape[0], u.shape[1])))
+        # TODO: I don't remember the exact shape `f` should have, so I'm not putting an assertion here.
+        #       If you get hit by a wrong shape of `f`, please open an issue with the necessary info!
+
         if f is None:
             self._this.setUnaryEnergy(eigen.c_matrixXf(u))
         else:

tests/issue29.py

+# probs of shape 3d image per class: Nb_classes x Height x Width x Depth
+# assume the image has shape (69, 51, 72)
+import numpy as np
+import pydensecrf.densecrf as dcrf
+from pydensecrf.utils import unary_from_softmax, create_pairwise_gaussian
+
+###
+
+#shape = (69, 51, 72)
+#probs = np.random.randn(5, 69, 51).astype(np.float32)
+#probs /= probs.sum(axis=0, keepdims=True)
+#
+#d = dcrf.DenseCRF(np.prod(shape), probs.shape[0])
+#U = unary_from_softmax(probs)
+#print(U.shape)
+#d.setUnaryEnergy(U)
+#feats = create_pairwise_gaussian(sdims=(1.0, 1.0, 1.0), shape=shape)
+#d.addPairwiseEnergy(feats, compat=3, kernel=dcrf.FULL_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC)
+#Q = d.inference(5)
+#new_image = np.argmax(Q, axis=0).reshape((shape[0], shape[1],shape[2]))
+
+
+###
+
+SHAPE, NLABELS = (69, 51, 72), 5
+probs = np.random.randn(NLABELS, 68, 50).astype(np.float32)  # WRONG shape here
+probs /= probs.sum(axis=0, keepdims=True)
+
+d = dcrf.DenseCRF(np.prod(SHAPE), NLABELS)
+
+d.setUnaryEnergy(unary_from_softmax(probs))  # THIS SHOULD THROW and not crash later
+feats = create_pairwise_gaussian(sdims=(1.0, 1.0, 1.0), shape=SHAPE)
+d.addPairwiseEnergy(feats, compat=3, kernel=dcrf.FULL_KERNEL, normalization=dcrf.NORMALIZE_SYMMETRIC)
+
+Q = d.inference(5)
+new_image = np.argmax(Q, axis=0).reshape(SHAPE)