X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=sky.py;h=040ec67e0ed194687d28cbb7f4d1d8c2808736b4;hb=eaed6307836d88abe7c0f4be733a38364ba20e2f;hp=e93c88a1a372c9f64aff923dfc063523ab38db5d;hpb=bfcef9a8c82ed45528601e85725166241bbee916;p=culture.git diff --git a/sky.py b/sky.py index e93c88a..040ec67 100755 --- a/sky.py +++ b/sky.py @@ -37,31 +37,31 @@ class Sky(problem.Problem): 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)]) + "><" ) - def __init__(self, height=6, width=8, nb_birds=3, speed=1, nb_iterations=4): + def __init__( + self, + height=6, + width=8, + nb_birds=3, + speed=2, + nb_iterations=2, + avoid_collision=True, + ): self.height = height self.width = width self.nb_birds = nb_birds self.speed = speed self.nb_iterations = nb_iterations - - def direction_tokens(self): - return self.token_forward, self.token_backward + self.avoid_collision = avoid_collision def generate_frame_sequences(self, nb): frame_sequences = [] for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): - result = torch.zeros( - self.nb_iterations, self.height, self.width, dtype=torch.int64 - ) - i, j, vi, vj = ( torch.empty(self.nb_birds, dtype=torch.int64), torch.empty(self.nb_birds, dtype=torch.int64), @@ -69,75 +69,91 @@ class Sky(problem.Problem): torch.empty(self.nb_birds, dtype=torch.int64), ) + def collision_okay(): + if not self.avoid_collision: + return True + + count = torch.zeros(self.height, self.width, dtype=torch.int64) + + for n in range(self.nb_birds): + count[i[n], j[n]] += 1 + count[i[n] - vi[n], j[n]] += 1 + count[i[n], j[n] - vj[n]] += 1 + + return count.max() <= 1 + col = ( torch.randperm(self.colors.size(0) - 1)[: self.nb_birds].sort().values + 1 ) - for n in range(self.nb_birds): + while True: while True: - i[n] = torch.randint(self.height, (1,)) - j[n] = torch.randint(self.width, (1,)) - vm = torch.randint(4, (1,)) - vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 - if ( - i[n] - vi[n] >= 0 - and i[n] - vi[n] < self.height - and j[n] - vj[n] >= 0 - and j[n] - vj[n] < self.width - ): + for n in range(self.nb_birds): + while True: + i[n] = torch.randint(self.height, (1,)) + j[n] = torch.randint(self.width, (1,)) + vm = torch.randint(4, (1,)) + vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 + if ( + i[n] - vi[n] >= 0 + and i[n] - vi[n] < self.height + and j[n] - vj[n] >= 0 + and j[n] - vj[n] < self.width + ): + break + + if collision_okay(): break - for l in range(self.nb_iterations): - for n in range(self.nb_birds): - c = col[n] - result[l, i[n], j[n]] = c - result[l, i[n] - vi[n], j[n]] = c - result[l, i[n], j[n] - vj[n]] = c + result = torch.zeros( + self.nb_iterations * self.speed, + self.height, + self.width, + dtype=torch.int64, + ) - if (i[n] == 0 and vi[n] == -1) or ( - i[n] == self.height - 1 and vi[n] == 1 - ): - vi[n] = -vi[n] + fine = torch.empty(self.nb_iterations * self.speed) - if (j[n] == 0 and vj[n] == -1) or ( - j[n] == self.width - 1 and vj[n] == 1 - ): - vj[n] = -vj[n] + t_to_keep = ( + torch.arange(self.nb_iterations, device=result.device) * self.speed + ) - i[n] += vi[n] - j[n] += vj[n] + for l in range(self.nb_iterations * self.speed): + fine[l] = collision_okay() + for n in range(self.nb_birds): + c = col[n] + result[l, i[n], j[n]] = c + result[l, i[n] - vi[n], j[n]] = c + result[l, i[n], j[n] - vj[n]] = c - frame_sequences.append(result) + if (i[n] == 0 and vi[n] == -1) or ( + i[n] == self.height - 1 and vi[n] == 1 + ): + vi[n] = -vi[n] - return frame_sequences + if (j[n] == 0 and vj[n] == -1) or ( + j[n] == self.width - 1 and vj[n] == 1 + ): + vj[n] = -vj[n] - def generate_token_sequences(self, nb): - frame_sequences = self.generate_frame_sequences(nb) + i[n] += vi[n] + j[n] += vj[n] - result = [] + result = result[t_to_keep] + fine = fine[t_to_keep] - for frame_sequence in frame_sequences: - a = [] - if torch.rand(1) < 0.5: - for frame in frame_sequence: - if len(a) > 0: - a.append(torch.tensor([self.token_forward])) - a.append(frame.flatten()) - else: - for frame in reversed(frame_sequence): - if len(a) > 0: - a.append(torch.tensor([self.token_backward])) - a.append(frame.flatten()) + if fine[-1]: + break - result.append(torch.cat(a, dim=0)[None, :]) + frame_sequences.append(result) - return torch.cat(result, dim=0) + return frame_sequences ###################################################################### def frame2img(self, x, scale=15): - x = x.reshape(-1, self.height, self.width) + x = x.reshape(x.size(0), self.height, -1) m = torch.logical_and( x >= 0, x < self.first_bird_token + self.nb_bird_tokens ).long() @@ -163,92 +179,94 @@ class Sky(problem.Problem): return x - def seq2img(self, seq, scale=15): - all = [ - self.frame2img( - seq[:, : self.height * self.width].reshape(-1, self.height, self.width), - scale, + def seq2str(self, seq): + result = [] + for s in seq: + result.append("".join([self.token2char[v] for v in s])) + return result + + def save_image( + self, + result_dir, + filename, + prompts, + answers, + predicted_prompts=None, + predicted_answers=None, + ): + if predicted_prompts is None: + predicted_prompts = 255 + + 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, ) - ] + 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) + y[...] = c[:, :, None, None] + y[:, :, margin:-margin, margin:-margin] = x + return y - separator = torch.full((seq.size(0), 3, self.height * scale - 1, 1), 0) + margin = 4 - t = self.height * self.width + img_prompts = add_frame(self.frame2img(prompts.to("cpu")), 0, 1) + img_answers = add_frame(self.frame2img(answers.to("cpu")), 0, 1) - while t < seq.size(1): - direction_tokens = seq[:, t] - t += 1 + # img_prompts = add_frame(img_prompts, 255, margin) + # img_answers = add_frame(img_answers, 255, margin) - direction_images = self.colors[ - torch.full( - (direction_tokens.size(0), self.height * scale - 1, scale), 0 - ) - ].permute(0, 3, 1, 2) - - for n in range(direction_tokens.size(0)): - if direction_tokens[n] == self.token_forward: - for k in range(scale): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - 3 + scale // 2 - abs(k - scale // 2), - ] = 0 - elif direction_tokens[n] == self.token_backward: - for k in range(scale): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - 3 + abs(k - scale // 2), - ] = 0 - else: - for k in range(2, scale - 2): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - k, - ] = 0 - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - scale - 1 - k, - ] = 0 - - all += [ - separator, - direction_images, - separator, - self.frame2img( - seq[:, t : t + self.height * self.width].reshape( - -1, self.height, self.width - ), - scale, - ), - ] - - t += self.height * self.width - - return torch.cat(all, dim=3) + img_prompts = add_frame(img_prompts, predicted_prompts, margin) + img_answers = add_frame(img_answers, predicted_answers, margin) - def seq2str(self, seq): - result = [] - for s in seq: - result.append("".join([self.token2char[v] for v in s])) - return result + separator = img_prompts.new_full( + (img_prompts.size(0), img_prompts.size(1), img_prompts.size(2), margin), 255 + ) + + img = torch.cat([img_prompts, img_answers], dim=3) - def save_image(self, input, result_dir, filename): - 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) + torchvision.utils.save_image( + img.float() / 255.0, image_name, nrow=6, padding=margin * 2, pad_value=1.0 + ) - def save_quizzes(self, input, result_dir, filename_prefix): - self.save_image(input, result_dir, filename_prefix + ".png") + ###################################################################### + + def nb_token_values(self): + return len(self.colors) + + 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) + return prompts, answers + + def save_quizzes( + self, + result_dir, + filename_prefix, + prompts, + answers, + predicted_prompts=None, + predicted_answers=None, + ): + self.save_image( + result_dir, + filename_prefix + ".png", + prompts, + answers, + predicted_prompts, + predicted_answers, + ) ###################################################################### @@ -258,10 +276,19 @@ if __name__ == "__main__": sky = Sky(height=6, width=8, speed=1, nb_iterations=4) - start_time = time.perf_counter() - seq = sky.generate_frame_sequences(nb=64) - delay = time.perf_counter() - start_time - print(f"{seq.size(0)/delay:02f} seq/s") + 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 + + sky.save_quizzes( + "/tmp", "test", prompts, answers, predicted_prompts, predicted_answers + ) + + # start_time = time.perf_counter() + # token_sequences = sky.generate_token_sequences(nb=64) + # delay = time.perf_counter() - start_time + # print(f"{token_sequences.size(0)/delay:02f} seq/s") # print(sky.seq2str(seq[:4])) @@ -278,10 +305,10 @@ if __name__ == "__main__": # m = (torch.rand(seq.size()) < 0.05).long() # seq = (1 - m) * seq + m * 23 - print(seq.size()) - img = sky.seq2img(seq) - print(img.size()) + # print(seq.size()) + # img = sky.seq2img(token_sequences) + # print(img.size()) - torchvision.utils.save_image( - img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0 - ) + # torchvision.utils.save_image( + # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0 + # )