core/deepinterp_engine.py

"""The 'deepinterpolation' engine: real Allen Institute DeepInterpolation.

This is the wired slot for the actual TensorFlow/Keras `deepinterpolation`
package (https://github.com/AllenInstitute/deepinterpolation). It is NOT part of
the light default image — install it via Dockerfile.deepinterp and point
DEEPINTERP_MODEL at a pretrained inference model (HDF5).

The learned network predicts frame t from a stack of its temporal neighbors
(center excluded) just like the fast engine, but with a deep U-Net that has been
trained on real two-photon / Ophys data, so it removes far more noise.
"""
from __future__ import annotations

import os

import numpy as np


def available() -> bool:
    try:
        import deepinterpolation  # noqa: F401
        return True
    except Exception:  # noqa: BLE001
        return False


def interpolate(movie: np.ndarray, pre: int = 30, post: int = 30,
                omit: int = 0) -> np.ndarray:
    """Run real DeepInterpolation inference over a (T, H, W) movie.

    Requires the `deepinterpolation` package (TensorFlow) and a pretrained model
    given by env DEEPINTERP_MODEL. Raises a clear error if either is missing so
    the UI/API can fall back or report it.
    """
    model_path = os.environ.get("DEEPINTERP_MODEL", "")
    if not available():
        raise RuntimeError(
            "The 'deepinterpolation' package is not installed. Build the image "
            "with Dockerfile.deepinterp to use this engine."
        )
    if not model_path or not os.path.exists(model_path):
        raise RuntimeError(
            "Set DEEPINTERP_MODEL to a pretrained DeepInterpolation inference "
            "model (HDF5). See https://github.com/AllenInstitute/deepinterpolation."
        )

    # Real inference path. Kept import-local so the light image never imports TF.
    import tempfile

    import tifffile
    from deepinterpolation.generic import ClassLoader

    movie = np.asarray(movie, dtype=np.float32)
    with tempfile.TemporaryDirectory() as td:
        in_tif = os.path.join(td, "input.tif")
        out_base = os.path.join(td, "di_out")
        tifffile.imwrite(in_tif, movie)

        generator_param = {
            "type": "generator",
            "name": "SingleTifGenerator",
            "pre_frame": int(pre),
            "post_frame": int(post),
            "pre_post_omission": int(omit),
            "train_path": in_tif,
            "batch_size": 1,
            "start_frame": int(pre),
            "end_frame": int(movie.shape[0] - post - 1),
        }
        inference_param = {
            "type": "inferrence",
            "name": "core_inferrence",
            "model_path": model_path,
            "output_file": out_base + ".h5",
        }

        data_generator = ClassLoader(generator_param).find_and_build()(generator_param)
        inferrence_class = ClassLoader(inference_param).find_and_build()(
            inference_param, data_generator
        )
        inferrence_class.run()

        import h5py
        with h5py.File(out_base + ".h5", "r") as f:
            pred = np.asarray(f["data"][:]).astype(np.float32)
    pred = np.squeeze(pred)

    # The network only produces frames [pre, T-post); pad edges with the input so
    # the output has the same shape (T, H, W) the contract expects.
    out = movie.copy()
    out[int(pre):int(pre) + pred.shape[0]] = pred
    return out