29 lines
898 B
Python
29 lines
898 B
Python
|
|
from ultralytics import YOLO
|
|
|
|
import os
|
|
import shutil
|
|
|
|
|
|
model_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/HelmetHeadBAC/trainHardHatsV1HelmetHeadShoe2/weights/best.pt'
|
|
image_path = '/home/admin-root/haotian/锻8/tensorrtx/yolov8/images_20250624105811'
|
|
output_path = '/home/admin-root/haotian/python哈汽锻8安全帽识别/export/duan8_11/images'
|
|
|
|
image_names = os.listdir(image_path)
|
|
target_list = [1]
|
|
|
|
model = YOLO(model_path)
|
|
|
|
all_image_path = [os.path.join(image_path, t) for t in image_names]
|
|
|
|
for i in range(len(all_image_path)):
|
|
results = model(all_image_path[i], conf=0.3, device=0)
|
|
if len(results[0].boxes.cls) > 0:
|
|
for c in results[0].boxes.cls:
|
|
if c.item() in target_list:
|
|
print(all_image_path[i])
|
|
shutil.copy(all_image_path[i], os.path.join(output_path, image_names[i]))
|
|
break
|
|
|
|
|
|
# |