better box-text visualization

tensorpack/utils/palette.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# File: palette.py
+
+import numpy as np
+
+__all__ = ['PALETTE_RGB']
+
+# copied from https://stackoverflow.com/questions/2328339/how-to-generate-n-different-colors-for-any-natural-number-n
+PALETTE_HEX = [
+    "#000000", "#FFFF00", "#1CE6FF", "#FF34FF", "#FF4A46", "#008941", "#006FA6", "#A30059",
+    "#FFDBE5", "#7A4900", "#0000A6", "#63FFAC", "#B79762", "#004D43", "#8FB0FF", "#997D87",
+    "#5A0007", "#809693", "#FEFFE6", "#1B4400", "#4FC601", "#3B5DFF", "#4A3B53", "#FF2F80",
+    "#61615A", "#BA0900", "#6B7900", "#00C2A0", "#FFAA92", "#FF90C9", "#B903AA", "#D16100",
+    "#DDEFFF", "#000035", "#7B4F4B", "#A1C299", "#300018", "#0AA6D8", "#013349", "#00846F",
+    "#372101", "#FFB500", "#C2FFED", "#A079BF", "#CC0744", "#C0B9B2", "#C2FF99", "#001E09",
+    "#00489C", "#6F0062", "#0CBD66", "#EEC3FF", "#456D75", "#B77B68", "#7A87A1", "#788D66",
+    "#885578", "#FAD09F", "#FF8A9A", "#D157A0", "#BEC459", "#456648", "#0086ED", "#886F4C",
+    "#34362D", "#B4A8BD", "#00A6AA", "#452C2C", "#636375", "#A3C8C9", "#FF913F", "#938A81",
+    "#575329", "#00FECF", "#B05B6F", "#8CD0FF", "#3B9700", "#04F757", "#C8A1A1", "#1E6E00",
+    "#7900D7", "#A77500", "#6367A9", "#A05837", "#6B002C", "#772600", "#D790FF", "#9B9700",
+    "#549E79", "#FFF69F", "#201625", "#72418F", "#BC23FF", "#99ADC0", "#3A2465", "#922329",
+    "#5B4534", "#FDE8DC", "#404E55", "#0089A3", "#CB7E98", "#A4E804", "#324E72", "#6A3A4C",
+    "#83AB58", "#001C1E", "#D1F7CE", "#004B28", "#C8D0F6", "#A3A489", "#806C66", "#222800",
+    "#BF5650", "#E83000", "#66796D", "#DA007C", "#FF1A59", "#8ADBB4", "#1E0200", "#5B4E51",
+    "#C895C5", "#320033", "#FF6832", "#66E1D3", "#CFCDAC", "#D0AC94",
+    "#7ED379", "#012C58"]
+
+
+def _parse_hex_color(s):
+    r = int(s[1:3], 16)
+    g = int(s[3:5], 16)
+    b = int(s[5:7], 16)
+    return (r, g, b)
+
+
+PALETTE_RGB = np.asarray(
+    list(map(_parse_hex_color, PALETTE_HEX)),
+    dtype='int32')

tensorpack/utils/rect.py

@@ -82,7 +82,7 @@ class IntBox(BoxBase):
         self.y2 = np.clip(self.y2, 0, shape[0] - 1)
 
     def roi(self, img):
-        assert self.validate(img.shape[:2]), "{} vs {}".format(self, img.shape[:2])
+        assert self.is_valid_box(img.shape[:2]), "{} vs {}".format(self, img.shape[:2])
         return img[self.y1:self.y2 + 1, self.x1:self.x2 + 1]
 
     # def expand(self, frac):

tensorpack/utils/viz.py

@@ -9,7 +9,8 @@ import sys
 import io
 from .fs import mkdir_p
 from .argtools import shape2d
-from .rect import BoxBase
+from .rect import BoxBase, IntBox
+from .palette import PALETTE_RGB
 
 try:
     import cv2
@@ -352,7 +353,7 @@ def intensity_to_rgb(intensity, cmap='cubehelix', normalize=False):
     return intensity.astype('float32') * 255.0
 
 
-def draw_boxes(im, boxes, labels=None, color=(218, 218, 218)):
+def draw_boxes(im, boxes, labels=None):
     """
     Args:
         im (np.ndarray): will not be modified
@@ -362,6 +363,8 @@ def draw_boxes(im, boxes, labels=None, color=(218, 218, 218)):
     Returns:
         np.ndarray
     """
+    FONT = cv2.FONT_HERSHEY_SIMPLEX
+    FONT_SCALE = 0.4
     if isinstance(boxes, list):
         arr = np.zeros((len(boxes), 4), dtype='int32')
         for idx, b in enumerate(boxes):
@@ -376,14 +379,38 @@ def draw_boxes(im, boxes, labels=None, color=(218, 218, 218)):
     sorted_inds = np.argsort(-areas)
 
     im = im.copy()
+    COLOR = (218, 218, 218)
+    COLOR_DIFF_WEIGHT = np.asarray((2, 4, 3), dtype='int32')    # https://www.wikiwand.com/en/Color_difference
+    COLOR_CANDIDATES = PALETTE_RGB[[0, 1, 2, 3, 18, 113], :]
+    if im.ndim == 2 or (im.ndim == 3 and im.shape[2] == 1):
+        im = cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
     for i in sorted_inds:
         box = boxes[i, :]
-        cv2.rectangle(im, (box[0], box[1]), (box[2], box[3]), color, thickness=1)
 
+        best_color = COLOR
         if labels is not None:
             label = labels[i]
-            cv2.putText(im, label, (box[0], box[1] - 3),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, lineType=cv2.LINE_AA)
+
+            # find the best placement for the text
+            ((linew, lineh), _) = cv2.getTextSize(label, FONT, FONT_SCALE, 1)
+            bottom_left = [box[0] + 1, box[1] - 0.3 * lineh]
+            top_left = [box[0] + 1, box[1] - 1.3 * lineh]
+            if top_left[1] < 0:     # out of image
+                top_left[1] = box[3] - 1.3 * lineh
+                bottom_left[1] = box[3] - 0.3 * lineh
+            textbox = IntBox(int(top_left[0]), int(top_left[1]),
+                             int(top_left[0] + linew), int(top_left[1] + lineh))
+            textbox.clip_by_shape(im.shape[:2])
+            # find the best color
+            mean_color = textbox.roi(im).mean(axis=(0, 1))
+            best_color_ind = (np.square(COLOR_CANDIDATES - mean_color) *
+                              COLOR_DIFF_WEIGHT).sum(axis=1).argmax()
+            best_color = COLOR_CANDIDATES[best_color_ind]
+
+            cv2.putText(im, label, (textbox.x1, textbox.y2),
+                        FONT, FONT_SCALE, color=best_color.tolist(), lineType=cv2.LINE_AA)
+        cv2.rectangle(im, (box[0], box[1]), (box[2], box[3]),
+                      color=best_color.tolist(), thickness=1)
     return im