fed_analysis: support new dataset with added masks (NVIDIA#633)

* support new dataset with added masks

* do not hard-code filesep

examples/federated_analysis/data/prepare_data.py

@@ -26,7 +26,7 @@ def create_datasets(root, subdirs, extension, shuffle, seed):
 
     data_lists = []
     for subdir in subdirs:
-        search_string = os.path.join(root, "**", subdir, "*" + extension)
+        search_string = os.path.join(root, "**", subdir, "images", "*" + extension)
         data_list = glob.glob(search_string, recursive=True)
 
         assert (
@@ -56,7 +56,7 @@ def save_data_list(data, data_list_file, data_root, key="data"):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input_dir", type=str, required=True, help="Location of image files")
-    parser.add_argument("--input_ext", type=str, default=".png", help="Search extions")
+    parser.add_argument("--input_ext", type=str, default=".png", help="Search extension")
     parser.add_argument("--output_dir", type=str, default="./data", help="Output location of data lists")
     parser.add_argument(
         "--subdirs",