Update.

[culture.git] / world.py

diff --git a/world.py b/world.py
index 89833e6..ab02c82 100755 (executable)
--- a/world.py
+++ b/world.py
@@ -18,7 +18,6 @@ from torch.nn import functional as F
 colors = torch.tensor(
     [
         [255, 255, 255],
 colors = torch.tensor(
     [
         [255, 255, 255],

-        [0, 0, 0],

         [255, 0, 0],
         [0, 128, 0],
         [0, 0, 255],
         [255, 0, 0],
         [0, 128, 0],
         [0, 0, 255],


@@ -27,7 +26,13 @@ colors = torch.tensor(
     ]
 )
 
     ]
 )
 

-token2char = "_X" + "".join([str(n) for n in range(len(colors) - 2)]) + ">"
+token_background = 0
+first_bird_token = 1
+nb_bird_tokens = len(colors) - 1
+token_forward = first_bird_token + nb_bird_tokens
+token_backward = token_forward + 1
+
+token2char = "_" + "".join([str(n) for n in range(len(colors) - 1)]) + "><"

 
 
 def generate(
 
 
 def generate(


@@ -37,13 +42,17 @@ def generate(
     max_nb_obj=2,
     nb_iterations=2,
 ):
     max_nb_obj=2,
     nb_iterations=2,
 ):

-    f_start = torch.zeros(nb, height, width, dtype=torch.int64)
-    f_end = torch.zeros(nb, height, width, dtype=torch.int64)
-    n = torch.arange(f_start.size(0))
+    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 n in range(nb):
-        nb_fish = torch.randint(max_nb_obj, (1,)).item() + 1
-        for c in torch.randperm(colors.size(0) - 2)[:nb_fish].sort().values:
+        nb_birds = torch.randint(max_nb_obj, (1,)).item() + 1
+        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,
             i, j = (
                 torch.randint(height - 2, (1,))[0] + 1,
                 torch.randint(width - 2, (1,))[0] + 1,


@@ -51,10 +60,10 @@ def generate(
             vm = torch.randint(4, (1,))[0]
             vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
 
             vm = torch.randint(4, (1,))[0]
             vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * (2 * (vm % 2) - 1)
 

-            f_start[n, i, j] = c + 2
-            f_start[n, i - vi, j - vj] = c + 2
-            f_start[n, i + vj, j - vi] = c + 2
-            f_start[n, i - vj, j + vi] = c + 2
+            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
 
             for l in range(nb_iterations):
                 i += vi


@@ -66,42 +75,97 @@ def generate(
                     i += vi
                     j += vj
 
                     i += vi
                     j += vj
 

-            f_end[n, i, j] = c + 2
-            f_end[n, i - vi, j - vj] = c + 2
-            f_end[n, i + vj, j - vi] = c + 2
-            f_end[n, i - vj, j + vi] = c + 2
-
-    return torch.cat(
-        [
-            f_end.flatten(1),
-            torch.full((f_end.size(0), 1), len(colors)),
-            f_start.flatten(1),
-        ],
-        dim=1,
-    )
+            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))

 
 

-def sample2img(seq, height, width):
-    f_start = seq[:, : height * width].reshape(-1, height, width)
-    f_start = (f_start >= len(colors)).long() + (f_start < len(colors)).long() * f_start
-    f_end = seq[:, height * width + 1 :].reshape(-1, height, width)
-    f_end = (f_end >= len(colors)).long() + (f_end < len(colors)).long() * 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, :]
+            )

 
 

-    img_f_start, img_f_end = colors[f_start], colors[f_end]
+    return torch.cat(result, dim=0)
+
+
+def sample2img(seq, height, width, upscale=15):
+    f_first = seq[:, : height * width].reshape(-1, height, width)
+    f_second = seq[:, height * width + 1 :].reshape(-1, height, width)
+    direction = seq[:, height * width]
+
+    def mosaic(x, upscale):
+        x = x.reshape(-1, height, width)
+        m = torch.logical_and(x >= 0, x < first_bird_token + nb_bird_tokens).long()
+        x = colors[x * m].permute(0, 3, 1, 2)
+        s = x.shape
+        x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale)
+        x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale)
+
+        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, upscale - 2):
+                            x[n, :, i * upscale + k, j * upscale + k] = 0
+                            x[n, :, i * upscale + upscale - 1 - k, j * upscale + k] = 0
+
+        return x
+
+    direction_symbol = torch.full((direction.size(0), height * upscale, upscale), 0)
+    direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2)
+    separator = torch.full((direction.size(0), 3, height * upscale, 1), 0)
+
+    for n in range(direction_symbol.size(0)):
+        if direction[n] == token_forward:
+            for k in range(upscale):
+                direction_symbol[
+                    n,
+                    :,
+                    (height * upscale) // 2 - upscale // 2 + k,
+                    3 + abs(k - upscale // 2),
+                ] = 0
+        elif direction[n] == token_backward:
+            for k in range(upscale):
+                direction_symbol[
+                    n,
+                    :,
+                    (height * upscale) // 2 - upscale // 2 + k,
+                    3 + upscale // 2 - abs(k - upscale // 2),
+                ] = 0
+        else:
+            for k in range(2, upscale - 2):
+                direction_symbol[
+                    n, :, (height * upscale) // 2 - upscale // 2 + k, k
+                ] = 0
+                direction_symbol[
+                    n, :, (height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k
+                ] = 0

 
 

-    img = torch.cat(
+    return torch.cat(

         [
         [

-            img_f_start,
-            torch.full(
-                (img_f_start.size(0), img_f_start.size(1), 1, img_f_start.size(3)), 1
-            ),
-            img_f_end,
+            mosaic(f_first, upscale),
+            separator,
+            direction_symbol,
+            separator,
+            mosaic(f_second, upscale),

         ],
         ],

-        dim=2,
+        dim=3,

     )
 
     )
 

-    return img.permute(0, 3, 1, 2)
-

 
 def seq2str(seq):
     result = []
 
 def seq2str(seq):
     result = []


@@ -117,13 +181,18 @@ if __name__ == "__main__":
 
     height, width = 6, 8
     start_time = time.perf_counter()
 
     height, width = 6, 8
     start_time = time.perf_counter()

-    seq = generate(nb=64, height=height, width=width, max_nb_obj=3)
+    seq = generate(nb=90, height=height, width=width, max_nb_obj=3)

     delay = time.perf_counter() - start_time
     print(f"{seq.size(0)/delay:02f} samples/s")
 
     print(seq2str(seq[:4]))
 
     delay = time.perf_counter() - start_time
     print(f"{seq.size(0)/delay:02f} samples/s")
 
     print(seq2str(seq[:4]))
 

+    # m = (torch.rand(seq.size()) < 0.05).long()
+    # seq = (1 - m) * seq + m * 23
+

     img = sample2img(seq, height, width)
     print(img.size())
 
     img = sample2img(seq, height, width)
     print(img.size())
 

-    torchvision.utils.save_image(img.float() / 255.0, "world.png", nrow=8, padding=2)
+    torchvision.utils.save_image(
+        img.float() / 255.0, "/tmp/world.png", nrow=6, padding=4
+    )