# Written by Francois Fleuret <francois@fleuret.org>
-import math, sys, tqdm, os
+import math, sys, tqdm, os, warnings
import torch, torchvision
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,
- )
+ def add_frame(x, c, margin, bottom=False):
+ if bottom:
+ h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
+ else:
+ h, w, di, dj = (
+ x.size(2) + 2 * margin,
+ x.size(3) + 2 * margin,
+ margin,
+ margin,
+ )
+
+ y = x.new_full((x.size(0), x.size(1), h, w), 0)
+
if type(c) is int:
y[...] = c
else:
c = c.long()[:, None]
- c = c * torch.tensor([192, 192, 192], device=c.device) + (
+ c = c * torch.tensor([0, 0, 0], device=c.device) + (
1 - c
) * torch.tensor([255, 255, 255], device=c.device)
y[...] = c[:, :, None, None]
- y[:, :, margin:-margin, margin:-margin] = x
+
+ y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
+
return y
margin = 4
- img_prompts = add_frame(self.frame2img(prompts.to("cpu")), 0, 1)
- img_answers = add_frame(self.frame2img(answers.to("cpu")), 0, 1)
+ img_prompts = add_frame(self.frame2img(prompts.to("cpu")), c=0, margin=1)
+ img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1)
- # img_prompts = add_frame(img_prompts, 255, margin)
- # img_answers = add_frame(img_answers, 255, margin)
+ img_prompts = add_frame(img_prompts, c=255, margin=margin, bottom=True)
+ img_answers = add_frame(img_answers, c=255, margin=margin, bottom=True)
- img_prompts = add_frame(img_prompts, predicted_prompts, margin)
- img_answers = add_frame(img_answers, predicted_answers, margin)
+ img_prompts = add_frame(
+ img_prompts, c=predicted_prompts, margin=margin, bottom=True
+ )
+ img_answers = add_frame(
+ img_answers, c=predicted_answers, margin=margin, bottom=True
+ )
+
+ marker_size = 8
separator = img_prompts.new_full(
- (img_prompts.size(0), img_prompts.size(1), img_prompts.size(2), margin), 255
+ (
+ img_prompts.size(0),
+ img_prompts.size(1),
+ img_prompts.size(2),
+ marker_size,
+ ),
+ 255,
)
- img = torch.cat([img_prompts, img_answers], dim=3)
+ for k in range(2, 2 * marker_size - 3):
+ i = k + 1 - marker_size
+ j = marker_size - 2 - abs(k - marker_size + 1)
+ separator[:, :, separator.size(2) // 2 + i, j] = 0
+ separator[:, :, separator.size(2) // 2 + i + 1, j] = 0
+
+ img = torch.cat([img_prompts, separator, img_answers], dim=3)
image_name = os.path.join(result_dir, filename)
torchvision.utils.save_image(
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)
+ # warnings.warn("dirty test with longer answer", RuntimeWarning)
+ # answers = torch.cat(
+ # [
+ # frame_sequences[:, frame_sequences.size(1) // 2 :],
+ # frame_sequences[:, frame_sequences.size(1) // 2 :],
+ # ],
+ # dim=3,
+ # ).flatten(1)
+
return prompts, answers
def save_quizzes(