Merge pull request #6 from Scyfer/util_functions

Util functions

util_inference_example.py

+"""
+Usage: python util_inference_example.py image annotations
+
+Adapted from the dense_inference.py to demonstate the usage of the util
+functions.
+"""
+
+import sys
+import numpy as np
+import cv2
+import densecrf as dcrf
+import matplotlib.pylab as plt
+from skimage.segmentation import relabel_sequential
+
+from utils import compute_unary, create_pairwise_bilateral, \
+    create_pairwise_gaussian
+
+fn_im = sys.argv[1]
+fn_anno = sys.argv[2]
+
+##################################
+### Read images and annotation ###
+##################################
+img = cv2.imread(fn_im)
+labels = relabel_sequential(cv2.imread(fn_anno, 0))[0].flatten()
+M = 21 # 21 Classes to match the C++ example
+
+###########################
+### Setup the CRF model ###
+###########################
+use_2d = False
+if use_2d:
+    # Example using the DenseCRF2D code
+    d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], M)
+
+    # get unary potentials (neg log probability)
+    U = compute_unary(labels, M)
+    d.setUnaryEnergy(U)
+
+    # This adds the color-independent term, features are the locations only.
+    d.addPairwiseGaussian(sxy=(3, 3), compat=3, kernel=dcrf.DIAG_KERNEL,
+                          normalization=dcrf.NORMALIZE_SYMMETRIC)
+
+    # This adds the color-dependent term, i.e. features are (x,y,r,g,b).
+    d.addPairwiseBilateral(sxy=(80, 80), srgb=(13, 13, 13), rgbim=img,
+                           compat=10,
+                           kernel=dcrf.DIAG_KERNEL,
+                           normalization=dcrf.NORMALIZE_SYMMETRIC)
+else:
+    # Example using the DenseCRF class and the util functions
+    d = dcrf.DenseCRF(img.shape[0] * img.shape[1], M)
+
+    # get unary potentials (neg log probability)
+    U = compute_unary(labels, M)
+    d.setUnaryEnergy(U)
+
+    # This creates the color-independent features and then add them to the CRF
+    feats = create_pairwise_gaussian(sdims=(3, 3), shape=img.shape[:2])
+    d.addPairwiseEnergy(feats, compat=3,
+                        kernel=dcrf.DIAG_KERNEL,
+                        normalization=dcrf.NORMALIZE_SYMMETRIC)
+
+    # This creates the color-dependent features and then add them to the CRF
+    feats = create_pairwise_bilateral(sdims=(80, 80), schan=(13, 13, 13),
+                                      img=img, chdim=2)
+    d.addPairwiseEnergy(feats, compat=10,
+                        kernel=dcrf.DIAG_KERNEL,
+                        normalization=dcrf.NORMALIZE_SYMMETRIC)
+
+
+####################################
+### Do inference and compute map ###
+####################################
+Q = d.inference(5)
+map = np.argmax(Q, axis=0).reshape(img.shape[:2])
+
+res = map.astype('float32') * 255 / map.max()
+plt.imshow(res)
+plt.show()
+
+
+# Manually inference
+Q, tmp1, tmp2 = d.startInference()
+for i in range(5):
+    print("KL-divergence at {}: {}".format(i, d.klDivergence(Q)))
+    d.stepInference(Q, tmp1, tmp2)
+

utils.py

