Diff of /utils/vis_utils.py [000000] .. [98e649]

Switch to side-by-side view

--- a
+++ b/utils/vis_utils.py
@@ -0,0 +1,175 @@
+import numpy as np
+import colorsys
+import random
+import cv2
+
+def get_n_hls_colors(num):
+    hls_colors = []
+    i = 0
+    step = 360.0 / num
+    while i < 360:
+        h = i
+        s = 90 + random.random() * 10
+        l = 50 + random.random() * 10
+        _hlsc = [h / 360.0, l / 100.0, s / 100.0]
+        hls_colors.append(_hlsc)
+        i += step
+
+    return hls_colors
+
+def ncolors(num):
+    rgb_colors = []
+    if num < 1:
+        return np.array(rgb_colors)
+    hls_colors = get_n_hls_colors(num)
+    for hlsc in hls_colors:
+        _r, _g, _b = colorsys.hls_to_rgb(hlsc[0], hlsc[1], hlsc[2])
+        rgb_colors.append([_r, _g, _b])
+        # r, g, b = [int(x * 255.0) for x in (_r, _g, _b)]
+        # rgb_colors.append([r, g, b])
+
+    return np.array(rgb_colors)
+
+def random_colors(N, bright=True):
+    """
+    Generate random colors.
+    To get visually distinct colors, generate them in HSV space then
+    convert to RGB.
+    """
+    brightness = 1.0 if bright else 0.7
+    hsv = [(i / N, 1, brightness) for i in range(N)]
+    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
+    # random.shuffle(colors)
+    return colors
+
+def mask2png(mask, file_name=None, suffix="png"):
+    """ mask: (w, h)
+        img_rgb: (w, h, rgb)
+    """
+    nums = np.unique(mask)[1:]
+    if len(nums) < 1:
+        colors = np.array([[0,0,0]])
+    else:
+        # colors = ncolors(len(nums))
+        colors = (np.array(random_colors(len(nums))) * 255).astype(int)
+        colors = np.insert(colors, 0, [0,0,0], 0)
+    
+    # 保证mask中的值为1-N连续
+    mask_ordered = np.zeros_like(mask)
+    for cnt, l in enumerate(nums, 1):
+        mask_ordered[mask==l] = cnt
+
+    im_rgb = colors[mask_ordered.astype(int)]
+    if file_name is not None:
+        cv2.imwrite(file_name+"."+suffix, im_rgb[:, :, ::-1])
+    return im_rgb
+
+def apply_mask(image, mask, color, alpha=0.5, scale=1):
+    """Apply the given mask to the image.
+    """
+    for c in range(3):
+        image[:, :, c] = np.where(mask == 1,
+                                  image[:, :, c] *
+                                  (1 - alpha) + alpha * color[c] * scale,
+                                  image[:, :, c])
+    return image
+
+def img_mask_png(image, mask, file_name=None, alpha=0.5, suffix="png"):
+    """ mask: (h, w)
+        image: (h, w, rgb)
+    """
+    nums = np.unique(mask)[1:]
+    if len(nums) < 1:
+        colors = np.array([[0,0,0]])
+    else:
+        colors = ncolors(len(nums))
+        colors = np.insert(colors, 0, [0,0,0], 0)
+    
+    # 保证mask中的值为1-N连续
+    mask_ordered = np.zeros_like(mask)
+    for cnt, l in enumerate(nums, 1):
+        mask_ordered[mask==l] = cnt
+    
+    # mask_rgb = colors[mask_ordered.astype(int)]
+    mix_im = image.copy()
+    for i in np.unique(mask_ordered)[1:]:
+        mix_im = apply_mask(mix_im, mask_ordered==i, colors[int(i)], alpha=alpha, scale=255)
+
+    if file_name is not None:
+        cv2.imwrite(file_name+"."+suffix, mix_im[:, :, ::-1])
+    return mix_im
+
+def _find_contour(mask):
+    # _, contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 顶点
+    _, contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
+
+    cont = np.zeros_like(mask)
+    for contour in contours:
+        cont[contour[:,:,1], contour[:,:,0]] = 1
+    return cont
+
+def masks_to_contours(masks):
+    # 包含多个区域
+    nums = np.unique(masks)[1:]
+    cont_mask = np.zeros_like(masks)
+    for i in nums:
+        cont_mask += _find_contour(masks==i)
+    return (cont_mask>0).astype(int)
+    
+def batchToColorImg(batch, minv=None, maxv=None, scale=255.):
+    """ batch: (N, H, W, C)
+    """
+    if batch.ndim == 3:
+        N, H, W = batch.shape
+    elif batch.ndim == 4:
+        N, H, W, _ = batch.shape
+    colorImg = np.zeros(shape=(N, H, W, 3))
+    for i in range(N):
+        if minv is None:
+            a = (batch[i] - batch[i].min()) / (batch[i].max() - batch[i].min()) * 255
+        else:
+            a = (batch[i] - minv) / (maxv - minv) * scale
+        a = cv2.applyColorMap(a.astype(np.uint8), cv2.COLORMAP_JET)
+        colorImg[i, ...] = a[..., ::-1] / 255.
+    return colorImg
+
+if __name__ == "__main__":
+    a = np.zeros((100, 100))
+    a[0:5, 3:8] = 1
+    a[75:85, 85:95] = 2
+
+    # colors = ncolors(2)[::-1]
+    colors = np.array(random_colors(2)) * 255
+    colors = np.insert(colors, 0, [0, 0, 0], 0)
+
+    b = colors[a.astype(int)].astype(np.uint8)
+    import cv2, skimage
+    # skimage.io.imsave("test_io.png", b)
+    # # cv2.imwrite("test.jpg", b[:, :, ::-1])
+    # print()
+
+    # mask2png(a, "test")
+
+    ################################################
+    # img_mask_png(b, a, "test")
+
+    #############################################
+    # cont_mask = find_contours(a)
+    # print()
+    # # skimage.io.imsave("test_cont.png", cont_mask) 
+    # b[cont_mask>0, :] = [255, 255, 255]
+    # skimage.io.imsave("test_cont.png", b)
+
+    gt0 = skimage.io.imread("gt0.png", as_gray=False)
+    print()
+    gt0[gt0==54] = 0
+    # cont_mask = find_contours(gt0==237)  # Array([ 18,  54,  73, 237], dtype=uint8)
+    # cont_mask += find_contours(gt0==18)
+    # cont_mask += find_contours(gt0==73)
+    cont_mask = masks_to_contours(gt0)
+    
+    colors = np.array(random_colors(1)) * 255
+    colors = np.insert(colors, 0, [0, 0, 0], 0)
+    cont_mask = colors[cont_mask.astype(int)].astype(np.uint8)
+
+    skimage.io.imsave("test_cont.png", cont_mask)