Clearer and more unified unary utilities.

Author: lucasb-eyer

README.md

@@ -72,10 +72,10 @@ don't know how to without introducing an explicit dependency on numpy.
 There's two common ways of getting unary potentials:
 
 1. From a hard labeling generated by a human or some other processing.
-   This case is covered by `from pydensecrf.utils import compute_unary`.
+   This case is covered by `from pydensecrf.utils import unary_from_labels`.
 
 2. From a probability distribution computed by, e.g. the softmax output of a
-   deep network. For this, see `from pydensecrf.utils import softmax_to_unary`.
+   deep network. For this, see `from pydensecrf.utils import unary_from_softmax`.
 
 For usage of both of these, please refer to their docstrings or have a look at [the example](examples/inference.py).
 

examples/inference.py

@@ -16,7 +16,7 @@
     imwrite = imsave
     # TODO: Use scipy instead.
 
-from pydensecrf.utils import compute_unary, create_pairwise_bilateral, create_pairwise_gaussian
+from pydensecrf.utils import unary_from_labels, create_pairwise_bilateral, create_pairwise_gaussian
 
 if len(sys.argv) != 4:
     print("Usage: python {} IMAGE ANNO OUTPUT".format(sys.argv[0]))
@@ -52,8 +52,8 @@
 # Compute the number of classes in the label image.
 # We subtract one because the number shouldn't include the value 0 which stands
 # for "unknown" or "unsure".
-M = len(set(labels.flat)) - 1
-print(M, " labels and \"unknown\" 0: ", set(labels.flat))
+n_labels = len(set(labels.flat)) - 1
+print(n_labels, " labels and \"unknown\" 0: ", set(labels.flat))
 
 ###########################
 ### Setup the CRF model ###
@@ -64,10 +64,10 @@
     print("Using 2D specialized functions")
 
     # Example using the DenseCRF2D code
-    d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], M)
+    d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], n_labels)
 
     # get unary potentials (neg log probability)
-    U = compute_unary(labels, M, GT_PROB=0.7)
+    U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=True)
     d.setUnaryEnergy(U)
 
     # This adds the color-independent term, features are the locations only.
@@ -83,10 +83,10 @@
     print("Using generic 2D functions")
 
     # Example using the DenseCRF class and the util functions
-    d = dcrf.DenseCRF(img.shape[1] * img.shape[0], M)
+    d = dcrf.DenseCRF(img.shape[1] * img.shape[0], n_labels)
 
     # get unary potentials (neg log probability)
-    U = compute_unary(labels, M, GT_PROB=0.7)
+    U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=True)
     d.setUnaryEnergy(U)
 
     # This creates the color-independent features and then add them to the CRF

pydensecrf/utils.py

@@ -1,61 +1,90 @@
 import numpy as np
+from logging import warning
 
 
-def compute_unary(labels, M, GT_PROB=0.5):
+def unary_from_labels(labels, n_labels, gt_prob, zero_unsure=True):
     """
     Simple classifier that is 50% certain that the annotation is correct.
     (same as in the inference example).
 
 
     Parameters
     ----------
-    labels: nummpy.array
-        The label-map. The label value `0` is not a label, but the special
-        value indicating that the location has no label/information and thus
-        every label is equally likely.
-    M: int
-        The number of labels there are, not including the special `0` value.
-    GT_PROB: float
+    labels: numpy.array
+        The label-map, i.e. an array of your data's shape where each unique
+        value corresponds to a label.
+    n_labels: int
+        The total number of labels there are.
+        If `zero_unsure` is True (the default), this number should not include
+        `0` in counting the labels, since `0` is not a label!
+    gt_prob: float
         The certainty of the ground-truth (must be within (0,1)).
+    zero_unsure: bool
+        If `True`, treat the label value `0` as meaning "could be anything",
+        i.e. entries with this value will get uniform unary probability.
+        If `False`, do not treat the value `0` specially, but just as any
+        other class.
     """
-    assert 0 < GT_PROB < 1, "`GT_PROB must be in (0,1)."
+    assert 0 < gt_prob < 1, "`gt_prob must be in (0,1)."
 
     labels = labels.flatten()
 
-    u_energy = -np.log(1.0 / M)
-    n_energy = -np.log((1.0 - GT_PROB) / (M - 1))
-    p_energy = -np.log(GT_PROB)
+    n_energy = -np.log((1.0 - gt_prob) / (n_labels - 1))
+    p_energy = -np.log(gt_prob)
 
     # Note that the order of the following operations is important.
     # That's because the later ones overwrite part of the former ones, and only
     # after all of them is `U` correct!
-    U = np.full((M, len(labels)), n_energy, dtype='float32')
-    U[labels - 1, np.arange(U.shape[1])] = p_energy
-    U[:, labels == 0] = u_energy
+    U = np.full((n_labels, len(labels)), n_energy, dtype='float32')
+    U[labels - 1 if zero_unsure else labels, np.arange(U.shape[1])] = p_energy
+
+    # Overwrite 0-labels using uniform probability, i.e. "unsure".
+    if zero_unsure:
+        U[:, labels == 0] = -np.log(1.0 / n_labels)
+
     return U
 
 
-def softmax_to_unary(sm, GT_PROB=1):
-    """
-    Util function that converts softmax scores (classwise probabilities) to
-    unary potentials (the negative log likelihood per node).
+def compute_unary(labels, M, GT_PROB=0.5):
+    """Deprecated, use `unary_from_labels` instead."""
+    warning("pydensecrf.compute_unary is deprecated, use unary_from_labels instead.")
+    return unary_from_labels(labels, M, GT_PROB)
+
+
+def unary_from_softmax(sm, scale=None, clip=1e-5):
+    """Converts softmax class-probabilities to unary potentials (NLL per node).
 
     Parameters
     ----------
-    sm: nummpy.array
-        Softmax input. The first dimension is expected to be the classes,
-        all others will be flattend.
-    GT_PROB: float
-        The certainty of the softmax output (default is 1).
-
+    sm: numpy.array
+        Output of a softmax where the first dimension is the classes,
+        all others will be flattend. This means `sm.shape[0] == n_classes`.
+    scale: float
+        The certainty of the softmax output (default is None).
+        If not None, the softmax outputs are scaled to range from uniform
+        probability for 0 outputs to `scale` probability for 1 outputs.
+    clip: float
+        Minimum value to which probability should be clipped.
+        This is because the unary is the negative log of the probability, and
+        log(0) = inf, so we need to clip 0 probabilities to a positive value.
     """
     num_cls = sm.shape[0]
-    if GT_PROB < 1:
+    if scale is not None:
+        assert 0 < scale <= 1, "`scale` needs to be in (0,1]"
         uniform = np.ones(sm.shape) / num_cls
-        sm = GT_PROB * sm + (1 - GT_PROB) * uniform
+        sm = scale * sm + (1 - scale) * uniform
+    if clip is not None:
+        sm = np.clip(sm, clip, 1.0)
     return -np.log(sm).reshape([num_cls, -1]).astype(np.float32)
 
 
+def softmax_to_unary(sm, GT_PROB=1):
+    """Deprecated, use `unary_from_softmax` instead."""
+    warning("pydensecrf.softmax_to_unary is deprecated, use unary_from_softmax instead.")
+    scale = None if GT_PROB == 1 else GT_PROB
+    return unary_from_softmax(sm, scale, clip=None)
+
+
 def create_pairwise_gaussian(sdims, shape):
     """
     Util function that create pairwise gaussian potentials. This works for all