Windows下LSTR测试的复现

LSTR项目地址：https://github.com/liuruijin17/LSTR

一、基础配置

本机配置为RTX3050，CUDA11.2，CUDNN8.1.0，Pytorch1.8.1（Pytorch没有CUDA11.2版本的，可以使用适配低于11.2版本的Pytorch）

官方要求Pytorch只要高于1.5就行了

1.创建虚拟环境

conda env create --n lstr python=3.8

激活虚拟环境

conda activate lstr

安装依赖项

pip install -r requirements.txt

可以测试一下看一下有没有缺什么包，缺什么包就补什么包

python test.py LSTR --testiter 500000 --modality eval --split testing --debug --debugDec

二、数据集下载

测试使用的数据集为TuSimple数据集，下载链接如下：
https://github.com/TuSimple/tusimple-benchmark/issues/3

下载train_set.zip（约10G）,test_set.zip（约10G）以及test_label.json

将数据集解压放在LSTR-main\cache\TuSimple\LaneDetection文件夹下，如下：

三、图片测试

在目录下创建一个新的文件夹命名为images，然后将要测试的图片放进该文件夹，执行命令

python test.py LSTR --testiter 500000 --modality images --image_root ./ --debug

在detections目录下就可以查看检测的结果啦

四、视频测试

新建video_test.py，将test.py里的代码复制过来，修改如下部分：
1.将image_root改为video_root（就是个参数名，改不改随便，主要是怕乱套了）

def parse_args():
    parser = argparse.ArgumentParser(description="Test CornerNet")
    parser.add_argument("cfg_file", help="config file", type=str)
    parser.add_argument("--testiter", dest="testiter",
                        help="test at iteration i",
                        default=None, type=int)
    parser.add_argument("--split", dest="split",
                        help="which split to use",
                        default="validation", type=str)
    parser.add_argument("--suffix", dest="suffix", default=None, type=str)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--modality", dest="modality",
                        default=None, type=str)
    parser.add_argument("--batch", dest='batch',
                        help="select a value to maximum your FPS",
                        default=1, type=int)
    parser.add_argument("--debugEnc", action="store_true")
    parser.add_argument("--debugDec", action="store_true")
    parser.add_argument("--video_root", dest="video_root",
                        default=None, type=str)
    args = parser.parse_args()
    return args

2.修改test函数

def test(db, testiter,
         debug=False, suffix=None, modality=None, batch=1,
         debugEnc=False, debugDec=False, video_root=None):

注释掉

# result_dir = os.path.join(result_dir, str(testiter), split)

# if suffix is not None:
# result_dir = os.path.join(result_dir, suffix)

# make_dirs([result_dir])

再赋值

result_dir = "./result\\LSTR\\500000\\validation"

添加

elif modality == 'video':
    video_root = args.video_root
    if video_root == None:
        raise ValueError('--video_root is not defined!')
    print("processing [video]...")
    test_file = "test.video"
    video_testing = importlib.import_module(test_file).testing
    video_testing(db, nnet, video_root, debug=debug, evaluator=None)

3.在\LSTR-main\test下新建video.py

import os
import torch
import cv2
import json
import time
import numpy as np
from torch.autograd import Variable
import torch.nn.functional as F

from torch import nn
import matplotlib.pyplot as plt
from copy import deepcopy
from tqdm import tqdm
from config import system_configs

from utils import crop_image, normalize_

from sample.vis import *


class PostProcess(nn.Module):
    @torch.no_grad()
    def forward(self, outputs, target_sizes):
        out_logits, out_curves = outputs['pred_logits'], outputs['pred_curves']
        assert len(out_logits) == len(target_sizes)
        assert target_sizes.shape[1] == 2
        prob = F.softmax(out_logits, -1)
        scores, labels = prob.max(-1)
        labels[labels != 1] = 0
        results = torch.cat([labels.unsqueeze(-1).float(), out_curves], dim=-1)

        return results

def kp_detection(db, nnet, video_root, debug=False, evaluator=None):
    input_size  = db.configs["input_size"]  # [h w]
    image_dir = video_root

    postprocessors = {'bbox': PostProcess()}
    cap = cv2.VideoCapture(image_dir)
    while (cap.isOpened()):
        ret, image = cap.read()
        height, width = image.shape[0:2]

        images = np.zeros((1, 3, input_size[0], input_size[1]), dtype=np.float32)
        masks = np.ones((1, 1, input_size[0], input_size[1]), dtype=np.float32)
        orig_target_sizes = torch.tensor(input_size).unsqueeze(0).cuda()
        pad_image     = image.copy()
        pad_mask      = np.zeros((height, width, 1), dtype=np.float32)
        resized_image = cv2.resize(pad_image, (input_size[1], input_size[0]))
        resized_mask  = cv2.resize(pad_mask, (input_size[1], input_size[0]))
        masks[0][0]   = resized_mask.squeeze()
        resized_image = resized_image / 255.
        normalize_(resized_image, db.mean, db.std)
        resized_image = resized_image.transpose(2, 0, 1)
        images[0]     = resized_image
        images        = torch.from_numpy(images).cuda(non_blocking=True)
        masks         = torch.from_numpy(masks).cuda(non_blocking=True)
        torch.cuda.synchronize(0)  # 0 is the GPU id
        t0            = time.time()
        outputs, weights      = nnet.test([images, masks])
        torch.cuda.synchronize(0)  # 0 is the GPU id
        t             = time.time() - t0
        results = postprocessors['bbox'](outputs, orig_target_sizes)
        if evaluator is not None:
            evaluator.add_prediction(1, results.cpu().numpy(), t)

        if debug:
            pred = results[0].cpu().numpy()
            img  = pad_image
            img_h, img_w, _ = img.shape
            pred = pred[pred[:, 0].astype(int) == 1]
            overlay = img.copy()
            color = (0, 255, 0)
            for i, lane in enumerate(pred):
                lane = lane[1:]  # remove conf
                lower, upper = lane[0], lane[1]
                lane = lane[2:]  # remove upper, lower positions

                # generate points from the polynomial
                ys = np.linspace(lower, upper, num=100)
                points = np.zeros((len(ys), 2), dtype=np.int32)
                points[:, 1] = (ys * img_h).astype(int)
                points[:, 0] = ((lane[0] / (ys - lane[1]) ** 2 + lane[2] / (ys - lane[1]) + lane[3] + lane[4] * ys -
                                    lane[5]) * img_w).astype(int)
                points = points[(points[:, 0] > 0) & (points[:, 0] < img_w)]

                # draw lane with a polyline on the overlay
                for current_point, next_point in zip(points[:-1], points[1:]):
                    overlay = cv2.line(overlay, tuple(current_point), tuple(next_point), color=color, thickness=15)
                    # draw lane ID
                if len(points) > 0:
                    cv2.putText(img, str(i), tuple(points[0]), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
                                color=color,
                                thickness=3)
            # Add lanes overlay
            w = 0.6
            img = ((1. - w) * img + w * overlay).astype(np.uint8)
            # cv2.namedWindow("resized", 0)  # 0可以改变窗口大小了
            cv2.imshow("resized", img)  # 显示视频
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

                # cv2.imwrite(os.path.join(result_dir, image_names[ind][:-4] + '.jpg'), img)
    cap.release
    cv2.destroyAllWindows()


    return 0

def testing(db, nnet, video_root, debug=False, evaluator=None):
    return globals()[system_configs.sampling_function](db, nnet, video_root, debug=debug, evaluator=evaluator)

4.测试

python videotest.py LSTR --testiter 500000 --modality video --video_root path\to\video --debug