- 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
+ 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))
+
+ 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:
+ i, j = (
+ torch.randint(height - 2, (1,))[0] + 1,
+ torch.randint(width - 2, (1,))[0] + 1,
+ )