X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=world.py;h=ab02c8240fcb70a28a7e112bcc3e29707dd1c69f;hb=f87a57354a1e575181e760fdaedbb2c2d5cf9fa0;hp=118a470b29b159d97826dc895362db8e8d673ded;hpb=8adf0586ee5aeb9fbdf81b78c7ff4b484a9b82ab;p=culture.git diff --git a/world.py b/world.py index 118a470..ab02c82 100755 --- a/world.py +++ b/world.py @@ -18,7 +18,6 @@ from torch.nn import functional as F colors = torch.tensor( [ [255, 255, 255], - [0, 0, 0], [255, 0, 0], [0, 128, 0], [0, 0, 255], @@ -27,7 +26,13 @@ colors = torch.tensor( ] ) -token2char = "_X" + "".join([str(n) for n in range(len(colors) - 2)]) + ">" +token_background = 0 +first_bird_token = 1 +nb_bird_tokens = len(colors) - 1 +token_forward = first_bird_token + nb_bird_tokens +token_backward = token_forward + 1 + +token2char = "_" + "".join([str(n) for n in range(len(colors) - 1)]) + "><" def generate( @@ -37,13 +42,17 @@ def generate( max_nb_obj=2, nb_iterations=2, ): - f_start = torch.zeros(nb, height, width, dtype=torch.int64) - f_end = torch.zeros(nb, height, width, dtype=torch.int64) - n = torch.arange(f_start.size(0)) + pairs = [] for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): - nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1 - for c in torch.randperm(colors.size(0) - 2)[:nb_fish].sort().values: + f_start = torch.zeros(height, width, dtype=torch.int64) + f_end = torch.zeros(height, width, dtype=torch.int64) + n = torch.arange(f_start.size(0)) + + nb_birds = torch.randint(max_nb_obj, (1,)).item() + 1 + for c in ( + (torch.randperm(nb_bird_tokens) + first_bird_token)[:nb_birds].sort().values + ): i, j = ( torch.randint(height - 2, (1,))[0] + 1, torch.randint(width - 2, (1,))[0] + 1, @@ -51,10 +60,10 @@ def generate( vm = torch.randint(4, (1,))[0] vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1) - f_start[n, i, j] = c + 2 - f_start[n, i - vi, j - vj] = c + 2 - f_start[n, i + vj, j - vi] = c + 2 - f_start[n, i - vj, j + vi] = c + 2 + f_start[i, j] = c + f_start[i - vi, j - vj] = c + f_start[i + vj, j - vi] = c + f_start[i - vj, j + vi] = c for l in range(nb_iterations): i += vi @@ -66,42 +75,97 @@ def generate( i += vi j += vj - f_end[n, i, j] = c + 2 - f_end[n, i - vi, j - vj] = c + 2 - f_end[n, i + vj, j - vi] = c + 2 - f_end[n, i - vj, j + vi] = c + 2 - - return torch.cat( - [ - f_end.flatten(1), - torch.full((f_end.size(0), 1), len(colors)), - f_start.flatten(1), - ], - dim=1, - ) + f_end[i, j] = c + f_end[i - vi, j - vj] = c + f_end[i + vj, j - vi] = c + f_end[i - vj, j + vi] = c + pairs.append((f_start, f_end)) -def sample2img(seq, height, width): - f_start = seq[:, : height * width].reshape(-1, height, width) - f_start = (f_start >= len(colors)).long() + (f_start < len(colors)).long() * f_start - f_end = seq[:, height * width + 1 :].reshape(-1, height, width) - f_end = (f_end >= len(colors)).long() + (f_end < len(colors)).long() * f_end + result = [] + for p in pairs: + if torch.rand(1) < 0.5: + result.append( + torch.cat( + [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()], + dim=0, + )[None, :] + ) + else: + result.append( + torch.cat( + [p[1].flatten(), torch.tensor([token_backward]), p[0].flatten()], + dim=0, + )[None, :] + ) - img_f_start, img_f_end = colors[f_start], colors[f_end] + return torch.cat(result, dim=0) + + +def sample2img(seq, height, width, upscale=15): + 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) + m = torch.logical_and(x >= 0, x < first_bird_token + nb_bird_tokens).long() + x = colors[x * m].permute(0, 3, 1, 2) + s = x.shape + x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale) + x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale) + + for n in range(m.size(0)): + for i in range(m.size(1)): + for j in range(m.size(2)): + if m[n, i, j] == 0: + for k in range(2, upscale - 2): + x[n, :, i * upscale + k, j * upscale + k] = 0 + x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0 + + return x + + 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 - img = torch.cat( + return torch.cat( [ - img_f_start, - torch.full( - (img_f_start.size(0), img_f_start.size(1), 1, img_f_start.size(3)), 1 - ), - img_f_end, + mosaic(f_first, upscale), + separator, + direction_symbol, + separator, + mosaic(f_second, upscale), ], - dim=2, + dim=3, ) - return img.permute(0, 3, 1, 2) - def seq2str(seq): result = [] @@ -117,13 +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 + img = sample2img(seq, height, width) print(img.size()) - torchvision.utils.save_image(img.float() / 255.0, "world.png", nrow=8, padding=2) + torchvision.utils.save_image( + img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4 + )