visualize.py 15.5 KB
Newer Older
1
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
Q
qingqing01 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division

17
import os
Q
qingqing01 已提交
18 19
import cv2
import numpy as np
F
Feng Ni 已提交
20 21
from PIL import Image, ImageDraw, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
22
import math
Q
qingqing01 已提交
23 24


G
Guanghua Yu 已提交
25
def visualize_box_mask(im, results, labels, threshold=0.5):
Q
qingqing01 已提交
26 27 28 29 30 31
    """
    Args:
        im (str/np.ndarray): path of image/np.ndarray read by cv2
        results (dict): include 'boxes': np.ndarray: shape:[N,6], N: number of box,
                        matix element:[class, score, x_min, y_min, x_max, y_max]
                        MaskRCNN's results include 'masks': np.ndarray:
G
Guanghua Yu 已提交
32
                        shape:[N, im_h, im_w]
Q
qingqing01 已提交
33 34 35 36 37 38 39
        labels (list): labels:['class1', ..., 'classn']
        threshold (float): Threshold of score.
    Returns:
        im (PIL.Image.Image): visualized image
    """
    if isinstance(im, str):
        im = Image.open(im).convert('RGB')
40
    elif isinstance(im, np.ndarray):
Q
qingqing01 已提交
41
        im = Image.fromarray(im)
42
    if 'masks' in results and 'boxes' in results and len(results['boxes']) > 0:
Q
qingqing01 已提交
43
        im = draw_mask(
G
Guanghua Yu 已提交
44
            im, results['boxes'], results['masks'], labels, threshold=threshold)
45
    if 'boxes' in results and len(results['boxes']) > 0:
G
Guanghua Yu 已提交
46
        im = draw_box(im, results['boxes'], labels, threshold=threshold)
Q
qingqing01 已提交
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
    if 'segm' in results:
        im = draw_segm(
            im,
            results['segm'],
            results['label'],
            results['score'],
            labels,
            threshold=threshold)
    return im


def get_color_map_list(num_classes):
    """
    Args:
        num_classes (int): number of class
    Returns:
        color_map (list): RGB color list
    """
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
    return color_map


G
Guanghua Yu 已提交
79
def draw_mask(im, np_boxes, np_masks, labels, threshold=0.5):
Q
qingqing01 已提交
80 81 82 83
    """
    Args:
        im (PIL.Image.Image): PIL image
        np_boxes (np.ndarray): shape:[N,6], N: number of box,
G
Guanghua Yu 已提交
84 85
            matix element:[class, score, x_min, y_min, x_max, y_max]
        np_masks (np.ndarray): shape:[N, im_h, im_w]
Q
qingqing01 已提交
86 87 88 89 90 91 92 93 94 95
        labels (list): labels:['class1', ..., 'classn']
        threshold (float): threshold of mask
    Returns:
        im (PIL.Image.Image): visualized image
    """
    color_list = get_color_map_list(len(labels))
    w_ratio = 0.4
    alpha = 0.7
    im = np.array(im).astype('float32')
    clsid2color = {}
G
Guanghua Yu 已提交
96 97 98
    expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
    np_boxes = np_boxes[expect_boxes, :]
    np_masks = np_masks[expect_boxes, :, :]
W
wangguanzhong 已提交
99 100
    im_h, im_w = im.shape[:2]
    np_masks = np_masks[:, :im_h, :im_w]
G
Guanghua Yu 已提交
101 102 103
    for i in range(len(np_masks)):
        clsid, score = int(np_boxes[i][0]), np_boxes[i][1]
        mask = np_masks[i]
Q
qingqing01 已提交
104 105 106 107 108
        if clsid not in clsid2color:
            clsid2color[clsid] = color_list[clsid]
        color_mask = clsid2color[clsid]
        for c in range(3):
            color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
G
Guanghua Yu 已提交
109
        idx = np.nonzero(mask)
Q
qingqing01 已提交
110 111 112 113 114 115
        color_mask = np.array(color_mask)
        im[idx[0], idx[1], :] *= 1.0 - alpha
        im[idx[0], idx[1], :] += alpha * color_mask
    return Image.fromarray(im.astype('uint8'))


