80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
import pandas as pd
|
|
import vaex as vx
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from skimage import exposure
|
|
from skimage.util import img_as_ubyte
|
|
import numpy as np
|
|
from matplotlib import cm
|
|
import sys
|
|
|
|
base_size = 1080, 1920
|
|
|
|
|
|
def scale_to(width=None, height=None):
|
|
isnone = (width is None, height is None)
|
|
ret = {
|
|
(False, False): lambda w, h: (w, h),
|
|
(True, True): lambda w, h: (width, height),
|
|
(False, True): lambda w, h: (width, width * (h / w)),
|
|
(True, False): lambda w, h: (height * (w / h), height),
|
|
}
|
|
return lambda *args: tuple(map(int, ret[isnone](*args)))
|
|
|
|
|
|
# xz -1 1
|
|
bining = {
|
|
("zx", -1, 1): scale_to(width=base_size[0]), # main view, top down
|
|
# ('yx',1,1): lambda size,w,h: (size,int(size*(w/h))), #
|
|
# ('zy',-1,1): lambda size,w,h: (int(size*(h/w)),size), #
|
|
}
|
|
|
|
print("Loading stars.csv")
|
|
stars = pd.read_csv(
|
|
"stars.csv",
|
|
names=["id", "name", "num_bodies", "has_scoopable", "mult", "x", "y", "z"],
|
|
usecols=["id", "num_bodies", "x", "y", "z", "mult"],
|
|
index_col=0,
|
|
)
|
|
stars = vx.from_pandas(stars, copy_index=False)
|
|
|
|
filename = "heuristic.png"
|
|
fnt = ImageFont.truetype(r"FiraCode-Regular", 40)
|
|
|
|
for (binby_key, m1, m2), calcshape in bining.items():
|
|
binby = [m1 * stars[binby_key[0]], m2 * stars[binby_key[1]]]
|
|
|
|
mm = [binby[0].minmax(), binby[1].minmax()]
|
|
|
|
w, h = [mm[0][1] - mm[0][0], mm[1][1] - mm[1][0]]
|
|
shape = calcshape(w, h)
|
|
hm_all = stars.sum("num_bodies", binby=binby, shape=shape, limits="minmax")
|
|
hm_all_mask = hm_all != 0
|
|
hm_all = exposure.equalize_hist(hm_all)
|
|
hm_all -= hm_all.min()
|
|
hm_all /= hm_all.max()
|
|
|
|
hm_boost = stars.sum(
|
|
"astype(mult>1.0,'int')", binby=binby, shape=shape, limits="minmax"
|
|
)
|
|
hm_boost_mask = hm_boost != 0
|
|
|
|
hm_boost = exposure.equalize_hist(hm_boost)
|
|
hm_boost -= hm_boost.min()
|
|
hm_boost /= hm_boost.max()
|
|
|
|
# R = cm.Reds_r()
|
|
G = cm.Greens_r(hm_all)
|
|
B = cm.Blues_r(hm_boost)
|
|
|
|
img = np.zeros((base_size[0], base_size[1], 4))
|
|
img[:, :, :] = 0.0
|
|
img[:, :, 3] = 1.0
|
|
canvas = img[: shape[0], : shape[1], :]
|
|
canvas[hm_all_mask] = G[hm_all_mask]
|
|
canvas[hm_boost_mask] = B[hm_boost_mask]
|
|
pil_img = Image.fromarray(img_as_ubyte(img))
|
|
draw = ImageDraw.Draw(pil_img)
|
|
messages = ["Hello World"]
|
|
draw.multiline_text((shape[0], 0), "\n".join(messages), font=fnt)
|
|
pil_img.save(filename)
|