render.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import numpy as np
  2. import cv2
  3. from math import sqrt
  4. _LINE_THICKNESS_SCALING = 500.0
  5. np.random.seed(0)
  6. RAND_COLORS = np.random.randint(50, 255, (64, 3), "int") # used for class visu
  7. RAND_COLORS[0] = [220, 220, 220]
  8. def render_box(img, box, color=(200, 200, 200)):
  9. """
  10. Render a box. Calculates scaling and thickness automatically.
  11. :param img: image to render into
  12. :param box: (x1, y1, x2, y2) - box coordinates
  13. :param color: (b, g, r) - box color
  14. :return: updated image
  15. """
  16. x1, y1, x2, y2 = box
  17. thickness = int(
  18. round(
  19. (img.shape[0] * img.shape[1])
  20. / (_LINE_THICKNESS_SCALING * _LINE_THICKNESS_SCALING)
  21. )
  22. )
  23. thickness = max(1, thickness)
  24. img = cv2.rectangle(
  25. img,
  26. (int(x1), int(y1)),
  27. (int(x2), int(y2)),
  28. color,
  29. thickness=thickness
  30. )
  31. return img
  32. def render_filled_box(img, box, color=(200, 200, 200)):
  33. """
  34. Render a box. Calculates scaling and thickness automatically.
  35. :param img: image to render into
  36. :param box: (x1, y1, x2, y2) - box coordinates
  37. :param color: (b, g, r) - box color
  38. :return: updated image
  39. """
  40. x1, y1, x2, y2 = box
  41. img = cv2.rectangle(
  42. img,
  43. (int(x1), int(y1)),
  44. (int(x2), int(y2)),
  45. color,
  46. thickness=cv2.FILLED
  47. )
  48. return img
  49. _TEXT_THICKNESS_SCALING = 700.0
  50. _TEXT_SCALING = 520.0
  51. def get_text_size(img, text, normalised_scaling=1.0):
  52. """
  53. Get calculated text size (as box width and height)
  54. :param img: image reference, used to determine appropriate text scaling
  55. :param text: text to display
  56. :param normalised_scaling: additional normalised scaling. Default 1.0.
  57. :return: (width, height) - width and height of text box
  58. """
  59. thickness = int(
  60. round(
  61. (img.shape[0] * img.shape[1])
  62. / (_TEXT_THICKNESS_SCALING * _TEXT_THICKNESS_SCALING)
  63. )
  64. * normalised_scaling
  65. )
  66. thickness = max(1, thickness)
  67. scaling = img.shape[0] / _TEXT_SCALING * normalised_scaling
  68. return cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, scaling, thickness)[0]
  69. def render_text(img, text, pos, color=(200, 200, 200), normalised_scaling=1.0):
  70. """
  71. Render a text into the image. Calculates scaling and thickness automatically.
  72. :param img: image to render into
  73. :param text: text to display
  74. :param pos: (x, y) - upper left coordinates of render position
  75. :param color: (b, g, r) - text color
  76. :param normalised_scaling: additional normalised scaling. Default 1.0.
  77. :return: updated image
  78. """
  79. x, y = pos
  80. thickness = int(
  81. round(
  82. (img.shape[0] * img.shape[1])
  83. / (_TEXT_THICKNESS_SCALING * _TEXT_THICKNESS_SCALING)
  84. )
  85. * normalised_scaling
  86. )
  87. thickness = max(1, thickness)
  88. scaling = img.shape[0] / _TEXT_SCALING * normalised_scaling
  89. size = get_text_size(img, text, normalised_scaling)
  90. cv2.putText(
  91. img,
  92. text,
  93. (int(x), int(y + size[1])),
  94. cv2.FONT_HERSHEY_SIMPLEX,
  95. scaling,
  96. color,
  97. thickness=thickness,
  98. )
  99. return img