Merge pull request lucasb-eyer#21 from lucasb-eyer/examples-fix

Examples fixFixes to the examples, it's only one now, but it's better!

Author: lucasb-eyer

README.md

@@ -77,7 +77,7 @@ There's two common ways of getting unary potentials:
 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`.
 
-For usage of both of these, please refer to their docstrings or have a look at [the example](examples/utils_example.py).
+For usage of both of these, please refer to their docstrings or have a look at [the example](examples/inference.py).
 
 Pairwise potentials
 -------------------

examples/inference.py

@@ -1,51 +1,124 @@
-#!/usr/bin/python
-
 """
-Adapted from the original C++ example: densecrf/examples/dense_inference.cpp
-http://www.philkr.net/home/densecrf Version 2.2
+Adapted from the inference.py to demonstate the usage of the util functions.
 """
 
+import sys
 import numpy as np
-import cv2
 import pydensecrf.densecrf as dcrf
-from skimage.segmentation import relabel_sequential
-import sys
+
+# Get im{read,write} from somewhere.
+try:
+    from cv2 import imread, imwrite
+except ImportError:
+    # Note that, sadly, skimage unconditionally import scipy and matplotlib,
+    # so you'll need them if you don't have OpenCV. But you probably have them.
+    from skimage.io import imread, imsave
+    imwrite = imsave
+    # TODO: Use scipy instead.
+
+from pydensecrf.utils import compute_unary, create_pairwise_bilateral, create_pairwise_gaussian
 
 if len(sys.argv) != 4:
     print("Usage: python {} IMAGE ANNO OUTPUT".format(sys.argv[0]))
     print("")
     print("IMAGE and ANNO are inputs and OUTPUT is where the result should be written.")
     sys.exit(1)
 
-img = cv2.imread(sys.argv[1], 1)
-labels = relabel_sequential(cv2.imread(sys.argv[2], 0))[0].flatten()
-output = sys.argv[3]
+fn_im = sys.argv[1]
+fn_anno = sys.argv[2]
+fn_output = sys.argv[3]
+
+##################################
+### Read images and annotation ###
+##################################
+img = imread(fn_im)
+
+# Convert the annotation's RGB color to a single 32-bit integer color 0xBBGGRR
+anno_rgb = imread(fn_anno).astype(np.uint32)
+anno_lbl = anno_rgb[:,:,0] + (anno_rgb[:,:,1] << 8) + (anno_rgb[:,:,2] << 16)
+
+# Convert the 32bit integer color to 1, 2, ... labels.
+# Note that all-black, i.e. the value 0 for background will stay 0.
+colors, labels = np.unique(anno_lbl, return_inverse=True)
+
+# And create a mapping back from the labels to 32bit integer colors.
+# But remove the all-0 black, that won't exist in the MAP!
+colors = colors[1:]
+colorize = np.empty((len(colors), 3), np.uint8)
+colorize[:,0] = (colors & 0x0000FF)
+colorize[:,1] = (colors & 0x00FF00) >> 8
+colorize[:,2] = (colors & 0xFF0000) >> 16
+
+# 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))
+
+###########################
+### Setup the CRF model ###
+###########################
+use_2d = False
+# use_2d = True
+if use_2d:
+    print("Using 2D specialized functions")
+
+    # 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, GT_PROB=0.7)
+    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:
+    print("Using generic 2D functions")
+
+    # Example using the DenseCRF class and the util functions
+    d = dcrf.DenseCRF(img.shape[1] * img.shape[0], M)
+
+    # get unary potentials (neg log probability)
+    U = compute_unary(labels, M, GT_PROB=0.7)
+    d.setUnaryEnergy(U)
 
-M = labels.max() + 1  # number of labels
+    # 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)
 
-# Setup the CRF model
-d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], M)
+    # 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)
 
-# Certainty that the ground truth is correct
-GT_PROB = 0.5
 
-# Simple classifier that is 50% certain that the annotation is correct
-u_energy = -np.log(1.0 / M)
-n_energy = -np.log((1.0 - GT_PROB) / (M - 1))
-p_energy = -np.log(GT_PROB)
+####################################
+### Do inference and compute MAP ###
+####################################
 
-U = np.zeros((M, img.shape[0] * img.shape[1]), dtype='float32')
-U[:, labels > 0] = n_energy
-U[labels, np.arange(U.shape[1])] = p_energy
-U[:, labels == 0] = u_energy
-d.setUnaryEnergy(U)
+# Run five inference steps.
+Q = d.inference(5)
 
-d.addPairwiseGaussian(sxy=3, compat=3)
-d.addPairwiseBilateral(sxy=80, srgb=13, rgbim=img, compat=10)
+# Find out the most probable class for each pixel.
+MAP = np.argmax(Q, axis=0)
 
-# Do the inference
-res = np.argmax(d.inference(5), axis=0).astype('float32')
+# Convert the MAP (labels) back to the corresponding colors and save the image.
+MAP = colorize[MAP,:]
+imsave(fn_output, MAP.reshape(img.shape))
 
-res *= 255 / res.max()
-res = res.reshape(img.shape[:2])
-cv2.imwrite(output, res.astype('uint8'))
+# Just randomly manually run inference iterations
+Q, tmp1, tmp2 = d.startInference()
+for i in range(5):
+    print("KL-divergence at {}: {}".format(i, d.klDivergence(Q)))
+    d.stepInference(Q, tmp1, tmp2)

examples/out1.png

@@ -26,8 +26,10 @@ def compute_unary(labels, M, GT_PROB=0.5):
     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
+    # 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
     return U