euler-inference 2.0.1__py3-none-any.whl

euler_inference/__init__.py

@@ -0,0 +1,8 @@
+"""Model inference pipeline using euler-loading."""
+
+# Note: We don't import submodules here to avoid RuntimeWarning when running
+# `python -m euler_inference`. Import directly from submodules instead:
+#   from euler_inference.config import InferenceConfig
+#   from euler_inference.inference import run_inference
+
+__all__ = ["config", "inference", "models"]

euler_inference/__main__.py

@@ -0,0 +1,5 @@
+"""Allow running as ``python -m euler_inference``."""
+
+from euler_inference.inference import main
+
+main()

euler_inference/_outputs.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import logging
+from pathlib import Path
+from typing import IO, Any
+
+import numpy as np
+from euler_loading import MultiModalDataset
+
+from euler_inference.config import InferenceConfig, OutputConfig
+
+
+def get_save_path(output_base: Path, full_id: str, suffix: str, extension: str) -> Path:
+    """
+    Get save path from full_id with configurable suffix and extension.
+
+    The full_id is "/" separated, representing the hierarchical structure
+    the file originates from (e.g., "Scene01/clone/Camera_0/00001").
+
+    We save as: output_base/Scene01/clone/Camera_0/00001{suffix}.{extension}
+    """
+    full_id = full_id.lstrip("/")
+    return output_base / f"{full_id}{suffix}.{extension}"
+
+
+def _normalize_image_array(data: np.ndarray) -> np.ndarray:
+    if data.dtype in (np.float32, np.float64):
+        return (np.clip(data, 0, 1) * 255).astype(np.uint8)
+    if data.dtype != np.uint8:
+        return data.astype(np.uint8)
+    return data
+
+
+def _build_image(data: np.ndarray):
+    from PIL import Image
+
+    normalized = _normalize_image_array(data)
+    if normalized.ndim == 2:
+        return Image.fromarray(normalized, mode="L")
+    if normalized.ndim == 3 and normalized.shape[2] == 3:
+        return Image.fromarray(normalized, mode="RGB")
+    if normalized.ndim == 3 and normalized.shape[2] == 4:
+        return Image.fromarray(normalized, mode="RGBA")
+    raise ValueError(f"Unsupported array shape for image: {data.shape}")
+
+
+def _save_image_output(data: np.ndarray, target: Path | IO[bytes], output_type: str) -> None:
+    format_name = {"png": "PNG", "jpg": "JPEG", "jpeg": "JPEG"}[output_type]
+    _build_image(data).save(target, format=format_name)
+
+
+def _save_exr_output(data: np.ndarray, target: Path) -> None:
+    try:
+        import Imath
+        import OpenEXR
+    except ImportError as exc:
+        raise ImportError(
+            "OpenEXR package required for EXR output. Install with: pip install OpenEXR"
+        ) from exc
+
+    half_chan = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT))
+    if data.ndim == 2:
+        height, width = data.shape
+        header = OpenEXR.Header(width, height)
+        header["channels"] = {"Y": half_chan}
+        exr = OpenEXR.OutputFile(str(target), header)
+        exr.writePixels({"Y": data.astype(np.float32).tobytes()})
+        exr.close()
+        return
+
+    if data.ndim == 3 and data.shape[2] == 3:
+        height, width, _ = data.shape
+        header = OpenEXR.Header(width, height)
+        header["channels"] = {"R": half_chan, "G": half_chan, "B": half_chan}
+        exr = OpenEXR.OutputFile(str(target), header)
+        exr.writePixels({
+            "R": data[:, :, 0].astype(np.float32).tobytes(),
+            "G": data[:, :, 1].astype(np.float32).tobytes(),
+            "B": data[:, :, 2].astype(np.float32).tobytes(),
+        })
+        exr.close()
+        return
+
+    raise ValueError(f"Unsupported array shape for EXR: {data.shape}")
+
+
+def save_output(data: np.ndarray, target: Path | IO[bytes], output_type: str) -> None:
+    """
+    Save output data to file in the specified format.
+
+    Args:
+        data: numpy array to save
+        target: Path to save to, or a writable file-like object (for zip streaming)
+        output_type: File format ("npy", "png", "jpg", "jpeg", "exr")
+
+    Raises:
+        ValueError: If output_type is not supported
+    """
+    if isinstance(target, Path):
+        target.parent.mkdir(parents=True, exist_ok=True)
+
+    if output_type == "npy":
+        np.save(target, data)
+        return
+
+    if output_type in ("png", "jpg", "jpeg"):
+        _save_image_output(data, target, output_type)
+        return
+
+    if output_type == "exr":
+        if not isinstance(target, Path):
+            raise ValueError("EXR output is not supported in zip mode")
+        _save_exr_output(data, target)
+        return
+
+    raise ValueError(f"Unsupported output type: {output_type}")
+
+
+def _get_writer_kwargs(
+    output_config: OutputConfig,
+    *,
+    strict: bool = True,
+    dataset_name: str | None = None,
+) -> dict[str, Any]:
+    """Derive DatasetWriter/ZipDatasetWriter kwargs from an OutputConfig."""
+    writer_meta = output_config.writer or {}
+    name = writer_meta.get("name", dataset_name or output_config.key)
+    type_ = writer_meta.get(
+        "type",
+        writer_meta.get("euler_train", {}).get("modality_type", output_config.key),
+    )
+    default_modality_type = output_config.key if strict else "other"
+    euler_train = writer_meta.get("euler_train", {
+        "used_as": "target",
+        "modality_type": default_modality_type,
+    })
+    kwargs = dict(name=name, type=type_, euler_train=euler_train, separator=None)
+
+    meta = writer_meta.get("meta")
+    if meta is not None:
+        kwargs["meta"] = meta
+
+    euler_loading = writer_meta.get("euler_loading")
+    if euler_loading is not None:
+        merged = {
+            key: value for key, value in euler_train.items()
+            if key in ("used_as", "modality_type")
+        }
+        merged.update(euler_loading)
+        kwargs["euler_loading"] = merged
+    return kwargs
+
+
+def create_writers(
+    config: InferenceConfig,
+    dataset_name: str | None = None,
+    outputs: list[OutputConfig] | None = None,
+) -> dict[str, Any]:
+    """Create one DatasetWriter or ZipDatasetWriter per output key."""
+    from ds_crawler import DatasetWriter, ZipDatasetWriter
+
+    output_base = Path(config.output_base_path)
+    writers: dict[str, Any] = {}
+    selected_outputs = outputs if outputs is not None else config.outputs
+    for output_config in selected_outputs:
+        kwargs = _get_writer_kwargs(
+            output_config,
+            strict=config.strict,
+            dataset_name=dataset_name,
+        )
+        if config.zip:
+            root = output_base / f"{output_config.key}.zip"
+            writers[output_config.key] = ZipDatasetWriter(root, **kwargs)
+        else:
+            root = output_base / output_config.key
+            writers[output_config.key] = DatasetWriter(root, **kwargs)
+    return writers
+
+
+@dataclass(frozen=True)
+class _WriterBackedOutputPlan:
+    output_config: OutputConfig
+    source_modality: str
+    output_root: Path
+
+
+def _resolve_source_modality(
+    output_config: OutputConfig,
+    modality_names: set[str],
+) -> tuple[str | None, bool]:
+    if output_config.source_modality:
+        return output_config.source_modality, True
+    if output_config.key in modality_names:
+        return output_config.key, False
+    return None, False
+
+
+def _output_root(base_path: str, output_key: str, *, zip_mode: bool) -> Path:
+    output_base = Path(base_path)
+    if zip_mode:
+        return output_base / f"{output_key}.zip"
+    return output_base / output_key
+
+
+def _build_output_plans(
+    config: InferenceConfig,
+    dataset: MultiModalDataset,
+    logger: logging.Logger,
+) -> tuple[list[_WriterBackedOutputPlan], list[OutputConfig]]:
+    """Split outputs into euler-loading-backed vs legacy serialization."""
+    modality_names = set(config.dataset.modalities)
+    writer_backed: list[_WriterBackedOutputPlan] = []
+    legacy: list[OutputConfig] = []
+
+    for output_config in config.outputs:
+        source_modality, explicit = _resolve_source_modality(
+            output_config,
+            modality_names,
+        )
+
+        if source_modality is None:
+            legacy.append(output_config)
+            continue
+
+        if source_modality not in modality_names:
+            if explicit:
+                raise ValueError(
+                    f"Output '{output_config.key}' references source_modality="
+                    f"{source_modality!r}, but that modality is not present in "
+                    "dataset.modalities."
+                )
+            legacy.append(output_config)
+            continue
+
+        try:
+            dataset.get_writer(source_modality)
+        except (KeyError, ValueError) as exc:
+            if explicit:
+                raise ValueError(
+                    f"Output '{output_config.key}' requested source_modality="
+                    f"{source_modality!r}, but no euler-loading writer is "
+                    f"available: {exc}"
+                ) from exc
+            logger.warning(
+                "Output '%s' matches modality '%s', but no euler-loading "
+                "writer is available (%s). Falling back to the legacy "
+                "serializer.",
+                output_config.key,
+                source_modality,
+                exc,
+            )
+            legacy.append(output_config)
+            continue
+
+        if output_config.writer:
+            logger.warning(
+                "Output '%s' is source-backed via modality '%s'; "
+                "outputs[].writer metadata is ignored for this mode.",
+                output_config.key,
+                source_modality,
+            )
+
+        output_root = _output_root(
+            config.output_base_path,
+            output_config.key,
+            zip_mode=config.zip,
+        )
+        logger.info(
+            "Output '%s': mirroring source modality '%s' into %s. Source "
+            "filenames and extensions are preserved.",
+            output_config.key,
+            source_modality,
+            output_root,
+        )
+        writer_backed.append(
+            _WriterBackedOutputPlan(
+                output_config=output_config,
+                source_modality=source_modality,
+                output_root=output_root,
+            )
+        )
+
+    for output_config in legacy:
+        root = _output_root(
+            config.output_base_path,
+            output_config.key,
+            zip_mode=config.zip,
+        )
+        logger.info(
+            "Output '%s': using the legacy serializer into %s.",
+            output_config.key,
+            root,
+        )
+
+    return writer_backed, legacy
+
+
+def _create_source_backed_writers(
+    plans: list[_WriterBackedOutputPlan],
+    *,
+    dataset: MultiModalDataset,
+) -> dict[str, Any]:
+    writers: dict[str, Any] = {}
+    for plan in plans:
+        writers[plan.output_config.key] = dataset.create_output_writer(
+            plan.source_modality,
+            plan.output_root,
+            zip=plan.output_root.suffix.lower() == ".zip",
+        )
+    return writers
+
+
+def _full_id_leaf(full_id: str) -> str:
+    parts = [part for part in full_id.split("/") if part]
+    return parts[-1] if parts else full_id
+
+
+@dataclass
+class PreparedOutputs:
+    dataset: MultiModalDataset
+    expected_output_keys: tuple[str, ...]
+    writer_backed_plans: list[_WriterBackedOutputPlan]
+    legacy_outputs: list[OutputConfig]
+    source_backed_writers: dict[str, Any]
+    legacy_writers: dict[str, Any]
+    zip_mode: bool
+
+    @classmethod
+    def prepare(
+        cls,
+        *,
+        config: InferenceConfig,
+        dataset: MultiModalDataset,
+        logger: logging.Logger,
+    ) -> "PreparedOutputs":
+        writer_backed_plans, legacy_outputs = _build_output_plans(
+            config,
+            dataset,
+            logger,
+        )
+        source_backed_writers = _create_source_backed_writers(
+            writer_backed_plans,
+            dataset=dataset,
+        ) if writer_backed_plans else {}
+        legacy_writers = create_writers(
+            config,
+            dataset_name=dataset.get_dataset_name(),
+            outputs=legacy_outputs,
+        ) if legacy_outputs else {}
+        return cls(
+            dataset=dataset,
+            expected_output_keys=tuple(output.key for output in config.outputs),
+            writer_backed_plans=writer_backed_plans,
+            legacy_outputs=legacy_outputs,
+            source_backed_writers=source_backed_writers,
+            legacy_writers=legacy_writers,
+            zip_mode=config.zip,
+        )
+
+    def validate_prediction_keys(self, outputs: dict[str, Any]) -> None:
+        missing_keys = [
+            output_key
+            for output_key in self.expected_output_keys
+            if output_key not in outputs
+        ]
+        if missing_keys:
+            raise ValueError(
+                "Model did not return the configured output keys "
+                f"{missing_keys}. Got: {list(outputs.keys())}"
+            )
+
+    def write_predictions(
+        self,
+        *,
+        sample_index: int,
+        full_id: str,
+        outputs: dict[str, Any],
+    ) -> None:
+        for plan in self.writer_backed_plans:
+            self.dataset.write_sample(
+                sample_index,
+                {plan.source_modality: outputs[plan.output_config.key]},
+                {plan.source_modality: self.source_backed_writers[plan.output_config.key]},
+            )
+
+        if not self.legacy_writers:
+            return
+
+        file_id = _full_id_leaf(full_id)
+        for output_config in self.legacy_outputs:
+            output_data = outputs[output_config.key]
+            basename = (
+                f"{file_id}{output_config.effective_suffix}.{output_config.type}"
+            )
+            writer = self.legacy_writers[output_config.key]
+            if self.zip_mode:
+                with writer.open(full_id, basename) as handle:
+                    save_output(output_data, handle, output_config.type)
+            else:
+                path = writer.get_path(full_id, basename)
+                save_output(output_data, path, output_config.type)
+
+    def finalize(self) -> None:
+        for writer in self.legacy_writers.values():
+            writer.save_index()
+
+        for writer in self.source_backed_writers.values():
+            writer.save_index()

