Ultra-Fast-Lane-Detection-v2-plus/demo.py at main · scott-mao/Ultra-Fast-Lane-Detection-v2-plus · 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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import torch, os, cv2
from utils.dist_utils import dist_print
import torch, os
from utils.common import merge_config, get_model
import tqdm
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset

def pred2coords(pred, row_anchor, col_anchor, local_width = 1, original_image_width = 1640, original_image_height = 590):
    batch_size, num_grid_row, num_cls_row, num_lane_row = pred['loc_row'].shape # 1, 200, 72, 4
    batch_size, num_grid_col, num_cls_col, num_lane_col = pred['loc_col'].shape # 1, 100, 81, 4

    max_indices_row = pred['loc_row'].argmax(1).cpu()
    # n , num_cls, num_lanes
    print(f'### max_indices_row: {max_indices_row.shape}; ') # 1, 200, 72, 4 => 1, 72, 4     200是anchor?
    print(f'    第一行：{max_indices_row[0][0][0]}：{pred["loc_row"][0][max_indices_row[0][0][0]][0][0]}、'
                '{max_indices_row[0][0][1]}：{pred["loc_row"][0][max_indices_row[0][0][1]][0][0]}、'
                '{max_indices_row[0][0][2]}：{pred["loc_row"][0][max_indices_row[0][0][2]][0][0]}、'
                '{max_indices_row[0][0][3]};：{pred["loc_row"][0][max_indices_row[0][0][3]][0][0]}')
    print(f'    第5行：{max_indices_row[0][5][0]}、{max_indices_row[0][5][1]}、{max_indices_row[0][5][2]}、{max_indices_row[0][5][3]};')
    print(f'    最后一行：{max_indices_row[0][71][0]}、{max_indices_row[0][71][1]}、{max_indices_row[0][71][2]}、{max_indices_row[0][71][3]};')
    valid_row = pred['exist_row'].argmax(1).cpu()
    # n, num_cls, num_lanes
    print(f'### valid_row: {valid_row.shape};') # 1, 2, 72, 4 => 1, 72, 4
    print(f'    第一行：{valid_row[0][0][0]}、{valid_row[0][0][1]}、{valid_row[0][0][2]}、{valid_row[0][0][3]};')
    print(f'    第5行：{valid_row[0][5][0]}、{valid_row[0][5][1]}、{valid_row[0][5][2]}、{valid_row[0][5][3]};')
    print(f'    最后一行：{valid_row[0][71][0]}、{valid_row[0][71][1]}、{valid_row[0][71][2]}、{valid_row[0][71][3]};')

    max_indices_col = pred['loc_col'].argmax(1).cpu()
    # n , num_cls, num_lanes
    valid_col = pred['exist_col'].argmax(1).cpu()
    # n, num_cls, num_lanes

    pred['loc_row'] = pred['loc_row'].cpu()
    pred['loc_col'] = pred['loc_col'].cpu()

    coords = []

    row_lane_idx = [1,2]
    col_lane_idx = [0,3]
    exit(0)

    for i in row_lane_idx:
        tmp = []
        if valid_row[0,:,i].sum() > num_cls_row / 2:
            for k in range(valid_row.shape[1]):
                if valid_row[0,k,i]:
                    all_ind = torch.tensor(list(range(max(0,max_indices_row[0,k,i] - local_width), min(num_grid_row-1, max_indices_row[0,k,i] + local_width) + 1)))

                    out_tmp = (pred['loc_row'][0,all_ind,k,i].softmax(0) * all_ind.float()).sum() + 0.5
                    out_tmp = out_tmp / (num_grid_row-1) * original_image_width
                    tmp.append((int(out_tmp), int(row_anchor[k] * original_image_height)))
            coords.append(tmp)

    for i in col_lane_idx:
        tmp = []
        if valid_col[0,:,i].sum() > num_cls_col / 4:
            for k in range(valid_col.shape[1]):
                if valid_col[0,k,i]:
                    all_ind = torch.tensor(list(range(max(0,max_indices_col[0,k,i] - local_width), min(num_grid_col-1, max_indices_col[0,k,i] + local_width) + 1)))

                    out_tmp = (pred['loc_col'][0,all_ind,k,i].softmax(0) * all_ind.float()).sum() + 0.5

                    out_tmp = out_tmp / (num_grid_col-1) * original_image_height
                    tmp.append((int(col_anchor[k] * original_image_width), int(out_tmp)))
            coords.append(tmp)

    return coords
if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True

    args, cfg = merge_config()
    cfg.batch_size = 1
    print('setting batch_size to 1 for demo generation')

    dist_print('start testing...')
    assert cfg.backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']

    if cfg.dataset == 'CULane':
        cls_num_per_lane = 18
    elif cfg.dataset == 'Tusimple':
        cls_num_per_lane = 56
    else:
        raise NotImplementedError

    net = get_model(cfg)

    state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
    compatible_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()

    export_onnx = False
    if export_onnx:
        # input torch.Size([1, 3, 320, 1600])
        print(f'### Generate onnx file ...')
        x = torch.randn(1, 3, 320, 1600).cuda()
        torch.onnx.export(net, x, "ufldv2.onnx", verbose=True, input_names=["input"], output_names=["output"],opset_version=12)
        i_debug = 1
        if 1 == i_debug:
            print(f'### onnx file generated! ^_^ ###')
            exit()


    img_transforms = transforms.Compose([
        transforms.Resize((int(cfg.train_height / cfg.crop_ratio), cfg.train_width)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    if cfg.dataset == 'CULane':
        #splits = ['test0_normal.txt', 'test1_crowd.txt', 'test2_hlight.txt', 'test3_shadow.txt', 'test4_noline.txt', 'test5_arrow.txt', 'test6_curve.txt', 'test7_cross.txt', 'test8_night.txt']
        splits = ['test6_curve.txt']
        datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, 'list/test_split/'+split),img_transform = img_transforms, crop_size = cfg.train_height) for split in splits]
        img_w, img_h = 1640, 590
    elif cfg.dataset == 'Tusimple':
        splits = ['test.txt']
        datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, split),img_transform = img_transforms, crop_size = cfg.train_height) for split in splits]
        img_w, img_h = 1280, 720
    else:
        raise NotImplementedError
    for split, dataset in zip(splits, datasets):
        loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle = False, num_workers=1)
        fourcc = cv2.VideoWriter_fourcc(*'MJPG')
        print(split[:-3]+'avi')
        vout = cv2.VideoWriter(split[:-3]+'avi', fourcc , 30.0, (img_w, img_h))
        for i, data in enumerate(tqdm.tqdm(loader)):
            imgs, names = data
            imgs = imgs.cuda()
            print(f'### imgs: {imgs.shape};')
            with torch.no_grad():
                pred = net(imgs)
            for k, v in pred.items():
                print(f'### {k}: {v.shape};')

            vis = cv2.imread(os.path.join(cfg.data_root,names[0]))
            coords = pred2coords(pred, cfg.row_anchor, cfg.col_anchor, original_image_width = img_w, original_image_height = img_h)
            for lane in coords:
                for coord in lane:
                    cv2.circle(vis,coord,5,(0,255,0),-1)
            vout.write(vis)

        vout.release()