+)
+
+token2char = "_X01234>"
+
+
+def generate(
+ nb,
+ height,
+ width,
+ obj_length=6,
+ mask_height=3,
+ mask_width=3,
+ nb_obj=3,
+):
+ 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
+ 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
+ i += vi
+ j += vj
+ if i < 0 or i >= height or j < 0 or j >= width or intact[n, i, j] != 0:
+ i -= vi
+ j -= vj
+ vi, vj = -vj, vi
+ 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,
+ )
+
+
+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]
+
+ img = torch.cat(
+ [
+ img_intact,
+ torch.full(
+ (img_intact.size(0), img_intact.size(1), 1, img_intact.size(3)), 1
+ ),
+ img_masked,
+ ],
+ dim=2,
+ )
+
+ return img.permute(0, 3, 1, 2)
+
+
+def seq2str(seq):
+ result = []
+ for s in seq:
+ result.append("".join([token2char[v] for v in s]))
+ return result