X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=world.py;h=36aa1e97638403195c0d711fde941be83ef17a7a;hb=1fb7ff66e07f267511f0f2592509a3f3212fc8f0;hp=2e13ec38c0c9c073fef95dc55f3bb5925d75cdcd;hpb=276696d1e7407a40d55c2a50b12839b52c0fed1e;p=culture.git diff --git a/world.py b/world.py index 2e13ec3..36aa1e9 100755 --- a/world.py +++ b/world.py @@ -19,27 +19,152 @@ colors = torch.tensor( [ [255, 255, 255], [255, 0, 0], - [0, 128, 0], + [0, 192, 0], [0, 0, 255], - [255, 200, 0], + [255, 192, 0], + [0, 255, 255], + [255, 0, 255], + [192, 255, 192], + [255, 192, 192], + [192, 192, 255], [192, 192, 192], ] ) token_background = 0 -first_fish_token = 1 -nb_fish_tokens = len(colors) - 1 -token_forward = first_fish_token + nb_fish_tokens +first_bird_token = 1 +nb_bird_tokens = colors.size(0) - 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)]) + "><" +token2char = "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" -def generate( +def generate_seq( + nb, height, width, nb_birds=3, nb_iterations=2, return_iterations=False +): + pairs = [] + kept_iterations = [] + + for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): + while True: + iterations = [] + + f_start = torch.zeros(height, width, dtype=torch.int64) + + i, j, vi, vj = ( + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + ) + + col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1 + + for n in range(nb_birds): + c = col[n] + + while True: + i[n], j[n] = ( + torch.randint(height, (1,))[0], + torch.randint(width, (1,))[0], + ) + vm = torch.randint(4, (1,))[0] + vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 + if ( + i[n] - vi[n] >= 0 + and i[n] - vi[n] < height + and j[n] - vj[n] >= 0 + and j[n] - vj[n] < width + and f_start[i[n], j[n]] == 0 + and f_start[i[n] - vi[n], j[n]] == 0 + and f_start[i[n], j[n] - vj[n]] == 0 + ): + break + + f_start[i[n], j[n]] = c + f_start[i[n] - vi[n], j[n]] = c + f_start[i[n], j[n] - vj[n]] = c + + f_end = f_start.clone() + + for l in range(nb_iterations): + iterations.append(f_end.clone()) + f_end[...] = 0 + nb_collisions = 0 + for n in range(nb_birds): + c = col[n] + + pi, pj, pvi, pvj = ( + i[n].item(), + j[n].item(), + vi[n].item(), + vj[n].item(), + ) + + if (i[n] == 0 and vi[n] == -1) or ( + i[n] == height - 1 and vi[n] == 1 + ): + vi[n] = -vi[n] + if (j[n] == 0 and vj[n] == -1) or ( + j[n] == width - 1 and vj[n] == 1 + ): + vj[n] = -vj[n] + + i[n] += vi[n] + j[n] += vj[n] + + if not ( + f_end[i[n], j[n]] == 0 + and f_end[i[n] - vi[n], j[n]] == 0 + and f_end[i[n], j[n] - vj[n]] == 0 + ): + nb_collisions += 1 + + f_end[i[n], j[n]] = c + f_end[i[n] - vi[n], j[n]] = c + f_end[i[n], j[n] - vj[n]] = c + + iterations.append(f_end.clone()) + + if nb_collisions == 0: + break + + kept_iterations.append(iterations) + pairs.append((f_start, 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, :] + ) + + if return_iterations: + # iterations = torch.cat([ torch.cat([ x[None, None] for x in l], dim = 1) for l in kept_iterations ], dim=0) + return torch.cat(result, dim=0), kept_iterations + else: + return torch.cat(result, dim=0) + + +###################################################################### + + +def generate_seq_old( nb, height, width, - max_nb_obj=2, + nb_birds=3, nb_iterations=2, ): pairs = [] @@ -49,9 +174,8 @@ def generate( f_end = torch.zeros(height, width, dtype=torch.int64) n = torch.arange(f_start.size(0)) - nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1 for c in ( - (torch.randperm(nb_fish_tokens) + first_fish_token)[:nb_fish].sort().values + (torch.randperm(nb_bird_tokens) + first_bird_token)[:nb_birds].sort().values ): i, j = ( torch.randint(height - 2, (1,))[0] + 1, @@ -102,32 +226,37 @@ def generate( return torch.cat(result, dim=0) -def sample2img(seq, height, width, upscale=15): +def frame2img(x, height, width, upscale=15): + 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) + + x[:, :, :, torch.arange(0, x.size(3), upscale)] = 0 + x[:, :, torch.arange(0, x.size(2), upscale), :] = 0 + x = x[:, :, 1:, 1:] + + 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 + + +def seq2img(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_fish_token + nb_fish_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 = torch.full((direction.size(0), height * upscale - 1, upscale), 0) direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2) - separator = torch.full((direction.size(0), 3, height * upscale, 1), 0) + separator = torch.full((direction.size(0), 3, height * upscale - 1, 1), 0) for n in range(direction_symbol.size(0)): if direction[n] == token_forward: @@ -136,7 +265,7 @@ def sample2img(seq, height, width, upscale=15): n, :, (height * upscale) // 2 - upscale // 2 + k, - 3 + abs(k - upscale // 2), + 3 + upscale // 2 - abs(k - upscale // 2), ] = 0 elif direction[n] == token_backward: for k in range(upscale): @@ -144,7 +273,7 @@ def sample2img(seq, height, width, upscale=15): n, :, (height * upscale) // 2 - upscale // 2 + k, - 3 + upscale // 2 - abs(k - upscale // 2), + 3 + abs(k - upscale // 2), ] = 0 else: for k in range(2, upscale - 2): @@ -157,11 +286,11 @@ def sample2img(seq, height, width, upscale=15): return torch.cat( [ - mosaic(f_first, upscale), + frame2img(f_first, height, width, upscale), separator, direction_symbol, separator, - mosaic(f_second, upscale), + frame2img(f_second, height, width, upscale), ], dim=3, ) @@ -181,18 +310,30 @@ if __name__ == "__main__": height, width = 6, 8 start_time = time.perf_counter() - seq = generate(nb=90, height=height, width=width, max_nb_obj=3) + seq, it = generate_seq( + nb=64, height=height, width=width, nb_iterations=100, return_iterations=True + ) 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 + for t in range(len(it[0])): + img = torch.cat([frame2img(f[t], height, width) for f in it], dim=0) + torchvision.utils.save_image( + img.float() / 255.0, + f"/tmp/frame_{t:03d}.png", + nrow=8, + padding=6, + pad_value=0, + ) + + # m = (torch.rand(seq.size()) < 0.05).long() + # seq = (1 - m) * seq + m * 23 - img = sample2img(seq, height, width) + img = seq2img(seq, height, width) print(img.size()) torchvision.utils.save_image( - img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4 + img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0 )