+import numpy as np
+
+
+def compute_unary(labels, M, GT_PROB=0.5):
+    """
+    Simple classifier that is 50% certain that the annotation is correct
+    (same as in the inference example).
+    """
+    u_energy = -np.log(1.0 / M)
+    n_energy = -np.log((1.0 - GT_PROB) / (M - 1))
+    p_energy = -np.log(GT_PROB)
+
+    U = np.zeros((M, len(labels)), dtype='float32')
+    U[:, labels > 0] = n_energy
+    U[labels, np.arange(U.shape[1])] = p_energy
+    U[:, labels == 0] = u_energy
+    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).
+
+    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).
+
+    """
+    num_cls = sm.shape[0]
+    if GT_PROB < 1:
+        uniform = np.ones(sm.shape) / num_cls
+        sm = GT_PROB * sm + (1 - GT_PROB) * uniform
+    return -np.log(sm).reshape([num_cls, -1]).astype(np.float32)
+
+
+def create_pairwise_gaussian(sdims, shape):
+    """
+    Util function that create pairwise gaussian potentials. This works for all
+    image dimensions. For the 2D case does the same as
+    `DenseCRF2D.addPairwiseGaussian`.
+
+    Parameters
+    ----------
+    sdims: list or tuple
+        The scaling factors per dimension. This is referred to `sxy` in
+        `DenseCRF2D.addPairwiseGaussian`.
+    shape: list or tuple
+        The shape the CRF has.
+
+    """
+    # create mesh
+    cord_range = [range(s) for s in shape]
+    mesh = np.array(np.meshgrid(*cord_range, indexing='ij'), dtype=np.float32)
+
+    # scale mesh accordingly
+    for i, s in enumerate(sdims):
+        mesh[i] /= s
+    return mesh.reshape([len(sdims), -1])
+
+
+def create_pairwise_bilateral(sdims, schan, img, chdim=-1):
+    """
+    Util function that create pairwise bilateral potentials. This works for
+    all image dimensions. For the 2D case does the same as
+    `DenseCRF2D.addPairwiseBilateral`.
+
+    Parameters
+    ----------
+    sdims: list or tuple
+        The scaling factors per dimension. This is referred to `sxy` in
+        `DenseCRF2D.addPairwiseBilateral`.
+    schan: list or tuple
+        The scaling factors per channel in the image. This is referred to
+        `srgb` in `DenseCRF2D.addPairwiseBilateral`.
+    img: numpy.array
+        The input image.
+    chdim: int, optional
+        This specifies where the channel dimension is in the image. For
+        example `chdim=2` for a RGB image of size (240, 300, 3). If the
+        image has no channel dimension (e.g. it has only one channel) use
+        `chdim=-1`.
+
+    """
+    # Put channel dim in right position
+    if chdim == -1:
+        # We don't have a channel, add a new axis
+        im_feat = img[np.newaxis].astype(np.float32)
+    else:
+        # Put the channel dim as axis 0, all others stay relatively the same
+        im_feat = np.rollaxis(img, chdim).astype(np.float32)
+
+    # scale image features per channel
+    for i, s in enumerate(schan):
+        im_feat[i] /= s
+
+    # create a mesh
+    cord_range = [range(s) for s in im_feat.shape[1:]]
+    mesh = np.array(np.meshgrid(*cord_range, indexing='ij'), dtype=np.float32)
+
+    # scale mesh accordingly
+    for i, s in enumerate(sdims):
+        mesh[i] /= s
+
+    feats = np.concatenate([mesh, im_feat])
+    return feats.reshape([feats.shape[0], -1])
+
+
+def _create_pairwise_gaussian_2d(sx, sy, shape):
+    """
+    A simple reference implementation for the 2D case. The ND implementation
+    is faster.
+    """
+    feat_size = 2
+    feats = np.zeros((feat_size, shape[0], shape[1]), dtype=np.float32)
+    for i in range(shape[0]):
+        for j in range(shape[1]):
+            feats[0, i, j] = i / sx
+            feats[1, i, j] = j / sy
+    return feats.reshape([feat_size, -1])
+
+
+def _create_pairwise_bilateral_2d(sx, sy, sr, sg, sb, img):
+    """
+    A simple reference implementation for the 2D case. The ND implementation
+    is faster.
+    """
+    feat_size = 5
+    feats = np.zeros((feat_size, img.shape[0], img.shape[1]), dtype=np.float32)
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            feats[0, i, j] = i / sx
+            feats[1, i, j] = j / sy
+            feats[2, i, j] = img[i, j, 0] / sr
+            feats[3, i, j] = img[i, j, 1] / sg
+            feats[4, i, j] = img[i, j, 2] / sb
+    return feats.reshape([feat_size, -1])
+