AAAI Version

AAAI Supplementary Material/ForNet Creation Code/get_in_indices.py

@@ -0,0 +1,42 @@
+import argparse
+import json
+
+from datadings.reader import MsgpackReader
+from datadings.torch import Dataset
+from tqdm.auto import tqdm
+
+parser = argparse.ArgumentParser(description="Foreground size ratio")
+parser.add_argument("-mode", choices=["train", "val"], default="train", help="Train or val data?")
+parser.add_argument("-ds", "--dataset", choices=["imagenet", "tinyimagenet"], required=True, help="Dataset to use")
+parser.add_argument("-r", "--dataset_root", required=True, type=str, help="Root directory of the dataset")
+parser.add_argument("-s", "--segment_root", required=True, type=str, help="Root directory of the segmentation dataset")
+args = parser.parse_args()
+
+
+if args.dataset == "imagenet":
+    reader = MsgpackReader(f"{args.dataset_root}imagenet/msgpack/{args.mode}.msgpack")
+    dataset = Dataset(reader)
+elif args.dataset == "tinyimagenet":
+    reader = MsgpackReader(f"{args.dataset_root}TinyINSegment/TinyIN_{args.mode}.msgpack")
+    dataset = Dataset(reader)
+else:
+    raise ValueError(f"Unknown dataset: {args.dataset}")
+
+if args.dataset.startswith("imagenet"):
+    with open("wordnet_data/imagenet1k_synsets.json", "r") as f:
+        id_to_synset = json.load(f)
+    id_to_synset = {int(k): v for k, v in id_to_synset.items()}
+elif args.dataset.startswith("tinyimagenet"):
+    with open("wordnet_data/tinyimagenet_synset_names.txt", "r") as f:
+        synsets = f.readlines()
+    id_to_synset = [int(synset.split(":")[0].strip()[1:]) for synset in synsets]
+    id_to_synset = sorted(id_to_synset)
+
+key_to_idx = {
+    f"n{id_to_synset[data['label']]:08d}/{data['key'].split('.')[0]}": i
+    for i, data in enumerate(tqdm(dataset, leave=True))
+}
+
+print(len(key_to_idx))
+with open(f"{args.segment_root}/{args.mode}_indices.json", "w") as f:
+    json.dump(key_to_idx, f)