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()
# mean=(0.3598, 0.3653, 0.3662)
# std=(0.2573, 0.2663, 0.2756)
# Imagenet mean, std
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
transform = Compose(Resize(resize_shape), ToTensor(),
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 = []
import argparse
import cv2
import torch
from model import SCNN
from utils.prob2lines import getLane
from utils.transforms import *
net = SCNN(input_size=(800, 288), pretrained=False)
mean = (0.3598, 0.3653, 0.3662) # CULane mean, std
std = (0.2573, 0.2663, 0.2756)
transform_img = Resize((800, 288))
transform_to_net = Compose(ToTensor(), Normalize(mean=mean, std=std))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--img_path",
'-i',
type=str,
default="demo/demo.jpg",
help="Path to demo img")
parser.add_argument("--weight_path",
'-w',
type=str,
help="Path to model weights")
parser.add_argument("--visualize",
'-v',
action="store_true",
default=False,
help="Visualize the result")
6].modules[1].modules[0].modules[396].modules[1].modules[2].modules[
0].modules[0].weight
model2['layer2.1.weight'] = model1.modules[42].modules[8].weight
model2['layer2.1.bias'] = model1.modules[42].modules[8].bias
model2['fc.0.weight'] = model1.modules[43].modules[1].modules[3].weight
model2['fc.0.bias'] = model1.modules[43].modules[1].modules[3].bias
model2['fc.2.weight'] = model1.modules[43].modules[1].modules[5].weight
model2['fc.2.bias'] = model1.modules[43].modules[1].modules[5].bias
save_name = os.path.join('experiments', 'vgg_SCNN_DULR_w9',
'vgg_SCNN_DULR_w9.pth')
torch.save(model2, save_name)
# load and save again
net = SCNN(pretrained=False)
d = torch.load(save_name)
net.load_state_dict(d, strict=False)
for m in net.backbone.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None:
m.bias.data.zero_()
save_dict = {
"epoch": 0,
"net": net.state_dict(),
"optim": None,
"lr_scheduler": None
}
if not os.path.exists(os.path.join('experiments', 'vgg_SCNN_DULR_w9')):
model2['message_passing.down_up.weight'] = model1.modules[42].modules[6].modules[0].modules[0].modules[140].modules[1].modules[2].modules[0].modules[0].weight
model2['message_passing.left_right.weight'] = model1.modules[42].modules[6].modules[1].modules[0].modules[2].modules[0].modules[1].modules[1].modules[0].weight
model2['message_passing.right_left.weight'] = model1.modules[42].modules[6].modules[1].modules[0].modules[396].modules[1].modules[2].modules[0].modules[0].weight
model2['layer2.1.weight'] = model1.modules[42].modules[8].weight
model2['layer2.1.bias'] = model1.modules[42].modules[8].bias
model2['fc.0.weight'] = model1.modules[43].modules[1].modules[3].weight
model2['fc.0.bias'] = model1.modules[43].modules[1].modules[3].bias
model2['fc.2.weight'] = model1.modules[43].modules[1].modules[5].weight
model2['fc.2.bias'] = model1.modules[43].modules[1].modules[5].bias
save_name = os.path.join('experiments', 'vgg_SCNN_DULR_w9', 'vgg_SCNN_DULR_w9.pth')
torch.save(model2, save_name)
# load and save again
net = SCNN(input_size=(800, 288), pretrained=False)
d = torch.load(save_name)
net.load_state_dict(d, strict=False)
for m in net.backbone.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None:
m.bias.data.zero_()
save_dict = {
"epoch": 0,
"net": net.state_dict(),
"optim": None,
"lr_scheduler": None
}
import argparse
import cv2
import torch
from model import SCNN
from utils.prob2lines import getLane
from utils.transforms import *
import pdb
image_resize_width = 800
net = SCNN(input_size=(image_resize_width, 288), pretrained=False)
mean = (0.3598, 0.3653, 0.3662) # CULane mean, std
std = (0.2573, 0.2663, 0.2756)
transform = Compose(Resize((image_resize_width, 288)), ToTensor(),
Normalize(mean=mean, std=std))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--img_path",
'-i',
type=str,
default="demo/demo.jpg",
help="Path to demo img")
parser.add_argument("--weight_path",
'-w',
type=str,
help="Path to model weights")
parser.add_argument("--visualize",
import argparse
import cv2
import torch
from model import SCNN
from utils.prob2lines import getLane
from utils.transforms import *
net = SCNN(pretrained=False)
mean=(0.3598, 0.3653, 0.3662) # CULane mean, std
std=(0.2573, 0.2663, 0.2756)
transform = Compose(Resize((800, 288)), ToTensor(),
Normalize(mean=mean, std=std))
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--img_path", '-i', type=str, default="demo/demo.jpg", help="Path to demo img")
parser.add_argument("--weight_path", '-w', type=str, help="Path to model weights")
parser.add_argument("--visualize", '-v', action="store_true", default=False, help="Visualize the result")
args = parser.parse_args()
return args
def main():
args = parse_args()
img_path = args.img_path
weight_path = args.weight_path
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
collate_fn=train_dataset.collate,
num_workers=6)
# ------------ val data ------------
transform_val_img = Resize(resize_shape)
transform_val_x = Compose(ToTensor(), Normalize(mean=mean, std=std))
transform_val = Compose(transform_val_img, transform_val_x)
val_dataset = Dataset_Type(Dataset_Path[dataset_name], "val", transform_val)
val_loader = DataLoader(val_dataset,
batch_size=8,
collate_fn=val_dataset.collate,
num_workers=4)
# ------------ preparation ------------
if exp_cfg['model'] == "scnn":
net = SCNN(resize_shape, pretrained=True)
elif exp_cfg['model'] == "enet_sad":
net = ENet_SAD(resize_shape, sad=True, dataset=dataset_name)
else:
raise Exception(
"Model not match. 'model' in 'cfg.json' should be 'scnn' or 'enet_sad'."
)
net = net.to(device)
net = torch.nn.DataParallel(net)
optimizer = optim.SGD(net.parameters(), **exp_cfg['optim'])
lr_scheduler = PolyLR(optimizer, 0.9, **exp_cfg['lr_scheduler'])
best_val_loss = 1e6
# ------------ data and model ------------
# CULane mean, std
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))
# val_dataset = CULane(CULane_path, "val", transform)
# val_loader = DataLoader(val_dataset, batch_size=8, collate_fn=val_dataset.collate, num_workers=4)
test_dataset = CULane(CULane_path, "test", transform)
test_loader = DataLoader(test_dataset,
batch_size=8,
collate_fn=test_dataset.collate,
num_workers=4)
net = SCNN(pretrained=False)
save_name = os.path.join(exp_dir, exp_dir.split('/')[-1] + '_best.pth')
save_name = "/home/lion/hanyibo/SCNN/experiments/vgg_SCNN_DULR_w9/vgg_SCNN_DULR_w9.pth"
save_dict = torch.load(save_name, map_location='cpu')
print("loading", save_name, "......")
net.load_state_dict(save_dict['net'])
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)
val_dataset = Dataset_Type(Dataset_Path[dataset_name], "val", transform_val)
# Horovod: use DistributedSampler to partition the test data.
val_sampler = torch.utils.data.distributed.DistributedSampler(
val_dataset, num_replicas=hvd.size(), rank=hvd.rank())
val_loader = DataLoader(val_dataset,
batch_size=val_batch_size,
collate_fn=val_dataset.collate,
sampler=val_sampler,
**kwargs)
# Horovod: print logs on the first worker.
verbose = 1 if hvd.rank() == 0 else 0
# ------------ preparation ------------
net = SCNN(resize_shape, pretrained=True)
lr_scaler = 1
if torch.cuda.is_available():
net.cuda()
# Horovod: Scale learning rate as per number of devices
if hvd.nccl_built():
lr_scaler = hvd.local_size()
net = torch.nn.DataParallel(net)
lr = exp_cfg['optim']['lr']
momentum = exp_cfg['optim']['momentum']
weight_decay = exp_cfg['optim']['weight_decay']
nesterov = exp_cfg['optim']['nesterov']
# Horovod: scale learning rate by lr_scaler.
optimizer = optim.SGD(net.parameters(),
batch_size=exp_cfg['dataset']['batch_size'],
shuffle=True,
collate_fn=dataset.collate,
num_workers=8)
val_loader = DataLoader(val_dataset,
batch_size=8,
collate_fn=dataset.collate,
num_workers=4)
# ------------ preparation ------------
seg_classes = 5
if hasattr(dataset, 'seg_classes'):
seg_classes = getattr(dataset, 'seg_classes')
net = SCNN(resize_shape,
pretrained=True,
seg_classes=seg_classes,
weights=Dataset_Type.get_weights(
exp_cfg['dataset']['other']['seg_mode']))
net = net.to(device)
#net = torch.nn.DataParallel(net)
optimizer = optim.Adam(net.parameters(), **exp_cfg['optim'])
lr_scheduler = PolyLR(optimizer, 0.9, **exp_cfg['lr_scheduler'])
best_val_loss = 1e6
def train(epoch):
print("Train Epoch: {}".format(epoch))
net.train()
train_loss = 0
train_loss_seg = 0
Normalize(mean=mean, std=std))
transform_val_img = Resize(resize_shape)
transform_val_x = Compose(ToTensor(), Normalize(mean=mean, std=std))
transform_val = Compose(transform_val_img, transform_val_x)
dataset_name = exp_cfg['dataset'].pop('dataset_name')
Dataset_Type = getattr(dataset, dataset_name)
dataset = Dataset_Type(Dataset_Path[dataset_name], transform_train,
**exp_cfg['dataset']['other'])
seg_classes = 5
if hasattr(dataset, 'seg_classes'):
seg_classes = getattr(dataset, 'seg_classes')
net = SCNN(resize_shape,
pretrained=True,
seg_classes=seg_classes,
weights=Dataset_Type.get_weights(
exp_cfg['dataset']['other']['seg_mode']))
def main():
print('Loading model')
net.load_state_dict(
torch.load(model_path, map_location=torch.device('cpu'))['net'])
dummy_input = torch.zeros((2, 3, resize_shape[0], resize_shape[1]))
#net(dummy_input)
torch.onnx.export(net, dummy_input, "model.onnx", opset_version=11)
if __name__ == "__main__":
main()
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)
