visualization.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import cv2
  15. import numpy as np
  16. def tlwhs_to_tlbrs(tlwhs):
  17. tlbrs = np.copy(tlwhs)
  18. if len(tlbrs) == 0:
  19. return tlbrs
  20. tlbrs[:, 2] += tlwhs[:, 0]
  21. tlbrs[:, 3] += tlwhs[:, 1]
  22. return tlbrs
  23. def get_color(idx):
  24. idx = idx * 3
  25. color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
  26. return color
  27. def resize_image(image, max_size=800):
  28. if max(image.shape[:2]) > max_size:
  29. scale = float(max_size) / max(image.shape[:2])
  30. image = cv2.resize(image, None, fx=scale, fy=scale)
  31. return image
  32. def plot_tracking(image,
  33. tlwhs,
  34. obj_ids,
  35. scores=None,
  36. frame_id=0,
  37. fps=0.,
  38. ids2=None):
  39. im = np.ascontiguousarray(np.copy(image))
  40. im_h, im_w = im.shape[:2]
  41. top_view = np.zeros([im_w, im_w, 3], dtype=np.uint8) + 255
  42. text_scale = max(1, image.shape[1] / 1600.)
  43. text_thickness = 2
  44. line_thickness = max(1, int(image.shape[1] / 500.))
  45. radius = max(5, int(im_w / 140.))
  46. cv2.putText(
  47. im,
  48. 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
  49. (0, int(15 * text_scale)),
  50. cv2.FONT_HERSHEY_PLAIN,
  51. text_scale, (0, 0, 255),
  52. thickness=2)
  53. for i, tlwh in enumerate(tlwhs):
  54. x1, y1, w, h = tlwh
  55. intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
  56. obj_id = int(obj_ids[i])
  57. id_text = '{}'.format(int(obj_id))
  58. if ids2 is not None:
  59. id_text = id_text + ', {}'.format(int(ids2[i]))
  60. _line_thickness = 1 if obj_id <= 0 else line_thickness
  61. color = get_color(abs(obj_id))
  62. cv2.rectangle(
  63. im,
  64. intbox[0:2],
  65. intbox[2:4],
  66. color=color,
  67. thickness=line_thickness)
  68. cv2.putText(
  69. im,
  70. id_text, (intbox[0], intbox[1] + 10),
  71. cv2.FONT_HERSHEY_PLAIN,
  72. text_scale, (0, 0, 255),
  73. thickness=text_thickness)
  74. if scores is not None:
  75. text = '{:.2f}'.format(float(scores[i]))
  76. cv2.putText(
  77. im,
  78. text, (intbox[0], intbox[1] - 10),
  79. cv2.FONT_HERSHEY_PLAIN,
  80. text_scale, (0, 255, 255),
  81. thickness=text_thickness)
  82. return im
  83. def plot_trajectory(image, tlwhs, track_ids):
  84. image = image.copy()
  85. for one_tlwhs, track_id in zip(tlwhs, track_ids):
  86. color = get_color(int(track_id))
  87. for tlwh in one_tlwhs:
  88. x1, y1, w, h = tuple(map(int, tlwh))
  89. cv2.circle(
  90. image, (int(x1 + 0.5 * w), int(y1 + h)), 2, color, thickness=2)
  91. return image
  92. def plot_detections(image, tlbrs, scores=None, color=(255, 0, 0), ids=None):
  93. im = np.copy(image)
  94. text_scale = max(1, image.shape[1] / 800.)
  95. thickness = 2 if text_scale > 1.3 else 1
  96. for i, det in enumerate(tlbrs):
  97. x1, y1, x2, y2 = np.asarray(det[:4], dtype=np.int)
  98. if len(det) >= 7:
  99. label = 'det' if det[5] > 0 else 'trk'
  100. if ids is not None:
  101. text = '{}# {:.2f}: {:d}'.format(label, det[6], ids[i])
  102. cv2.putText(
  103. im,
  104. text, (x1, y1 + 30),
  105. cv2.FONT_HERSHEY_PLAIN,
  106. text_scale, (0, 255, 255),
  107. thickness=thickness)
  108. else:
  109. text = '{}# {:.2f}'.format(label, det[6])
  110. if scores is not None:
  111. text = '{:.2f}'.format(scores[i])
  112. cv2.putText(
  113. im,
  114. text, (x1, y1 + 30),
  115. cv2.FONT_HERSHEY_PLAIN,
  116. text_scale, (0, 255, 255),
  117. thickness=thickness)
  118. cv2.rectangle(im, (x1, y1), (x2, y2), color, 2)
  119. return im