Merge pull request #304 from JdeRobot/issue-303

Update examples

Author: dpascualhe

examples/convert_dataset_ontology.py

@@ -30,7 +30,16 @@ def parse_args():
         required=True,
         help="Directory to save the new dataset",
     )
-    return parser.parse_args()
+    parser.add_argument(
+        "--resize",
+        type=str,
+        required=False,
+        help="Resize images to a specific size (e.g. 512x512)",
+    )
+
+    args = parser.parse_args()
+    args.resize = tuple(map(int, args.resize.split("x"))) if args.resize else None
+    return args
 
 
 def main():
@@ -44,7 +53,12 @@ def main():
     with open(args.ontology_translation, "r", encoding="utf-8") as f:
         ontology_translation = json.load(f)
 
-    dataset.export(args.outdir, new_ontology, ontology_translation)
+    dataset.export(
+        args.outdir,
+        resize=args.resize,
+        new_ontology=new_ontology,
+        ontology_translation=ontology_translation,
+    )
 
 
 if __name__ == "__main__":

examples/gaia_image.py

@@ -25,6 +25,12 @@ def parse_args() -> argparse.Namespace:
         required=False,
         help="Resize images to a specific size (e.g. 512x512)",
     )
+    parser.add_argument(
+        "--split",
+        type=str,
+        required=False,
+        help="What split to export (train, val, test). Export all splits if not specified",
+    )
 
     args = parser.parse_args()
     args.resize = tuple(map(int, args.resize.split("x"))) if args.resize else None
@@ -36,6 +42,9 @@ def main():
     args = parse_args()
 
     dataset = GaiaImageSegmentationDataset(dataset_fname=args.dataset)
+    if args.split:
+        dataset.dataset = dataset.dataset[dataset.dataset["split"] == args.split]
+        dataset.has_label_count = False
     dataset.export(outdir=args.outdir, resize=args.resize)
 
 

examples/goose_image.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 
 from detectionmetrics.datasets.goose import GOOSEImageSegmentationDataset
 
@@ -32,10 +33,19 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Directory where dataset will be stored in common format",
     )
+    parser.add_argument(
+        "--new_ontology",
+        type=str,
+        help="New ontology JSON file name",
+    )
+    parser.add_argument(
+        "--ontology_translation",
+        type=str,
+        help="Ontology translation JSON file name",
+    )
     parser.add_argument(
         "--resize",
         type=str,
-        required=False,
         help="Resize images to a specific size (e.g. 512x512)",
     )
 
@@ -48,12 +58,26 @@ def main():
     """Main function"""
     args = parse_args()
 
+    new_ontology, ontology_translation = None, None
+    if args.new_ontology is not None:
+        with open(args.new_ontology, "r", encoding="utf-8") as f:
+            new_ontology = json.load(f)
+
+    if args.ontology_translation is not None:
+        with open(args.ontology_translation, "r", encoding="utf-8") as f:
+            ontology_translation = json.load(f)
+
     dataset = GOOSEImageSegmentationDataset(
         train_dataset_dir=args.train_dataset_dir,
         val_dataset_dir=args.val_dataset_dir,
         test_dataset_dir=args.test_dataset_dir,
     )
-    dataset.export(args.outdir, resize=args.resize)
+    dataset.export(
+        args.outdir,
+        resize=args.resize,
+        new_ontology=new_ontology,
+        ontology_translation=ontology_translation,
+    )
 
 
 if __name__ == "__main__":

examples/rellis3d_image.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 
 from detectionmetrics.datasets.rellis3d import Rellis3DImageSegmentationDataset
 
@@ -34,10 +35,19 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Directory where dataset will be stored in common format",
     )
+    parser.add_argument(
+        "--new_ontology",
+        type=str,
+        help="New ontology JSON file name",
+    )
+    parser.add_argument(
+        "--ontology_translation",
+        type=str,
+        help="Ontology translation JSON file name",
+    )
     parser.add_argument(
         "--resize",
         type=str,
-        required=False,
         help="Resize images to a specific size (e.g. 512x512)",
     )
 
@@ -50,12 +60,26 @@ def main():
     """Main function"""
     args = parse_args()
 
+    new_ontology, ontology_translation = None, None
+    if args.new_ontology is not None:
+        with open(args.new_ontology, "r", encoding="utf-8") as f:
+            new_ontology = json.load(f)
+
+    if args.ontology_translation is not None:
+        with open(args.ontology_translation, "r", encoding="utf-8") as f:
+            ontology_translation = json.load(f)
+
     dataset = Rellis3DImageSegmentationDataset(
         dataset_dir=args.dataset_dir,
         split_dir=args.split_dir,
         ontology_fname=args.ontology_fname,
     )
-    dataset.export(args.outdir, resize=args.resize)
+    dataset.export(
+        args.outdir,
+        resize=args.resize,
+        new_ontology=new_ontology,
+        ontology_translation=ontology_translation,
+    )
 
 
 if __name__ == "__main__":