X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=sky.py;h=6ef8a3af2184777c223dfd1803647a49bd3dd54d;hb=ff436fbfeea7830a600c81175b5376d02925a1a0;hp=040ec67e0ed194687d28cbb7f4d1d8c2808736b4;hpb=64abc9f3a07a8211f308271fde7d8f876a968ab5;p=culture.git

diff --git a/sky.py b/sky.py
index 040ec67..6ef8a3a 100755
--- a/sky.py
+++ b/sky.py
@@ -5,7 +5,7 @@
 
 # Written by Francois Fleuret <francois@fleuret.org>
 
-import math, sys, tqdm, os
+import math, sys, tqdm, os, warnings
 
 import torch, torchvision
 
@@ -200,42 +200,74 @@ class Sky(problem.Problem):
         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)
+        h = img_prompts.size(2)
+        img_answers = add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1)
+
+        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, 255, margin)
-        # img_answers = add_frame(img_answers, 255, 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
+        )
 
-        img_prompts = add_frame(img_prompts, predicted_prompts, margin)
-        img_answers = add_frame(img_answers, predicted_answers, margin)
+        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
         )
 
     ######################################################################
@@ -246,8 +278,20 @@ class Sky(problem.Problem):
     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(