1 files changed, 151 insertions, 0 deletions
diff --git a/display_animation_tkinter.py b/display_animation_tkinter.py
new file mode 100644
index 0000000..f26202c
--- /dev/null
+++ b/display_animation_tkinter.py
@@ -0,0 +1,151 @@
+import argparse
+import json
+import math
+import time
+import tkinter as tk
+
+import numpy as np
+
+from utils import rgb_to_hex
+
+
+def parse_args():
+    return {
+        "cell_size": 8,  # in pixels
+        "num_cells_w": 100,  # TODO: w/h or r/c?
+        "num_cells_h": 100,
+        "colors_path": "configs/colors_64_v0.json",
+        "corners_path": "configs/corners_hollow4x4_v0.json",
+        "frame_delay": 1000 // 15,  # time b/w frames, in millis
+    }
+
+
+def seq_to_frames(bin_seq: np.ndarray, args: dict) -> np.ndarray:  # TODO: Namespace
+    """
+    Converts a binary sequence to an array of frames.
+    TODO: doc, expecting size of exactly one frame in the future?
+
+    np.packbits seems relevant but limited to 8 bits
+
+    Args:
+        bin_seq: a 1D array of binary values.
+        args: config parameters
+
+    Returns:
+        array of shape (num_frames, args["num_cells_h"], args["num_cells_w"]), where each element is
+        a hex string for the color of that cell
+    """
+    bin_seq = bin_seq.copy()  # so that we don't mutate bin_seq
+
+    with open(args["colors_path"], "r") as f:
+        colors = json.load(f)
+
+    with open(args["corners_path"], "r") as f:
+        corners = json.load(f)
+
+    assert len(corners["corner_colors"]) == 4, "Hardcoded for 4 corners"
+    corner_width = corners["corner_width"]
+    corner_height = corners["corner_height"]
+
+    num_colors = len(colors)
+    bits_per_cell = int(math.log2(num_colors))
+    assert 2**bits_per_cell == num_colors, "Assumed the number of colors is a power of 2."
+
+    # bits_per_frame = bits_per_cell * args["num_cells_w"] * args["num_cells_h"]
+    # num_frames = int(math.ceil(len(bin_seq) / bits_per_frame))
+    # bin_seq.resize(bits_per_frame * num_frames)
+    # # frames_bits = bin_seq.reshape(num_frames, args["num_cells_h"], args["num_cells_w"], bits_per_cell)
+    #
+    # pows = 2 ** np.arange(bits_per_cell)
+    # frames_vals = (frames_bits * pows).sum(axis=-1)  # low to high bit order
+
+    cells_per_frame = args["num_cells_w"] * args["num_cells_h"] - 4 * corner_width * corner_height
+    bits_per_frame = bits_per_cell * cells_per_frame
+    num_frames = int(math.ceil(len(bin_seq) / bits_per_frame))
+    bin_seq.resize(bits_per_frame * num_frames)
+    frames_bits = bin_seq.reshape(num_frames, cells_per_frame, bits_per_cell)
+    pows = 2 ** np.arange(bits_per_cell)
+    frames_vals = (frames_bits * pows).sum(axis=-1)  # low to high bit order
+
+    # Efficiently map frame_vals to the corresponding hex colors (https://stackoverflow.com/a/55950051)
+    color_mapping_arr = np.empty(num_colors, dtype="<U7")
+    for i in range(num_colors):
+        color_mapping_arr[i] = rgb_to_hex(colors[str(i)])  # JSON keys stored as str
+
+    frames_colors = color_mapping_arr[frames_vals]
+
+    num_top_cells = (args["num_cells_w"] - 2 * corner_width) * corner_height  # TODO: explain
+    top_cells = frames_colors[:, :num_top_cells].reshape(num_frames, corner_height, -1)
+    center_cells = frames_colors[:, num_top_cells:-num_top_cells].reshape(num_frames, -1, args["num_cells_w"])
+    bottom_cells = frames_colors[:, -num_top_cells:].reshape(num_frames, corner_height, -1)
+
+    corners_numpy_dict = {key: np.broadcast_to(val, (num_frames, corner_height, corner_width))
+                          for key, val in corners["corner_colors"].items()}
+    top_rows = np.concatenate([corners_numpy_dict["0"], top_cells, corners_numpy_dict["1"]], axis=2)
+    bottom_rows = np.concatenate([corners_numpy_dict["2"], bottom_cells, corners_numpy_dict["3"]], axis=2)
+
+    return np.concatenate([top_rows, center_cells, bottom_rows], axis=1)
+
+
+class AnimatedFrames:
+    def __init__(self, frames: np.ndarray, args: dict):
+        self.frames = frames
+        self.args = args
+        self.width_pixels = args["cell_size"] * args["num_cells_w"]
+        self.height_pixels = args["cell_size"] * args["num_cells_h"]
+
+        self.root = tk.Tk()
+        self.root.columnconfigure(0, weight=1)
+        self.root.rowconfigure(0, weight=1)
+        self.canvas = tk.Canvas(self.root, width=self.width_pixels, height=self.height_pixels,
+                                borderwidth=0, highlightthickness=0)  # https://stackoverflow.com/a/63220348
+        self.canvas.grid(column=0, row=0)
+
+        # We will create each cell (rectangle) now, and then update their colors in different frames
+        self.inds_to_id = dict()  # map from (row_ind, col_ind) to the id of the rectangle at that cell position
+        for i in range(args["num_cells_h"]):
+            for j in range(args["num_cells_w"]):
+                self.inds_to_id[(i, j)] = self.canvas.create_rectangle(
+                    i * args["cell_size"], j * args["cell_size"], (i + 1) * args["cell_size"],
+                    (j + 1) * args["cell_size"], width=0
+                )
+
+        self.frame_ind = 0  # the index in `self.frames` of the next frame to display
+        self.start_time = None
+        self._debug_text_id = self.canvas.create_text((args["num_cells_w"] - 2) * args["cell_size"], (args["num_cells_h"] - 2) * args["cell_size"])
+
+    def display_frame(self):
+        func_start_time = time.time_ns()
+        for inds_tuple, color_hex in np.ndenumerate(self.frames[self.frame_ind % len(self.frames)]):
+            self.canvas.itemconfigure(self.inds_to_id[inds_tuple], fill=color_hex)
+
+        self.canvas.itemconfigure(self._debug_text_id, text=self.frame_ind)
+        self.frame_ind += 1
+
+        adjusted_delay = round((self.start_time + self.args["frame_delay"] * int(1e6)
+                                * self.frame_ind - func_start_time) / 1e6)
+        assert adjusted_delay > 1, adjusted_delay  # o/w we lagged too far behind, assuming 1 ms for this instruction
+        self.canvas.after(adjusted_delay, self.display_frame)  # TODO: check delay is exact
+        # TODO: set self time and just sleep until then
+
+    def animate(self):
+        epilepsy_warning_id = self.canvas.create_text(self.width_pixels / 2, self.height_pixels / 2,
+                                                      text="Warning: Epilepsy")
+
+        def delete_and_animate():
+            self.canvas.delete(epilepsy_warning_id)
+            self.start_time = time.time_ns()
+            self.display_frame()
+
+        self.canvas.after(5000, delete_and_animate)
+        # self.display_frame()
+        self.root.mainloop()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    rand_bin_seq = np.random.randint(2, size=1000000)
+    frames = seq_to_frames(rand_bin_seq, args)
+
+    AnimatedFrames(frames, args).animate()