AAAI Version

AAAI Supplementary Material/ForNet Creation Code/inspect_img_versions.py

@@ -0,0 +1,65 @@
+import argparse
+import os
+from random import shuffle
+
+import numpy as np
+from matplotlib import pyplot as plt
+from PIL import Image
+
+parser = argparse.ArgumentParser(description="Inspect image versions")
+parser.add_argument("-f", "--base_folder", type=str, required=True, help="Base folder to inspect")
+parser.add_argument("-opt_fg_ratio", type=float, default=0.3, help="Optimal foreground ratio")
+
+args = parser.parse_args()
+
+bg_folder = os.path.join(args.base_folder, "backgrounds")
+fg_folder = os.path.join(args.base_folder, "foregrounds")
+
+classes = os.listdir(fg_folder)
+classes = sorted(classes, key=lambda x: int(x[1:]))
+assert len(classes) in [200, 1_000], f"Expected 200 or 1_000 classes, got {len(classes)}"
+
+total_images = set()
+for in_cls in classes:
+    cls_images = {
+        os.path.join(in_cls, "_".join(img.split(".")[0].split("_")[:3]))
+        for img in os.listdir(os.path.join(fg_folder, in_cls))
+    }
+    total_images.update(cls_images)
+total_images = list(total_images)
+shuffle(total_images)
+print(total_images[:5], "...")
+
+# in_cls to print name/lemma
+with open("wordnet_data/tinyimagenet_synset_names.txt", "r") as f:
+    in_cls_to_name = {line.split(":")[0].strip(): line.split(":")[1].strip() for line in f.readlines() if len(line) > 2}
+
+for image_name in total_images:
+    in_cls, img_name = image_name.split("/")
+    versions = set()
+    for img in os.listdir(os.path.join(fg_folder, in_cls)):
+        if img.startswith(img_name):
+            versions.add(img)
+    if len(versions) <= 1:
+        print(f"Image {image_name} has only one version")
+        continue
+    versions = sorted(list(versions))
+
+    bg_img = Image.open(os.path.join(bg_folder, in_cls, f"{versions[0].split('.')[0]}.JPEG"))
+
+    # plot all versions at once
+    fig, axs = plt.subplots(1, len(versions) + 1, figsize=(15, 5))
+    axs[0].imshow(bg_img)
+    axs[0].axis("off")
+    axs[0].set_title("Background")
+    for version, ax in zip(versions, axs[1:]):
+        img = Image.open(os.path.join(fg_folder, in_cls, version))
+        img_mask = np.array(img.convert("RGBA").split()[-1])
+        fg_dens_fact = np.sum(img_mask) / (255 * img_mask.size)
+        fg_ratio = np.sum(img_mask) / (255 * bg_img.size[0] * bg_img.size[1])
+        ax.imshow(img)
+        ax.axis("off")
+        ax.set_title(f"{version}\nfg_rat: {fg_ratio:.2f} dist to optimal: {abs(fg_ratio - args.opt_fg_ratio):.2f}")
+    fig.suptitle(in_cls_to_name[in_cls])
+    plt.show()
+    plt.close()