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(
31 nb_fish_tokens = len(colors) - 1
32 token_forward = first_fish_token + nb_fish_tokens
33 token_backward = token_forward + 1
35 token2char = "_" + "".join([str(n) for n in range(len(colors) - 1)]) + "><"
47 for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"):
48 f_start = torch.zeros(height, width, dtype=torch.int64)
49 f_end = torch.zeros(height, width, dtype=torch.int64)
50 n = torch.arange(f_start.size(0))
52 nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1
54 (torch.randperm(nb_fish_tokens) + first_fish_token)[:nb_fish].sort().values
57 torch.randint(height - 2, (1,))[0] + 1,
58 torch.randint(width - 2, (1,))[0] + 1,
60 vm = torch.randint(4, (1,))[0]
61 vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
64 f_start[i - vi, j - vj] = c
65 f_start[i + vj, j - vi] = c
66 f_start[i - vj, j + vi] = c
68 for l in range(nb_iterations):
71 if i < 0 or i >= height or j < 0 or j >= width:
79 f_end[i - vi, j - vj] = c
80 f_end[i + vj, j - vi] = c
81 f_end[i - vj, j + vi] = c
83 pairs.append((f_start, f_end))
87 if torch.rand(1) < 0.5:
90 [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()],
97 [p[1].flatten(), torch.tensor([token_backward]), p[0].flatten()],
102 return torch.cat(result, dim=0)
105 def sample2img(seq, height, width, upscale=15):
106 f_first = seq[:, : height * width].reshape(-1, height, width)
107 f_second = seq[:, height * width + 1 :].reshape(-1, height, width)
108 direction = seq[:, height * width]
110 def mosaic(x, upscale):
111 x = x.reshape(-1, height, width)
112 m = torch.logical_and(x >= 0, x < first_fish_token + nb_fish_tokens).long()
113 x = colors[x * m].permute(0, 3, 1, 2)
115 x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale)
116 x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale)
118 for n in range(m.size(0)):
119 for i in range(m.size(1)):
120 for j in range(m.size(2)):
122 for k in range(2, upscale - 2):
123 x[n, :, i * upscale + k, j * upscale + k] = 0
124 x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0
128 direction_symbol = torch.full((direction.size(0), height * upscale, upscale), 0)
129 direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2)
130 separator = torch.full((direction.size(0), 3, height * upscale, 1), 0)
132 for n in range(direction_symbol.size(0)):
133 if direction[n] == token_forward:
134 for k in range(upscale):
138 (height * upscale) // 2 - upscale // 2 + k,
139 3 + abs(k - upscale // 2),
141 elif direction[n] == token_backward:
142 for k in range(upscale):
146 (height * upscale) // 2 - upscale // 2 + k,
147 3 + upscale // 2 - abs(k - upscale // 2),
150 for k in range(2, upscale - 2):
152 n, :, (height * upscale) // 2 - upscale // 2 + k, k
155 n, :, (height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k
160 mosaic(f_first, upscale),
164 mosaic(f_second, upscale),
173 result.append("".join([token2char[v] for v in s]))
177 ######################################################################
179 if __name__ == "__main__":
183 start_time = time.perf_counter()
184 seq = generate(nb=90, height=height, width=width, max_nb_obj=3)
185 delay = time.perf_counter() - start_time
186 print(f"{seq.size(0)/delay:02f} samples/s")
188 print(seq2str(seq[:4]))
190 m = (torch.rand(seq.size()) < 0.05).long()
191 seq = (1 - m) * seq + m * 23
193 img = sample2img(seq, height, width)
196 torchvision.utils.save_image(
197 img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4