Update tools/ai/demo_utils.py

tools/ai/demo_utils.py

@@ -1,111 +1,111 @@
-import cv2
-import random
-import numpy as np
-
-from PIL import Image
-
-def get_strided_size(orig_size, stride):
-    return ((orig_size[0]-1)//stride+1, (orig_size[1]-1)//stride+1)
-
-def get_strided_up_size(orig_size, stride):
-    strided_size = get_strided_size(orig_size, stride)
-    return strided_size[0]*stride, strided_size[1]*stride
-
-def imshow(image, delay=0, mode='RGB', title='show'):
-    if mode == 'RGB':
-        demo_image = image[..., ::-1]
-    else:
-        demo_image = image
-
-    cv2.imshow(title, demo_image)
-    if delay >= 0:
-        cv2.waitKey(delay)
-
-def transpose(image):
-    return image.transpose((1, 2, 0))
-
-def denormalize(image, mean=None, std=None, dtype=np.uint8, tp=True):
-    if tp:
-        image = transpose(image)
-        
-    if mean is not None:
-        image = (image * std) + mean
-    
-    if dtype == np.uint8:
-        image *= 255.
-        return image.astype(np.uint8)
-    else:
-        return image
-
-def colormap(cam, shape=None, mode=cv2.COLORMAP_JET):
-    if shape is not None:
-        h, w, c = shape
-        cam = cv2.resize(cam, (w, h))
-    cam = cv2.applyColorMap(cam, mode)
-    return cam
-
-def decode_from_colormap(data, colors):
-    ignore = (data == 255).astype(np.int32)
-
-    mask = 1 - ignore
-    data *= mask
-
-    h, w = data.shape
-    image = colors[data.reshape((h * w))].reshape((h, w, 3))
-
-    ignore = np.concatenate([ignore[..., np.newaxis], ignore[..., np.newaxis], ignore[..., np.newaxis]], axis=-1)
-    image[ignore.astype(np.bool)] = 255
-    return image
-
-def normalize(cam, epsilon=1e-5):
-    cam = np.maximum(cam, 0)
-    max_value = np.max(cam, axis=(0, 1), keepdims=True)
-    return np.maximum(cam - epsilon, 0) / (max_value + epsilon)
-
-def crf_inference(img, probs, t=10, scale_factor=1, labels=21):
-    import pydensecrf.densecrf as dcrf
-    from pydensecrf.utils import unary_from_softmax
-
-    h, w = img.shape[:2]
-    n_labels = labels
-
-    d = dcrf.DenseCRF2D(w, h, n_labels)
-
-    unary = unary_from_softmax(probs)
-    unary = np.ascontiguousarray(unary)
-
-    d.setUnaryEnergy(unary)
-    d.addPairwiseGaussian(sxy=3/scale_factor, compat=3)
-    d.addPairwiseBilateral(sxy=80/scale_factor, srgb=13, rgbim=np.copy(img), compat=10)
-    Q = d.inference(t)
-
-    return np.array(Q).reshape((n_labels, h, w))
-
-def crf_with_alpha(ori_image, cams, alpha):
-    # h, w, c -> c, h, w
-    # cams = cams.transpose((2, 0, 1))
-
-    bg_score = np.power(1 - np.max(cams, axis=0, keepdims=True), alpha)
-    bgcam_score = np.concatenate((bg_score, cams), axis=0)
-    
-    cams_with_crf = crf_inference(ori_image, bgcam_score, labels=bgcam_score.shape[0])
-    # return cams_with_crf.transpose((1, 2, 0))
-    return cams_with_crf
-
-def crf_inference_label(img, labels, t=10, n_labels=21, gt_prob=0.7):
-    import pydensecrf.densecrf as dcrf
-    from pydensecrf.utils import unary_from_labels
-
-    h, w = img.shape[:2]
-
-    d = dcrf.DenseCRF2D(w, h, n_labels)
-    
-    unary = unary_from_labels(labels, n_labels, gt_prob=gt_prob, zero_unsure=False)
-
-    d.setUnaryEnergy(unary)
-    d.addPairwiseGaussian(sxy=3, compat=3)
-    d.addPairwiseBilateral(sxy=50, srgb=5, rgbim=np.ascontiguousarray(np.copy(img)), compat=10)
-
-    q = d.inference(t)
-
+import cv2
+import random
+import numpy as np
+
+from PIL import Image
+
+def get_strided_size(orig_size, stride):
+    return ((orig_size[0]-1)//stride+1, (orig_size[1]-1)//stride+1)
+
+def get_strided_up_size(orig_size, stride):
+    strided_size = get_strided_size(orig_size, stride)
+    return strided_size[0]*stride, strided_size[1]*stride
+
+def imshow(image, delay=0, mode='RGB', title='show'):
+    if mode == 'RGB':
+        demo_image = image[..., ::-1]
+    else:
+        demo_image = image
+
+    cv2.imshow(title, demo_image)
+    if delay >= 0:
+        cv2.waitKey(delay)
+
+def transpose(image):
+    return image.transpose((1, 2, 0))
+
+def denormalize(image, mean=None, std=None, dtype=np.uint8, tp=True):
+    if tp:
+        image = transpose(image)
+        
+    if mean is not None:
+        image = (image * std) + mean
+    
+    if dtype == np.uint8:
+        image *= 255.
+        return image.astype(np.uint8)
+    else:
+        return image
+
+def colormap(cam, shape=None, mode=cv2.COLORMAP_JET):
+    if shape is not None:
+        h, w, c = shape
+        cam = cv2.resize(cam, (w, h))
+    cam = cv2.applyColorMap(cam, mode)
+    return cam
+
+def decode_from_colormap(data, colors):
+    ignore = (data == 255).astype(np.int32)
+
+    mask = 1 - ignore
+    data *= mask
+
+    h, w = data.shape
+    image = colors[data.reshape((h * w))].reshape((h, w, 3))
+
+    ignore = np.concatenate([ignore[..., np.newaxis], ignore[..., np.newaxis], ignore[..., np.newaxis]], axis=-1)
+    image[ignore.astype(bool)] = 255
+    return image
+
+def normalize(cam, epsilon=1e-5):
+    cam = np.maximum(cam, 0)
+    max_value = np.max(cam, axis=(0, 1), keepdims=True)
+    return np.maximum(cam - epsilon, 0) / (max_value + epsilon)
+
+def crf_inference(img, probs, t=10, scale_factor=1, labels=21):
+    import pydensecrf.densecrf as dcrf
+    from pydensecrf.utils import unary_from_softmax
+
+    h, w = img.shape[:2]
+    n_labels = labels
+
+    d = dcrf.DenseCRF2D(w, h, n_labels)
+
+    unary = unary_from_softmax(probs)
+    unary = np.ascontiguousarray(unary)
+
+    d.setUnaryEnergy(unary)
+    d.addPairwiseGaussian(sxy=3/scale_factor, compat=3)
+    d.addPairwiseBilateral(sxy=80/scale_factor, srgb=13, rgbim=np.copy(img), compat=10)
+    Q = d.inference(t)
+
+    return np.array(Q).reshape((n_labels, h, w))
+
+def crf_with_alpha(ori_image, cams, alpha):
+    # h, w, c -> c, h, w
+    # cams = cams.transpose((2, 0, 1))
+
+    bg_score = np.power(1 - np.max(cams, axis=0, keepdims=True), alpha)
+    bgcam_score = np.concatenate((bg_score, cams), axis=0)
+    
+    cams_with_crf = crf_inference(ori_image, bgcam_score, labels=bgcam_score.shape[0])
+    # return cams_with_crf.transpose((1, 2, 0))
+    return cams_with_crf
+
+def crf_inference_label(img, labels, t=10, n_labels=21, gt_prob=0.7):
+    import pydensecrf.densecrf as dcrf
+    from pydensecrf.utils import unary_from_labels
+
+    h, w = img.shape[:2]
+
+    d = dcrf.DenseCRF2D(w, h, n_labels)
+    
+    unary = unary_from_labels(labels, n_labels, gt_prob=gt_prob, zero_unsure=False)
+
+    d.setUnaryEnergy(unary)
+    d.addPairwiseGaussian(sxy=3, compat=3)
+    d.addPairwiseBilateral(sxy=50, srgb=5, rgbim=np.ascontiguousarray(np.copy(img)), compat=10)
+
+    q = d.inference(t)
+
     return np.argmax(np.array(q).reshape((n_labels, h, w)), axis=0)