from ultralytics import YOLO
import os
def transform_yolo(imgs_path, save_path, width=640, height=640):
img_list = model(imgs_path, stream=True)
# img_path_list
img_path_list = os.listdir(imgs_path)
loc_list = []
save_txt = []
len_list = 0
for loc in img_path_list:
path = imgs_path + loc
txt_path = save_path + loc.split(".")[0] + ".txt"
loc_list.append(path)
save_txt.append(txt_path)
for (num_img, path) in zip(img_list, save_txt):
mid = len(num_img)
with open(path, "a", encoding="utf-8") as f:
for i in num_img.boxes.data.tolist():
len_list += 1
x = "{:.6f}".format((i[0] + i[2]) / 2 / width)
y = "{:.6f}".format((i[1] + i[3]) / 2 / height)
w = "{:.6f}".format((i[2] - i[0]) / width)
h = "{:.6f}".format((i[3] - i[1]) / height)
class_id = int(i[5])
f.write(f"{class_id} {x} {y} {w} {h}\n")
if len_list == mid:
len_list = 0
continue
# 模型路径
model_path = r""
# 图片所在文件夹路径
imgs_path = r""
# yolo格式的txt文件保存路径
save_path = r""
# 加载模型
model = YOLO(model_path)
transform_yolo(imgs_path, save_path, width=640, height=640)