X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=sky.py;h=ec476a629fb34b5c1d1d7565ca952d1fb26eca34;hb=e2c3b8046c3fddef8aacb74cf5f848d42044897e;hp=3458d853110096b041ed9f64340a9d1a48f7777f;hpb=6dbc18a5db82b12b06212841426896412e8bd6de;p=culture.git diff --git a/sky.py b/sky.py index 3458d85..ec476a6 100755 --- a/sky.py +++ b/sky.py @@ -5,7 +5,7 @@ # Written by Francois Fleuret -import math, sys, tqdm +import math, sys, tqdm, os import torch, torchvision @@ -14,40 +14,46 @@ from torch.nn import functional as F ###################################################################### +import problem + + +class Sky(problem.Problem): + colors = torch.tensor( + [ + [255, 255, 255], + [255, 0, 0], + [0, 192, 0], + [0, 0, 255], + [255, 192, 0], + [0, 255, 255], + [255, 0, 255], + [192, 255, 192], + [255, 192, 192], + [192, 192, 255], + [192, 192, 192], + ] + ) + + token_background = 0 + first_bird_token = 1 + nb_bird_tokens = colors.size(0) - 1 + token_forward = first_bird_token + nb_bird_tokens + token_backward = token_forward + 1 + + token2char = ( + "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" + ) -colors = torch.tensor( - [ - [255, 255, 255], - [255, 0, 0], - [0, 192, 0], - [0, 0, 255], - [255, 192, 0], - [0, 255, 255], - [255, 0, 255], - [192, 255, 192], - [255, 192, 192], - [192, 192, 255], - [192, 192, 192], - ] -) - -token_background = 0 -first_bird_token = 1 -nb_bird_tokens = colors.size(0) - 1 -token_forward = first_bird_token + nb_bird_tokens -token_backward = token_forward + 1 - -token2char = "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" - - -class Sky: - def __init__(self, height, width): - self.heigh = heigh + def __init__(self, height=6, width=8, nb_birds=3, nb_iterations=2): + self.height = height self.width = width + self.nb_birds = nb_birds + self.nb_iterations = nb_iterations - def generate_seq( - nb, height, width, nb_birds=3, nb_iterations=2, return_iterations=False - ): + def direction_tokens(self): + return self.token_forward, self.token_backward + + def generate_seq(self, nb, return_iterations=False): pairs = [] kept_iterations = [] @@ -55,32 +61,37 @@ class Sky: while True: iterations = [] - f_start = torch.zeros(height, width, dtype=torch.int64) + f_start = torch.zeros(self.height, self.width, dtype=torch.int64) i, j, vi, vj = ( - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), ) - col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1 + col = ( + torch.randperm(self.colors.size(0) - 1)[: self.nb_birds] + .sort() + .values + + 1 + ) - for n in range(nb_birds): + for n in range(self.nb_birds): c = col[n] while True: i[n], j[n] = ( - torch.randint(height, (1,))[0], - torch.randint(width, (1,))[0], + torch.randint(self.height, (1,))[0], + torch.randint(self.width, (1,))[0], ) vm = torch.randint(4, (1,))[0] vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 if ( i[n] - vi[n] >= 0 - and i[n] - vi[n] < height + and i[n] - vi[n] < self.height and j[n] - vj[n] >= 0 - and j[n] - vj[n] < width + and j[n] - vj[n] < self.width and f_start[i[n], j[n]] == 0 and f_start[i[n] - vi[n], j[n]] == 0 and f_start[i[n], j[n] - vj[n]] == 0 @@ -93,11 +104,11 @@ class Sky: f_end = f_start.clone() - for l in range(nb_iterations): + for l in range(self.nb_iterations): iterations.append(f_end.clone()) f_end[...] = 0 nb_collisions = 0 - for n in range(nb_birds): + for n in range(self.nb_birds): c = col[n] pi, pj, pvi, pvj = ( @@ -108,11 +119,11 @@ class Sky: ) if (i[n] == 0 and vi[n] == -1) or ( - i[n] == height - 1 and vi[n] == 1 + i[n] == self.height - 1 and vi[n] == 1 ): vi[n] = -vi[n] if (j[n] == 0 and vj[n] == -1) or ( - j[n] == width - 1 and vj[n] == 1 + j[n] == self.width - 1 and vj[n] == 1 ): vj[n] = -vj[n] @@ -143,7 +154,11 @@ class Sky: if torch.rand(1) < 0.5: result.append( torch.cat( - [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()], + [ + p[0].flatten(), + torch.tensor([self.token_forward]), + p[1].flatten(), + ], dim=0, )[None, :] ) @@ -152,7 +167,7 @@ class Sky: torch.cat( [ p[1].flatten(), - torch.tensor([token_backward]), + torch.tensor([self.token_backward]), p[0].flatten(), ], dim=0, @@ -168,27 +183,26 @@ class Sky: ###################################################################### def generate_seq_old( + self, nb, - height, - width, - nb_birds=3, - nb_iterations=2, ): 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) + f_start = torch.zeros(self.height, self.width, dtype=torch.int64) + f_end = torch.zeros(self.height, self.width, dtype=torch.int64) n = torch.arange(f_start.size(0)) for c in ( - (torch.randperm(nb_bird_tokens) + first_bird_token)[:nb_birds] + (torch.randperm(self.nb_bird_tokens) + self.first_bird_token)[ + : self.nb_birds + ] .sort() .values ): i, j = ( - torch.randint(height - 2, (1,))[0] + 1, - torch.randint(width - 2, (1,))[0] + 1, + torch.randint(self.height - 2, (1,))[0] + 1, + torch.randint(self.width - 2, (1,))[0] + 1, ) vm = torch.randint(4, (1,))[0] vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * ( @@ -200,10 +214,10 @@ class Sky: f_start[i + vj, j - vi] = c f_start[i - vj, j + vi] = c - for l in range(nb_iterations): + for l in range(self.nb_iterations): i += vi j += vj - if i < 0 or i >= height or j < 0 or j >= width: + if i < 0 or i >= self.height or j < 0 or j >= self.width: i -= vi j -= vj vi, vj = -vi, -vj @@ -222,7 +236,11 @@ class Sky: if torch.rand(1) < 0.5: result.append( torch.cat( - [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()], + [ + p[0].flatten(), + torch.tensor([self.token_forward]), + p[1].flatten(), + ], dim=0, )[None, :] ) @@ -231,7 +249,7 @@ class Sky: torch.cat( [ p[1].flatten(), - torch.tensor([token_backward]), + torch.tensor([self.token_backward]), p[0].flatten(), ], dim=0, @@ -240,10 +258,12 @@ class Sky: return torch.cat(result, dim=0) - def frame2img(x, height, width, upscale=15): - 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) + def frame2img(self, x, upscale=15): + x = x.reshape(-1, self.height, self.width) + m = torch.logical_and( + x >= 0, x < self.first_bird_token + self.nb_bird_tokens + ).long() + x = self.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) @@ -262,81 +282,93 @@ class Sky: return x - def seq2img(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 seq2img(self, seq, upscale=15): + f_first = seq[:, : self.height * self.width].reshape( + -1, self.height, self.width + ) + f_second = seq[:, self.height * self.width + 1 :].reshape( + -1, self.height, self.width + ) + direction = seq[:, self.height * self.width] direction_symbol = torch.full( - (direction.size(0), height * upscale - 1, upscale), 0 + (direction.size(0), self.height * upscale - 1, upscale), 0 ) - direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2) - separator = torch.full((direction.size(0), 3, height * upscale - 1, 1), 0) + direction_symbol = self.colors[direction_symbol].permute(0, 3, 1, 2) + separator = torch.full((direction.size(0), 3, self.height * upscale - 1, 1), 0) for n in range(direction_symbol.size(0)): - if direction[n] == token_forward: + if direction[n] == self.token_forward: for k in range(upscale): direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, 3 + upscale // 2 - abs(k - upscale // 2), ] = 0 - elif direction[n] == token_backward: + elif direction[n] == self.token_backward: for k in range(upscale): direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, 3 + abs(k - upscale // 2), ] = 0 else: for k in range(2, upscale - 2): direction_symbol[ - n, :, (height * upscale) // 2 - upscale // 2 + k, k + n, :, (self.height * upscale) // 2 - upscale // 2 + k, k ] = 0 direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k, ] = 0 return torch.cat( [ - frame2img(f_first, height, width, upscale), + self.frame2img(f_first, upscale), separator, direction_symbol, separator, - frame2img(f_second, height, width, upscale), + self.frame2img(f_second, upscale), ], dim=3, ) - def seq2str(seq): + def seq2str(self, seq): result = [] for s in seq: - result.append("".join([token2char[v] for v in s])) + result.append("".join([self.token2char[v] for v in s])) return result + def save_image(self, input, result_dir, filename, logger): + img = self.seq2img(input.to("cpu")) + image_name = os.path.join(result_dir, filename) + torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=6, padding=4) + logger(f"wrote {image_name}") + + def save_quizzes(self, input, result_dir, filename_prefix, logger): + self.save_image(input, result_dir, filename_prefix + ".png", logger) + ###################################################################### if __name__ == "__main__": import time - height, width = 6, 8 + sky = Sky(height=6, width=8, nb_iterations=100) + start_time = time.perf_counter() - seq, it = generate_seq( - nb=64, height=height, width=width, nb_iterations=100, return_iterations=True - ) + seq, it = sky.generate_seq(nb=64, return_iterations=True) delay = time.perf_counter() - start_time print(f"{seq.size(0)/delay:02f} samples/s") - print(seq2str(seq[:4])) + print(sky.seq2str(seq[:4])) for t in range(len(it[0])): - img = torch.cat([frame2img(f[t], height, width) for f in it], dim=0) + img = torch.cat([sky.frame2img(f[t]) for f in it], dim=0) torchvision.utils.save_image( img.float() / 255.0, f"/tmp/frame_{t:03d}.png", @@ -348,7 +380,7 @@ if __name__ == "__main__": # m = (torch.rand(seq.size()) < 0.05).long() # seq = (1 - m) * seq + m * 23 - img = seq2img(seq, height, width) + img = sky.seq2img(seq) print(img.size()) torchvision.utils.save_image(