X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=world.py;h=ab02c8240fcb70a28a7e112bcc3e29707dd1c69f;hb=f87a57354a1e575181e760fdaedbb2c2d5cf9fa0;hp=0392940d22b6af336e2f18103d18fa09aac62f6a;hpb=f08c6c01dcc03b727c69478c3a1de7ebf9facd95;p=culture.git diff --git a/world.py b/world.py index 0392940..ab02c82 100755 --- a/world.py +++ b/world.py @@ -18,91 +18,154 @@ 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], - [255, 255, 0], + [255, 200, 0], [192, 192, 192], ] ) -token2char = "_X01234>" +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( nb, height, width, - obj_length=6, - mask_height=3, - mask_width=3, - nb_obj=3, + max_nb_obj=2, + nb_iterations=2, ): - intact = torch.zeros(nb, height, width, dtype=torch.int64) - n = torch.arange(intact.size(0)) - - for n in range(nb): - for c in torch.randperm(colors.size(0) - 2)[:nb_obj] + 2: - z = intact[n].flatten() - m = (torch.rand(z.size()) * (z == 0)).argmax(dim=0) - i, j = m // width, m % width + pairs = [] + + for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): + 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, + ) vm = torch.randint(4, (1,))[0] vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1) - for l in range(obj_length): - intact[n, i, j] = c + + 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 j += vj - if i < 0 or i >= height or j < 0 or j >= width or intact[n, i, j] != 0: + if i < 0 or i >= height or j < 0 or j >= width: i -= vi j -= vj - vi, vj = -vj, vi + vi, vj = -vi, -vj i += vi j += vj - if ( - i < 0 - or i >= height - or j < 0 - or j >= width - or intact[n, i, j] != 0 - ): - break - - masked = intact.clone() - - for n in range(nb): - i = torch.randint(height - mask_height + 1, (1,))[0] - j = torch.randint(width - mask_width + 1, (1,))[0] - masked[n, i : i + mask_height, j : j + mask_width] = 1 - return torch.cat( - [ - masked.flatten(1), - torch.full((masked.size(0), 1), len(colors)), - intact.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): - intact = seq[:, : height * width].reshape(-1, height, width) - masked = seq[:, height * width + 1 :].reshape(-1, height, width) - img_intact, img_masked = colors[intact], colors[masked] + 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, :] + ) + + 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_intact, - torch.full( - (img_intact.size(0), img_intact.size(1), 1, img_intact.size(3)), 1 - ), - img_masked, + 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 = [] @@ -118,13 +181,18 @@ if __name__ == "__main__": height, width = 6, 8 start_time = time.perf_counter() - seq = generate(nb=64, height=height, width=width) + 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 + )