Merge pull request #300 from JdeRobot/issue-299

Update computational cost estimation and add CLI support

Author: dpascualhe

detectionmetrics/cli/__init__.py

@@ -1,5 +1,129 @@
+from detectionmetrics import datasets
+from detectionmetrics import models
 from detectionmetrics.cli.evaluate import evaluate
+from detectionmetrics.cli.computational_cost import computational_cost
 
 REGISTRY = {
     "evaluate": evaluate,
+    "computational_cost": computational_cost,
 }
+
+
+def get_model(task, input_type, model_format, model, ontology, model_cfg):
+    # Init model from registry
+    model_name = f"{model_format}_{input_type}_{task}"
+    if model_name not in models.REGISTRY:
+        raise ValueError(
+            f"Model format not supported: {model_format}. "
+            f"Must be one of {models.REGISTRY.keys()}",
+        )
+    return models.REGISTRY[model_name](model, model_cfg, ontology)
+
+
+def get_dataset(
+    task,
+    input_type,
+    dataset_format,
+    dataset_fname,
+    dataset_dir,
+    split_dir,
+    train_dataset_dir,
+    val_dataset_dir,
+    test_dataset_dir,
+    images_dir,
+    labels_dir,
+    data_suffix,
+    label_suffix,
+    ontology,
+    split,
+):
+    # Check if required data is available
+    if dataset_format == "gaia":
+        if dataset_fname is None:
+            raise ValueError("--dataset is required for 'gaia' format")
+
+    elif dataset_format in ["rellis3d", "wildscenes"]:
+        if dataset_dir is None:
+            raise ValueError(
+                "--dataset_dir is required for 'rellis3d' and 'wildscenes' formats"
+            )
+        if split_dir is None:
+            raise ValueError(
+                "--split_dir is required for 'rellis3d' and 'wildscenes' formats"
+            )
+
+        if dataset_format == "rellis3d" and ontology is None:
+            raise ValueError("--dataset_ontology is required for 'rellis3d' format")
+
+    elif dataset_format in ["goose", "generic"]:
+        if "train" in split and train_dataset_dir is None:
+            raise ValueError(
+                "--train_dataset_dir is required for 'train' split in 'goose' and 'generic' formats"
+            )
+        elif "val" in split and val_dataset_dir is None:
+            raise ValueError(
+                "--val_dataset_dir is required for 'val' split in 'goose' and 'generic' formats"
+            )
+        elif "test" in split and test_dataset_dir is None:
+            raise ValueError(
+                "--test_dataset_dir is required for 'test' split in 'goose' and 'generic' formats"
+            )
+
+        if dataset_format == "generic":
+            if data_suffix is None:
+                raise ValueError("--data_suffix is required for 'generic' format")
+            if label_suffix is None:
+                raise ValueError("--label_suffix is required for 'generic' format")
+            if ontology is None:
+                raise ValueError("--dataset_ontology is required for 'generic' format")
+
+    elif dataset_format == "rugd":
+        if images_dir is None:
+            raise ValueError("--images_dir is required for 'rugd' format")
+        if labels_dir is None:
+            raise ValueError("--labels_dir is required for 'rugd' format")
+
+    else:
+        raise ValueError(f"Dataset format not supported: {dataset_format}")
+
+    # Get arguments to init dataset
+    if dataset_format == "gaia":
+        dataset_args = {"dataset_fname": dataset_fname}
+    elif dataset_format == "rellis3d":
+        dataset_args = {
+            "dataset_dir": dataset_dir,
+            "split_dir": split_dir,
+            "ontology_fname": ontology,
+        }
+    elif dataset_format == "goose":
+        dataset_args = {
+            "train_dataset_dir": train_dataset_dir,
+            "val_dataset_dir": val_dataset_dir,
+            "test_dataset_dir": test_dataset_dir,
+        }
+    elif dataset_format == "generic":
+        dataset_args = {
+            "data_suffix": data_suffix,
+            "label_suffix": label_suffix,
+            "ontology_fname": ontology,
+            "train_dataset_dir": train_dataset_dir,
+            "val_dataset_dir": val_dataset_dir,
+            "test_dataset_dir": test_dataset_dir,
+        }
+    elif dataset_format == "rugd":
+        dataset_args = {
+            "images_dir": images_dir,
+            "labels_dir": labels_dir,
+            "ontology_fname": ontology,
+        }
+    else:
+        raise ValueError(f"Dataset format not supported: {dataset_format}")
+
+    # Init dataset from registry
+    dataset_name = f"{dataset_format}_{input_type}_{task}"
+    if dataset_name not in datasets.REGISTRY:
+        raise ValueError(
+            f"Dataset format not supported: {dataset_format}. "
+            f"Must be one of {datasets.REGISTRY.keys()}",
+        )
+    return datasets.REGISTRY[dataset_name](**dataset_args)

