--- a +++ b/FastRCNN/utils/visualization.py @@ -0,0 +1,108 @@ +import numpy as np + + +def vis_image(img, ax=None): + """Visualize a color image. + + Args: + img (~numpy.ndarray): An array of shape :math:`(3, height, width)`. + This is in RGB format and the range of its value is + :math:`[0, 255]`. + ax (matplotlib.axes.Axis): The visualization is displayed on this + axis. If this is :obj:`None` (default), a new axis is created. + + Returns: + ~matploblib.axes.Axes: + Returns the Axes object with the plot for further tweaking. + + """ + from matplotlib import pyplot as plot + if ax is None: + fig = plot.figure() + ax = fig.add_subplot(1, 1, 1) + # CHW -> HWC + img = img.transpose((1, 2, 0))[:, :, 0] + ax.imshow(img.astype(np.uint8), cmap='gray') + return ax + + +def vis_bbox(img, bbox, label=None, score=None, label_names=None, ax=None): + """Visualize bounding boxes inside image. + + Example: + + >>> from chainercv.datasets import VOCDetectionDataset + >>> from chainercv.datasets import voc_bbox_label_names + >>> from chainercv.visualizations import vis_bbox + >>> import matplotlib.pyplot as plot + >>> dataset = VOCDetectionDataset() + >>> img, bbox, label = dataset[60] + >>> vis_bbox(img, bbox, label, + ... label_names=voc_bbox_label_names) + >>> plot.show() + + Args: + img (~numpy.ndarray): An array of shape :math:`(3, height, width)`. + This is in RGB format and the range of its value is + :math:`[0, 255]`. + bbox (~numpy.ndarray): An array of shape :math:`(R, 4)`, where + :math:`R` is the number of bounding boxes in the image. + Each element is organized + by :math:`(y_{min}, x_{min}, y_{max}, x_{max})` in the second axis. + label (~numpy.ndarray): An integer array of shape :math:`(R,)`. + The values correspond to id for label names stored in + :obj:`label_names`. This is optional. + score (~numpy.ndarray): A float array of shape :math:`(R,)`. + Each value indicates how confident the prediction is. + This is optional. + label_names (iterable of strings): Name of labels ordered according + to label ids. If this is :obj:`None`, labels will be skipped. + ax (matplotlib.axes.Axis): The visualization is displayed on this + axis. If this is :obj:`None` (default), a new axis is created. + + Returns: + ~matploblib.axes.Axes: + Returns the Axes object with the plot for further tweaking. + + """ + from matplotlib import pyplot as plot + + #if label is not None and not len(bbox) == len(label): + # raise ValueError('The length of label must be same as that of bbox') + #if score is not None and not len(bbox) == len(score): + # raise ValueError('The length of score must be same as that of bbox') + + # Returns newly instantiated matplotlib.axes.Axes object if ax is None + ax = vis_image(img, ax=ax) + + # If there is no bounding box to display, visualize the image and exit. + if len(bbox) == 0: + return ax + + if len(bbox) != 0: + #print(bbox) + bb = bbox[0] + i = 0 + xy = (bb[1], bb[0]) + height = bb[2] - bb[0] + width = bb[3] - bb[1] + ax.add_patch(plot.Rectangle( + xy, width, height, fill=False, edgecolor='red', linewidth=3)) + + caption = list() + + if label is not None and label_names is not None: + lb = label[i] + if not (0 <= lb < len(label_names)): + raise ValueError('No corresponding name is given') + caption.append(label_names[lb]) + if score is not None: + sc = score[i] + caption.append('{:.2f}'.format(sc)) + + if len(caption) > 0: + ax.text(bb[1], bb[0], + ': '.join(caption), + style='italic') +# bbox={'facecolor': 'white', 'alpha': 0.7, 'pad': 5}) + return ax \ No newline at end of file