get_point.py 848 B

1234567891011121314151617181920212223242526
  1. import cv2
  2. import numpy as np
  3. import pyyolo
  4. names_filepath = os.path.sep.join(['yolo-config', "bot.names"])
  5. weights_filepath = os.path.sep.join(['yolo-config', "yolov3-tiny-bot_2000.weights"])
  6. cfg_filepath = os.path.sep.join(['yolo-config', "yolov3-tiny-bot.cfg"])
  7. image_filepath = './1.png'
  8. meta = pyyolo.load_names(names_filepath)
  9. net = pyyolo.load_net(cfg_filepath, weights_filepath, False)
  10. im = cv2.imread(image_filepath)
  11. yolo_img = pyyolo.array_to_image(im)
  12. res = pyyolo.detect(net, meta, yolo_img)
  13. colors = np.random.rand(meta.classes, 3) * 255
  14. for r in res:
  15. cv2.rectangle(im, r.bbox.get_point(pyyolo.BBox.Location.TOP_LEFT, is_int=True),
  16. r.bbox.get_point(pyyolo.BBox.Location.BOTTOM_RIGHT, is_int=True), tuple(colors[r.id].tolist()), 2)
  17. cv2.imshow('Frame', im)
  18. cv2.waitKey(0)