Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import, division | |
| import numpy as np | |
| from PIL import ImageDraw | |
| def _pascal_color_map(N=256, normalized=False): | |
| """ | |
| Python implementation of the color map function for the PASCAL VOC data set. | |
| Official Matlab version can be found in the PASCAL VOC devkit | |
| http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html#devkit | |
| """ | |
| def bitget(byteval, idx): | |
| return (byteval & (1 << idx)) != 0 | |
| dtype = 'float32' if normalized else 'uint8' | |
| cmap = np.zeros((N, 3), dtype=dtype) | |
| for i in range(N): | |
| r = g = b = 0 | |
| c = i | |
| for j in range(8): | |
| r = r | (bitget(c, 0) << 7 - j) | |
| g = g | (bitget(c, 1) << 7 - j) | |
| b = b | (bitget(c, 2) << 7 - j) | |
| c = c >> 3 | |
| cmap[i] = np.array([r, g, b]) | |
| cmap = cmap / 255 if normalized else cmap | |
| return cmap | |
| def plot_scribble(ax, scribble, frame, output_size=None, **kwargs): | |
| """ Plot scribbles into an axis. | |
| # Arguments | |
| ax: Matplotlib Axis. Axis where to plot the scribble lines. | |
| scribbles: Scribble. Scribble to plot. | |
| frame: Integer. Frame of the scribble to plot. | |
| output_size: Tuple. Image size to scale the scribble points `(H, W)`. | |
| **kwargs: Dictionary. Additional parameters to pass at the | |
| `ax.plot(**kwargs)` method. | |
| # Returns | |
| matplotlib.axis: Returns the given axis with the scribbles plotted on | |
| it. | |
| """ | |
| scribbles = scribble['scribbles'] | |
| if frame >= len(scribbles): | |
| raise ValueError('Frame value not valid') | |
| cmap = _pascal_color_map(normalized=True) | |
| frame_scribbles = scribbles[frame] | |
| for line in frame_scribbles: | |
| path, obj_id = line['path'], line['object_id'] | |
| path = np.asarray(path, dtype=np.float32) | |
| color = cmap[obj_id] | |
| if output_size: | |
| img_size = np.asarray(output_size, dtype=np.float32) | |
| img_size -= 1 | |
| path *= img_size | |
| ax.plot(*path.T, color=color, **kwargs) | |
| return ax | |
| def draw_scribble(img, scribble, frame, output_size=None, width=5): | |
| """ Draw scribbles into a PIL Image. | |
| # Arguments | |
| img: PIL Image. Image where to draw the scribbles. | |
| scribbles: Scribble. Scribble to plot. | |
| frame: Integer. Frame of the scribble to plot. | |
| output_size: Tuple. Image size to scale the scribble points `(H, W)`. | |
| width: Integer. Width of the drawed lines. | |
| # Returns | |
| PIL Image: Returns the original image `img` with the scribble draw on | |
| it. | |
| """ | |
| scribbles = scribble['scribbles'] | |
| if frame >= len(scribbles): | |
| raise ValueError('Frame value not valid') | |
| cmap = _pascal_color_map(normalized=False) | |
| frame_scribbles = scribbles[frame] | |
| draw = ImageDraw.Draw(img) | |
| for line in frame_scribbles: | |
| path, obj_id = line['path'], line['object_id'] | |
| path = np.asarray(path, dtype=np.float32) | |
| color = cmap[obj_id] | |
| if output_size: | |
| img_size = np.asarray(output_size, dtype=np.float32) | |
| img_size -= 1 | |
| path *= img_size | |
| path = path.ravel().tolist() | |
| draw.line(path, fill=tuple(color), width=width) | |
| del draw | |
| return img | |
| def overlay_mask(im, ann, alpha=0.5, colors=None, contour_thickness=None): | |
| """ Overlay mask over image. | |
| This function allows you to overlay a mask over an image with some | |
| transparency. | |
| # Arguments | |
| im: Numpy Array. Array with the image. The shape must be (H, W, 3) and | |
| the pixels must be represented as `np.uint8` data type. | |
| ann: Numpy Array. Array with the mask. The shape must be (H, W) and the | |
| values must be intergers | |
| alpha: Float. Proportion of alpha to apply at the overlaid mask. | |
| colors: Numpy Array. Optional custom colormap. It must have shape (N, 3) | |
| being N the maximum number of colors to represent. | |
| contour_thickness: Integer. Thickness of each object index contour draw | |
| over the overlay. This function requires to have installed the | |
| package `opencv-python`. | |
| # Returns | |
| Numpy Array: Image of the overlay with shape (H, W, 3) and data type | |
| `np.uint8`. | |
| """ | |
| im, ann = np.asarray(im, dtype=np.uint8), np.asarray(ann, dtype=np.int) | |
| if im.shape[:-1] != ann.shape: | |
| raise ValueError('First two dimensions of `im` and `ann` must match') | |
| if im.shape[-1] != 3: | |
| raise ValueError('im must have three channels at the 3 dimension') | |
| colors = colors or _pascal_color_map() | |
| colors = np.asarray(colors, dtype=np.uint8) | |
| mask = colors[ann] | |
| fg = im * alpha + (1 - alpha) * mask | |
| img = im.copy() | |
| img[ann > 0] = fg[ann > 0] | |
| if contour_thickness: # pragma: no cover | |
| import cv2 | |
| for obj_id in np.unique(ann[ann > 0]): | |
| contours = cv2.findContours((ann == obj_id).astype( | |
| np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:] | |
| cv2.drawContours(img, contours[0], -1, colors[obj_id].tolist(), | |
| contour_thickness) | |
| return img | |