|
|
@ -10,7 +10,7 @@ from model import ColorTransformerModel |
|
|
|
# import matplotlib.colors as mcolors |
|
|
|
|
|
|
|
|
|
|
|
def make_image(ckpt: str, fname: str, color=True): |
|
|
|
def make_image(ckpt: str, fname: str, color=True, **kwargs): |
|
|
|
M = ColorTransformerModel.load_from_checkpoint(ckpt) |
|
|
|
|
|
|
|
# preds = M(rgb_tensor) |
|
|
@ -36,10 +36,10 @@ def make_image(ckpt: str, fname: str, color=True): |
|
|
|
ax.axis("off") |
|
|
|
# ax.axis("square") |
|
|
|
|
|
|
|
plt.savefig(f"{fname}.png", dpi=300) |
|
|
|
plt.savefig(f"{fname}.png", **kwargs) |
|
|
|
|
|
|
|
|
|
|
|
def create_circle(ckpt: str, fname: str, dpi: int = 150, skip: bool = True): |
|
|
|
def create_circle(ckpt: str, fname: str, skip: bool = True, **kwargs): |
|
|
|
if isinstance(ckpt, str): |
|
|
|
M = ColorTransformerModel.load_from_checkpoint(ckpt) |
|
|
|
else: |
|
|
@ -48,7 +48,7 @@ def create_circle(ckpt: str, fname: str, dpi: int = 150, skip: bool = True): |
|
|
|
rgb_tensor, _ = extract_colors() |
|
|
|
rgb_tensor = preprocess_data(rgb_tensor) |
|
|
|
preds = M(rgb_tensor.to(M.device)) |
|
|
|
plot_preds(preds, rgb_tensor.detach().cpu().numpy(), fname=fname, dpi=dpi) |
|
|
|
plot_preds(preds, rgb_tensor.detach().cpu().numpy(), fname=fname, **kwargs) |
|
|
|
|
|
|
|
|
|
|
|
def plot_preds( |
|
|
@ -116,10 +116,11 @@ if __name__ == "__main__": |
|
|
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
|
|
# make the following accept a list of arguments |
|
|
|
parser.add_argument("-v", "--version", type=int, nargs="+", default=[0, 1]) |
|
|
|
parser.add_argument("-v", "--version", type=int, nargs="+", default=[0]) |
|
|
|
parser.add_argument( |
|
|
|
"--dpi", type=int, default=150, help="Resolution for saved image." |
|
|
|
) |
|
|
|
parser.add_argument("--figsize", type=int, default=3, help="Figure size") |
|
|
|
args = parser.parse_args() |
|
|
|
versions = args.version |
|
|
|
for v in versions: |
|
|
@ -130,7 +131,7 @@ if __name__ == "__main__": |
|
|
|
if len(ckpt) > 0: |
|
|
|
ckpt = ckpt[-1] |
|
|
|
print(f"Generating image for checkpoint: {ckpt}") |
|
|
|
create_circle(ckpt, fname=name, dpi=args.dpi) |
|
|
|
create_circle(ckpt, fname=name, dpi=args.dpi, figsize=[args.figsize] * 2) |
|
|
|
else: |
|
|
|
print(f"No checkpoint found for version {v}") |
|
|
|
# make_image(ckpt, fname=name + "b", color=False) |
|
|
|
# make_image(ckpt, fname=name + "b", color=False, dpi=args.dpi,) |
|
|
|