cryo-EM inpainting

cryo-em-inpainting.py

from typing import Tuple

import einops
import imageio
import numpy as np
from scipy.interpolate import LSQBivariateSpline
import napari

image = imageio.imread('image/EMPIAR-10164_TS_01_000_0.0.tif')
mask = imageio.imread('mask/EMPIAR-10164_TS_01_000_0.0.tif')

viewer = napari.Viewer()
image_layer = viewer.add_image(image)
mask_layer = viewer.add_labels(mask)


def central_crop_2d(image, percentage=25):
    h, w = image.shape[-2], image.shape[-1]
    mh, mw = h // 2, w // 2
    dh, dw = int(h * (percentage / 100 / 2)), int(w * (percentage / 100 / 2))
    fh, fw = mh - dh, mw - dw
    ch, cw = mh + dh, mw + dw
    return image[..., fh:ch, fw:cw]


def estimate_background(
        image, mask, background_model_resolution: Tuple[int, int] = (5, 5),
        n_samples_for_fit: int = 20000
):
    """Estimate image background with a bivariate cubic spline."""
    if mask is not None:
        idx_background = np.argwhere(mask == 0)
    else:
        idx_background = einops.rearrange(np.indices(dimensions=image.shape), 'yx h w -> (h w) yx')
    sample_idx = np.random.choice(idx_background.shape[0], size=n_samples_for_fit, replace=False)
    idx_background = idx_background[sample_idx]
    y, x = idx_background[:, 0], idx_background[:, 1]
    z = image[(y, x)]
    ty = np.linspace(0, image.shape[0], num=background_model_resolution[0])
    tx = np.linspace(0, image.shape[1], num=background_model_resolution[1])

    spline = LSQBivariateSpline(x, y, z, tx, ty)
    y = np.arange(image.shape[0])
    x = np.arange(image.shape[1])
    return spline(y, x, grid=True)


def estimate_background_std(image, mask):
    """Estimate the standard deviation of the background from a central crop."""
    image = central_crop_2d(image, percentage=25)
    mask = central_crop_2d(mask, percentage=25)
    return np.std(central_crop_2d(image)[central_crop_2d(mask) == 0])


def inpaint(image, mask, background_model_resolution: Tuple[int, int] = (5, 5),
            add_noise: bool = True):
    background = estimate_background(image, mask, background_model_resolution)
    idx_mask = np.argwhere(mask == 1)
    idx_mask = (idx_mask[:, 0], idx_mask[:, 1])
    inpainted = np.copy(image)
    inpainted[idx_mask] = background[idx_mask]
    if add_noise is True:
        background_std = estimate_background_std(image, mask)
        n_pixels_to_inpaint = idx_mask[0].shape[0]
        inpainted[idx_mask] += np.random.normal(loc=0, scale=background_std,
                                                size=n_pixels_to_inpaint)
    return inpainted


background = estimate_background(
    image, mask, background_model_resolution=(5, 5), n_samples_for_fit=20000
)
inpainted_flat = inpaint(image, mask, background_model_resolution=(5, 5), add_noise=False)
inpainted_noisy = inpaint(image, mask, background_model_resolution=(5, 5), add_noise=True)

viewer.add_image(background, contrast_limits=image_layer.contrast_limits, visible=False)
viewer.add_image(inpainted_flat, contrast_limits=image_layer.contrast_limits, visible=False)
viewer.add_image(inpainted_noisy, contrast_limits=image_layer.contrast_limits, visible=False)

napari.run()