Internal change

PiperOrigin-RevId: 380650647

official/vision/beta/data/process_coco_few_shot.sh

@@ -13,7 +13,7 @@ done
 
 cocosplit_url="dl.yf.io/fs-det/datasets/cocosplit"
 wget --recursive --no-parent -q --show-progress --progress=bar:force:noscroll \
-    -P "${tmp_dir}" -A "5k.json,*10shot*.json,*30shot*.json" \
+    -P "${tmp_dir}" -A "trainvalno5k.json,5k.json,*10shot*.json,*30shot*.json" \
     "http://${cocosplit_url}/"
 mv "${tmp_dir}/${cocosplit_url}/"* "${tmp_dir}"
 rm -rf "${tmp_dir}/${cocosplit_url}/"
@@ -45,4 +45,14 @@ python create_coco_tf_record.py \
     --output_file_prefix="${output_dir}/5k" \
     --num_shards=10
 
+python create_coco_tf_record.py \
+    --logtostderr \
+    --image_dir=/namespace/vale-project/datasets/mscoco_raw/images/train2014 \
+    --image_dir=/namespace/vale-project/datasets/mscoco_raw/images/val2014 \
+    --image_info_file="${tmp_dir}/datasplit/trainvalno5k_base.json" \
+    --object_annotations_file="${tmp_dir}/datasplit/trainvalno5k_base.json" \
+    --caption_annotations_file="" \
+    --output_file_prefix="${output_dir}/trainvalno5k_base" \
+    --num_shards=200
+
 rm -rf "${tmp_dir}"

official/vision/beta/data/process_coco_few_shot_json_files.py

@@ -76,10 +76,28 @@ for _seed, _shots in itertools.product(SEEDS, SHOTS):
             _shots,
             _category))
 
+# Base class IDs, as defined in
+# https://github.com/ucbdrive/few-shot-object-detection/blob/master/fsdet/evaluation/coco_evaluation.py#L60-L65
+BASE_CLASS_IDS = [8, 10, 11, 13, 14, 15, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
+                  35, 36, 37, 38, 39, 40, 41, 42, 43, 46, 47, 48, 49, 50, 51,
+                  52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 65, 70, 73, 74, 75,
+                  76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
+
 
 def main(unused_argv):
   workdir = FLAGS.workdir
 
+  # Filter novel class annotations from the training set.
+  file_path = os.path.join(workdir, 'datasplit', 'trainvalno5k.json')
+  with tf.io.gfile.GFile(file_path, 'r') as f:
+    json_dict = json.load(f)
+
+  json_dict['annotations'] = [a for a in json_dict['annotations']
+                              if a['category_id'] in BASE_CLASS_IDS]
+  output_path = os.path.join(workdir, 'datasplit', 'trainvalno5k_base.json')
+  with tf.io.gfile.GFile(output_path, 'w') as f:
+    json.dump(json_dict, f)
+
   for seed, shots in itertools.product(SEEDS, SHOTS):
     # Retrieve all examples for a given seed and shots setting.
     file_paths = [os.path.join(workdir, suffix)