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: