Add MODNet checkpoint asset utilities (#377)

This allows you to now call run_reconstruction with
download_masking_model set to False. This will skip the download of the
MODNet weights, instead requiring the user to somehow ensure the model
weights are installed. They can be installed from file using the newly
added openlifu.util.assets.install_modnet_from_file.

Author: ebrahimebrahim

src/openlifu/nav/photoscan.py

@@ -18,13 +18,13 @@
 import numpy as np
 import onnxruntime as ort
 import OpenEXR
-import requests
 import trimesh
 import vtk
 from PIL import Image
 from vtk.util.numpy_support import numpy_to_vtk
 
 from openlifu.util.annotations import OpenLIFUFieldData
+from openlifu.util.assets import download_and_install_modnet, get_modnet_path
 
 logger_meshrecon = logging.getLogger("MeshRecon")
 logger_meshroom = logging.getLogger("Meshroom")
@@ -300,6 +300,7 @@ def run_reconstruction(
     locations: List[Tuple[float, float, float]] | None = None,
     return_durations: bool = False,
     progress_callback : Callable[[int,str],None] | None = None,
+    download_masking_model: bool = True,
 ) -> Tuple[Photoscan, Path] | Tuple[Photoscan, Path, Dict[str, float]]:
     """Run Meshroom with the given images and pipeline.
     Args:
@@ -322,6 +323,8 @@ def run_reconstruction(
         return_durations (bool): If True, also return a dictionary mapping node names to durations in seconds.
         progress_callback: An optional function that will be called to report progress. The function should accept two arguments:
             an integer progress value from 0 to 100 followed by a string message describing the step currently being worked on.
+        download_masking_model: Whether to auto-download the masking model weights if they are not present;
+            only relevant if use_masks is enabled.
 
     Returns:
         Union[Tuple[Photoscan, Path], Tuple[Photoscan, Path, Dict[str, float]]]:
@@ -432,7 +435,7 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
         start_time = time.perf_counter()
         masks_dir = temp_dir / "masks"
         masks_dir.mkdir(parents=True, exist_ok=True)
-        make_masks(new_paths, masks_dir)
+        make_masks(new_paths, masks_dir, download_model=download_masking_model)
         command.append( f"PrepareDenseScene_1.masksFolders=['{masks_dir.as_posix()}']" )
         durations["MaskCreation"] = time.perf_counter() - start_time
 
@@ -597,38 +600,6 @@ def inverse_transform(img):
 
     return inverse_transform(image) if inverse else transform(image)
 
-
-def get_modnet_path() -> Path:
-    """Get the MODNet checkpoint path. Download it if not present.
-    """
-    package = "openlifu.nav.modnet_checkpoints"
-    filename = "modnet_photographic_portrait_matting.onnx"
-    url = "https://data.kitware.com/api/v1/file/67feb2cb31a330568827ab32/download"
-    try:
-        # Try to find the checkpoint in the package
-        resource_path = importlib.resources.files(package) / filename
-        if resource_path.is_file():
-            logger_meshrecon.info(f"Found existing MODNet checkpoint at {resource_path}")
-            return resource_path
-    except (FileNotFoundError, ModuleNotFoundError):
-        pass
-
-    # Fallback: Download the checkpoint
-    base_dir = Path(importlib.resources.files(package))
-    full_path = base_dir / filename
-    logger_meshrecon.info(f"MODNet checkpoint not found. Downloading from {url}...")
-    response = requests.get(url, stream=True, timeout=(10, 300))
-    if response.status_code == 200:
-        with open(full_path, 'wb') as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                if chunk:
-                    f.write(chunk)
-        logger_meshrecon.info(f"Downloaded MODNet checkpoint to {full_path}")
-    else:
-        raise RuntimeError(f"Failed to download MODNet checkpoint: {response.status_code} - {response.text}")
-
-    return full_path
-
 def preprocess_image_modnet(image: np.ndarray, ref_size: int = 512) -> np.ndarray:
     """
     Preprocess an input image for MODNet inference.
@@ -675,7 +646,7 @@ def preprocess_image_modnet(image: np.ndarray, ref_size: int = 512) -> np.ndarra
     return image
 
 
-def make_masks(image_paths: list[Path], output_dir: Path, threshold: float = 0.01) -> None:
+def make_masks(image_paths: list[Path], output_dir: Path, threshold: float = 0.01, download_model=True) -> None:
     """
     Runs MODNet on a list of image paths and saves the output masks.
 
@@ -687,10 +658,23 @@ def make_masks(image_paths: list[Path], output_dir: Path, threshold: float = 0.0
         image_paths (List[str]): List of input image file paths.
         output_dir (str): Directory where the output masks will be saved.
         threshold (float): Threshold to binarize the soft segmentation output.
+        download_model (bool): Whether to auto-download the model weights if they are not present.
     """
+
     # Load the ONNX model
+
     ckpt_path = get_modnet_path()
+    if not ckpt_path.exists():
+        if download_model:
+            logger_meshrecon.info(f"Downloading MODNet checkpoint to {ckpt_path}")
+            download_and_install_modnet()
+        else:
+            raise FileNotFoundError(f"MODNet checkpoint not found at {ckpt_path}. Install it using an appropirate utility in openlifu.util.assets.")
+    else:
+        logger_meshrecon.info(f"Found existing MODNet checkpoint at {ckpt_path}")
+
     session = ort.InferenceSession(ckpt_path, providers=["CPUExecutionProvider"])  # or CUDAExecutionProvider
+
     for image_path in image_paths:
         image = Image.open(image_path)
         exif = image.getexif()

src/openlifu/util/assets.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import importlib
 import shutil
 import tempfile
 from pathlib import Path
@@ -50,3 +51,28 @@ def install_asset(destination:PathLike, path_to_asset:PathLike|None, url_to_asse
                 Path(temp_file_path).unlink()
     else:
         raise ValueError("Either path_to_asset or url_to_asset must be provided.")
+
+
+def get_modnet_path() -> Path:
+    """Get the MODNet checkpoint path.
+    It may or may not exist; see `download_and_install_modnet` and `install_modnet_from_file`.
+    If `get_modnet_path().exists()` is False, then use one of those two options to install.
+    """
+    package = "openlifu.nav.modnet_checkpoints"
+    filename = "modnet_photographic_portrait_matting.onnx"
+    base_dir = Path(importlib.resources.files(package))
+    return  base_dir / filename
+
+def download_and_install_modnet() -> Path:
+    """Download and install the MODNet checkpoint. Returns path to installed MODNet checkpoint."""
+    url = "https://data.kitware.com/api/v1/file/67feb2cb31a330568827ab32/download"
+    modnet_path = get_modnet_path()
+    install_asset(modnet_path, url_to_asset=url)
+    return modnet_path
+
+def install_modnet_from_file(path_to_modnet_file:PathLike) -> Path:
+    """Copy MODNet checkpoint to the appropriate place for openlifu to use it.
+    Returns path to installed MODNet checkpoint."""
+    modnet_path = get_modnet_path()
+    install_asset(modnet_path, path_to_asset=path_to_modnet_file)
+    return modnet_path