395 lines
14 KiB
Python
395 lines
14 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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
|
|
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
import PIL
|
|
from PIL import Image, ImageDraw, ImageFile
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
|
from collections import deque
|
|
|
|
def imagedraw_textsize_c(draw, text):
|
|
if int(PIL.__version__.split('.')[0]) < 10:
|
|
tw, th = draw.textsize(text)
|
|
else:
|
|
left, top, right, bottom = draw.textbbox((0, 0), text)
|
|
tw, th = right - left, bottom - top
|
|
|
|
return tw, th
|
|
|
|
def imagedraw_textsize_c(draw, text):
|
|
if int(PIL.__version__.split('.')[0]) < 10:
|
|
tw, th = draw.textsize(text)
|
|
else:
|
|
left, top, right, bottom = draw.textbbox((0, 0), text)
|
|
tw, th = right - left, bottom - top
|
|
|
|
return tw, th
|
|
|
|
|
|
def visualize_box_mask(im, results, labels, threshold=0.5):
|
|
"""
|
|
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]
|
|
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')
|
|
else:
|
|
im = Image.fromarray(im)
|
|
if 'boxes' in results and len(results['boxes']) > 0:
|
|
im = draw_box(im, results['boxes'], 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
|
|
|
|
|
|
def draw_box(im, np_boxes, labels, threshold=0.5):
|
|
"""
|
|
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']
|
|
threshold (float): threshold of box
|
|
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))
|
|
expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
|
|
np_boxes = np_boxes[expect_boxes, :]
|
|
|
|
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])
|
|
|
|
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)
|
|
|
|
# draw label
|
|
text = "{} {:.4f}".format(labels[clsid], score)
|
|
tw, th = imagedraw_textsize_c(draw, 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 get_color(idx):
|
|
idx = idx * 3
|
|
color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
|
|
return color
|
|
|
|
|
|
def plot_tracking(image,
|
|
tlwhs,
|
|
obj_ids,
|
|
scores=None,
|
|
frame_id=0,
|
|
fps=0.,
|
|
ids2names=[],
|
|
do_entrance_counting=False,
|
|
entrance=None):
|
|
im = np.ascontiguousarray(np.copy(image))
|
|
im_h, im_w = im.shape[:2]
|
|
|
|
text_scale = max(0.5, image.shape[1] / 3000.)
|
|
text_thickness = 2
|
|
line_thickness = max(1, int(image.shape[1] / 500.))
|
|
|
|
cv2.putText(
|
|
im,
|
|
'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
|
|
(0, int(15 * text_scale) + 5),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
for i, tlwh in enumerate(tlwhs):
|
|
x1, y1, w, h = tlwh
|
|
intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
|
|
obj_id = int(obj_ids[i])
|
|
id_text = 'ID: {}'.format(int(obj_id))
|
|
if ids2names != []:
|
|
assert len(
|
|
ids2names) == 1, "plot_tracking only supports single classes."
|
|
id_text = 'ID: {}_'.format(ids2names[0]) + id_text
|
|
_line_thickness = 1 if obj_id <= 0 else line_thickness
|
|
color = get_color(abs(obj_id))
|
|
cv2.rectangle(
|
|
im, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness)
|
|
cv2.putText(
|
|
im,
|
|
id_text, (intbox[0], intbox[1] - 25),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 255, 255),
|
|
thickness=text_thickness)
|
|
|
|
if scores is not None:
|
|
text = 'score: {:.2f}'.format(float(scores[i]))
|
|
cv2.putText(
|
|
im,
|
|
text, (intbox[0], intbox[1] - 6),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 255, 0),
|
|
thickness=text_thickness)
|
|
if do_entrance_counting:
|
|
entrance_line = tuple(map(int, entrance))
|
|
cv2.rectangle(
|
|
im,
|
|
entrance_line[0:2],
|
|
entrance_line[2:4],
|
|
color=(0, 255, 255),
|
|
thickness=line_thickness)
|
|
return im
|
|
|
|
|
|
def plot_tracking_dict(image,
|
|
num_classes,
|
|
tlwhs_dict,
|
|
obj_ids_dict,
|
|
scores_dict,
|
|
frame_id=0,
|
|
fps=0.,
|
|
ids2names=[],
|
|
do_entrance_counting=False,
|
|
do_break_in_counting=False,
|
|
do_illegal_parking_recognition=False,
|
|
illegal_parking_dict=None,
|
|
entrance=None,
|
|
records=None,
|
|
center_traj=None):
|
|
im = np.ascontiguousarray(np.copy(image))
|
|
im_h, im_w = im.shape[:2]
|
|
if do_break_in_counting or do_illegal_parking_recognition:
|
|
entrance = np.array(entrance[:-1]) # last pair is [im_w, im_h]
|
|
|
|
text_scale = max(0.5, image.shape[1] / 3000.)
|
|
text_thickness = 2
|
|
line_thickness = max(1, int(image.shape[1] / 500.))
|
|
|
|
if num_classes == 1:
|
|
if records is not None:
|
|
start = records[-1].find('Total')
|
|
end = records[-1].find('In')
|
|
cv2.putText(
|
|
im,
|
|
records[-1][start:end], (0, int(40 * text_scale) + 10),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
if num_classes == 1 and do_entrance_counting:
|
|
entrance_line = tuple(map(int, entrance))
|
|
cv2.rectangle(
|
|
im,
|
|
entrance_line[0:2],
|
|
entrance_line[2:4],
|
|
color=(0, 255, 255),
|
|
thickness=line_thickness)
|
|
# find start location for entrance counting data
|
|
start = records[-1].find('In')
|
|
cv2.putText(
|
|
im,
|
|
records[-1][start:-1], (0, int(60 * text_scale) + 10),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
if num_classes == 1 and (do_break_in_counting or
|
|
do_illegal_parking_recognition):
|
|
np_masks = np.zeros((im_h, im_w, 1), np.uint8)
|
|
cv2.fillPoly(np_masks, [entrance], 255)
|
|
|
|
# Draw region mask
|
|
alpha = 0.3
|
|
im = np.array(im).astype('float32')
|
|
mask = np_masks[:, :, 0]
|
|
color_mask = [0, 0, 255]
|
|
idx = np.nonzero(mask)
|
|
color_mask = np.array(color_mask)
|
|
im[idx[0], idx[1], :] *= 1.0 - alpha
|
|
im[idx[0], idx[1], :] += alpha * color_mask
|
|
im = np.array(im).astype('uint8')
|
|
|
|
if do_break_in_counting:
|
|
# find start location for break in counting data
|
|
start = records[-1].find('Break_in')
|
|
cv2.putText(
|
|
im,
|
|
records[-1][start:-1],
|
|
(entrance[0][0] - 10, entrance[0][1] - 10),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
if illegal_parking_dict is not None and len(illegal_parking_dict) != 0:
|
|
for key, value in illegal_parking_dict.items():
|
|
x1, y1, w, h = value['bbox']
|
|
plate = value['plate']
|
|
if plate is None:
|
|
plate = ""
|
|
|
|
# red box
|
|
cv2.rectangle(im, (int(x1), int(y1)),
|
|
(int(x1 + w), int(y1 + h)), (0, 0, 255), 2)
|
|
|
|
cv2.putText(
|
|
im,
|
|
"illegal_parking:" + plate,
|
|
(int(x1) + 5, int(16 * text_scale + y1 + 15)),
|
|
cv2.FONT_ITALIC,
|
|
text_scale * 1.5, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
for cls_id in range(num_classes):
|
|
tlwhs = tlwhs_dict[cls_id]
|
|
obj_ids = obj_ids_dict[cls_id]
|
|
scores = scores_dict[cls_id]
|
|
cv2.putText(
|
|
im,
|
|
'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
|
|
(0, int(15 * text_scale) + 5),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
record_id = set()
|
|
for i, tlwh in enumerate(tlwhs):
|
|
x1, y1, w, h = tlwh
|
|
intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
|
|
center = tuple(map(int, (x1 + w / 2., y1 + h / 2.)))
|
|
obj_id = int(obj_ids[i])
|
|
if center_traj is not None:
|
|
record_id.add(obj_id)
|
|
if obj_id not in center_traj[cls_id]:
|
|
center_traj[cls_id][obj_id] = deque(maxlen=30)
|
|
center_traj[cls_id][obj_id].append(center)
|
|
|
|
id_text = '{}'.format(int(obj_id))
|
|
if ids2names != []:
|
|
id_text = '{}_{}'.format(ids2names[cls_id], id_text)
|
|
else:
|
|
id_text = 'class{}_{}'.format(cls_id, id_text)
|
|
|
|
_line_thickness = 1 if obj_id <= 0 else line_thickness
|
|
|
|
in_region = False
|
|
if do_break_in_counting:
|
|
center_x = min(x1 + w / 2., im_w - 1)
|
|
center_down_y = min(y1 + h, im_h - 1)
|
|
if in_quadrangle([center_x, center_down_y], entrance, im_h,
|
|
im_w):
|
|
in_region = True
|
|
|
|
color = get_color(abs(obj_id)) if in_region == False else (0, 0,
|
|
255)
|
|
cv2.rectangle(
|
|
im,
|
|
intbox[0:2],
|
|
intbox[2:4],
|
|
color=color,
|
|
thickness=line_thickness)
|
|
cv2.putText(
|
|
im,
|
|
id_text, (intbox[0], intbox[1] - 25),
|
|
cv2.FONT_ITALIC,
|
|
text_scale,
|
|
color,
|
|
thickness=text_thickness)
|
|
|
|
if do_break_in_counting and in_region:
|
|
cv2.putText(
|
|
im,
|
|
'Break in now.', (intbox[0], intbox[1] - 50),
|
|
cv2.FONT_ITALIC,
|
|
text_scale, (0, 0, 255),
|
|
thickness=text_thickness)
|
|
|
|
if scores is not None:
|
|
text = 'score: {:.2f}'.format(float(scores[i]))
|
|
cv2.putText(
|
|
im,
|
|
text, (intbox[0], intbox[1] - 6),
|
|
cv2.FONT_ITALIC,
|
|
text_scale,
|
|
color,
|
|
thickness=text_thickness)
|
|
if center_traj is not None:
|
|
for traj in center_traj:
|
|
for i in traj.keys():
|
|
if i not in record_id:
|
|
continue
|
|
for point in traj[i]:
|
|
cv2.circle(im, point, 3, (0, 0, 255), -1)
|
|
return im
|
|
|
|
|
|
def in_quadrangle(point, entrance, im_h, im_w):
|
|
mask = np.zeros((im_h, im_w, 1), np.uint8)
|
|
cv2.fillPoly(mask, [entrance], 255)
|
|
p = tuple(map(int, point))
|
|
if mask[p[1], p[0], :] > 0:
|
|
return True
|
|
else:
|
|
return False
|