clean up plot

check.py

@@ -1,3 +1,4 @@
+# import matplotlib.patches as patches
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -22,11 +23,13 @@ def make_image(ckpt: str, fname: str, color=True):
     preds = M(rgb_tensor)
     sorted_inds = np.argsort(preds.detach().numpy().ravel())
 
-    fig, ax = plt.subplots(figsize=(20, 5))
+    fig, ax = plt.subplots()
     for i in range(len(sorted_inds)):
         idx = sorted_inds[i]
         color = rgb_values[idx]
-        ax.plot([i, i],[0, 5], lw=0.5, c=color, antialiased=False, alpha=1)
+        ax.plot([i + 0.5, i + 0.5], [0, 1], lw=1, c=color, antialiased=True, alpha=1)
+        # rect = patches.Rectangle((i, 0), 1, 5, linewidth=0.1, edgecolor=None, facecolor=None, alpha=1)
+        # ax.add_patch(rect)
         ax.axis("off")
         # ax.axis("square")
 
@@ -35,7 +38,13 @@ def make_image(ckpt: str, fname: str, color=True):
 
 if __name__ == "__main__":
     # name = "color_128_0.3_1.00e-06"
-    name = "color_64_1_1.0e-3.png"
+    import glob
+
+    v = 29
+    name = f"v{v}"
     # ckpt = f"/teamspace/jobs/{name}/work/colors/lightning_logs/version_2/checkpoints/epoch=999-step=8000.ckpt"
-    ckpt = "/teamspace/studios/this_studio/colors/lightning_logs/version_26/checkpoints/epoch=99-step=1500.ckpt"
+    ckpt = glob.glob(
+        f"/teamspace/studios/this_studio/colors/lightning_logs/version_{v}/checkpoints/*.ckpt"
+    )[-1]
     make_image(ckpt, fname=name)
+    make_image(ckpt, fname=name + "b", color=False)