From 276696d1e7407a40d55c2a50b12839b52c0fed1e Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sat, 22 Jun 2024 10:37:14 +0200 Subject: [PATCH] Update. --- world.py | 54 +++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 7 deletions(-) diff --git a/world.py b/world.py index a5e9601..2e13ec3 100755 --- a/world.py +++ b/world.py @@ -103,8 +103,9 @@ def generate( def sample2img(seq, height, width, upscale=15): - f_start = seq[:, : height * width].reshape(-1, height, width) - f_end = seq[:, height * width + 1 :].reshape(-1, height, width) + f_first = seq[:, : height * width].reshape(-1, height, width) + f_second = seq[:, height * width + 1 :].reshape(-1, height, width) + direction = seq[:, height * width] def mosaic(x, upscale): x = x.reshape(-1, height, width) @@ -124,7 +125,46 @@ def sample2img(seq, height, width, upscale=15): return x - return torch.cat([mosaic(f_start, upscale), mosaic(f_end, upscale)], dim=3) + direction_symbol = torch.full((direction.size(0), height * upscale, upscale), 0) + direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2) + separator = torch.full((direction.size(0), 3, height * upscale, 1), 0) + + for n in range(direction_symbol.size(0)): + if direction[n] == token_forward: + for k in range(upscale): + direction_symbol[ + n, + :, + (height * upscale) // 2 - upscale // 2 + k, + 3 + abs(k - upscale // 2), + ] = 0 + elif direction[n] == token_backward: + for k in range(upscale): + direction_symbol[ + n, + :, + (height * upscale) // 2 - upscale // 2 + k, + 3 + upscale // 2 - abs(k - upscale // 2), + ] = 0 + else: + for k in range(2, upscale - 2): + direction_symbol[ + n, :, (height * upscale) // 2 - upscale // 2 + k, k + ] = 0 + direction_symbol[ + n, :, (height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k + ] = 0 + + return torch.cat( + [ + mosaic(f_first, upscale), + separator, + direction_symbol, + separator, + mosaic(f_second, upscale), + ], + dim=3, + ) def seq2str(seq): @@ -141,18 +181,18 @@ if __name__ == "__main__": height, width = 6, 8 start_time = time.perf_counter() - seq = generate(nb=64, height=height, width=width, max_nb_obj=3) + seq = generate(nb=90, height=height, width=width, max_nb_obj=3) delay = time.perf_counter() - start_time print(f"{seq.size(0)/delay:02f} samples/s") print(seq2str(seq[:4])) - # m = (torch.rand(seq.size()) < 0.05).long() - # seq = (1 - m) * seq + m * 23 + m = (torch.rand(seq.size()) < 0.05).long() + seq = (1 - m) * seq + m * 23 img = sample2img(seq, height, width) print(img.size()) torchvision.utils.save_image( - img.float() / 255.0, "/tmp/world.png", nrow=8, padding=2 + img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4 ) -- 2.39.5