def for_loop(net, data_loader, train_optimizer):
is_train = train_optimizer is not None
net.train() if is_train else net.eval()
total_loss, total_time, total_num, preds, targets = 0.0, 0.0, 0, [], []
data_bar = tqdm(data_loader, dynamic_ncols=True)
with (torch.enable_grad() if is_train else torch.no_grad()):
for data, target, grad, boundary, name in data_bar:
data, target, grad, boundary = data.cuda(), target.cuda(), grad.cuda(), boundary.cuda()
torch.cuda.synchronize()
start_time = time.time()
seg, edge = net(data, grad)
prediction = torch.argmax(seg.detach(), dim=1)
torch.cuda.synchronize()
end_time = time.time()
semantic_loss = semantic_criterion(seg, target)
edge_loss = edge_criterion(edge, target, boundary)
task_loss = task_criterion(seg, edge, target)
loss = semantic_loss + 20 * edge_loss + task_loss
if is_train:
train_optimizer.zero_grad()
loss.backward()
train_optimizer.step()
total_num += data.size(0)
total_time += end_time - start_time
total_loss += loss.item() * data.size(0)
preds.append(prediction.cpu())
targets.append(target.cpu())
if not is_train:
if data_loader.dataset.split == 'test':
# revert train id to regular id
for key in sorted(trainId2label.keys(), reverse=True):
prediction[prediction == key] = trainId2label[key].id
# save pred images
save_root = '{}/{}_{}_{}/{}'.format(save_path, backbone_type, crop_h, crop_w, data_loader.dataset.split)
if not os.path.exists(save_root):
os.makedirs(save_root)
for pred_tensor, pred_name in zip(prediction, name):
pred_img = ToPILImage()(pred_tensor.unsqueeze(dim=0).byte().cpu())
if data_loader.dataset.split == 'val':
pred_img.putpalette(get_palette())
pred_name = pred_name.replace('leftImg8bit', 'color')
path = '{}/{}'.format(save_root, pred_name)
pred_img.save(path)
data_bar.set_description('{} Epoch: [{}/{}] Loss: {:.4f} FPS: {:.0f}'
.format(data_loader.dataset.split.capitalize(), epoch, epochs,
total_loss / total_num, total_num / total_time))
# compute metrics
preds = torch.cat(preds, dim=0)
targets = torch.cat(targets, dim=0)
pa, mpa, class_iou, category_iou = compute_metrics(preds, targets)
print('{} Epoch: [{}/{}] PA: {:.2f}% mPA: {:.2f}% Class_mIOU: {:.2f}% Category_mIOU: {:.2f}%'
.format(data_loader.dataset.split.capitalize(), epoch, epochs,
pa * 100, mpa * 100, class_iou * 100, category_iou * 100))
return total_loss / total_num, pa * 100, mpa * 100, class_iou * 100, category_iou * 100
def train_val(net, data_loader, train_optimizer):
is_train = train_optimizer is not None
net.train() if is_train else net.eval()
total_loss, total_correct, total_time, total_num, data_bar = 0.0, 0.0, 0.0, 0, tqdm(
data_loader)
with (torch.enable_grad() if is_train else torch.no_grad()):
for data, target, name in data_bar:
data, target = data.cuda(), target.cuda()
torch.cuda.synchronize()
start_time = time.time()
out = net(data)
prediction = torch.argmax(out, dim=1)
torch.cuda.synchronize()
end_time = time.time()
loss = loss_criterion(out, target)
if is_train:
train_optimizer.zero_grad()
loss.backward()
train_optimizer.step()
total_num += data.size(0)
total_time += end_time - start_time
total_loss += loss.item() * data.size(0)
total_correct += torch.sum(
prediction == target).item() / target.numel() * data.size(0)
if not is_train and epoch % save_step == 0:
# revert train id to regular id
for key in trainId2label.keys():
prediction[prediction == key] = trainId2label[key].id
# save pred images
for pred_tensor, pred_name in zip(prediction, name):
pred_img = ToPILImage()(
pred_tensor.unsqueeze(dim=0).byte().cpu())
pred_img.putpalette(palette)
pred_img.save('results/{}'.format(
pred_name.replace('leftImg8bit', 'color')))
data_bar.set_description(
'{} Epoch: [{}/{}] Loss: {:.4f} mPA: {:.2f}% FPS: {:.0f}'.
format('Train' if is_train else 'Val', epoch, epochs,
total_loss / total_num, total_correct / total_num * 100,
total_num / total_time))
return total_loss / total_num, total_correct / total_num * 100
def savelabel(label, name):
l = ToPILImage()(label[0].unsqueeze(0).byte()).convert('P')
l.putpalette(lab_palette)
l.save(name)
target = Image.new('RGB', (image_width * num_width, image_height))
images = [image]
image = transform(image).unsqueeze(dim=0).cuda()
# model load
model = FastSCNN(in_channels=3, num_classes=19)
model.load_state_dict(
torch.load(model_weight, map_location=torch.device('cpu')))
model = model.cuda()
model.eval()
# predict and save image
with torch.no_grad():
output = model(image)
pred = torch.argmax(output, dim=1)
pred_image = ToPILImage()(pred.byte().cpu())
pred_image.putpalette(palette)
if 'test' not in input_pic:
gt_image = Image.open('{}/gtFine/{}'.format(
data_path, input_pic.replace('leftImg8bit', 'gtFine_color')))
images.append(gt_image)
images.append(pred_image)
# concat images
for i in range(len(images)):
left, top, right, bottom = image_width * i, 0, image_width * (
i + 1), image_height
target.paste(images[i], (left, top, right, bottom))
target.save(
os.path.split(input_pic)[-1].replace('leftImg8bit', 'result'))
grad = torch.from_numpy(
np.expand_dims(np.asarray(grad, np.float32) / 255.0,
axis=0).copy()).unsqueeze(dim=0).cuda()
# model load
model = GatedSCNN(model_weight.split('/')[-1].split('_')[0],
num_classes=19)
model.load_state_dict(
torch.load(model_weight, map_location=torch.device('cpu')))
model = model.cuda()
model.eval()
# predict and save image
with torch.no_grad():
output, _ = model(image, grad)
pred = torch.argmax(output, dim=1)
pred_image = ToPILImage()(pred.byte().cpu())
pred_image.putpalette(get_palette())
if 'test' not in input_pic:
gt_image = Image.open('{}/gtFine/{}'.format(
data_path, input_pic.replace('leftImg8bit', 'gtFine_color')))
images.append(gt_image)
images.append(pred_image)
# concat images
for i in range(len(images)):
left, top, right, bottom = image_width * i, 0, image_width * (
i + 1), image_height
target.paste(images[i], (left, top, right, bottom))
target.save(
os.path.split(input_pic)[-1].replace('leftImg8bit', 'result'))
