Ultra fast lane detection 批量测试自己的图片

这部分我在网上找了很久，都没能找到，无奈只能自己改一下了。主要是修改demo.py的内容。没有直接修改源代码，而是修改了网上的测试自己的视频的代码
其实demo.py的思想很简单 就是先将视频截取成图片，然后进行预测，将预测出来是图片再合成视频。所以修改的思路就是，直接输入图片，然后直接保存预测后的图片。
修改 test_path(需要测试的图片的路径)和path(图片保存路径)就可以用了。
删除线格式
经过实践我发现 根据网上的过程修改完出来的效果是错的，事实证明不能投机取巧，所以我自己修改了源代码，效果很好。
其实认真看demo.py 会发现它的逻辑是从data_root的test.txt读取需要测试的图片的数据 然后测试图片 最后把所有图片合成视频
所以测试自己的图片的步骤很简单 将configs/tusimple.py(基于tusimple,我觉得基于那个模型都可以主要是看自己的图片尺寸更符合哪个模型) 的dataroot修改成自己的数据集

dataset='Tusimple'
data_root = 'img_highspeed/train_data'
# data_root = r'TUSIMPLEROOT'

train_data下的文件夹
别的文件是训练用的 我们主要在意的是test文件夹 和test.txt test文件夹里面存放要训练的文件 test.txt里存放图片的路径
test.txt的内容
这样修改后demo.py中的测试图片就可以由自己指定了

接下来就是生成图片 这一步就比较简单 还是之前的想法 在结果图片生成视频之前将图片保存下来
只需要加一个保存图片的语句

vis = cv2.imread(os.path.join(cfg.data_root,names[0]))
            for i in range(out_j.shape[1]):
                if np.sum(out_j[:, i] != 0) > 2:
                    for k in range(out_j.shape[0]):
                        if out_j[k, i] > 0:
                            ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1, int(img_h * (row_anchor[cls_num_per_lane-1-k]/288)) - 1 )
                            # print(ppp)
                            cv2.circle(vis,ppp,5,(0,255,0),-1)
            cv2.imwrite(r'D:\Desktop\result/' + names[0].split('/')[1],vis)

在图片完成之后将图片保存下来就可以 位置可以修改为自己的路径
效果还是很不错的哈
不过要达到这个效果还是要训练自己的模型

底下是完整代码

后面的部分 是我尝试得到mask图片的过程 可以参考以下 解除注释是能得到mask结果的 我尝试了将点连成线但是效果不佳 如果只检测一条车道线的话可以用
mask的结果

import torch, os, cv2
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print
import torch
import scipy.special, tqdm
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from data.constant import culane_row_anchor, tusimple_row_anchor

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

    args, cfg = merge_config()


    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 = parsingNet(pretrained = False, backbone=cfg.backbone,cls_dim = (cfg.griding_num+1,cls_num_per_lane,4),
                    use_aux=False).cuda() # we dont need auxiliary segmentation in testing

    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()

    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        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']
        datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, 'list/test_split/'+split),img_transform = img_transforms) for split in splits]
        img_w, img_h = 1640, 590
        row_anchor = culane_row_anchor
    elif cfg.dataset == 'Tusimple':
        splits = ['test.txt']
        datasets = [LaneTestDataset(cfg.data_root,os.path.join(cfg.data_root, split),img_transform = img_transforms) for split in splits]
        img_w, img_h = 1280, 720
        row_anchor = tusimple_row_anchor
    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()
            with torch.no_grad():
                out = net(imgs)

            col_sample = np.linspace(0, 800 - 1, cfg.griding_num)
            col_sample_w = col_sample[1] - col_sample[0]


            out_j = out[0].data.cpu().numpy()
            out_j = out_j[:, ::-1, :]
            prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
            idx = np.arange(cfg.griding_num) + 1
            idx = idx.reshape(-1, 1, 1)
            loc = np.sum(prob * idx, axis=0)
            out_j = np.argmax(out_j, axis=0)
            loc[out_j == cfg.griding_num] = 0
            out_j = loc

            # import pdb; pdb.set_trace()
            vis = cv2.imread(os.path.join(cfg.data_root,names[0]))
            for i in range(out_j.shape[1]):
                if np.sum(out_j[:, i] != 0) > 2:
                    for k in range(out_j.shape[0]):
                        if out_j[k, i] > 0:
                            ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1, int(img_h * (row_anchor[cls_num_per_lane-1-k]/288)) - 1 )
                            # print(ppp)
                            cv2.circle(vis,ppp,5,(0,255,0),-1)
            cv2.imwrite(r'D:\Desktop\result/' + names[0].split('/')[1],vis)

            # mask = np.zeros((img_h,img_w,3),dtype=np.uint8)
            # list = []
            # for i in range(out_j.shape[1]):
            #     if np.sum(out_j[:, i] != 0) > 2:
            #         for k in range(out_j.shape[0]):
            #             if out_j[k, i] > 0:
            #                 ppp =[int(out_j[k, i] * col_sample_w * img_w / 800) - 1, int(img_h * (row_anchor[cls_num_per_lane-1-k]/288)) - 1]
            # #                 list.append(ppp)
            # # pts = np.array(list,np.int32)
            # # pts.reshape((-1,1,2))
            # # cv2.polylines(mask,[pts],False,(255,255,255),5)
            #                 cv2.circle(mask, ppp, 5, (0, 255, 0), -1)
            # #
            # #                 # cv2.polylines(mask,ppp,False,(255,255,255),5)
            # cv2.imwrite(r'D:\Desktop\lane_mask/' + names[0].split('/')[1],mask)


            vout.write(vis)
        vout.release()

参考：
测试自己的视频（别人写的）