You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

pretrained_yolov4.py 931 B

4 years ago
12345678910111213141516171819202122232425262728
  1. import numpy as np
  2. import cv2
  3. from PIL import Image
  4. from examples.model_zoo.common import yolo4_input_processing, yolo4_output_processing, \
  5. result_to_json, read_class_names, draw_boxes_and_labels_to_image_with_json
  6. from examples.model_zoo.yolo import YOLOv4
  7. import tensorlayer as tl
  8. tl.logging.set_verbosity(tl.logging.DEBUG)
  9. INPUT_SIZE = 416
  10. image_path = './data/kite.jpg'
  11. class_names = read_class_names('./model/coco.names')
  12. original_image = cv2.imread(image_path)
  13. image = cv2.cvtColor(np.array(original_image), cv2.COLOR_BGR2RGB)
  14. model = YOLOv4(NUM_CLASS=80, pretrained=True)
  15. model.set_eval()
  16. batch_data = yolo4_input_processing(original_image)
  17. feature_maps = model(batch_data)
  18. pred_bbox = yolo4_output_processing(feature_maps)
  19. json_result = result_to_json(image, pred_bbox)
  20. image = draw_boxes_and_labels_to_image_with_json(image, json_result, class_names)
  21. image = Image.fromarray(image.astype(np.uint8))
  22. image.show()

TensorLayer3.0 是一款兼容多种深度学习框架为计算后端的深度学习库。计划兼容TensorFlow, Pytorch, MindSpore, Paddle.