from ultralytics import YOLO

import glob
import os.path as osp
import argparse
import cv2

def pltoresult(jsonfilePath):
    # Load a model

    model = YOLO('best.pt')  # load a pretrained model (recommended for training)
    # Train the model

        # model.train(data='NEU_train.yaml', epochs=150, imgsz=224, batch=64, workers=0, optimizer='s6D', device='0')# model_name = uuid.uuid1().-_str__()
    # result = model.predict(source="3333.png",save_txt=True, device='0', save=True)

    #方法1
    jsonfileList = glob.glob(osp.join(jsonfilePath, "*.png"))
    for jsonfile in jsonfileList:
        result = model.predict(source=jsonfile, save_txt=True, device='0', save=True)

    # 方法2
    # files = os.listdir(jsonfilePath)
    # for file in files:
    #     result = model.predict(source=file, save_txt=True, device='0', save=True)

if __name__ == "__main__":
    # jsonfilePath = "../test/"
    # pltoresult(jsonfilePath)

    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default='yolov5s.pt', help='input file')
    args = parser.parse_args()
    pltoresult(args.input)
