12345678910111213141516171819202122232425262728 |
- import os
- import sys
- from glob import glob
- import cv2
- import torch
- sys.path.append(os.path.abspath(os.path.dirname(__file__)) + '/../../')
- from botr.yolov8.model import Predictor
- ROOT = os.path.abspath(os.path.dirname(__file__)) + '/../../'
- model_path = ROOT + 'botr/yolov8/weights.pt'
- image_size = 640
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- def detect(image_np, predictor=None, show=False):
- if predictor is None:
- predictor = Predictor(image_size, device, model_path)
- result_list = predictor.predict(image_np, show=show)
- return result_list
- if __name__ == '__main__':
- p = r'C:\Users\Administrator\Desktop\test_b_table\real2.png'
- paths = glob(r'C:\Users\Administrator\Desktop\test_b_table\error10.png')
- for p in paths:
- img = cv2.imread(p)
- detect(img, show=True)
|