detectionmetrics/cli/batch.py

@@ -21,7 +21,8 @@ def batch(command, jobs_cfg):
         jobs_cfg["model"] = [jobs_cfg["model"]]
 
     # Same for dataset
-    if not isinstance(jobs_cfg["dataset"], list):
+    has_dataset = "dataset" in jobs_cfg
+    if has_dataset and not isinstance(jobs_cfg["dataset"], list):
         jobs_cfg["dataset"] = [jobs_cfg["dataset"]]
 
     # Build list of model configurations
@@ -54,11 +55,16 @@ def batch(command, jobs_cfg):
 
     # Build list of jobs (IDs must be unique)
     all_jobs = {}
-    for model_cfg, dataset_cfg in product(model_cfgs, jobs_cfg["dataset"]):
-        job_id = f"{model_cfg['id']}-{dataset_cfg['id']}"
+    job_iter = product(model_cfgs, jobs_cfg["dataset"]) if has_dataset else model_cfgs
+    for job_components in job_iter:
+        if not isinstance(job_components, tuple):
+            job_components = (job_components,)
+
+        job_id = "-".join([str(jc["id"]) for jc in job_components])
         if job_id in all_jobs:
             raise ValueError(f"Job ID {job_id} is not unique")
-        all_jobs[job_id] = (model_cfg, dataset_cfg)
+
+        all_jobs[job_id] = job_components
 
     print("\n" + "-" * 80)
     print(f"{len(all_jobs)} job(s) will be executed:")
@@ -69,7 +75,7 @@ def batch(command, jobs_cfg):
     # Start processing jobs
     pbar = tqdm(all_jobs.items(), total=len(all_jobs), leave=True)
     preds_outdir = None
-    for job_id, (model_cfg, dataset_cfg) in pbar:
+    for job_id, job_components in pbar:
         job_out_fname = os.path.join(jobs_cfg["outdir"], f"{job_id}.csv")
         if jobs_cfg.get("store_results_per_sample", False):
             preds_outdir = os.path.join(jobs_cfg["outdir"], f"preds-{job_id}")
@@ -84,31 +90,50 @@ def batch(command, jobs_cfg):
 
         ctx = click.get_current_context()
         try:
