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 = "_X" + "".join([str(n) for n in range(len(colors) - 2)]) + ">"
40 f_start = torch.zeros(nb, height, width, dtype=torch.int64)
41 f_end = torch.zeros(nb, height, width, dtype=torch.int64)
42 n = torch.arange(f_start.size(0))
44 for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
45 nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1
46 for c in torch.randperm(colors.size(0) - 2)[:nb_fish].sort().values:
48 torch.randint(height - 2, (1,))[0] + 1,
49 torch.randint(width - 2, (1,))[0] + 1,
51 vm = torch.randint(4, (1,))[0]
52 vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
54 f_start[n, i, j] = c + 2
55 f_start[n, i - vi, j - vj] = c + 2
56 f_start[n, i + vj, j - vi] = c + 2
57 f_start[n, i - vj, j + vi] = c + 2
59 for l in range(nb_iterations):
62 if i < 0 or i >= height or j < 0 or j >= width:
69 f_end[n, i, j] = c + 2
70 f_end[n, i - vi, j - vj] = c + 2
71 f_end[n, i + vj, j - vi] = c + 2
72 f_end[n, i - vj, j + vi] = c + 2
77 torch.full((f_end.size(0), 1), len(colors)),
84 def sample2img(seq, height, width):
85 f_start = seq[:, : height * width].reshape(-1, height, width)
86 f_start = (f_start >= len(colors)).long() + (f_start < len(colors)).long() * f_start
87 f_end = seq[:, height * width + 1 :].reshape(-1, height, width)
88 f_end = (f_end >= len(colors)).long() + (f_end < len(colors)).long() * f_end
90 img_f_start, img_f_end = colors[f_start], colors[f_end]
96 (img_f_start.size(0), img_f_start.size(1), 1, img_f_start.size(3)), 1
103 return img.permute(0, 3, 1, 2)
109 result.append("".join([token2char[v] for v in s]))
113 ######################################################################
115 if __name__ == "__main__":
119 start_time = time.perf_counter()
120 seq = generate(nb=64, height=height, width=width, max_nb_obj=3)
121 delay = time.perf_counter() - start_time
122 print(f"{seq.size(0)/delay:02f} samples/s")
124 print(seq2str(seq[:4]))
126 img = sample2img(seq, height, width)
129 torchvision.utils.save_image(img.float() / 255.0, "world.png", nrow=8, padding=2)