File size: 3,378 Bytes
8f6f117 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """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
|