visualization.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  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 get_color(idx):
  17. idx = idx * 3
  18. color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
  19. return color
  20. def plot_tracking(image,
  21. tlwhs,
  22. obj_ids,
  23. scores=None,
  24. frame_id=0,
  25. fps=0.,
  26. ids2names=[]):
  27. im = np.ascontiguousarray(np.copy(image))
  28. im_h, im_w = im.shape[:2]
  29. top_view = np.zeros([im_w, im_w, 3], dtype=np.uint8) + 255
  30. text_scale = max(1, image.shape[1] / 1600.)
  31. text_thickness = 2
  32. line_thickness = max(1, int(image.shape[1] / 500.))
  33. radius = max(5, int(im_w / 140.))
  34. cv2.putText(
  35. im,
  36. 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
  37. (0, int(15 * text_scale)),
  38. cv2.FONT_HERSHEY_PLAIN,
  39. text_scale, (0, 0, 255),
  40. thickness=2)
  41. for i, tlwh in enumerate(tlwhs):
  42. x1, y1, w, h = tlwh
  43. intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
  44. obj_id = int(obj_ids[i])
  45. id_text = '{}'.format(int(obj_id))
  46. if ids2names != []:
  47. assert len(
  48. ids2names) == 1, "plot_tracking only supports single classes."
  49. id_text = '{}_'.format(ids2names[0]) + id_text
  50. _line_thickness = 1 if obj_id <= 0 else line_thickness
  51. color = get_color(abs(obj_id))
  52. cv2.rectangle(
  53. im,
  54. intbox[0:2],
  55. intbox[2:4],
  56. color=color,
  57. thickness=line_thickness)
  58. cv2.putText(
  59. im,
  60. id_text, (intbox[0], intbox[1] - 10),
  61. cv2.FONT_HERSHEY_PLAIN,
  62. text_scale, (0, 0, 255),
  63. thickness=text_thickness)
  64. if scores is not None:
  65. text = '{:.2f}'.format(float(scores[i]))
  66. cv2.putText(
  67. im,
  68. text, (intbox[0], intbox[1] + 10),
  69. cv2.FONT_HERSHEY_PLAIN,
  70. text_scale, (0, 255, 255),
  71. thickness=text_thickness)
  72. return im
  73. def plot_tracking_dict(image,
  74. num_classes,
  75. tlwhs_dict,
  76. obj_ids_dict,
  77. scores_dict,
  78. frame_id=0,
  79. fps=0.,
  80. ids2names=[]):
  81. im = np.ascontiguousarray(np.copy(image))
  82. im_h, im_w = im.shape[:2]
  83. top_view = np.zeros([im_w, im_w, 3], dtype=np.uint8) + 255
  84. text_scale = max(1, image.shape[1] / 1600.)
  85. text_thickness = 2
  86. line_thickness = max(1, int(image.shape[1] / 500.))
  87. radius = max(5, int(im_w / 140.))
  88. for cls_id in range(num_classes):
  89. tlwhs = tlwhs_dict[cls_id]
  90. obj_ids = obj_ids_dict[cls_id]
  91. scores = scores_dict[cls_id]
  92. cv2.putText(
  93. im,
  94. 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)),
  95. (0, int(15 * text_scale)),
  96. cv2.FONT_HERSHEY_PLAIN,
  97. text_scale, (0, 0, 255),
  98. thickness=2)
  99. for i, tlwh in enumerate(tlwhs):
  100. x1, y1, w, h = tlwh
  101. intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
  102. obj_id = int(obj_ids[i])
  103. id_text = '{}'.format(int(obj_id))
  104. if ids2names != []:
  105. id_text = '{}_{}'.format(ids2names[cls_id], id_text)
  106. else:
  107. id_text = 'class{}_{}'.format(cls_id, id_text)
  108. _line_thickness = 1 if obj_id <= 0 else line_thickness
  109. color = get_color(abs(obj_id))
  110. cv2.rectangle(
  111. im,
  112. intbox[0:2],
  113. intbox[2:4],
  114. color=color,
  115. thickness=line_thickness)
  116. cv2.putText(
  117. im,
  118. id_text, (intbox[0], intbox[1] - 10),
  119. cv2.FONT_HERSHEY_PLAIN,
  120. text_scale, (0, 0, 255),
  121. thickness=text_thickness)
  122. if scores is not None:
  123. text = '{:.2f}'.format(float(scores[i]))
  124. cv2.putText(
  125. im,
  126. text, (intbox[0], intbox[1] + 10),
  127. cv2.FONT_HERSHEY_PLAIN,
  128. text_scale, (0, 255, 255),
  129. thickness=text_thickness)
  130. return im