def single():
print('Mode: Single')
img = Image.open('test_content/example_01.png').convert('RGB')
class_encoding = color_encoding = OrderedDict([
('unlabeled', (0, 0, 0)), ('road', (128, 64, 128)),
('sidewalk', (244, 35, 232)), ('building', (70, 70, 70)),
('wall', (102, 102, 156)), ('fence', (190, 153, 153)),
('pole', (153, 153, 153)), ('traffic_light', (250, 170, 30)),
('traffic_sign', (220, 220, 0)), ('vegetation', (107, 142, 35)),
('terrain', (152, 251, 152)), ('sky', (70, 130, 180)),
('person', (220, 20, 60)), ('rider', (255, 0, 0)),
('car', (0, 0, 142)), ('truck', (0, 0, 70)), ('bus', (0, 60, 100)),
('train', (0, 80, 100)), ('motorcycle', (0, 0, 230)),
('bicycle', (119, 11, 32))
])
num_classes = len(class_encoding)
model = ERFNet(num_classes)
model_path = os.path.join(args.save_dir, args.name)
print('Loading model at:', model_path)
checkpoint = torch.load(model_path)
# model = ENet(num_classes)
model = model.cuda()
model.load_state_dict(checkpoint['state_dict'])
img = img.resize((args.width, args.height), Image.BILINEAR)
start = time.time()
images = transforms.ToTensor()(img)
torch.reshape(images, (1, 3, args.width, args.height))
images = images.unsqueeze(0)
with torch.no_grad():
images = images.cuda()
predictions = model(images)
end = time.time()
print('model speed:', int(1 / (end - start)), "FPS")
_, predictions = torch.max(predictions.data, 1)
label_to_rgb = transforms.Compose(
[utils.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()])
color_predictions = utils.batch_transform(predictions.cpu(),
label_to_rgb)
end = time.time()
print('model+transform:', int(1 / (end - start)), "FPS")
utils.imshow_batch(images.data.cpu(), color_predictions)
def train(train_loader, val_loader, class_weights, class_encoding):
print("Training...")
num_classes = len(class_encoding)
model = ERFNet(num_classes)
criterion = nn.CrossEntropyLoss(weight=class_weights)
optimizer = optim.Adam(model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
if use_cuda:
model = model.cuda()
criterion = criterion.cuda()
# Optionally resume from a checkpoint
if args.resume:
model, optimizer, start_epoch, best_miou, val_miou, train_miou, val_loss, train_loss = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name, True)
print(
"Resuming from model: Start epoch = {0} | Best mean IoU = {1:.4f}".
format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
val_miou = []
train_miou = []
val_loss = []
train_loss = []
# Start Training
train = Train(model, train_loader, optimizer, criterion, metric, use_cuda)
val = Test(model, val_loader, criterion, metric, use_cuda)
for epoch in range(start_epoch, args.epochs):
print(">> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step()
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
print(
">> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".format(
epoch, epoch_loss, miou))
train_loss.append(epoch_loss)
train_miou.append(miou)
#preform a validation test
if (epoch + 1) % 10 == 0 or epoch + 1 == args.epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
val_loss.append(loss)
val_miou.append(miou)
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("Best model thus far. Saving...")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
val_miou, train_miou, val_loss,
train_loss, args)
return model, train_loss, train_miou, val_loss, val_miou
def video():
print('testing from video')
cameraWidth = 1920
cameraHeight = 1080
cameraMatrix = np.matrix([[1.3878727764994030e+03, 0, cameraWidth / 2],
[0, 1.7987055172413220e+03, cameraHeight / 2],
[0, 0, 1]])
distCoeffs = np.matrix([
-5.8881725390917083e-01, 5.8472404395779809e-01,
-2.8299599929891900e-01, 0
])
vidcap = cv2.VideoCapture('test_content/massachusetts.mp4')
success = True
i = 0
while success:
success, img = vidcap.read()
if i % 1000 == 0:
print("frame: ", i)
if args.rmdistort:
P = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(
cameraMatrix, distCoeffs, (cameraWidth, cameraHeight),
None)
map1, map2 = cv2.fisheye.initUndistortRectifyMap(
cameraMatrix, distCoeffs, np.eye(3), P, (1920, 1080),
cv2.CV_16SC2)
img = cv2.remap(img, map1, map2, cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
# img = img.convert('RGB')
# cv2.imshow('',img)
# cv2.waitKey(0)
# img2 = Image.open(filename).convert('RGB')
class_encoding = color_encoding = OrderedDict([
('unlabeled', (0, 0, 0)), ('road', (128, 64, 128)),
('sidewalk', (244, 35, 232)), ('building', (70, 70, 70)),
('wall', (102, 102, 156)), ('fence', (190, 153, 153)),
('pole', (153, 153, 153)), ('traffic_light', (250, 170, 30)),
('traffic_sign', (220, 220, 0)),
('vegetation', (107, 142, 35)), ('terrain', (152, 251, 152)),
('sky', (70, 130, 180)), ('person', (220, 20, 60)),
('rider', (255, 0, 0)), ('car', (0, 0, 142)),
('truck', (0, 0, 70)), ('bus', (0, 60, 100)),
('train', (0, 80, 100)), ('motorcycle', (0, 0, 230)),
('bicycle', (119, 11, 32))
])
num_classes = len(class_encoding)
model_path = os.path.join(args.save_dir, args.name)
checkpoint = torch.load(model_path)
model = ERFNet(num_classes)
model = model.cuda()
model.load_state_dict(checkpoint['state_dict'])
img = img.resize((args.width, args.height), Image.BILINEAR)
start = time.time()
images = transforms.ToTensor()(img)
torch.reshape(images, (1, 3, args.width, args.height))
images = images.unsqueeze(0)
with torch.no_grad():
images = images.cuda()
predictions = model(images)
end = time.time()
print('model speed:', int(1 / (end - start)), "FPS")
_, predictions = torch.max(predictions.data, 1)
label_to_rgb = transforms.Compose([
utils.LongTensorToRGBPIL(class_encoding),
transforms.ToTensor()
])
color_predictions = utils.batch_transform(
predictions.cpu(), label_to_rgb)
end = time.time()
print('model+transform:', int(1 / (end - start)), "FPS")
utils.imshow_batch(images.data.cpu(), color_predictions)
i += 1
loaders, w_class, class_encoding = load_dataset(dataset)
train_loader, val_loader, test_loader = loaders
if args.mode.lower() in {'train'}:
model, tl, tmiou, vl, vmiou = train(train_loader, val_loader,
w_class, class_encoding)
plt.plot(tl, label="train loss")
plt.plot(tmiou, label="train miou")
plt.plot(vl, label="val loss")
plt.plot(vmiou, label="val miou")
plt.legend()
plt.xlabel("Epoch")
plt.ylabel("loss/accuracy")
plt.grid(True)
plt.xticks()
plt.savefig('./plots/train.png')
elif args.mode.lower() == 'test':
num_classes = len(class_encoding)
#model = ENet(num_classes)
model = ERFNet(num_classes)
if use_cuda:
model = model.cuda()
optimizer = optim.Adam(model.parameters())
model = utils.load_checkpoint(model, optimizer, args.save_dir,
args.name)[0]
test(model, test_loader, w_class, class_encoding)
else:
raise RuntimeError(
"\"{0}\" is not a valid choice for execution mode.".format(
args.mode))