def init_model(self):
torch.backends.cudnn.benchmark = True
self.args, self.cfg = merge_config()
assert self.cfg.backbone in ['18','34','50','101','152','50next','101next','50wide','101wide']
if self.cfg.dataset == 'CULane':
cls_num_per_lane = 18
elif self.cfg.dataset == 'Tusimple':
cls_num_per_lane = 56
else:
raise NotImplementedError
self.net = parsingNet(pretrained = False, backbone=self.cfg.backbone,cls_dim = (self.cfg.griding_num+1,cls_num_per_lane,4),
use_aux=False).cuda() # we dont need auxiliary segmentation in testing
state_dict = torch.load(self.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
self.net.load_state_dict(compatible_state_dict, strict=False)
self.net.eval()
self.img_transforms = transforms.Compose([
transforms.Resize((288, 800)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
def __init__(self):
torch.backends.cudnn.benchmark = True
self.args, self.cfg = merge_config()
self.cls_num_per_lane = 56
self.row_anchor = tusimple_row_anchor
self.net = parsingNet(pretrained=False,
backbone=self.cfg.backbone,
cls_dim=(self.cfg.griding_num + 1,
self.cls_num_per_lane,
self.cfg.num_lanes),
use_aux=False).cuda()
state_dict = torch.load(self.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
self.net.load_state_dict(compatible_state_dict, strict=False)
#not recommend to uncommen this line
self.net.eval()
self.img_transforms = transforms.Compose([
transforms.Resize((288, 800)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
self.img_w = 960
self.img_h = 480
self.scale_factor = 1
self.color = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
self.idx = np.arange(self.cfg.griding_num) + 1
self.idx = self.idx.reshape(-1, 1, 1)
self.cpu_img = None
self.gpu_img = None
self.type = None
self.gpu_output = None
self.cpu_output = None
col_sample = np.linspace(0, 800 - 1, self.cfg.griding_num)
self.col_sample_w = col_sample[1] - col_sample[0]
from utils.common import merge_config
from utils.dist_utils import dist_print
import scipy.special, tqdm
import numpy as np
import torch.nn.functional as F
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from sklearn.linear_model import RANSACRegressor
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
if __name__ == "__main__":
torch.backends.cudnn.benchmark = True
args, cfg = merge_config()
dist_print('start testing...')
net = E2ENet(Channels=96,
nums_lane=4,
culomn_channels=cfg.griding_num,
row_channels=cfg.row_num,
initialed=True).cuda()
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:
