add --count and --batch args for data_export.py

Author: jiangyinzuo

ann_benchmarks/results.py

@@ -84,7 +84,7 @@ def load_all_results(dataset: Optional[str] = None,
     Yields:
         tuple: A tuple containing properties as a dictionary and an h5py file object.
     """
-    for root, _, files in os.walk(build_result_filepath(dataset, count)):
+    for root, _, files in os.walk(build_result_filepath(dataset, count, batch_mode=batch_mode)):
         for filename in files:
             if os.path.splitext(filename)[-1] != ".hdf5":
                 continue
@@ -110,4 +110,4 @@ def get_unique_algorithms() -> Set[str]:
     for batch_mode in [False, True]:
         for properties, _ in load_all_results(batch_mode=batch_mode):
             algorithms.add(properties["algo"])
-    return algorithms
+    return algorithms

data_export.py

@@ -9,14 +9,21 @@
     parser = argparse.ArgumentParser()
     parser.add_argument("--output", help="Path to the output file", required=True)
     parser.add_argument("--recompute", action="store_true", help="Recompute metrics")
+    parser.add_argument(
+        "-k", "--count", default=10, type=int, help="The number of near neighbours to search for"
+    )
+    parser.add_argument("--batch", action="store_true", help="Batch mode")
     args = parser.parse_args()
 
     datasets = DATASETS.keys()
     dfs = []
     for dataset_name in datasets:
         print("Looking at dataset", dataset_name)
-        if len(list(load_all_results(dataset_name))) > 0:
-            results = load_all_results(dataset_name)
+        if len(list(load_all_results(dataset_name,
+                                        count=args.count,
+                                        batch_mode=args.batch
+                                     ))) > 0:
+            results = load_all_results(dataset_name, count=args.count, batch_mode=args.batch)
             dataset, _ = get_dataset(dataset_name)
             results = compute_metrics_all_runs(dataset, results, args.recompute)
             for res in results: