clean up inference pipeline

Author: geospatial-jeff

claymodel/cli.py

@@ -1,24 +1,14 @@
 """Clay model CLI."""
 
-import importlib.util
-
 import click
 
 from claymodel.api import load_metadata, load_model
-
-
-def _check_inference_deps() -> None:
-    """Check that inference optional dependencies are installed."""
-    missing = [
-        pkg
-        for pkg in ("rustac", "lazycogs", "geopandas", "shapely")
-        if importlib.util.find_spec(pkg) is None
-    ]
-    if missing:
-        raise click.ClickException(
-            f"Missing inference dependencies: {', '.join(missing)}\n"
-            "Install them with: pip install claymodel[inference]"
-        )
+from claymodel.inference.pipeline import (
+    generate_embeddings,
+    load_scene,
+    save_embeddings_geoparquet,
+    search_scene,
+)
 
 
 @click.group()
@@ -97,17 +87,7 @@ def embed(
     Searches Earth Search, loads COG bands via lazycogs, chips into patches,
     runs the Clay encoder, and saves embeddings as GeoParquet.
 
-    Requires: pip install claymodel[inference]
     """
-    _check_inference_deps()
-
-    from claymodel.inference.pipeline import (
-        generate_embeddings,
-        load_scene,
-        save_embeddings_geoparquet,
-        search_scene,
-    )
-
     click.echo(f"Scene: {scene_id}")
     click.echo(f"Model: {size} (checkpoint: {ckpt})")
     click.echo(f"Device: {device}")

claymodel/inference/pipeline.py

@@ -3,8 +3,6 @@
 Provides functions to search Earth Search STAC, load Sentinel-2 COGs via
 lazycogs, chip into model-ready patches, and generate embeddings with
 spatial metadata. Usable both programmatically and via the CLI.
-
-Requires optional dependencies: ``pip install claymodel[inference]``
 """
 
 __all__ = [
@@ -25,8 +23,12 @@
 from datetime import datetime
 from pathlib import Path
 
+import geopandas as gpd
+import lazycogs
 import numpy as np
+import rustac
 import torch
+from shapely.geometry import box
 
 from claymodel.api import load_metadata, load_model, normalize
 from claymodel.metadata import PlatformMetadata
@@ -130,8 +132,6 @@ async def _search_scene_async(
     Returns:
         Number of items found.
     """
-    import rustac
-
     return await rustac.search_to(
         str(parquet_path),
         EARTH_SEARCH_URL,
@@ -158,14 +158,6 @@ def search_scene(
         ValueError: If the scene is not found.
         ImportError: If rustac is not installed.
     """
-    try:
-        import rustac  # noqa: F401
-    except ImportError as exc:
-        raise ImportError(
-            "rustac is required for the inference pipeline. "
-            "Install it with: pip install claymodel[inference]"
-        ) from exc
-
     if output_dir is None:
         output_dir = Path(tempfile.mkdtemp(prefix="clay_"))
     else:
@@ -191,8 +183,6 @@ def search_scene(
         )
 
     # Read back the parquet to extract metadata
-    import geopandas as gpd
-
     gdf = gpd.read_parquet(parquet_path)
     row = gdf.iloc[0]
 
@@ -238,14 +228,6 @@ def load_scene(
     Raises:
         ImportError: If lazycogs is not installed.
     """
-    try:
-        import lazycogs
-    except ImportError as exc:
-        raise ImportError(
-            "lazycogs is required for the inference pipeline. "
-            "Install it with: pip install claymodel[inference]"
-        ) from exc
-
     parquet_str = str(search_result.parquet_path)
     store = lazycogs.store_for(parquet_str, skip_signature=True)
 
@@ -469,9 +451,6 @@ def save_embeddings_geoparquet(
     Raises:
         ImportError: If geopandas is not installed.
     """
-    import geopandas as gpd
-    from shapely.geometry import box
-
     output_path = Path(output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 

pyproject.toml

@@ -29,9 +29,14 @@ license-files = ["LICENSE"]
 dependencies = [
     "click>=8.1",
     "einops~=0.7.0",
+    "geopandas>=0.14",
+    "lazycogs",
     "lightning>=2.0.0",
     "pydantic>=2.0",
+    "pyarrow",
     "pyyaml",
+    "rustac",
+    "shapely>=2.0",
     "timm>=0.6.0",
     "torch>=2.4.0",
     "torchvision>=0.19.0",
@@ -51,14 +56,8 @@ finetune = [
     "tifffile",
     "rasterio",
 ]
-inference = [
-    "geopandas>=0.14",
-    "lazycogs",
-    "rustac",
-    "shapely>=2.0",
-]
 all = [
-    "claymodel[train,finetune,inference]",
+    "claymodel[train,finetune]",
 ]
 
 [project.scripts]
@@ -149,8 +148,8 @@ known-first-party = ["claymodel", "training", "finetune"]
 "claymodel/api.py" = ["PLR0912", "PLR0915", "PLR2004"]
 "claymodel/module.py" = ["PLR2004"]
 "claymodel/inference/masking.py" = ["PLR2004"]
-"claymodel/inference/pipeline.py" = ["PLC0415", "PLR0913", "PLR2004"]
-"claymodel/cli.py" = ["PLC0415", "PLR0913"]
+"claymodel/inference/pipeline.py" = ["PLR0913", "PLR2004"]
+"claymodel/cli.py" = ["PLR0913"]
 "tests/*.py" = [
     "ANN001",
     "ANN003",