- img = torch.cat((img, predicted_paths.unsqueeze(1)), 1)
-
- img = img.reshape((-1,) + img.size()[2:]).float() / 255.0
-
- torchvision.utils.save_image(img, name, padding=1, pad_value=0.85, nrow=6)
+ imgs = torch.cat((imgs, predicted_paths.unsqueeze(1)), 1)
+
+ # NxKxCxHxW
+ if path_correct is None:
+ path_correct = torch.zeros(imgs.size(0)) <= 1
+ path_correct = path_correct.cpu().long().view(-1, 1, 1, 1)
+ img = torch.tensor([224, 224, 224]).view(1, -1, 1, 1) * path_correct + torch.tensor(
+ [255, 0, 0]
+ ).view(1, -1, 1, 1) * (1 - path_correct)
+ img = img.expand(
+ -1, -1, imgs.size(3) + 2, 1 + imgs.size(1) * (1 + imgs.size(4))
+ ).clone()
+ for k in range(imgs.size(1)):
+ img[
+ :,
+ :,
+ 1 : 1 + imgs.size(3),
+ 1 + k * (1 + imgs.size(4)) : 1 + k * (1 + imgs.size(4)) + imgs.size(4),
+ ] = imgs[:, k]
+
+ img = img.float() / 255.0
+
+ torchvision.utils.save_image(img, name, nrow=4, padding=1, pad_value=224.0 / 256)