k20hcmus/FishEye8K

1import argparse2import platform3import sys4import time5from pathlib import Path6 7import pandas as pd8 9FILE = Path(__file__).resolve()10ROOT = FILE.parents[0]  # YOLO root directory11if str(ROOT) not in sys.path:12    sys.path.append(str(ROOT))  # add ROOT to PATH13# ROOT = ROOT.relative_to(Path.cwd())  # relative14 15import export16from models.experimental import attempt_load17from models.yolo import SegmentationModel18from segment.val import run as val_seg19from utils import notebook_init20from utils.general import LOGGER, check_yaml, file_size, print_args21from utils.torch_utils import select_device22from val import run as val_det23 24 25def run(26        weights=ROOT / 'yolo.pt',  # weights path27        imgsz=640,  # inference size (pixels)28        batch_size=1,  # batch size29        data=ROOT / 'data/coco.yaml',  # dataset.yaml path30        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu31        half=False,  # use FP16 half-precision inference32        test=False,  # test exports only33        pt_only=False,  # test PyTorch only34        hard_fail=False,  # throw error on benchmark failure35):36    y, t = [], time.time()37    device = select_device(device)38    model_type = type(attempt_load(weights, fuse=False))  # DetectionModel, SegmentationModel, etc.39    for i, (name, f, suffix, cpu, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, CPU, GPU)40        try:41            assert i not in (9, 10), 'inference not supported'  # Edge TPU and TF.js are unsupported42            assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13'  # CoreML43            if 'cpu' in device.type:44                assert cpu, 'inference not supported on CPU'45            if 'cuda' in device.type:46                assert gpu, 'inference not supported on GPU'47 48            # Export49            if f == '-':50                w = weights  # PyTorch format51            else:52                w = export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]  # all others53            assert suffix in str(w), 'export failed'54 55            # Validate56            if model_type == SegmentationModel:57                result = val_seg(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)58                metric = result[0][7]  # (box(p, r, map50, map), mask(p, r, map50, map), *loss(box, obj, cls))59            else:  # DetectionModel:60                result = val_det(data, w, batch_size, imgsz, plots=False, device=device, task='speed', half=half)61                metric = result[0][3]  # (p, r, map50, map, *loss(box, obj, cls))62            speed = result[2][1]  # times (preprocess, inference, postprocess)63            y.append([name, round(file_size(w), 1), round(metric, 4), round(speed, 2)])  # MB, mAP, t_inference64        except Exception as e:65            if hard_fail:66                assert type(e) is AssertionError, f'Benchmark --hard-fail for {name}: {e}'67            LOGGER.warning(f'WARNING ⚠️ Benchmark failure for {name}: {e}')68            y.append([name, None, None, None])  # mAP, t_inference69        if pt_only and i == 0:70            break  # break after PyTorch71 72    # Print results73    LOGGER.info('\n')74    parse_opt()75    notebook_init()  # print system info76    c = ['Format', 'Size (MB)', 'mAP50-95', 'Inference time (ms)'] if map else ['Format', 'Export', '', '']77    py = pd.DataFrame(y, columns=c)78    LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)')79    LOGGER.info(str(py if map else py.iloc[:, :2]))80    if hard_fail and isinstance(hard_fail, str):81        metrics = py['mAP50-95'].array  # values to compare to floor82        floor = eval(hard_fail)  # minimum metric floor to pass83        assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: mAP50-95 < floor {floor}'84    return py85 86 87def test(88        weights=ROOT / 'yolo.pt',  # weights path89        imgsz=640,  # inference size (pixels)90        batch_size=1,  # batch size91        data=ROOT / 'data/coco128.yaml',  # dataset.yaml path92        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu93        half=False,  # use FP16 half-precision inference94        test=False,  # test exports only95        pt_only=False,  # test PyTorch only96        hard_fail=False,  # throw error on benchmark failure97):98    y, t = [], time.time()99    device = select_device(device)100    for i, (name, f, suffix, gpu) in export.export_formats().iterrows():  # index, (name, file, suffix, gpu-capable)101        try:102            w = weights if f == '-' else \103                export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1]  # weights104            assert suffix in str(w), 'export failed'105            y.append([name, True])106        except Exception:107            y.append([name, False])  # mAP, t_inference108 109    # Print results110    LOGGER.info('\n')111    parse_opt()112    notebook_init()  # print system info113    py = pd.DataFrame(y, columns=['Format', 'Export'])114    LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)')115    LOGGER.info(str(py))116    return py117 118 119def parse_opt():120    parser = argparse.ArgumentParser()121    parser.add_argument('--weights', type=str, default=ROOT / 'yolo.pt', help='weights path')122    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')123    parser.add_argument('--batch-size', type=int, default=1, help='batch size')124    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')125    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')126    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')127    parser.add_argument('--test', action='store_true', help='test exports only')128    parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')129    parser.add_argument('--hard-fail', nargs='?', const=True, default=False, help='Exception on error or < min metric')130    opt = parser.parse_args()131    opt.data = check_yaml(opt.data)  # check YAML132    print_args(vars(opt))133    return opt134 135 136def main(opt):137    test(**vars(opt)) if opt.test else run(**vars(opt))138 139 140if __name__ == "__main__":141    opt = parse_opt()142    main(opt)143