From ffc0a92cddefa9fda9f25468c25ae1365b52be47 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 17 Mar 2023 17:40:57 +0100 Subject: [PATCH] Update --- maze.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/maze.py b/maze.py index 6e8e179..754cdea 100755 --- a/maze.py +++ b/maze.py @@ -207,34 +207,43 @@ def save_image( mazes = mazes.cpu() - mazes = colors[mazes.reshape(-1)].reshape(mazes.size() + (-1,)).permute(0, 3, 1, 2) + c_mazes = ( + colors[mazes.reshape(-1)].reshape(mazes.size() + (-1,)).permute(0, 3, 1, 2) + ) - imgs = mazes.unsqueeze(1) + imgs = c_mazes.unsqueeze(1) if target_paths is not None: target_paths = target_paths.cpu() - target_paths = ( + c_target_paths = ( colors[target_paths.reshape(-1)] .reshape(target_paths.size() + (-1,)) .permute(0, 3, 1, 2) ) - imgs = torch.cat((imgs, target_paths.unsqueeze(1)), 1) + imgs = torch.cat((imgs, c_target_paths.unsqueeze(1)), 1) if predicted_paths is not None: predicted_paths = predicted_paths.cpu() - predicted_paths = ( + c_predicted_paths = ( colors[predicted_paths.reshape(-1)] .reshape(predicted_paths.size() + (-1,)) .permute(0, 3, 1, 2) ) - imgs = torch.cat((imgs, predicted_paths.unsqueeze(1)), 1) + imgs = torch.cat((imgs, c_predicted_paths.unsqueeze(1)), 1) if score_paths is not None: - score_paths = (score_paths.cpu() * 255.0).long() - score_paths = score_paths.unsqueeze(1).expand(-1, 3, -1, -1) - imgs = torch.cat((imgs, score_paths.unsqueeze(1)), 1) + score_paths = score_paths.cpu() + c_score_paths = score_paths.unsqueeze(1).expand(-1, 3, -1, -1) + c_score_paths = ( + c_score_paths * colors[4].reshape(1, 3, 1, 1) + + (1 - c_score_paths) * colors[3].reshape(1, 3, 1, 1) + ).long() + c_score_paths = c_score_paths * (mazes.unsqueeze(1) == v_empty) + c_mazes * ( + mazes.unsqueeze(1) != v_empty + ) + imgs = torch.cat((imgs, c_score_paths.unsqueeze(1)), 1) # NxKxCxHxW if path_correct is None: -- 2.39.5