umap_visualizer.py

from vispy.visuals.transforms import STTransform
from vispy.scene import visuals
from vispy import scene
import vispy

from math import ceil
import pandas as pd
import numpy as np
import imageio
import os
import sys
from tqdm import tqdm


def load_data(file_path):
    if file_path.endswith(".npy"):
        data = np.load(file_path, allow_pickle=True)

    if file_path.endswith(".csv"):
        data = pd.read_csv(file_path).iloc[:, 3:].values

    return data


def median(arr):
    m, n = arr.shape
    middle = np.arange((m - 1) >> 1, (m >> 1) + 1)
    arr = np.partition(arr, middle, axis=0)
    return arr[middle].mean(axis=0)


def remove_outliers(arr, threshold=7):
    m = median(arr)
    s = np.abs(arr - m)
    return arr[(s < median(s) * threshold).all(axis=1)]


def prepare_data(data):
    # COLOR
    colors = data[:, :3]
    normalized = (colors - colors.min()) / (colors.max() - colors.min())

    color_map = np.ones((data.shape[0], 4), dtype=np.float32)
    color_map[:, :3] = normalized
    color_map[:, 3] = 1

    # DATA
    embedding = data[:, :3]
    return embedding, color_map


def create_canvas(datasets, animation=False):
    canvas = vispy.scene.SceneCanvas(keys="interactive", show=True, size=(960, 960))

    view_boxes = [
        scene.widgets.ViewBox(border_color="black", parent=canvas.scene)
        for _ in range(len(datasets))
    ]

    # Adding a view to the canvas.
    # view = canvas.central_widget.add_view()
    grid = canvas.central_widget.add_grid()
    grid.padding = 0

    ticker = 0
    for x in range(ceil(len(datasets) / 2)):
        for y in range(2):
            try:
                grid.add_widget(view_boxes[ticker], x, y)
                ticker += 1
            except:
                pass

    t1 = visuals.Text("Refik Anadol Studio", parent=canvas.scene, color="white")
    t1.font_size = 24
    t1.pos = canvas.size[0] // 2, canvas.size[1] // 8

    # t2 = visuals.Text(file_path, parent=canvas.scene, color="white")
    # t2.font_size = 12
    # t2.pos = canvas.size[0] // 2, canvas.size[1] // 1.1

    embeddings_sets = [prepare_data(data) for data in datasets]

    scatters = [visuals.Markers(symbol="o") for _ in range(len(datasets))]

    for scatter, embedding in zip(scatters, embeddings_sets):
        values, color_map = embedding
        scatter.set_data(
            values,
            edge_color=None,
            face_color=color_map,
            size=2,
            edge_width=0.01,
        )
        scatter.transform = STTransform(translate=[0, 0, 0])

    print("putting the data points in place 🧐")
    for view_box, scatter in zip(view_boxes, scatters):
        view_box.add(scatter)
        view_box.camera = "turntable"

    print("we are trying to set up the cameras 🎥")
    for idx, view_box in enumerate(view_boxes):
        if idx < len(view_boxes) - 1:
            view_boxes[idx].camera.link(view_boxes[idx + 1].camera)

    view_box.camera.set_range()
    print("everything seems as it should 🙏")

    if animation:
        ## Animation
        output_filename = "animation3.mp4"
        axis = [1, -0.5, -1]
        writer = imageio.get_writer(output_filename)
        print("warming up the processor, I'll put it in the oven 🔥")
        for i in tqdm(
            range(36),
            desc=f"Frames rendering.",
        ):
            im = canvas.render(alpha=True)
            writer.append_data(im)
            for view_box in view_boxes:
                view_box.camera.transform.rotate(10, axis)
        writer.close()
        print("Saved", os.path.join(os.getcwd(), output_filename))


# "/Users/cobanov/RAS/moma_umap/400k/8min01/moma_umap_n12_s2_m0.npy",
# "/Users/cobanov/RAS/moma_umap/400k/13a289/moma_umap_n15_s2_m0.npy",
# "/Users/cobanov/RAS/moma_umap/400k/14local25/moma_umap_n12_s2_m0.npy",

if __name__ == "__main__":
    file_paths = [
        # "stockholm_umap_1_n12.npy",
        "stockholm_umap_1_n15.npy",
        # "ddp_umap.csv",
        # "stockholm_umap_2_n12.npy",
        # "stockholm_umap_2_n15.npy",
        # "stockholm_umap_2_n12.npy"
        # "/Users/cobanov/RAS/moma_umap/400k/12min001/moma_umap_n20_s2_m0.npy",
        # "/Users/cobanov/RAS/moma_umap/400k/12min001/moma_umap_n30_s2_m0.npy",
    ]

    datasets = []

    for path in tqdm(file_paths, desc="Files are loading."):
        data = load_data(path)
        data = remove_outliers(data, 7)
        datasets.append(data)
        print(data.shape)

    create_canvas(datasets, animation=False)

    if sys.flags.interactive != 1:
        vispy.app.run()