G
Guanghua Yu 已提交
116
def draw_box(im, np_boxes, labels, threshold=0.5):
Q
qingqing01 已提交
117 118 119 120 121 122
    """
    Args:
        im (PIL.Image.Image): PIL image
        np_boxes (np.ndarray): shape:[N,6], N: number of box,
                               matix element:[class, score, x_min, y_min, x_max, y_max]
        labels (list): labels:['class1', ..., 'classn']
G
Guanghua Yu 已提交
123
        threshold (float): threshold of box
Q
qingqing01 已提交
124 125 126 127 128 129 130
    Returns:
        im (PIL.Image.Image): visualized image
    """
    draw_thickness = min(im.size) // 320
    draw = ImageDraw.Draw(im)
    clsid2color = {}
    color_list = get_color_map_list(len(labels))
G
Guanghua Yu 已提交
131 132
    expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
    np_boxes = np_boxes[expect_boxes, :]
Q
qingqing01 已提交
133 134 135 136 137 138 139

    for dt in np_boxes:
        clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
        if clsid not in clsid2color:
            clsid2color[clsid] = color_list[clsid]
        color = tuple(clsid2color[clsid])

C
cnn 已提交
140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
        if len(bbox) == 4:
            xmin, ymin, xmax, ymax = bbox
            print('class_id:{:d}, confidence:{:.4f}, left_top:[{:.2f},{:.2f}],'
                  'right_bottom:[{:.2f},{:.2f}]'.format(
                      int(clsid), score, xmin, ymin, xmax, ymax))
            # draw bbox
            draw.line(
                [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
                 (xmin, ymin)],
                width=draw_thickness,
                fill=color)
        elif len(bbox) == 8:
            x1, y1, x2, y2, x3, y3, x4, y4 = bbox
            draw.line(
                [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)],
                width=2,
                fill=color)
            xmin = min(x1, x2, x3, x4)
            ymin = min(y1, y2, y3, y4)
