draw.py 1.1 KB

123456789101112131415161718192021222324252627282930313233343536
  1. import numpy as np
  2. import cv2
  3. palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
  4. def compute_color_for_labels(label):
  5. """
  6. Simple function that adds fixed color depending on the class
  7. """
  8. color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
  9. return tuple(color)
  10. def draw_boxes(img, bbox, identities=None, offset=(0,0)):
  11. for i,box in enumerate(bbox):
  12. x1,y1,x2,y2 = [int(i) for i in box]
  13. x1 += offset[0]
  14. x2 += offset[0]
  15. y1 += offset[1]
  16. y2 += offset[1]
  17. # box text and bar
  18. id = int(identities[i]) if identities is not None else 0
  19. color = compute_color_for_labels(id)
  20. label = '{}{:d}'.format("", id)
  21. t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 2 , 2)[0]
  22. cv2.rectangle(img,(x1, y1),(x2,y2),color,3)
  23. cv2.rectangle(img,(x1, y1),(x1+t_size[0]+3,y1+t_size[1]+4), color,-1)
  24. cv2.putText(img,label,(x1,y1+t_size[1]+4), cv2.FONT_HERSHEY_PLAIN, 2, [255,255,255], 2)
  25. return img
  26. if __name__ == '__main__':
  27. for i in range(82):
  28. print(compute_color_for_labels(i))