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