Q
qingqing01 已提交
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185

        # draw label
        text = "{} {:.4f}".format(labels[clsid], score)
        tw, th = draw.textsize(text)
        draw.rectangle(
            [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
        draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
    return im


def draw_segm(im,
              np_segms,
              np_label,
              np_score,
              labels,
              threshold=0.5,
              alpha=0.7):
    """
    Draw segmentation on image
    """
    mask_color_id = 0
    w_ratio = .4
    color_list = get_color_map_list(len(labels))
    im = np.array(im).astype('float32')
    clsid2color = {}
    np_segms = np_segms.astype(np.uint8)
    for i in range(np_segms.shape[0]):
G
Guanghua Yu 已提交
186
        mask, score, clsid = np_segms[i], np_score[i], np_label[i]
Q
qingqing01 已提交
187 188 189 190 191 192 193 194 195 196
        if score < threshold:
            continue

        if clsid not in clsid2color:
            clsid2color[clsid] = color_list[clsid]
        color_mask = clsid2color[clsid]
        for c in range(3):
            color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
        idx = np.nonzero(mask)
        color_mask = np.array(color_mask)
C
cnn 已提交
197 198 199 200
        idx0 = np.minimum(idx[0], im.shape[0] - 1)
        idx1 = np.minimum(idx[1], im.shape[1] - 1)
        im[idx0, idx1, :] *= 1.0 - alpha
        im[idx0, idx1, :] += alpha * color_mask
Q
qingqing01 已提交
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219
        sum_x = np.sum(mask, axis=0)
        x = np.where(sum_x > 0.5)[0]
        sum_y = np.sum(mask, axis=1)
        y = np.where(sum_y > 0.5)[0]
        x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1]
        cv2.rectangle(im, (x0, y0), (x1, y1),
                      tuple(color_mask.astype('int32').tolist()), 1)
        bbox_text = '%s %.2f' % (labels[clsid], score)
        t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
        cv2.rectangle(im, (x0, y0), (x0 + t_size[0], y0 - t_size[1] - 3),
                      tuple(color_mask.astype('int32').tolist()), -1)
        cv2.putText(
            im,
            bbox_text, (x0, y0 - 2),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.3, (0, 0, 0),
            1,
            lineType=cv2.LINE_AA)
    return Image.fromarray(im.astype('uint8'))
220 221 222 223 224 225 226 227


def get_color(idx):
    idx = idx * 3
    color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
    return color


W
wangguanzhong 已提交
228 229 230 231 232 233 234
def visualize_pose(imgfile,
                   results,
                   visual_thresh=0.6,
                   save_name='pose.jpg',
                   save_dir='output',
                   returnimg=False,
                   ids=None):
235 236 237 238 239 240 241 242 243
    try:
        import matplotlib.pyplot as plt
        import matplotlib
        plt.switch_backend('agg')
    except Exception as e:
        logger.error('Matplotlib not found, please install matplotlib.'
                     'for example: `pip install matplotlib`.')
        raise e
    skeletons, scores = results['keypoint']
244
    skeletons = np.array(skeletons)
Z
zhiboniu 已提交
245 246 247
    kpt_nums = 17
    if len(skeletons) > 0:
        kpt_nums = skeletons.shape[1]
248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
    if kpt_nums == 17:  #plot coco keypoint
        EDGES = [(0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 7), (6, 8),
                 (7, 9), (8, 10), (5, 11), (6, 12), (11, 13), (12, 14),
                 (13, 15), (14, 16), (11, 12)]
    else:  #plot mpii keypoint
        EDGES = [(0, 1), (1, 2), (3, 4), (4, 5), (2, 6), (3, 6), (6, 7), (7, 8),
                 (8, 9), (10, 11), (11, 12), (13, 14), (14, 15), (8, 12),
                 (8, 13)]
    NUM_EDGES = len(EDGES)

    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
            [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
            [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
    cmap = matplotlib.cm.get_cmap('hsv')
    plt.figure()

    img = cv2.imread(imgfile) if type(imgfile) == str else imgfile

    color_set = results['colors'] if 'colors' in results else None

    if 'bbox' in results and ids is None:
        bboxs = results['bbox']
        for j, rect in enumerate(bboxs):
            xmin, ymin, xmax, ymax = rect
            color = colors[0] if color_set is None else colors[color_set[j] %
                                                               len(colors)]
            cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 1)

    canvas = img.copy()
    for i in range(kpt_nums):
        for j in range(len(skeletons)):
W
wangguanzhong 已提交
279
            if skeletons[j][i, 2] < visual_thresh:
280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302
                continue
            if ids is None:
                color = colors[i] if color_set is None else colors[color_set[j]
                                                                   %
                                                                   len(colors)]
            else:
                color = get_color(ids[j])

            cv2.circle(
                canvas,
                tuple(skeletons[j][i, 0:2].astype('int32')),
                2,
                color,
                thickness=-1)

    to_plot = cv2.addWeighted(img, 0.3, canvas, 0.7, 0)
    fig = matplotlib.pyplot.gcf()

    stickwidth = 2

    for i in range(NUM_EDGES):
        for j in range(len(skeletons)):
            edge = EDGES[i]
W
wangguanzhong 已提交
303 304
            if skeletons[j][edge[0], 2] < visual_thresh or skeletons[j][edge[
                    1], 2] < visual_thresh:
305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331
                continue

            cur_canvas = canvas.copy()
            X = [skeletons[j][edge[0], 1], skeletons[j][edge[1], 1]]
            Y = [skeletons[j][edge[0], 0], skeletons[j][edge[1], 0]]
            mX = np.mean(X)
            mY = np.mean(Y)
            length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
            polygon = cv2.ellipse2Poly((int(mY), int(mX)),
                                       (int(length / 2), stickwidth),
                                       int(angle), 0, 360, 1)
            if ids is None:
                color = colors[i] if color_set is None else colors[color_set[j]
                                                                   %
                                                                   len(colors)]
            else:
                color = get_color(ids[j])
            cv2.fillConvexPoly(cur_canvas, polygon, color)
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
    if returnimg:
        return canvas
    save_name = os.path.join(
        save_dir, os.path.splitext(os.path.basename(imgfile))[0] + '_vis.jpg')
    plt.imsave(save_name, canvas[:, :, ::-1])
    print("keypoint visualize image saved to: " + save_name)
    plt.close()
332 333 334 335


def visualize_attr(im, results, boxes=None):
    if isinstance(im, str):
336 337 338 339 340
        im = Image.open(im)
        im = np.ascontiguousarray(np.copy(im))
        im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
    else:
        im = np.ascontiguousarray(np.copy(im))
341

342
    im_h, im_w = im.shape[:2]
W
wangguanzhong 已提交
343 344
    text_scale = max(0.5, im.shape[0] / 3000.)
    text_thickness = 1
345

W
wangguanzhong 已提交
346
    line_inter = im.shape[0] / 40.
347 348
    for i, res in enumerate(results):
        if boxes is None:
W
wangguanzhong 已提交
349
            text_w = 3
350
            text_h = 1
351 352
        else:
            box = boxes[i]
W
wangguanzhong 已提交
353
            text_w = int(box[2]) + 3
354 355 356 357 358 359 360 361
            text_h = int(box[3])
        for text in res:
            text_h += int(line_inter)
            text_loc = (text_w, text_h)
            cv2.putText(
                im,
                text,
                text_loc,
W
wangguanzhong 已提交
362 363
                cv2.FONT_ITALIC,
                text_scale, (0, 255, 255),
364
                thickness=text_thickness)
365
    return im
J
JYChen 已提交
366 367


368 369 370 371 372 373
def visualize_action(im,
                     mot_boxes,
                     action_visual_collector=None,
                     action_text="",
                     video_action_score=None,
                     video_action_text=""):
J
JYChen 已提交
374
    im = cv2.imread(im) if isinstance(im, str) else im
375 376
    im_h, im_w = im.shape[:2]

J
JYChen 已提交
377
    text_scale = max(1, im.shape[1] / 1600.)
378 379 380
    text_thickness = 2

    if action_visual_collector:
J
JYChen 已提交
381 382 383 384 385 386
        id_action_dict = {}
        for collector, action_type in zip(action_visual_collector, action_text):
            id_detected = collector.get_visualize_ids()
            for pid in id_detected:
                id_action_dict[pid] = id_action_dict.get(pid, [])
                id_action_dict[pid].append(action_type)
387 388
        for mot_box in mot_boxes:
            # mot_box is a format with [mot_id, class, score, xmin, ymin, w, h] 
J
JYChen 已提交
389
            if mot_box[0] in id_action_dict:
390 391
                text_position = (int(mot_box[3] + mot_box[5] * 0.75),
                                 int(mot_box[4] - 10))
J
JYChen 已提交
392 393
                display_text = ', '.join(id_action_dict[mot_box[0]])
                cv2.putText(im, display_text, text_position,
394 395 396 397 398 399 400 401 402 403 404
                            cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), 2)

    if video_action_score:
        cv2.putText(
            im,
            video_action_text + ': %.2f' % video_action_score,
            (int(im_w / 2), int(15 * text_scale) + 5),
            cv2.FONT_ITALIC,
            text_scale, (0, 0, 255),
            thickness=text_thickness)

J
JYChen 已提交
405
    return im
Z
zhiboniu 已提交
406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440


def visualize_vehicleplate(im, results, boxes=None):
    if isinstance(im, str):
        im = Image.open(im)
        im = np.ascontiguousarray(np.copy(im))
        im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
    else:
        im = np.ascontiguousarray(np.copy(im))

    im_h, im_w = im.shape[:2]
    text_scale = max(0.5, im.shape[0] / 3000.)
    text_thickness = 1

    line_inter = im.shape[0] / 40.
    for i, res in enumerate(results):
        if boxes is None:
            text_w = 3
            text_h = 1
        else:
            box = boxes[i]
            text = res
            if text == "":
                continue
            text_w = int(box[2])
            text_h = int(box[5])
            text_loc = (text_w, text_h)
            cv2.putText(
                im,
                text,
                text_loc,
                cv2.FONT_ITALIC,
                text_scale, (0, 255, 255),
                thickness=text_thickness)
    return im