add draw_text and draw_bounding_boxes

README.md

@@ -14,4 +14,4 @@ pip install -e .
 - Python 3.5+
 - NumPy 1.11+
 - OpenCV 2.0+
-
+- Pillow

khandy/__init__.py

@@ -5,6 +5,7 @@ from .utils_fs import *
 from .utils_hash import *
 from .utils_list import *
 from .utils_numpy import *
+from .utils_draw import *
 from .utils_others import *
 
 from .boxes import *

khandy/utils_draw.py

@@ -0,0 +1,137 @@
+import numpy as np
+import PIL
+from PIL import Image
+from PIL import ImageDraw
+from PIL import ImageFont
+from PIL import ImageColor
+
+
+def _is_legal_color(color):
+    if color is None:
+        return True
+    if isinstance(color, str):
+        return True
+    return isinstance(color, (tuple, list)) and len(color) == 3
+        
+
+def _normalize_color(color, pil_mode, swap_rgb=False):
+    if color is None:
+        return color
+    if isinstance(color, str):
+        color = ImageColor.getrgb(color)
+    gray = color[0]
+    if swap_rgb:
+        color = (color[2], color[1], color[0])
+    if pil_mode == 'L':
+        color = gray
+    return color
+    
+    
+def draw_text(image, text, position, color=(255,0,0), font=None, font_size=15):
+    """Draws text on given image.
+    
+    Args:
+        image (ndarray).
+        text (str): text to be drawn.
+        position (Tuple[int, int]): position where to be drawn.
+        color (List[Union[str, Tuple[int, int, int]]]): text color.
+        font (str):  A filename or file-like object containing a TrueType font. If the file is not found in this 
+            filename, the loader may also search in other directories, such as the `fonts/` directory on Windows
+            or `/Library/Fonts/`, `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
+        font_size (int): The requested font size in points.
+
+    References:
+        torchvision.utils.draw_bounding_boxes
+    """
+    if isinstance(image, np.ndarray):
+        pil_image = Image.fromarray(image)
+    elif isinstance(image, PIL.Image.Image):
+        pil_image = image
+    else:
+        raise ValueError('Unsupported image type!')
+    assert pil_image.mode in ['L', 'RGB', 'RGBA']
+    
+    assert _is_legal_color(color)
+    color = _normalize_color(color, pil_image.mode, isinstance(image, np.ndarray))
+    
+    if font is None:
+        font_object = ImageFont.load_default()
+    else:
+        font_object = ImageFont.truetype(font, size=font_size)
+    
+    draw = ImageDraw.Draw(pil_image)
+    draw.text((position[0], position[1]), text, 
+              fill=color, font=font_object)
+
+    if isinstance(image, np.ndarray):
+        return np.asarray(pil_image)
+    return pil_image
+
+
+def draw_bounding_boxes(image, boxes, labels=None, colors=None,
+                        fill=False, width=1, font=None, font_size=15):
+    """Draws bounding boxes on given image.
+
+    Args:
+        image (ndarray).
+        boxes (ndarray): ndarray of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format.
+        labels (List[str]): List containing the labels of bounding boxes.
+        colors (List[Union[str, Tuple[int, int, int]]]): List containing the colors of bounding boxes or labels.
+        fill (bool): If `True` fills the bounding box with specified color.
+        width (int): Width of bounding box.
+        font (str):  A filename or file-like object containing a TrueType font. If the file is not found in this 
+            filename, the loader may also search in other directories, such as the `fonts/` directory on Windows
+            or `/Library/Fonts/`, `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
+        font_size (int): The requested font size in points.
+
+    References:
+        torchvision.utils.draw_bounding_boxes
+    """
+    if isinstance(image, np.ndarray):
+        pil_image = Image.fromarray(image)
+    elif isinstance(image, PIL.Image.Image):
+        pil_image = image
+    else:
+        raise ValueError('Unsupported image type!')
+    pil_image = pil_image.convert('RGB')
+    
+    if font is None:
+        font_object = ImageFont.load_default()
+    else:
+        font_object = ImageFont.truetype(font, size=font_size)
+
+    if fill:
+        draw = ImageDraw.Draw(pil_image, "RGBA")
+    else:
+        draw = ImageDraw.Draw(pil_image)
+
+    for i, bbox in enumerate(boxes):
+        if colors is None:
+            color = None
+        else:
+            color = colors[i]
+            
+        assert _is_legal_color(color)
+        color = _normalize_color(color, pil_image.mode, isinstance(image, np.ndarray))
+        
+        if fill:
+            if color is None:
+                fill_color = (255, 255, 255, 100)
+            elif isinstance(color, str):
+                # This will automatically raise Error if rgb cannot be parsed.
+                fill_color = ImageColor.getrgb(color) + (100,)
+            elif isinstance(color, tuple):
+                fill_color = color + (100,)
+            draw.rectangle(bbox, width=width, outline=color, fill=fill_color)
+        else:
+            draw.rectangle(bbox, width=width, outline=color)
+
+        if labels is not None:
+            margin = width + 1
+            draw.text((bbox[0] + margin, bbox[1] + margin), labels[i], fill=color, font=font_object)
+
+    if isinstance(image, np.ndarray):
+        return np.asarray(pil_image)
+    return pil_image
+    
+    

khandy/version.py

@@ -1,3 +1,3 @@
-__version__ = '0.1.2'
+__version__ = '0.1.3'
 
 __all__ = ['__version__']

requirements.txt

@@ -1,2 +1,3 @@
 numpy>=1.11.1
 opencv-python
+pillow

setup.py

@@ -1,11 +1,11 @@
 import sys
 from setuptools import find_packages, setup
 
-install_requires = ['numpy>=1.11.1', 'opencv-python']
+install_requires = ['numpy>=1.11.1', 'opencv-python', 'pillow']
 
 setup(
     name='KHandy',
-    version='0.1.2',
+    version='0.1.3',
     description='Handy Utilities for Computer Vision',
     long_description='Handy Utilities for Computer Vision',
     keywords='computer vision',