You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
from add_xml import add_xml
|
|
from create_xml import create_xml
|
|
from xz_yolov8_atm.yolov8_det import analysis_yolov8
|
|
import glob
|
|
import cv2
|
|
import os
|
|
from ultralytics import YOLO
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
imgpath = 'E:/BANK_xz_all/imgpath'
|
|
xmlpath = 'E:/BANK_xz_all/xz_yolov8_atm/output'
|
|
|
|
imglist = os.listdir(imgpath)
|
|
xmllist = os.listdir(xmlpath)
|
|
|
|
model_all = YOLO('xz_yolov8_atm/model_files/best(3)_all.pt')
|
|
|
|
for img in imglist:
|
|
|
|
images_path = os.path.join(imgpath,img)
|
|
print(images_path)
|
|
xml_path = os.path.join(xmlpath,img.split('.')[0] + '.xml')
|
|
print(xml_path)
|
|
|
|
images = cv2.imread(images_path)
|
|
# print(images)
|
|
|
|
results = analysis_yolov8(images=images,
|
|
model_coco=model_all,
|
|
confidence=0.5)
|
|
|
|
w,h,d = images.shape
|
|
img_shape = (w,h,d,img)
|
|
|
|
if os.path.isfile(xml_path):
|
|
|
|
add_labels = add_xml(inforsDict=results,
|
|
xmlFilePath=xml_path)
|
|
else:
|
|
create_new = create_xml(boxs=results,
|
|
img_shape=img_shape,
|
|
xml_path=xmlpath) |