euler_inference/api.py

@@ -0,0 +1,165 @@
+"""
+Python API for programmatic inference.
+
+Provides a simple `infer()` function so pipelines can call inference
+from Python without constructing JSON configs manually.
+
+Usage (model card):
+    from euler_inference.api import infer
+
+    # With a card file and placeholder bindings
+    infer(
+        model_card="model_card.json",
+        bindings={"weights": "/path/to/checkpoint.pt"},
+        data={"rgb": "/data/rgb"},
+        output_base_path="/output",
+    )
+
+    # With an already-resolved card dict (from server)
+    infer(
+        model_card={"model": "/abs/path/model.py", ...},
+        data={"rgb": "/data/rgb"},
+        output_base_path="/output",
+    )
+
+Usage (direct):
+    from euler_inference.api import infer
+
+    infer(
+        model_path="/path/to/model.py",
+        output_base_path="/output",
+        dataset_modalities={"rgb": "/data/rgb"},
+    )
+"""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from euler_inference.config import (
+    DatasetConfig,
+    ExternalModelConfig,
+    InferenceConfig,
+    OutputConfig,
+)
+from euler_inference.inference import run_inference
+
+
+def infer(
+    # Direct args (optional when model_card is used)
+    model_path: Optional[str] = None,
+    output_base_path: str = "",
+    dataset_modalities: Optional[dict] = None,
+    *,
+    # Model card args
+    model_card: Optional[Union[str, dict]] = None,
+    bindings: Optional[dict[str, str]] = None,
+    data: Optional[dict[str, str]] = None,
+    hierarchical_data: Optional[dict[str, str]] = None,
+    # Shared kwargs
+    model_config: Optional[dict] = None,
+    outputs: Optional[list[dict]] = None,
+    dataset_hierarchical_modalities: Optional[dict] = None,
+    device: Optional[str] = None,
+    max_samples: Optional[int] = None,
+    zip: bool = False,
+    strict: bool = True,
+    verbose: bool = False,
+) -> None:
+    """
+    Run inference with an external model.
+
+    Supports two modes:
+
+    1. **Model card mode** (preferred): Pass `model_card` as a path to a
+       model_card.json file or an already-resolved dict. Placeholder bindings
+       are provided via `bindings`, `data`, and `hierarchical_data`.
+
+    2. **Direct mode**: Pass `model_path`, `output_base_path`, and
+       `dataset_modalities` directly.
+
+    Args:
+        model_path: (Direct) Absolute path to the model.py file
+        output_base_path: Directory to save predictions
+        dataset_modalities: (Direct) Mapping of modality names to paths
+        model_card: Path to model_card.json (str) or already-resolved card
+            dict. When provided, direct args are ignored.
+        bindings: Placeholder bindings for model card resolution
+            (e.g., {"weights": "/path/to/checkpoint.pt"})
+        data: Input modality path bindings
+            (e.g., {"rgb": "/data/rgb"})
+        hierarchical_data: Hierarchical input path bindings
+            (e.g., {"textgt": "/data/textgt"})
+        model_config: Model-specific config dict
+        outputs: List of output dicts
+        dataset_hierarchical_modalities: Hierarchical modality paths
+        device: Device string ("cuda", "cpu", "mps"). Auto-detected if None.
+        max_samples: Max samples to process. None for all.
+        verbose: Enable verbose logging.
+    """
+    if model_card is not None:
+        from euler_inference.model_card import (
+            load_model_card,
+            model_card_to_config,
+            resolve_placeholders,
+        )
+
+        if isinstance(model_card, str):
+            # Path to card file — load and resolve placeholders
+            card_dict, card_dir = load_model_card(model_card)
+            all_bindings: dict[str, str] = {}
+            all_bindings.update(bindings or {})
+            all_bindings.update(data or {})
+            all_bindings.update(hierarchical_data or {})
+            card_dict = resolve_placeholders(card_dict, all_bindings)
+        elif isinstance(model_card, dict):
+            # Already-resolved dict (from server)
+            card_dict = model_card
+            card_dir = Path(".")
+        else:
+            raise TypeError(
+                f"model_card must be str or dict, got {type(model_card).__name__}"
+            )
+
+        config = model_card_to_config(
+            card_dict,
+            card_dir,
+            output_base_path=output_base_path,
+            device=device,
+            max_samples=max_samples,
+            zip=zip,
+            strict=strict,
+        )
+        run_inference(config, verbose=verbose)
+        return
+
+    # Direct code path
+    if model_path is None:
+        raise ValueError("Either model_card or model_path must be provided")
+    if dataset_modalities is None:
+        raise ValueError("dataset_modalities is required when not using model_card")
+
+    parsed_outputs = None
+    outputs_from_default = True
+    if outputs is not None:
+        parsed_outputs = [OutputConfig.from_dict(o) for o in outputs]
+        outputs_from_default = False
+
+    config = InferenceConfig(
+        external_model=ExternalModelConfig(
+            model_path=model_path,
+            model_config=model_config,
+        ),
+        dataset=DatasetConfig(
+            modalities=dataset_modalities,
+            hierarchical_modalities=dataset_hierarchical_modalities,
+        ),
+        output_base_path=output_base_path,
+        device=device,
+        max_samples=max_samples,
+        zip=zip,
+        strict=strict,
+        _outputs_from_default=outputs_from_default,
+        **({"outputs": parsed_outputs} if parsed_outputs else {}),
+    )
+
+    run_inference(config, verbose=verbose)