from pathlib import Path
from typing import Dict, List
from uuid import uuid4
import numpy as np
import pandas as pd
from skimage.measure import label, regionprops_table
from skimage.morphology import dilation
from sklearn.model_selection import train_test_split
from extensions.chaeo.accessors import MonoPatchStack
from extensions.chaeo.annotators import draw_boxes_on_3d_image
from extensions.chaeo.models import PatchStackObjectClassifier
from extensions.chaeo.process import mask_largest_object
from extensions.chaeo.products import export_patches_from_zstack, export_patch_masks_from_zstack, export_multichannel_patches_from_zstack, get_patches_from_zmask_meta, get_patch_masks_from_zmask_meta
from extensions.chaeo.zmask import build_zmask_from_object_mask, project_stack_from_focal_points
from extensions.ilastik.models import IlastikPixelClassifierModel
from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
from model_server.models import Model
from model_server.process import rescale
from model_server.workflows import Timer
def get_zmask_meta(
input_file_path: str,
ilastik_pixel_classifier: IlastikPixelClassifierModel,
segmentation_channel: int,
pxmap_threshold: float,
pxmap_foreground_channel: int = 0,
zmask_zindex: int = None,
zmask_clip: int = None,
zmask_filters: Dict = None,
zmask_type: str = 'boxes',
**kwargs,
) -> tuple:
ti = Timer()
stack = generate_file_accessor(Path(input_file_path))
fstem = Path(input_file_path).stem
ti.click('file_input')
# MIP if no zmask z-index is given, then classify pixels
if isinstance(zmask_zindex, int):
assert 0 < zmask_zindex < stack.nz
zmask_data = stack.get_one_channel_data(channel=segmentation_channel).data[:, :, :, zmask_zindex]
else:
zmask_data = stack.get_one_channel_data(channel=segmentation_channel).data.max(axis=-1, keepdims=True)
if zmask_clip:
zmask_data = rescale(zmask_data, zmask_clip)
mip = InMemoryDataAccessor(
zmask_data,
)
pxmap, _ = ilastik_pixel_classifier.infer(mip)
ti.click('infer_pixel_probability')
obmask = InMemoryDataAccessor(
pxmap.data > pxmap_threshold
)
ti.click('threshold_pixel_mask')
# make zmask
zmask, zmask_meta, df, interm = build_zmask_from_object_mask(
obmask.get_one_channel_data(pxmap_foreground_channel),
stack.get_one_channel_data(segmentation_channel),
mask_type=zmask_type,
filters=zmask_filters,
expand_box_by=kwargs['zmask_expand_box_by'],
)
ti.click('generate_zmasks')
# record pixel scale
df['pixel_scale_in_micrometers'] = float(stack.pixel_scale_in_micrometers.get('X'))
return ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm
# TODO: unpack and validate inputs
def export_patches_from_multichannel_zstack(
input_file_path: str,
output_folder_path: str,
models: List[Model],
pxmap_threshold: float,
pxmap_foreground_channel: int,
segmentation_channel: int,
patches_channel: int,
zmask_zindex: int = None, # None for MIP,
zmask_clip: int = None,
zmask_type: str = 'boxes',
zmask_filters: Dict = None,
zmask_expand_box_by: int = None,
export_pixel_probabilities=True,
export_2d_patches_for_training=True,
export_2d_patches_for_annotation=True,
draw_bounding_box_on_2d_patch=True,
draw_contour_on_2d_patch=False,
draw_mask_on_2d_patch=False,
export_3d_patches=True,
export_annotated_zstack=True,
draw_label_on_zstack=False,
export_patch_masks=True,
rgb_overlay_channels=(None, None, None),
rgb_overlay_weights=(1.0, 1.0, 1.0),
) -> Dict:
pixel_classifier = models[0]
ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm = get_zmask_meta(
input_file_path,
pixel_classifier,
segmentation_channel,
pxmap_threshold,
pxmap_foreground_channel=pxmap_foreground_channel,
zmask_zindex=zmask_zindex,
zmask_clip=zmask_clip,
zmask_expand_box_by=zmask_expand_box_by,
zmask_filters=zmask_filters,
zmask_type=zmask_type,
)
if export_pixel_probabilities:
write_accessor_data_to_file(
Path(output_folder_path) / 'pixel_probabilities' / (fstem + '.tif'),
pxmap
)
ti.click('export_pixel_probability')
if export_3d_patches and len(zmask_meta) > 0:
files = export_patches_from_zstack(
Path(output_folder_path) / '3d_patches',
stack.get_one_channel_data(patches_channel),
zmask_meta,
prefix=fstem,
draw_bounding_box=False,
rescale_clip=0.001,
make_3d=True,
)
ti.click('export_3d_patches')
if export_2d_patches_for_annotation and len(zmask_meta) > 0:
files = export_multichannel_patches_from_zstack(
Path(output_folder_path) / '2d_patches_annotation',
stack,
zmask_meta,
prefix=fstem,
rescale_clip=0.001,
make_3d=False,
focus_metric='max_sobel',
ch_white=patches_channel,
ch_rgb_overlay=rgb_overlay_channels,
draw_bounding_box=draw_bounding_box_on_2d_patch,
bounding_box_channel=1,
bounding_box_linewidth=2,
draw_contour=draw_contour_on_2d_patch,
draw_mask=draw_mask_on_2d_patch,
overlay_gain=rgb_overlay_weights,
)
df_patches = pd.DataFrame(files)
ti.click('export_2d_patches')
# associate 2d patches, dropping labeled objects that were not exported as patches
df = pd.merge(df, df_patches, left_index=True, right_on='df_index').drop(columns='df_index')
# prepopulate patch UUID
df['patch_id'] = df.apply(lambda _: uuid4(), axis=1)
if export_2d_patches_for_training and len(zmask_meta) > 0:
files = export_multichannel_patches_from_zstack(
Path(output_folder_path) / '2d_patches_training',
stack.get_one_channel_data(patches_channel),
zmask_meta,
prefix=fstem,
rescale_clip=0.001,
make_3d=False,
focus_metric='max_sobel',
)
ti.click('export_2d_patches')
if export_patch_masks and len(zmask_meta) > 0:
files = export_patch_masks_from_zstack(
Path(output_folder_path) / 'patch_masks',
stack.get_one_channel_data(patches_channel),
zmask_meta,
prefix=fstem,
)
if export_annotated_zstack:
annotated = InMemoryDataAccessor(
draw_boxes_on_3d_image(
stack.get_one_channel_data(patches_channel).data,
zmask_meta,
add_label=draw_label_on_zstack,
)
)
write_accessor_data_to_file(
Path(output_folder_path) / 'annotated_zstacks' / (fstem + '.tif'),
annotated
)
ti.click('export_annotated_zstack')
# generate multichannel projection from label centroids
dff = df[df['keeper']]
if len(zmask_meta) > 0:
interm['projected'] = project_stack_from_focal_points(
dff['centroid-0'].to_numpy(),
dff['centroid-1'].to_numpy(),
dff['zi'].to_numpy(),
stack,
degree=4,
)
else: # else just return MIP
interm['projected'] = stack.data.max(axis=-1)
return {
'pixel_model_id': pixel_classifier.model_id,
'input_filepath': input_file_path,
'number_of_objects': len(zmask_meta),
'pixeL_scale_in_micrometers': stack.pixel_scale_in_micrometers,
'success': True,
'timer_results': ti.events,
'dataframe': df[df['keeper'] == True],
'interm': interm,
}
def infer_object_map_from_zstack(
input_file_path: str,
output_folder_path: str,
models: List[Model],
pxmap_threshold: float,
pxmap_foreground_channel: int,
segmentation_channel: int,
patches_channel: int,
zmask_zindex: int = None, # None for MIP,
zmask_clip: int = None,
zmask_type: str = 'boxes',
zmask_filters: Dict = None,
# zmask_expand_box_by: int = None,
**kwargs,
) -> Dict:
assert len(models) == 2
pixel_classifier = models[0]
assert isinstance(pixel_classifier, IlastikPixelClassifierModel)
object_classifier = models[1]
assert isinstance(object_classifier, PatchStackObjectClassifier)
ti, stack, fstem, obmask, pxmap, zmask, zmask_meta, df, interm = get_zmask_meta(
input_file_path,
pixel_classifier,
segmentation_channel,
pxmap_threshold,
pxmap_foreground_channel=pxmap_foreground_channel,
zmask_zindex=zmask_zindex,
zmask_clip=zmask_clip,
# zmask_expand_box_by=zmask_expand_box_by,
zmask_filters=zmask_filters,
zmask_type=zmask_type,
**kwargs
)
# extract patches to accessor
patches_acc = get_patches_from_zmask_meta(
stack.get_one_channel_data(patches_channel),
zmask_meta,
rescale_clip=zmask_clip,
make_3d=False,
focus_metric='max_sobel',
**kwargs
)
# extract masks
patch_masks_acc = get_patch_masks_from_zmask_meta(
stack,
zmask_meta,
**kwargs
)
# send patches and mask stacks to object classifier
result_acc, _ = object_classifier.infer(patches_acc, patch_masks_acc)
labels_map = interm['label_map']
output_map = np.zeros(labels_map.shape, dtype=labels_map.dtype)
assert labels_map.shape == interm['label_map'].shape
assert labels_map.dtype == interm['label_map'].dtype
# assign labels to object map:
meta = []
for ii in range(0, len(zmask_meta)):
object_id = zmask_meta[ii]['info'].label
result_patch = mask_largest_object(result_acc.iat(ii))
object_class = np.unique(result_patch)[1]
output_map[labels_map == object_id] = object_class
meta.append({'object_id': ii, 'object_class': object_id})
output_path = Path(output_folder_path) / ('obj_classes_' + (fstem + '.tif'))
write_accessor_data_to_file(
output_path,
InMemoryDataAccessor(output_map)
)
ti.click('export_object_classes')
return {
'timer_results': ti.events,
'dataframe': pd.DataFrame(meta),
'interm': {},
'output_path': output_path.__str__(),
}
def transfer_ecotaxa_labels_to_patch_stacks(
where_masks: str,
where_patches: str,
object_csv: str,
ecotaxa_tsv: str,
where_output: str,
patch_size: tuple = (256, 256),
tr_split=0.6,
dilate_label_mask: bool = True, # to mitigate connected components error in ilastik
allow_multiple_objects: bool = False,
) -> Dict:
assert tr_split > 0.5 # reduce chance that low-probability objects are omitted from training
# read patch metadata
df_obj = pd.read_csv(
object_csv,
)
df_ecotaxa = pd.read_csv(
ecotaxa_tsv,
sep='\t',
header=[0],
dtype={
('object_annotation_date', '[t]'): str,
('object_annotation_time', '[t]'): str,
('object_annotation_category_id', '[t]'): str,
}
)
df_merge = pd.merge(df_obj, df_ecotaxa, left_on='patch_id', right_on='object_id')
# assign each unique lowest-level annotation to a class index
se_unique = pd.Series(
df_merge.object_annotation_hierarchy.unique()
)
df_split = (
se_unique.str.rsplit(
pat='>', n=1, expand=True
)
)
df_labels = pd.DataFrame({
'annotation_class_id': df_split.index + 1,
'hierarchy': se_unique,
'annotation_class': df_split.loc[:, 1].str.lower()
})
# join patch filenames and annotation classes
df_pf = pd.merge(
df_merge[['patch_filename', 'object_annotation_hierarchy']],
df_labels,
left_on='object_annotation_hierarchy',
right_on='hierarchy',
)
df_pl = df_pf[df_pf['object_annotation_hierarchy'].notnull()]
# export annotation classes and their summary stats
df_tr, df_te = train_test_split(df_pl, train_size=tr_split)
# df_labels['counts'] = df_pl['annotation_class_id'].value_counts()
df_labels = pd.merge(
df_labels,
pd.DataFrame(
[df_pl.annotation_class_id.value_counts(), df_tr.annotation_class_id.value_counts(), df_te.annotation_class_id.value_counts()],
index=['total', 'to_train', 'to_test']
).T,
left_on='annotation_class_id',
right_index=True,
how='outer'
)
df_labels.loc[df_labels.to_train.isna(), 'to_train'] = 0
df_labels.loc[df_labels.to_test.isna(), 'to_test'] = 0
for col in ['total', 'to_train', 'to_test']:
df_labels.loc[df_labels[col].isna(), col] = 0
df_labels.to_csv(Path(where_output) / 'labels_key.csv', index=False)
# export patches as z-stacks
for (dfk, dfv) in {'train': df_tr, 'test': df_te}.items():
zstack_keys = ['mask', 'label', 'raw']
zstacks = {f'{dfk}_{zsk}': np.zeros((*patch_size, 1, len(dfv)), dtype='uint8') for zsk in zstack_keys}
stack_meta = []
for fi, pl in enumerate(dfv.itertuples(name='PatchFile')):
fn = pl._asdict()['patch_filename']
ac = pl._asdict()['annotation_class']
aci = pl._asdict()['annotation_class_id']
stack_meta.append({'zi': fi, 'patch_filename': fn, 'annotation_class': ac, 'annotation_class_id': aci})
acc_bm = generate_file_accessor(Path(where_masks) / fn)
assert acc_bm.is_mask()
assert acc_bm.hw == patch_size, f'Unexpected patch size {patch_size}'
assert acc_bm.chroma == 1
assert acc_bm.nz == 1
mask = acc_bm.data[:, :, 0, 0]
if dilate_label_mask:
mask = dilation(mask)
if not allow_multiple_objects:
ob_id = label(acc_bm.data[:, :, 0, 0])
mask = mask_largest_object(ob_id)
zstacks[dfk + '_mask'][:, :, 0, fi] = mask
zstacks[dfk + '_label'][:, :, 0, fi] = (mask == 255) * aci
acc_pa = generate_file_accessor(Path(where_patches) / fn)
zstacks[dfk + '_raw'][:, :, :, fi] = acc_pa.data[:, :, :, 0]
for k in zstacks.keys():
write_accessor_data_to_file(Path(where_output) / f'zstack_{k}.tif', InMemoryDataAccessor(zstacks[k]))
pd.DataFrame(stack_meta).to_csv(Path(where_output) / f'{dfk}_stack.csv', index=False)