-def generate_seq_old(
- nb,
- height,
- width,
- nb_birds=3,
- nb_iterations=2,
-):
- 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))
-
- 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)
-
- 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:
- i -= vi
- j -= vj
- vi, vj = -vi, -vj
- i += vi
- j += vj
-
- 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))
-
- 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, :]
+ result = torch.zeros(
+ self.nb_iterations * self.speed,
+ self.height,
+ self.width,
+ dtype=torch.int64,
+ )
+
+ fine = torch.empty(self.nb_iterations * self.speed)
+
+ t_to_keep = (
+ torch.arange(self.nb_iterations, device=result.device) * self.speed
+ )
+
+ for l in range(self.nb_iterations * self.speed):
+ fine[l] = collision_okay()
+ for n in range(self.nb_birds):
+ c = col[n]
+ result[l, i[n], j[n]] = c
+ result[l, i[n] - vi[n], j[n]] = c
+ result[l, i[n], j[n] - vj[n]] = c
+
+ if (i[n] == 0 and vi[n] == -1) or (
+ i[n] == self.height - 1 and vi[n] == 1
+ ):
+ vi[n] = -vi[n]
+
+ if (j[n] == 0 and vj[n] == -1) or (
+ j[n] == self.width - 1 and vj[n] == 1
+ ):
+ vj[n] = -vj[n]
+
+ i[n] += vi[n]
+ j[n] += vj[n]
+
+ result = result[t_to_keep]
+ fine = fine[t_to_keep]
+
+ if fine[-1]:
+ break
+
+ frame_sequences.append(result)
+
+ return frame_sequences
+
+ ######################################################################
+
+ def generate_prompts_and_answers(self, nb):
+ frame_sequences = self.generate_frame_sequences(nb)
+ frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0)
+ prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1)
+ answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1)
+ return prompts, answers
+
+ ######################################################################
+
+ def frame2img(self, x, scale=15):
+ x = x.reshape(-1, self.height, self.width)
+ m = torch.logical_and(
+ x >= 0, x < self.first_bird_token + self.nb_bird_tokens
+ ).long()
+ x = self.colors[x * m].permute(0, 3, 1, 2)
+ s = x.shape
+ x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
+ x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
+
+ x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
+ x[:, :, torch.arange(0, x.size(2), scale), :] = 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, scale - 2):
+ for l in [0, 1]:
+ x[n, :, i * scale + k, j * scale + k - l] = 0
+ x[
+ n, :, i * scale + scale - 1 - k, j * scale + k - l
+ ] = 0
+
+ return x
+
+ def seq2str(self, seq):
+ result = []
+ for s in seq:
+ result.append("".join([self.token2char[v] for v in s]))
+ return result
+
+ def save_image(
+ self,
+ result_dir,
+ filename,
+ prompts,
+ answers,
+ predicted_prompts=None,
+ predicted_answers=None,
+ ):
+ if predicted_prompts is None:
+ predicted_prompts = 255
+
+ if predicted_answers is None:
+ predicted_answers = 255
+
+ def add_frame(x, c, margin):
+ y = x.new_full(
+ (x.size(0), x.size(1), x.size(2) + 2 * margin, x.size(3) + 2 * margin),
+ 0,