From 0f617cb50266fece9a1fa5a240365cca421f506d Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Mon, 6 Jul 2026 09:24:37 -0300 Subject: [PATCH] add datamint-inference --- datamint/__main__.py | 1 + .../client_cmd_tools/datamint_inference.py | 226 ++++++++++++++++++ docs/source/command_line_tools.rst | 32 ++- pyproject.toml | 1 + 4 files changed, 259 insertions(+), 1 deletion(-) create mode 100644 datamint/client_cmd_tools/datamint_inference.py diff --git a/datamint/__main__.py b/datamint/__main__.py index f6f5294e..64328745 100644 --- a/datamint/__main__.py +++ b/datamint/__main__.py @@ -10,6 +10,7 @@ "upload": "datamint.client_cmd_tools.datamint_upload", "init": "datamint.client_cmd_tools.datamint_init", "train": "datamint.client_cmd_tools.datamint_train", + "inference": "datamint.client_cmd_tools.datamint_inference", } diff --git a/datamint/client_cmd_tools/datamint_inference.py b/datamint/client_cmd_tools/datamint_inference.py new file mode 100644 index 00000000..d88db030 --- /dev/null +++ b/datamint/client_cmd_tools/datamint_inference.py @@ -0,0 +1,226 @@ +"""datamint-inference command-line tool. + +Run local inference with a registered Datamint model against a local file, without +writing any Python. Loads the model via MLflow (`models://latest`), builds a +`LocalResource` from the given file, and prints the resulting predictions. +""" +from __future__ import annotations + +import argparse +import logging +import sys +from typing import Any + +from rich.console import Console +from rich.table import Table + +from datamint.client_cmd_tools.datamint_upload import _is_valid_path_argparse, handle_api_key +from datamint.exceptions import DatamintException, ItemNotFoundError +from datamint.utils.logging_utils import ConsoleWrapperHandler, load_cmdline_logging_config + +_LOGGER = logging.getLogger(__name__) +_USER_LOGGER = logging.getLogger('user_logger') +CONSOLE: Console + +_OVERLAY_COLORS = ['red', 'lime', 'cyan', 'yellow', 'magenta', 'orange'] + + +class DatamintInferenceCliError(Exception): + """Expected, user-facing CLI error (bad input, model/project not found, etc.).""" + + +def _resolve_project_name(args: argparse.Namespace) -> str: + return args.project if args.project is not None else args.model_name + + +def _load_model(project_name: str, model_name: str): + from mlflow.exceptions import MlflowException + + from datamint.mlflow import set_project + from datamint.mlflow.flavors import load_model + + try: + set_project(project_name) + except ItemNotFoundError as e: + raise DatamintInferenceCliError( + f"Project '{project_name}' not found. Models are looked up by project; if " + f"'{model_name}' was registered under a different project name, pass --project " + "explicitly." + ) from e + + try: + return load_model(f'models:/{model_name}/latest') + except MlflowException as e: + raise DatamintInferenceCliError( + f"Could not load model '{model_name}' (project '{project_name}'): {e}" + ) from e + + +def _run_inference(model, file_path: str): + from datamint.entities.resource import LocalResource + + resource = LocalResource(local_filepath=file_path) + predictions = model.predict([resource]) + return resource, (predictions[0] if predictions else []) + + +def _print_predictions(console: Console, predictions: list) -> None: + if not predictions: + console.print("[warning]No predictions returned.[/warning]") + return + + table = Table(title="Predictions") + table.add_column("Type", style="key") + table.add_column("Identifier") + table.add_column("Confidence") + for ann in predictions: + confidence = getattr(ann, 'confiability', None) + conf_str = f"{confidence:.3f}" if isinstance(confidence, (int, float)) else "-" + identifier = getattr(ann, 'identifier', None) or getattr(ann, 'text_value', None) or '-' + table.add_row(ann.annotation_type, str(identifier), conf_str) + console.print(table) + + +def _save_overlay(console: Console, resource: Any, predictions: list, output_path: str) -> None: + import matplotlib + matplotlib.use('Agg') + import matplotlib.patches as patches + import matplotlib.pyplot as plt + import numpy as np + + img_np = np.array(resource.fetch_file_data(auto_convert=True)) + if img_np.ndim == 3 and img_np.shape[0] in (1, 3) and img_np.shape[0] != img_np.shape[-1]: + img_np = np.transpose(img_np, (1, 2, 0)) # (C, H, W) -> (H, W, C) + + fig, ax = plt.subplots(figsize=(8, 8)) + ax.imshow(img_np, cmap='gray' if img_np.ndim == 2 else None) + + drew_geometry = False + labels_only: list[str] = [] + + for i, ann in enumerate(predictions): + color = _OVERLAY_COLORS[i % len(_OVERLAY_COLORS)] + + if ann.annotation_type == 'segmentation' and getattr(ann, 'mask', None) is not None: + drew_geometry = True + mask_np = np.asarray(ann.mask) > 0 + rgba = np.zeros((*mask_np.shape, 4)) + rgba[..., :3] = matplotlib.colors.to_rgb(color) + rgba[..., 3] = 0.4 * mask_np + ax.imshow(rgba) + elif getattr(ann, 'geometry', None) is not None and hasattr(ann.geometry, 'points'): + drew_geometry = True + xs = [p[0] for p in ann.geometry.points] + ys = [p[1] for p in ann.geometry.points] + x_min, y_min = min(xs), min(ys) + rect = patches.Rectangle( + (x_min, y_min), max(xs) - x_min, max(ys) - y_min, + linewidth=2, edgecolor=color, facecolor='none', + ) + ax.add_patch(rect) + ax.text(x_min, y_min - 5, ann.identifier or '', color=color, fontsize=9) + else: + label = getattr(ann, 'identifier', None) or getattr(ann, 'text_value', None) + if label: + labels_only.append(str(label)) + + if labels_only: + ax.set_title(', '.join(labels_only), fontsize=11) + elif not drew_geometry: + ax.set_title('No predictions', fontsize=11) + + ax.axis('off') + fig.tight_layout() + fig.savefig(output_path, dpi=150) + plt.close(fig) + console.print(f"[success]Saved visualization to {output_path}[/success]") + + +def _execute(args: argparse.Namespace, console: Console) -> int: + project_name = _resolve_project_name(args) + + with console.status("[accent]Loading model...[/accent]"): + model = _load_model(project_name, args.model_name) + + with console.status("[accent]Running inference...[/accent]"): + resource, predictions = _run_inference(model, args.file) + + _print_predictions(console, predictions) + + if args.output: + _save_overlay(console, resource, predictions, args.output) + + return 0 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description='Run local inference with a registered Datamint model against a local file.', + epilog=""" +Examples: + datamint-inference file.png --model-name MyModel + # Predict using the model registered as 'MyModel' + datamint-inference file.png --model-name my-model-alias --project MyProject + # Model registered under a different name than its project + datamint-inference file.png --model-name MyModel --output result.png + # Also save a visualization of the predictions + +More Documentation: https://sonanceai.github.io/datamint-python-api/command_line_tools.html + """, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument('file', type=_is_valid_path_argparse, metavar='FILE', + help='Path to the local file to run inference on.') + parser.add_argument('--model-name', type=str, required=True, + help='Name of the registered model (loaded as models:/NAME/latest).') + parser.add_argument('--project', type=str, default=None, + help="Project the model was registered under. Defaults to --model-name " + "itself (the default for every trainer unless model_name was set " + "explicitly to something else).") + parser.add_argument('--output', type=str, default=None, metavar='PATH', + help='Save a rendered visualization of the predictions to this path.') + parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.') + + return parser.parse_args() + + +def main() -> None: + global CONSOLE + load_cmdline_logging_config() + CONSOLE = [h for h in _USER_LOGGER.handlers if isinstance(h, ConsoleWrapperHandler)][0].console + + args = _parse_args() + + if args.verbose: + logging.getLogger().handlers[0].setLevel(logging.DEBUG) + logging.getLogger('datamint').setLevel(logging.DEBUG) + _LOGGER.setLevel(logging.DEBUG) + _USER_LOGGER.setLevel(logging.DEBUG) + + try: + api_key = handle_api_key() + if api_key is None: + _USER_LOGGER.error("API key not provided. Aborting.") + sys.exit(1) + import os + + from datamint import configs + os.environ[configs.ENV_VARS[configs.APIKEY_KEY]] = api_key + + from datamint.mlflow.env_utils import ensure_mlflow_configured + ensure_mlflow_configured() + + sys.exit(_execute(args, CONSOLE)) + except DatamintInferenceCliError as e: + _USER_LOGGER.error(f'❌ {e}') + sys.exit(1) + except DatamintException as e: + _USER_LOGGER.error(f'❌ {e}') + sys.exit(1) + except KeyboardInterrupt: + CONSOLE.print("\nInference cancelled by user.", style='warning') + sys.exit(1) + + +if __name__ == '__main__': + main() diff --git a/docs/source/command_line_tools.rst b/docs/source/command_line_tools.rst index 6cf8f16d..7d29cc5b 100644 --- a/docs/source/command_line_tools.rst +++ b/docs/source/command_line_tools.rst @@ -27,6 +27,7 @@ There are three command-line tools available: - ``datamint-config``: To configure the Datamint API key and URL. - ``datamint-upload``: To upload DICOM, NIfTI, video, image, and segmentation files to the Datamint server. - ``datamint-train``: To train a model on a Datamint project using a built-in one-line trainer. +- ``datamint-inference``: To run local inference with a registered Datamint model against a local file. Configuring the Datamint settings --------------------------------- @@ -293,4 +294,33 @@ Advanced training options (custom losses, transforms, encoders, ``trainer_kwargs are intentionally not exposed here — use the Python SDK instead, see :doc:`Training your Model `. -See all available options by running ``datamint-train --help``. \ No newline at end of file +See all available options by running ``datamint-train --help``. + +Running local inference +------------------------ + +The ``datamint-inference`` command-line tool runs a registered Datamint model against a +local file, without writing any Python. It loads the model via MLflow +(``models://latest``), runs it against the given file, and prints the resulting +predictions: + +.. code-block:: bash + + datamint-inference file.png --model-name MyModel + +Models are looked up by project. By default, the project name is assumed to be the same +as ``--model-name``. If the model was registered under a different project, pass +``--project`` explicitly: + +.. code-block:: bash + + datamint-inference file.png --model-name my-model-alias --project MyProject + +To also save a visualization of the predictions overlaid on the input file, use +``--output``: + +.. code-block:: bash + + datamint-inference file.png --model-name MyModel --output result.png + +See all available options by running ``datamint-inference --help``. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index eb4a3255..138e1dc8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ datamint-upload = 'datamint.client_cmd_tools.datamint_upload:main' datamint-config = 'datamint.client_cmd_tools.datamint_config:main' datamint-init = 'datamint.client_cmd_tools.datamint_init:main' datamint-train = 'datamint.client_cmd_tools.datamint_train:main' +datamint-inference = 'datamint.client_cmd_tools.datamint_inference:main' [tool.poetry] # license = "Proprietary" # https://python-poetry.org/docs/pyproject/