def preload_lanemarking(weights_path):
net = SCNN(pretrained=False)
mean = (0.3598, 0.3653, 0.3662)
std = (0.2573, 0.2663, 0.2756)
transform = Compose(Resize((800, 288)), ToTensor(),
Normalize(mean=mean, std=std))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
save_dict = torch.load(weights_path, map_location='cpu')
net.load_state_dict(save_dict['net'])
net.eval()
net.to(device)
return {'net': net, 'transform': transform, 'device': device}
def main():
args = parse_args()
img_path = args.img_path
weight_path = args.weight_path
mode = args.mode
colors = Phoenix.get_colors(mode)
colors = np.array(colors)
num_classes = colors.shape[0]
print(num_classes)
net = SCNN(input_size=(512, 384), pretrained=False, seg_classes=num_classes, weights=Phoenix.get_weights(mode))
mean=(0.485, 0.456, 0.406)
std=(0.229, 0.224, 0.225)
transform_img = Compose(Resize((512, 384)), Rotation(2))
transform_to_net = Compose(ToTensor(), Normalize(mean=mean, std=std))
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = transform_img({'img': img})['img']
x = transform_to_net({'img': img})['img']
x.unsqueeze_(0)
save_dict = torch.load(weight_path, map_location='cpu')
net.load_state_dict(save_dict['net'])
net.eval()
seg_pred, exist_pred = net(x)[:2]
seg_pred = seg_pred.detach().cpu().numpy()
exist_pred = exist_pred.detach().cpu().numpy()
seg_pred = seg_pred[0]
exist = [1 if exist_pred[0, i] > 0.5 else 0 for i in range(num_classes)]
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
lane_img = np.zeros_like(img)
coord_mask = np.argmax(seg_pred, axis=0)
for i in range(0, num_classes):
#if exist_pred[0, i] > 0.5:
lane_img[coord_mask == (i + 1)] = colors[i]
img = cv2.addWeighted(src1=lane_img, alpha=0.8, src2=img, beta=1., gamma=0.)
cv2.imwrite("demo/demo_result.jpg", img)
"""for x in getLane.prob2lines_CULane(seg_pred, exist):
print(x)"""
if args.visualize:
print([1 if exist_pred[0, i] > 0.5 else 0 for i in range(num_classes)])
cv2.imshow("", img)
cv2.waitKey(0)
cv2.destroyAllWindows()
Normalize(mean=mean, std=std))
dataset_name = exp_cfg['dataset'].pop('dataset_name')
Dataset_Type = getattr(dataset, dataset_name)
test_dataset = Dataset_Type(Dataset_Path['Tusimple'], "test", transform)
test_loader = DataLoader(test_dataset,
batch_size=32,
collate_fn=test_dataset.collate,
num_workers=4)
net = SCNN(input_size=resize_shape, pretrained=False)
save_name = os.path.join(exp_dir, exp_dir.split('/')[-1] + '_best.pth')
save_dict = torch.load(save_name, map_location='cpu')
print("\nloading", save_name, "...... From Epoch: ", save_dict['epoch'])
net.load_state_dict(save_dict['net'])
net = torch.nn.DataParallel(net.to(device))
net.eval()
# ------------ test ------------
out_path = os.path.join(exp_dir, "coord_output")
evaluation_path = os.path.join(exp_dir, "evaluate")
if not os.path.exists(out_path):
os.mkdir(out_path)
if not os.path.exists(evaluation_path):
os.mkdir(evaluation_path)
dump_to_json = []
progressbar = tqdm(range(len(test_loader)))
with torch.no_grad():
for batch_idx, sample in enumerate(test_loader):
img = sample['img'].to(device)
img_name = sample['img_name']
import torch
from model import SCNN
import onnx
input_size = (800, 288)
FILE_NAME = 'scnn_pytorch_half.onnx'
net = SCNN(input_size=input_size, pretrained=False, test=True)
state = torch.load('D:/DHH/Downloads/exp10_best.pth')
net.load_state_dict(state['net'])
net.eval().cuda().half()
input_size = torch.randn(1, 3, 288, 800, device='cuda').half()
torch.onnx.export(net,
input_size,
FILE_NAME,
export_params=True,
opset_version=11,
do_constant_folding=True)
model = onnx.load(FILE_NAME)
onnx.checker.check_model(model, full_check=True)
print("==> Passed")
output = [node.name for node in model.graph.output]
input_all = [node.name for node in model.graph.input]
input_initializer = [node.name for node in model.graph.initializer]
net_feed_input = list(set(input_all) - set(input_initializer))
print('Inputs: ', net_feed_input)
print('Outputs: ', output)