from math import floor, sqrt
from pathlib import Path
import numpy as np
import pandas as pd
from scipy.stats import moment
from skimage.filters import sobel
from skimage.io import imsave
from skimage.measure import shannon_entropy
from tifffile import imwrite
from extensions.chaeo.annotators import draw_box_on_patch
from model_server.accessors import GenericImageDataAccessor, InMemoryDataAccessor
from model_server.process import pad, rescale, resample_to_8bit
def _make_rgb(zs):
h, w, c, nz = zs.shape
assert c <= 3
outdata = np.zeros((h, w, 3, nz), dtype=zs.dtype)
outdata[:, :, 0:c, :] = zs[:, :, :, :]
return outdata
def _focus_metrics():
return {
'max_intensity': lambda x: np.max(x),
'stdev': lambda x: np.std(x),
'max_sobel': lambda x: np.max(sobel(x)),
'rms_sobel': lambda x: sqrt(np.mean(sobel(x) ** 2)),
'entropy': lambda x: shannon_entropy(x),
'moment': lambda x: moment(x.flatten(), moment=2),
}
def _write_patch_to_file(where, fname, data):
ext = fname.split('.')[-1].upper()
where.mkdir(parents=True, exist_ok=True)
if ext == 'PNG':
assert data.dtype == 'uint8', f'Invalid data type {data.dtype}'
assert data.shape[2] <= 3, f'Cannot export images with more than 3 channels as PNGs'
assert data.shape[3] == 1, f'Cannot export z-stacks as PNGs'
if data.shape[2] == 1:
outdata = data[:, :, 0, 0]
elif data.shape[2] == 2: # add a blank blue channel
outdata = _make_rgb(data)
else: # preserve RGB order
outdata = data[:, :, :, 0]
imsave(where / fname, outdata, check_contrast=False)
return True
elif ext in ['TIF', 'TIFF']:
zcyx = np.moveaxis(data, [3, 2, 0, 1], [0, 1, 2, 3])
imwrite(where / fname, zcyx, imagej=True)
return True
else:
raise Exception(f'Unsupported file extension: {ext}')
def export_patch_masks_from_zstack(
where: Path,
zmask_meta: list,
pad_to: int = 256,
prefix='mask',
):
exported = []
for mi in zmask_meta:
obj = mi['info']
mask = np.expand_dims(mi['mask'], (2, 3))
if pad_to:
mask = pad(mask, pad_to)
ext = 'png'
fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}.{ext}'
mask8bit = 255 * mask.astype('uint8')
_write_patch_to_file(where, fname, mask8bit)
exported.append(fname)
return exported
def export_patches_from_zstack(
where: Path,
stack: GenericImageDataAccessor,
zmask_meta: list,
rescale_clip: float = 0.0,
pad_to: int = 256,
make_3d: bool = False,
prefix='patch',
focus_metric: str = None,
**kwargs
):
# assert stack.chroma == 1, 'Expecting monochromatic image data'
assert stack.nz > 1, 'Expecting z-stack'
exported = []
for mi in zmask_meta:
obj = mi['info']
sl = mi['slice']
rbb = mi['relative_bounding_box']
x0 = rbb['x0']
y0 = rbb['y0']
x1 = rbb['x1']
y1 = rbb['y1']
patch3d = stack.data[sl]
ph, pw, pc, pz = patch3d.shape
# make a 3d patch
if make_3d:
patch = patch3d
# make a 2d patch, find optimal z-position determined by focus_metric function
elif focus_metric is not None:
foc = _focus_metrics()[focus_metric]
sp_sl = np.s_[y0: y1, x0: x1, :, :]
subpatch = patch3d[sp_sl]
patch = np.zeros([ph, pw, pc, 1], dtype=patch3d.dtype)
for ci in range(0, pc):
me = [foc(subpatch[:, :, ci, zi]) for zi in range(0, pz)]
zif = np.argmax(me)
patch[:, :, ci, 0] = patch3d[:, :, ci, zif]
# make a 2d patch from middle of z-stack
else:
zim = floor(pz / 2)
patch = patch3d[:, :, :, [zim]]
assert len(patch.shape) == 4
assert patch.shape[2] == stack.chroma
if rescale_clip is not None:
patch = rescale(patch, rescale_clip)
if kwargs.get('draw_bounding_box') is True:
bci = kwargs.get('bounding_box_channel', 0)
assert bci < 3
if bci > 0:
patch = _make_rgb(patch)
for zi in range(0, patch.shape[3]):
patch[:, :, bci, zi] = draw_box_on_patch(
patch[:, :, bci, zi],
((x0, y0), (x1, y1)),
)
if pad_to:
patch = pad(patch, pad_to)
ext = 'tif' if make_3d else 'png'
fname = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}.{ext}'
_write_patch_to_file(where, fname, resample_to_8bit(patch))
exported.append(fname)
return exported
def export_3d_patches_with_focus_metrics(
where: Path,
stack: GenericImageDataAccessor,
zmask_meta: list,
rescale_clip: float = 0.0,
pad_to: int = 256,
prefix='patch',
**kwargs
):
"""
Export 3D patches as multi-level z-stacks, along with CSV of various focus methods for each z-position
:param kwargs:
annotate_focus_metric: name focus metric to use when drawing bounding box at optimal focus z-position
:return:
list of exported files
"""
assert stack.chroma == 1, 'Expecting monochromatic image data'
assert stack.nz > 1, 'Expecting z-stack'
def get_zstack_focus_metrics(zs):
nz = zs.shape[3]
me = _focus_metrics()
dd = {}
for zi in range(0, nz):
spf = zs[:, :, :, zi]
dd[zi] = {k: me[k](spf) for k in me.keys()}
return dd
exported = []
patch_meta = []
for mi in zmask_meta:
obj = mi['info']
sl = mi['slice']
rbb = mi['relative_bounding_box']
patch = stack.data[sl]
assert len(patch.shape) == 4
assert patch.shape[2] == stack.chroma
if rescale_clip is not None:
patch = rescale(patch, rescale_clip)
# unpack relative bounding box and define subset of patch data
x0 = rbb['x0']
y0 = rbb['y0']
x1 = rbb['x1']
y1 = rbb['y1']
sp_sl = np.s_[y0: y1, x0: x1, :, :]
subpatch = patch[sp_sl]
# compute focus metrics for all z-levels
me_dict = get_zstack_focus_metrics(subpatch)
patch_meta.append({'label': obj.label, 'zi': obj.zi, 'metrics': me_dict})
me_df = pd.DataFrame(me_dict).T
# drawing bounding box only on focused slice
ak = kwargs.get('annotate_focus_metric')
if ak and ak in me_df.columns:
zi_foc = me_df.idxmax().to_dict()[ak]
patch[:, :, 0, zi_foc] = draw_box_on_patch(
patch[:, :, 0, zi_foc],
((x0, y0), (x1, y1)),
)
if pad_to:
patch = pad(patch, pad_to)
fstem = f'{prefix}-la{obj.label:04d}-zi{obj.zi:04d}'
_write_patch_to_file(where, fstem + '.tif', resample_to_8bit(patch))
exported.append(fstem + '.tif')
me_df.to_csv(where / (fstem + '.csv'))
exported.append(fstem + '.csv')
return exported
def export_multichannel_patches_from_zstack(
where: Path,
stack: GenericImageDataAccessor,
zmask_meta: list,
ch_rgb_overlay: tuple = None,
overlay_gain: tuple = (1.0, 1.0, 1.0),
ch_white: int = None,
**kwargs
):
def _safe_add(a, g, b):
assert a.dtype == b.dtype
assert a.shape == b.shape
assert g >= 0.0
return np.clip(
a.astype('uint32') + g * b.astype('uint32'),
0,
np.iinfo(a.dtype).max
).astype(a.dtype)
idata = stack.data
if ch_white:
assert ch_white < stack.chroma
mdata = idata[:, :, [ch_white, ch_white, ch_white], :]
else:
mdata = idata
if ch_rgb_overlay:
assert len(ch_rgb_overlay) == 3
assert len(overlay_gain) == 3
for ii, ci in enumerate(ch_rgb_overlay):
if ci is None:
continue
assert isinstance(ci, int)
assert ci < stack.chroma
mdata[:, :, ii, :] = _safe_add(
mdata[:, :, ii, :],
overlay_gain[ii],
idata[:, :, ci, :]
)
mstack = InMemoryDataAccessor(mdata)
return export_patches_from_zstack(
where, mstack, zmask_meta, **kwargs
)