CLRNet 车道线检测

import torch
import cv2
from clrnet.engine.runner import Runner
from clrnet.utils.config import Config
import cv2
import os
import time
import torch.backends.cudnn as cudnn

cfg = Config.fromfile("configs/clrnet/clr_resnet18_culane.py")
cfg.gpus = 1
cfg.load_from = 'culane_r18.pth'
cfg.resume_from = ''
cfg.finetune_from = ''
cfg.seed = 0

def pre_process_image(path):
    img_w = 800
    img_h = 320
    cut_height = 270 
    image = cv2.imread(path)
    image = image[cut_height:, :, :]
    image = cv2.resize(image, (img_w, img_h))
    image = torch.Tensor([image])
    image = image.permute(0, 3, 1, 2)
    return image

def view(predictions):
        for lanes in predictions:
            img = cv2.imread("/root/CLRNet/v2-b40621ccca1d9a432adf1442dcc69540_b.jpg")
            out_file = "1.jpg"
            lanes = [lane.to_array(cfg) for lane in lanes]
            imshow_lanes(img, lanes, out_file=out_file)




COLORS = [
    (255, 0, 0),(0, 255, 0), (0, 0, 255),(255, 255, 0),(255, 0, 255),(0, 255, 255),(128, 255, 0),(255, 128, 0),(128, 0, 255),
    (255, 0, 128),(0, 128, 255),(0, 255, 128),(128, 255, 255),(255, 128, 255),(255, 255, 128),(60, 180, 0),(180, 60, 0),(0, 60, 180),
    (0, 180, 60),(60, 0, 180),(180, 0, 60),(255, 0, 0),(0, 255, 0),(0, 0, 255),(255, 255, 0),(255, 0, 255),(0, 255, 255),
    (128, 255, 0),(255, 128, 0),(128, 0, 255)]


def imshow_lanes(img, lanes, show=False, out_file=None, width=4):
    lanes_xys = []
    for _, lane in enumerate(lanes):
        xys = []
        for x, y in lane:
            if x <= 0 or y <= 0:
                continue
            x, y = int(x), int(y)
            xys.append((x, y))
        lanes_xys.append(xys)
    lanes_xys.sort(key=lambda xys : xys[0][0])

    for idx, xys in enumerate(lanes_xys):
        for i in range(1, len(xys)):
            cv2.line(img, xys[i - 1], xys[i], COLORS[idx], thickness=width)


    if show:
        cv2.imshow('view', img)
        cv2.waitKey(0)

    if out_file:
        cv2.imwrite(out_file, img)

t0=time.time()
cudnn.benchmark=True 
runner = Runner(cfg)
runner.net.eval()
predictions = []
image = pre_process_image("/root/CLRNet/v2-b40621ccca1d9a432adf1442dcc69540_b.jpg")
data = image.cuda()
with torch.no_grad():
    output = runner.net(data)
    output = runner.net.module.heads.get_lanes(output)
    print(output)
    t1=time.time()
    view(output)
    print((t1-t0))

想用最近的车道线检测算法 CLRNet 测试自己的本地图片,这是根据源码整合的简单的推理代码,但是output的输出是个空列表,有没有大佬知道这是怎么回事?是不是我忘记了处理什么地方。

请各位大佬给点意见!!!

你可能感兴趣的:(大数据,目标检测,计算机视觉,pytorch)