3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
10 import torch, torchvision
13 from torch.nn import functional as F
15 ######################################################################
18 colors = torch.tensor(
30 token2char = "_X01234>"
42 intact = torch.zeros(nb, height, width, dtype=torch.int64)
43 n = torch.arange(intact.size(0))
46 for c in torch.randperm(colors.size(0) - 2)[:nb_obj] + 2:
47 z = intact[n].flatten()
48 m = (torch.rand(z.size()) * (z == 0)).argmax(dim=0)
49 i, j = m // width, m % width
50 vm = torch.randint(4, (1,))[0]
51 vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
52 for l in range(obj_length):
56 if i < 0 or i >= height or j < 0 or j >= width or intact[n, i, j] != 0:
67 or intact[n, i, j] != 0
71 masked = intact.clone()
74 i = torch.randint(height - mask_height + 1, (1,))[0]
75 j = torch.randint(width - mask_width + 1, (1,))[0]
76 masked[n, i : i + mask_height, j : j + mask_width] = 1
81 torch.full((masked.size(0), 1), len(colors)),
88 def sample2img(seq, height, width):
89 intact = seq[:, : height * width].reshape(-1, height, width)
90 masked = seq[:, height * width + 1 :].reshape(-1, height, width)
91 img_intact, img_masked = colors[intact], colors[masked]
97 (img_intact.size(0), img_intact.size(1), 1, img_intact.size(3)), 1
104 return img.permute(0, 3, 1, 2)
110 result.append("".join([token2char[v] for v in s]))
114 ######################################################################
116 if __name__ == "__main__":
120 start_time = time.perf_counter()
121 seq = generate(nb=64, height=height, width=width)
122 delay = time.perf_counter() - start_time
123 print(f"{seq.size(0)/delay:02f} samples/s")
125 print(seq2str(seq[:4]))
127 img = sample2img(seq, height, width)
130 torchvision.utils.save_image(img.float() / 255.0, "world.png", nrow=8, padding=2)