demo.py

"""
测试小图 并显示出来 按q键结束显示
"""
import torch
import argparse
from utils import read_yml
from road_extractor import build_road_extractor


def parse_args():
    parser = argparse.ArgumentParser(description='Analysis road extractor model')
    parser.add_argument('--config', default='configs/LRDNet_RNBD.yml', help='train config file path')
    parser.add_argument('--checkpoints', type=str, default='./work_dir/2022-10-12-19_50_10/model99_resume.pth', help='the dir of checkpoints')
    parser.add_argument('--path_img', type=str, default='data\\RNBD\\test\img\\4_2.png', help='the path of image')

    args = parser.parse_args()

    return args


def main():
    args = parse_args()
    cfg = read_yml(args.config)
    model = build_road_extractor(cfg_model=cfg['model']).cuda()
    model.load_state_dict(torch.load(args.checkpoints, map_location='cuda:0')['state_dict'])

    model.forward_demo(args.path_img)


if __name__ == '__main__':
    main()