import json
import time
from ultralytics import YOLO
import os
import shutil

def inference(model_path, path_file):

    model = YOLO(model_path)
    model.predict(source="./data/datasets/train/images/",line_thickness=1, save_txt=True, save=True)
    path_file = './runs/segment/predict/labels/'
    label_path = os.listdir(path_file)
    print(label_path)
    res = []
    for img_name in label_path:
        img = img_name.split('.')
        txt_path = path_file + img[0] + '.txt'
        with open(txt_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in lines:
                val = line.split(' ')
                res.append(int(val[0]))
    res_json = {'results': res}
    return res_json

res_json = inference('./runs/segment/train/weights/best.pt','./data/datasets/train/images')
json_data = json.dumps(res_json)
file = open('./result.json', 'w', encoding='utf-8')
file.write(json_data)
file.close()
print("json文件写入完毕")

shutil.rmtree('./runs/segment/predict')