predict.py 844 B

12345678910111213141516171819202122232425262728
  1. import os
  2. import sys
  3. from glob import glob
  4. import cv2
  5. import torch
  6. sys.path.append(os.path.abspath(os.path.dirname(__file__)) + '/../../')
  7. from botr.yolov8.model import Predictor
  8. ROOT = os.path.abspath(os.path.dirname(__file__)) + '/../../'
  9. model_path = ROOT + 'botr/yolov8/weights.pt'
  10. image_size = 640
  11. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  12. def detect(image_np, predictor=None, show=False):
  13. if predictor is None:
  14. predictor = Predictor(image_size, device, model_path)
  15. result_list = predictor.predict(image_np, show=show)
  16. return result_list
  17. if __name__ == '__main__':
  18. p = r'C:\Users\Administrator\Desktop\test_b_table\real2.png'
  19. paths = glob(r'C:\Users\Administrator\Desktop\test_b_table\error10.png')
  20. for p in paths:
  21. img = cv2.imread(p)
  22. detect(img, show=True)