world.py

   1 #!/usr/bin/env python
   2
   3 import math
   4
   5 import torch, torchvision
   6
   7 from torch import nn
   8 from torch.nn import functional as F
   9 import cairo
  10
  11
  12 class Box:
  13     def __init__(self, x, y, w, h, r, g, b):
  14         self.x = x
  15         self.y = y
  16         self.w = w
  17         self.h = h
  18         self.r = r
  19         self.g = g
  20         self.b = b
  21
  22     def collision(self, scene):
  23         for c in scene:
  24             if (
  25                 self is not c
  26                 and max(self.x, c.x) <= min(self.x + self.w, c.x + c.w)
  27                 and max(self.y, c.y) <= min(self.y + self.h, c.y + c.h)
  28             ):
  29                 return True
  30         return False
  31
  32
  33 def scene2tensor(xh, yh, scene, size=64):
  34     width, height = size, size
  35     pixel_map = torch.ByteTensor(width, height, 4).fill_(255)
  36     data = pixel_map.numpy()
  37     surface = cairo.ImageSurface.create_for_data(
  38         data, cairo.FORMAT_ARGB32, width, height
  39     )
  40
  41     ctx = cairo.Context(surface)
  42     ctx.set_fill_rule(cairo.FILL_RULE_EVEN_ODD)
  43
  44     for b in scene:
  45         ctx.move_to(b.x * size, b.y * size)
  46         ctx.rel_line_to(b.w * size, 0)
  47         ctx.rel_line_to(0, b.h * size)
  48         ctx.rel_line_to(-b.w * size, 0)
  49         ctx.close_path()
  50         ctx.set_source_rgba(b.r, b.g, b.b, 1.0)
  51         ctx.fill()
  52
  53     hs = size * 0.1
  54     ctx.set_source_rgba(0.0, 0.0, 0.0, 1.0)
  55     ctx.move_to(xh * size - hs / 2, yh * size - hs / 2)
  56     ctx.rel_line_to(hs, 0)
  57     ctx.rel_line_to(0, hs)
  58     ctx.rel_line_to(-hs, 0)
  59     ctx.close_path()
  60     ctx.fill()
  61
  62     return pixel_map[None, :, :, :3].flip(-1).permute(0, 3, 1, 2).float() / 255
  63
  64
  65 def random_scene():
  66     scene = []
  67     colors = [
  68         (1.00, 0.00, 0.00),
  69         (0.00, 1.00, 0.00),
  70         (0.60, 0.60, 1.00),
  71         (1.00, 1.00, 0.00),
  72         (0.75, 0.75, 0.75),
  73     ]
  74
  75     for k in range(10):
  76         wh = torch.rand(2) * 0.2 + 0.2
  77         xy = torch.rand(2) * (1 - wh)
  78         c = colors[torch.randint(len(colors), (1,))]
  79         b = Box(
  80             xy[0].item(), xy[1].item(), wh[0].item(), wh[1].item(), c[0], c[1], c[2]
  81         )
  82         if not b.collision(scene):
  83             scene.append(b)
  84
  85     return scene
  86
  87
  88 def sequence(nb_steps=10, all_frames=False):
  89     delta = 0.1
  90     effects = [
  91         (False, 0, 0),
  92         (False, delta, 0),
  93         (False, 0, delta),
  94         (False, -delta, 0),
  95         (False, 0, -delta),
  96         (True, delta, 0),
  97         (True, 0, delta),
  98         (True, -delta, 0),
  99         (True, 0, -delta),
 100     ]
 101
 102     while True:
 103         frames = []
 104
 105         scene = random_scene()
 106         xh, yh = tuple(x.item() for x in torch.rand(2))
 107
 108         frames.append(scene2tensor(xh, yh, scene))
 109
 110         actions = torch.randint(len(effects), (nb_steps,))
 111         change = False
 112
 113         for a in actions:
 114             g, dx, dy = effects[a]
 115             if g:
 116                 for b in scene:
 117                     if b.x <= xh and b.x + b.w >= xh and b.y <= yh and b.y + b.h >= yh:
 118                         x, y = b.x, b.y
 119                         b.x += dx
 120                         b.y += dy
 121                         if (
 122                             b.x < 0
 123                             or b.y < 0
 124                             or b.x + b.w > 1
 125                             or b.y + b.h > 1
 126                             or b.collision(scene)
 127                         ):
 128                             b.x, b.y = x, y
 129                         else:
 130                             xh += dx
 131                             yh += dy
 132                             change = True
 133             else:
 134                 x, y = xh, yh
 135                 xh += dx
 136                 yh += dy
 137                 if xh < 0 or xh > 1 or yh < 0 or yh > 1:
 138                     xh, yh = x, y
 139
 140             if all_frames:
 141                 frames.append(scene2tensor(xh, yh, scene))
 142
 143         if not all_frames:
 144             frames.append(scene2tensor(xh, yh, scene))
 145
 146         if change:
 147             break
 148
 149     return frames, actions
 150
 151
 152 ######################################################################
 153
 154
 155 # ||x_i - c_j||^2 = ||x_i||^2 + ||c_j||^2 - 2<x_i, c_j>
 156 def sq2matrix(x, c):
 157     nx = x.pow(2).sum(1)
 158     nc = c.pow(2).sum(1)
 159     return nx[:, None] + nc[None, :] - 2 * x @ c.t()
 160
 161
 162 def update_centroids(x, c, nb_min=1):
 163     _, b = sq2matrix(x, c).min(1)
 164     b.squeeze_()
 165     nb_resets = 0
 166
 167     for k in range(0, c.size(0)):
 168         i = b.eq(k).nonzero(as_tuple=False).squeeze()
 169         if i.numel() >= nb_min:
 170             c[k] = x.index_select(0, i).mean(0)
 171         else:
 172             n = torch.randint(x.size(0), (1,))
 173             nb_resets += 1
 174             c[k] = x[n]
 175
 176     return c, b, nb_resets
 177
 178
 179 def kmeans(x, nb_centroids, nb_min=1):
 180     if x.size(0) < nb_centroids * nb_min:
 181         print("Not enough points!")
 182         exit(1)
 183
 184     c = x[torch.randperm(x.size(0))[:nb_centroids]]
 185     t = torch.full((x.size(0),), -1)
 186     n = 0
 187
 188     while True:
 189         c, u, nb_resets = update_centroids(x, c, nb_min)
 190         n = n + 1
 191         nb_changes = (u - t).sign().abs().sum() + nb_resets
 192         t = u
 193         if nb_changes == 0:
 194             break
 195
 196     return c, t
 197
 198
 199 ######################################################################
 200
 201
 202 def patchify(x, factor, invert_size=None):
 203     if invert_size is None:
 204         return (
 205             x.reshape(
 206                 x.size(0),  # 0
 207                 x.size(1),  # 1
 208                 factor,  # 2
 209                 x.size(2) // factor,  # 3
 210                 factor,  # 4
 211                 x.size(3) // factor,  # 5
 212             )
 213             .permute(0, 2, 4, 1, 3, 5)
 214             .reshape(-1, x.size(1), x.size(2) // factor, x.size(3) // factor)
 215         )
 216     else:
 217         return (
 218             x.reshape(
 219                 invert_size[0],  # 0
 220                 factor,  # 1
 221                 factor,  # 2
 222                 invert_size[1],  # 3
 223                 invert_size[2] // factor,  # 4
 224                 invert_size[3] // factor,  # 5
 225             )
 226             .permute(0, 3, 1, 4, 2, 5)
 227             .reshape(invert_size)
 228         )
 229
 230
 231 def train_encoder(input, device=torch.device("cpu")):
 232     class SomeLeNet(nn.Module):
 233         def __init__(self):
 234             super().__init__()
 235             self.conv1 = nn.Conv2d(1, 32, kernel_size=5)
 236             self.conv2 = nn.Conv2d(32, 64, kernel_size=5)
 237             self.fc1 = nn.Linear(256, 200)
 238             self.fc2 = nn.Linear(200, 10)
 239
 240         def forward(self, x):
 241             x = F.relu(F.max_pool2d(self.conv1(x), kernel_size=3))
 242             x = F.relu(F.max_pool2d(self.conv2(x), kernel_size=2))
 243             x = x.view(x.size(0), -1)
 244             x = F.relu(self.fc1(x))
 245             x = self.fc2(x)
 246             return x
 247
 248     ######################################################################
 249
 250     model = SomeLeNet()
 251
 252     nb_parameters = sum(p.numel() for p in model.parameters())
 253
 254     print(f"nb_parameters {nb_parameters}")
 255
 256     optimizer = torch.optim.SGD(model.parameters(), lr=lr)
 257     criterion = nn.CrossEntropyLoss()
 258
 259     model.to(device)
 260     criterion.to(device)
 261
 262     train_input, train_targets = train_input.to(device), train_targets.to(device)
 263     test_input, test_targets = test_input.to(device), test_targets.to(device)
 264
 265     mu, std = train_input.mean(), train_input.std()
 266     train_input.sub_(mu).div_(std)
 267     test_input.sub_(mu).div_(std)
 268
 269     start_time = time.perf_counter()
 270
 271     for k in range(nb_epochs):
 272         acc_loss = 0.0
 273
 274         for input, targets in zip(
 275             train_input.split(batch_size), train_targets.split(batch_size)
 276         ):
 277             output = model(input)
 278             loss = criterion(output, targets)
 279             acc_loss += loss.item()
 280
 281             optimizer.zero_grad()
 282             loss.backward()
 283             optimizer.step()
 284
 285         nb_test_errors = 0
 286         for input, targets in zip(
 287             test_input.split(batch_size), test_targets.split(batch_size)
 288         ):
 289             wta = model(input).argmax(1)
 290             nb_test_errors += (wta != targets).long().sum()
 291         test_error = nb_test_errors / test_input.size(0)
 292         duration = time.perf_counter() - start_time
 293
 294         print(f"loss {k} {duration:.02f}s {acc_loss:.02f} {test_error*100:.02f}%")
 295
 296
 297 ######################################################################
 298
 299 if __name__ == "__main__":
 300     import time
 301
 302     all_frames = []
 303     nb = 1000
 304     start_time = time.perf_counter()
 305     for n in range(nb):
 306         frames, actions = sequence(nb_steps=31)
 307         all_frames += frames
 308     end_time = time.perf_counter()
 309     print(f"{nb / (end_time - start_time):.02f} samples per second")
 310
 311     input = torch.cat(all_frames, 0)
 312
 313     # x = patchify(input, 8)
 314     # y = x.reshape(x.size(0), -1)
 315     # print(f"{x.size()=} {y.size()=}")
 316     # centroids, t = kmeans(y, 4096)
 317     # results = centroids[t]
 318     # results = results.reshape(x.size())
 319     # results = patchify(results, 8, input.size())
 320
 321     print(f"{input.size()=} {results.size()=}")
 322
 323     torchvision.utils.save_image(input[:64], "orig.png", nrow=8)
 324     torchvision.utils.save_image(results[:64], "qtiz.png", nrow=8)
 325
 326     # frames, actions = sequence(nb_steps=31, all_frames=True)
 327     # frames = torch.cat(frames, 0)
 328     # torchvision.utils.save_image(frames, "seq.png", nrow=8)