yolo26-face/test_widerface.py at main · derronqi/yolo26-face · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import os
import argparse
from ultralytics import YOLO

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/pose/yolo26s-face/weights/best.pt', help='model.pt path(s)')
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.01, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
    parser.add_argument('--device', type=str, default='cuda:0', help='augmented inference')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--e2e', action='store_true', help='NMS-free: use the one2one head (default is the standard NMS protocol, comparable with published baselines)')
    parser.add_argument('--save_folder', default='./widerface_evaluate/widerface_txt/', type=str, help='Dir to save txt results')
    parser.add_argument('--dataset_folder', default='./data/widerface/val/images/', type=str, help='dataset path')
    opt = parser.parse_args()
    print(opt)

    model = YOLO(opt.weights)
    if opt.e2e:
        model.model.end2end = True  # route dual-head models through the one2one (NMS-free) path

    # testing dataset
    testset_folder = opt.dataset_folder
    # locate wider_val.txt next to the image folder (supports both "<val>/images/" and flat "<val>/")
    folder = testset_folder.rstrip('/')
    base = os.path.dirname(folder) if os.path.basename(folder) == 'images' else folder
    testset_list = os.path.join(base, "wider_val.txt")
    with open(testset_list, 'r') as fr:
        test_dataset = fr.read().split()
        num_images = len(test_dataset)
    for img_name in test_dataset:
        # val images are stored flat (event prefix kept in img_name for the prediction subfolder)
        image_path = os.path.join(testset_folder, os.path.basename(img_name))
        # 8.4.143: predict defaults to NMS-free for dual-head models; nms=True restores the standard NMS protocol
        results = model.predict(source=image_path, imgsz=opt.img_size, conf=opt.conf_thres, iou=opt.iou_thres, augment=opt.augment, device=opt.device, max_det=2048, nms=not opt.e2e)

        save_name = opt.save_folder + img_name[:-4] + ".txt"
        dirname = os.path.dirname(save_name)
        if not os.path.isdir(dirname):
            os.makedirs(dirname)
        with open(save_name, "w") as fd:
            result = results[0].cpu().numpy()
            file_name = os.path.basename(save_name)[:-4] + "\n"
            bboxs_num = str(result.boxes.shape[0]) + '\n'
            fd.write(file_name)
            fd.write(bboxs_num)
            for box in result.boxes:
                conf = box.conf[0]
                cls  = box.cls[0]
                xyxy = box.xyxy[0]
                x1 = int(xyxy[0] + 0.5)
                y1 = int(xyxy[1] + 0.5)
                x2 = int(xyxy[2] + 0.5)
                y2 = int(xyxy[3] + 0.5)
                fd.write('%d %d %d %d %.04f' % (x1, y1, x2-x1, y2-y1, conf if conf <= 1 else 1) + '\n')