|
- # Copyright 2022 Huawei Technologies Co., Ltd
- #
- # 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.
- # ============================================================================
- """Images visualization script."""
- import numpy as np
- import cv2
-
-
- def get_color(idx):
- """
- Set the color for unique pedestrian.
- """
- idx = idx * 3
- color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
-
- return color
-
-
- def plot_tracking(image, tlwhs, obj_ids, frame_id=0, fps=0., ids2=None):
- """
- Show tracking results.
- """
- im = np.ascontiguousarray(np.copy(image))
-
- text_scale = max(1, image.shape[1] / 1600.)
- text_thickness = 1 if text_scale > 1.1 else 1
- line_thickness = max(1, int(image.shape[1] / 500.))
-
- cv2.putText(im, f'frame: {frame_id} fps: {fps:.2f} num: {len(tlwhs)}',
- (0, int(15 * text_scale)), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255), thickness=2)
-
- 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 = f'{int(obj_id)}'
- if ids2 is not None:
- id_text = id_text + f', {int(ids2[i])}'
- 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] + 30), cv2.FONT_HERSHEY_PLAIN, text_scale, (0, 0, 255),
- thickness=text_thickness)
- return im
|