From 483494965a97df9a09cd572246d1f3a4dc7248de Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sat, 11 Mar 2023 22:00:43 +0100 Subject: [PATCH] Update. --- beaver.py | 20 +++++++++++++------- maze.py | 32 +++++++++++++++++++++++++------- 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/beaver.py b/beaver.py index 920a446..a289867 100755 --- a/beaver.py +++ b/beaver.py @@ -75,11 +75,11 @@ parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") ############################## # maze options -parser.add_argument("--world_height", type=int, default=13) +parser.add_argument("--maze_height", type=int, default=13) -parser.add_argument("--world_width", type=int, default=21) +parser.add_argument("--maze_width", type=int, default=21) -parser.add_argument("--world_nb_walls", type=int, default=15) +parser.add_argument("--maze_nb_walls", type=int, default=15) ###################################################################### @@ -262,7 +262,13 @@ class TaskMaze(Task): mazes, paths = self.seq2map(input) _, predicted_paths = self.seq2map(result) - maze.save_image(f"result_{n_epoch:04d}.png", mazes, paths, predicted_paths) + maze.save_image( + f"result_{n_epoch:04d}.png", + mazes, + paths, + predicted_paths, + maze.path_correctness(mazes, predicted_paths), + ) model.train(t) @@ -276,9 +282,9 @@ task = TaskMaze( nb_train_samples=args.nb_train_samples, nb_test_samples=args.nb_test_samples, batch_size=args.batch_size, - height=args.world_height, - width=args.world_width, - nb_walls=args.world_nb_walls, + height=args.maze_height, + width=args.maze_width, + nb_walls=args.maze_nb_walls, device=device, ) diff --git a/maze.py b/maze.py index f4a4840..e377d2f 100755 --- a/maze.py +++ b/maze.py @@ -185,7 +185,7 @@ def create_maze_data( ###################################################################### -def save_image(name, mazes, target_paths, predicted_paths=None): +def save_image(name, mazes, target_paths, predicted_paths=None, path_correct=None): mazes, target_paths = mazes.cpu(), target_paths.cpu() colors = torch.tensor( @@ -204,7 +204,7 @@ def save_image(name, mazes, target_paths, predicted_paths=None): .reshape(target_paths.size() + (-1,)) .permute(0, 3, 1, 2) ) - img = torch.cat((mazes.unsqueeze(1), target_paths.unsqueeze(1)), 1) + imgs = torch.cat((mazes.unsqueeze(1), target_paths.unsqueeze(1)), 1) if predicted_paths is not None: predicted_paths = predicted_paths.cpu() @@ -213,11 +213,29 @@ def save_image(name, mazes, target_paths, predicted_paths=None): .reshape(predicted_paths.size() + (-1,)) .permute(0, 3, 1, 2) ) - 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) ###################################################################### -- 2.39.5