Esempio n. 1
0
 def _draw_bbox(self, img, bbox):
     bbox = bbox[0]
     left_top = (bbox[0], bbox[1])
     right_bottom = (bbox[2], bbox[3])
     cv2.rectangle(img,
                   left_top,
                   right_bottom,
                   color_val('red'),
                   thickness=self.thickness)
     cv2.putText(img, 'no wear helmet', (int(bbox[0]), int(bbox[1]) - 5),
                 cv2.FONT_HERSHEY_SIMPLEX, 0.8, color_val('red'), 2)
     return img
Esempio n. 2
0
 def _draw_bbox(self, img, bbox):
     bbox = bbox[0]
     left_top = (bbox[0], bbox[1])
     right_bottom = (bbox[2], bbox[3])
     cv2.rectangle(img,
                   left_top,
                   right_bottom,
                   color_val('red'),
                   thickness=self.thickness)
     return img
Esempio n. 3
0
def imshow_det_bboxes(img,
                      bboxes,
                      labels,
                      class_names=None,
                      score_thr=0,
                      bbox_color='green',
                      text_color='green',
                      thickness=1,
                      font_scale=0.5,
                      show=True,
                      win_name='',
                      wait_time=0,
                      out_file=None,
                      color_dist=None):
    """Draw bboxes and class labels (with scores) on an image.

    Args:
        img (str or ndarray): The image to be displayed.
        bboxes (ndarray): Bounding boxes (with scores), shaped (n, 4) or
            (n, 5).
        labels (ndarray): Labels of bboxes.
        class_names (list[str]): Names of each classes.
        score_thr (float): Minimum score of bboxes to be shown.
        bbox_color (str or tuple or :obj:`Color`): Color of bbox lines.
        text_color (str or tuple or :obj:`Color`): Color of texts.
        thickness (int): Thickness of lines.
        font_scale (float): Font scales of texts.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
        out_file (str or None): The filename to write the image.
        color_dist(dist):the bbox color dist
    """
    assert bboxes.ndim == 2
    assert labels.ndim == 1
    assert bboxes.shape[0] == labels.shape[0]
    assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
    img = mmcv.imread(img)
    if score_thr > 0:
        assert bboxes.shape[1] == 5
        scores = bboxes[:, -1]
        inds = scores > score_thr
        bboxes = bboxes[inds, :]
        labels = labels[inds]

    bbox_color = color_val(bbox_color)

    text_color = color_val(text_color)
    for bbox, label in zip(bboxes, labels):
        bbox_int = bbox.astype(np.int32)
        left_top = (bbox_int[0], bbox_int[1])
        right_bottom = (bbox_int[2], bbox_int[3])
        label_text = class_names[
            label] if class_names is not None else 'cls {}'.format(label)
        bbox_color = color_val(color_dist[label_text])
        text_color = color_val(color_dist[label_text])
        cv2.rectangle(img,
                      left_top,
                      right_bottom,
                      bbox_color,
                      thickness=thickness)
        if len(bbox) > 4:
            label_text += '|{:.02f}'.format(bbox[-1])
        cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2),
                    cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
Esempio n. 4
0
def imshow_tracklets(img,
                     bboxes,
                     labels=None,
                     ids=None,
                     thickness=2,
                     font_scale=0.4,
                     show=False,
                     win_name='',
                     color=None,
                     out_file=None):
    assert bboxes.ndim == 2
    assert labels.ndim == 1
    assert bboxes.shape[0] == labels.shape[0]
    # assert bboxes.shape[1] == 4 or bboxes.shape[1] == 5
    if isinstance(img, str):
        img = imread(img)
    i = 0
    if bboxes.shape[0] == 0:
        if out_file is not None:
            imwrite(img, out_file)
        return img
    if isinstance(bboxes, torch.Tensor):
        bboxes = bboxes.numpy()
        labels = labels.numpy()
        ids = ids.numpy()
    for bbox, label in zip(bboxes, labels):
        x1, y1, x2, y2, _ = bbox.astype(np.int32)
        if ids is not None:
            if color is None:
                bbox_color = random_color(ids[i])
                bbox_color = [int(255 * _c) for _c in bbox_color][::-1]
            else:
                bbox_color = mmcv.color_val(color)
            img[y1:y1 + 12, x1:x1 + 20, :] = bbox_color
            cv2.putText(img,
                        str(ids[i]), (x1, y1 + 10),
                        cv2.FONT_HERSHEY_COMPLEX,
                        font_scale,
                        color=color_val('black'))
        else:
            if color is None:
                bbox_color = color_val('green')
            else:
                bbox_color = mmcv.color_val(color)

        cv2.rectangle(img, (x1, y1), (x2, y2), bbox_color, thickness=thickness)

        if bbox[-1] < 0:
            bbox[-1] = np.nan
        # label_text = '{:.02f}'.format(bbox[-1])
        # img[y1 - 12:y1, x1:x1 + 30, :] = bbox_color
        # cv2.putText(
        #     img,
        #     label_text, (x1, y1 - 2),
        #     cv2.FONT_HERSHEY_COMPLEX,
        #     font_scale,
        #     color=color_val('black'))

        i += 1

    if show:
        imshow(img, win_name)
    if out_file is not None:
        imwrite(img, out_file)

    return img
Esempio n. 5
0
        imgs = data.get_data(images)
        print("start model recongnize")
        with torch.no_grad():
            output = model(return_loss=False, **imgs)
            outputs.append(output)

        use_softmax = True
        if use_softmax is True:
            print("Averaging score over {} clips with softmax".format(
                outputs[0].shape[0]))
            results = [softmax(res, dim=1).mean(axis=0) for res in outputs]
        else:
            print("Averaging score over {} clips without softmax (ie, raw)".format(
                outputs[0].shape[0]))
            results = [res.mean(axis=0) for res in outputs]
        pred = int(np.argmax(results, axis=1))
        print(pred)
        print(results[0][pred])
        print(gt_labels[pred])
        text_color = color_val('green')
        for image in images:
            cv2.putText(image, str(gt_labels[pred]), (100, 100),
                        cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color)
            cv2.imshow('image', image)
        cv2.waitKey(0)

    cap.release()
    cv2.destroyAllWindows()