|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
import numpy as np
|
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
from dataloader import extract_colors
|
|
|
|
|
from model import ColorTransformerModel
|
|
|
|
|
|
|
|
|
|
def make_image(ckpt: str, fname: str, color=True):
|
|
|
|
|
M = ColorTransformerModel.load_from_checkpoint(ckpt)
|
|
|
|
|
|
|
|
|
|
# preds = M(rgb_tensor)
|
|
|
|
|
if not color:
|
|
|
|
|
N = 949
|
|
|
|
|
linear_space = torch.linspace(0, 1, N)
|
|
|
|
|
rgb_tensor = linear_space.unsqueeze(1).repeat(1, 3)
|
|
|
|
|
else:
|
|
|
|
|
rgb_tensor, names = extract_colors()
|
|
|
|
|
|
|
|
|
|
rgb_values = rgb_tensor.detach().numpy()
|
|
|
|
|
preds = M(rgb_tensor)
|
|
|
|
|
sorted_inds = np.argsort(preds.detach().numpy().ravel())
|
|
|
|
|
|
|
|
|
|
fig, ax = plt.subplots(figsize=(10, 5))
|
|
|
|
|
for i in range(len(sorted_inds)):
|
|
|
|
|
idx = sorted_inds[i]
|
|
|
|
|
color = rgb_values[idx]
|
|
|
|
|
ax.vlines(i, ymin=0, ymax=1, lw=0.1, colors=color, antialiased=False, alpha=0.5)
|
|
|
|
|
ax.axis("off")
|
|
|
|
|
# ax.axis("square")
|
|
|
|
|
|
|
|
|
|
plt.savefig(f"{fname}.png", dpi=300)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
|
|
|
|
|
|
name = "color_128_0.3_1.00e-06"
|
|
|
|
|
ckpt = f"/teamspace/jobs/{name}/work/colors/lightning_logs/version_2/checkpoints/epoch=999-step=8000.ckpt"
|
|
|
|
|
make_image(ckpt, fname=name)
|