-            result = ctx.invoke(
-                cli_registry[command],
-                task=jobs_cfg["task"],
-                input_type=jobs_cfg["input_type"],
-                model_format=model_cfg["format"],
-                model=model_cfg["path"],
-                model_ontology=model_cfg["ontology"],
-                model_cfg=model_cfg["cfg"],
-                dataset_format=dataset_cfg["format"],
-                dataset_fname=dataset_cfg.get("fname", None),
-                dataset_dir=dataset_cfg.get("dir", None),
-                split_dir=dataset_cfg.get("split_dir", None),
-                train_dataset_dir=dataset_cfg.get("train_dir", None),
-                val_dataset_dir=dataset_cfg.get("val_dir", None),
-                test_dataset_dir=dataset_cfg.get("test_dir", None),
-                images_dir=dataset_cfg.get("data_dir", None),
-                labels_dir=dataset_cfg.get("labels_dir", None),
-                data_suffix=dataset_cfg.get("data_suffix", None),
-                label_suffix=dataset_cfg.get("label_suffix", None),
-                dataset_ontology=dataset_cfg.get("ontology", None),
-                split=dataset_cfg["split"],
-                ontology_translation=jobs_cfg.get("ontology_translation", None),
-                out_fname=job_out_fname,
-                predictions_outdir=preds_outdir,
+            params = {
+                "task": jobs_cfg["task"],
+                "input_type": jobs_cfg["input_type"],
+            }
+
+            model_cfg = job_components[0]
+            params.update(
+                {
+                    "model_format": model_cfg["format"],
+                    "model": model_cfg["path"],
+                    "model_ontology": model_cfg["ontology"],
+                    "model_cfg": model_cfg["cfg"],
+                    # "image_size": model_cfg.get("image_size", None),
+                }
             )
+            if has_dataset:
+                dataset_cfg = job_components[1]
+                params.update(
+                    {
+                        "dataset_format": dataset_cfg.get("format", None),
+                        "dataset_fname": dataset_cfg.get("fname", None),
+                        "dataset_dir": dataset_cfg.get("dir", None),
+                        "split_dir": dataset_cfg.get("split_dir", None),
+                        "train_dataset_dir": dataset_cfg.get("train_dir", None),
+                        "val_dataset_dir": dataset_cfg.get("val_dir", None),
+                        "test_dataset_dir": dataset_cfg.get("test_dir", None),
+                        "images_dir": dataset_cfg.get("data_dir", None),
+                        "labels_dir": dataset_cfg.get("labels_dir", None),
+                        "data_suffix": dataset_cfg.get("data_suffix", None),
+                        "label_suffix": dataset_cfg.get("label_suffix", None),
+                        "dataset_ontology": dataset_cfg.get("ontology", None),
+                        "split": dataset_cfg["split"],
+                        "ontology_translation": jobs_cfg.get(
+                            "ontology_translation", None
+                        ),
+                    }
+                )
+
+            params.update({"out_fname": job_out_fname})
+            if preds_outdir is not None:
+                params.update({"predictions_outdir": preds_outdir})
+
+            result = ctx.invoke(cli_registry[command], **params)
+
         except Exception as e:
             print(f"Error processing job {job_id}: {e}")
             continue

detectionmetrics/cli/computational_cost.py

@@ -0,0 +1,79 @@
+import click
+
+from detectionmetrics import cli
+from detectionmetrics.utils.io import read_json
+
+
+@click.command(name="computational_cost", help="Estimate model computational cost")
+@click.argument("task", type=click.Choice(["segmentation"], case_sensitive=False))
+@click.argument(
+    "input_type", type=click.Choice(["image", "lidar"], case_sensitive=False)
+)
+# model
+@click.option(
+    "--model_format",
+    type=click.Choice(
+        ["torch", "tensorflow", "tensorflow_explicit"], case_sensitive=False
+    ),
+    show_default=True,
+    default="torch",
+    help="Trained model format",
+)
+@click.option(
+    "--model",
+    type=click.Path(exists=True),
+    required=True,
+    help="Trained model filename (TorchScript) or directory (TensorFlow SavedModel)",
+)
+@click.option(
+    "--model_ontology",
+    type=click.Path(exists=True, dir_okay=False),
+    required=True,
+    help="JSON file containing model output ontology",
+)
+@click.option(
+    "--model_cfg",
+    type=click.Path(exists=True, dir_okay=False),
+    required=True,
+    help="JSON file with model configuration (norm. parameters, image size, etc.)",
+)
+@click.option(
+    "--image_size",
+    type=(int, int),
+    required=False,
+    help="Dummy image size used for computational cost estimation",
+)
+# output
+@click.option(
+    "--out_fname",
+    type=click.Path(writable=True),
+    help="CSV file where the computational cost estimation results will be stored",
+)
+def computational_cost(
+    task,
+    input_type,
+    model_format,
+    model,
+    model_ontology,
+    model_cfg,
+    image_size,
+    out_fname,
+):
+    """Estimate model computational cost"""
+
+    if image_size is None:
+        parsed_model_cfg = read_json(model_cfg)
+        if "image_size" in parsed_model_cfg:
+            image_size = parsed_model_cfg["image_size"]
+        else:
+            raise ValueError(
+                "Image size must be provided either as an argument or in the model configuration file"
+            )
+
+    model = cli.get_model(
+        task, input_type, model_format, model, model_ontology, model_cfg
+    )
+    results = model.get_computational_cost(image_size)
+    results.to_csv(out_fname)
+
+    return results