# 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) + (
- 1 - c
- ) * torch.tensor([255, 255, 255], device=c.device)
+ c = (
+ (c == 1).long() * torch.tensor([0, 255, 0], device=c.device)
+ + (c == 0).long() * torch.tensor([255, 255, 255], device=c.device)
+ + (c == -1).long() * torch.tensor([255, 0, 0], 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)
+ h = img_prompts.size(2)
+ 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 = 16
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)
+ separator[:, :, 0] = 0
+ separator[:, :, h - 1] = 0
+
+ for k in range(1, 2 * marker_size - 8):
+ i = k - (marker_size - 4)
+ j = marker_size - 5 - abs(i)
+ separator[:, :, h // 2 - 1 + i, 2 + j] = 0
+ separator[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
+
+ img = torch.cat([img_prompts, separator, img_answers], dim=3)
image_name = os.path.join(result_dir, filename)
torchvision.utils.save_image(
- img.float() / 255.0, image_name, nrow=6, padding=margin * 2, pad_value=1.0
+ img.float() / 255.0, image_name, nrow=6, padding=margin * 4, pad_value=1.0
)
######################################################################
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(
prompts, answers = sky.generate_prompts_and_answers(4)
- predicted_prompts = torch.rand(prompts.size(0)) < 0.5
- predicted_answers = torch.rand(answers.size(0)) < 0.5
+ predicted_prompts = torch.randint(3, (prompts.size(0),)) - 1
+ predicted_answers = torch.randint(3, (prompts.size(0),)) - 1
sky.save_quizzes(
"/tmp", "test", prompts, answers, predicted_prompts, predicted_answers