<commit_before>"""Code for segmenting 3D images."""

import SimpleITK as sitk

from jicbioimage.core.image import Image3D
from jicbioimage.core.transform import transformation

from utils import ColorImage3D


@transformation
def identity(im3d):
    return im3d


@transformation
def filter_median(im3d):
    itk_im = sitk.GetImageFromArray(im3d)
    median_filter = sitk.MedianImageFilter()
    itk_im = median_filter.Execute(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def gradient_magnitude(im3d):
    itk_im = sitk.GetImageFromArray(im3d)
    itk_im = sitk.GradientMagnitude(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def discrete_gaussian_filter(im3d, variance):
    itk_im = sitk.GetImageFromArray(im3d)
    gaussian_filter = sitk.DiscreteGaussianImageFilter()
    gaussian_filter.SetVariance(2.0)
    itk_im = gaussian_filter.Execute(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def morphological_watershed(im3d, level):
    itk_im = sitk.GetImageFromArray(im3d)
    itk_im = sitk.MorphologicalWatershed(itk_im, level=level)
    return ColorImage3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


def segment(stack):
    """Segment the stack into 3D regions representing cells."""
    stack = identity(stack)
    stack = filter_median(stack)
    stack = gradient_magnitude(stack)
    stack = discrete_gaussian_filter(stack, 2.0)
    stack = morphological_watershed(stack, 0.664)
    return stack<commit_msg>Implemented masking of root<commit_after>"""Code for segmenting 3D images."""

import SimpleITK as sitk

import numpy as np
import skimage.morphology

from jicbioimage.core.image import Image3D
from jicbioimage.core.transform import transformation
from jicbioimage.core.io import AutoWrite

from jicbioimage.transform import (
    threshold_otsu,
    remove_small_objects,
)

from utils import ColorImage3D


@transformation
def identity(im3d):
    return im3d


@transformation
def pseudo_convex_hull(im3d):
    stack = []
    ydim, xdim, zdim = im3d.shape
    for i in range(zdim):
        slice_hull = skimage.morphology.convex_hull_image(im3d[:, :, i])
        stack.append(slice_hull)
    return np.dstack(stack).view(Image3D)


@transformation
def remove_small_objects_in_plane(im3d, min_size):
    tmp_auto_write_on = AutoWrite.on
    AutoWrite.on = False
    stack = []
    ydim, xdim, zdim = im3d.shape
    for i in range(zdim):
        slice_hull = remove_small_objects(im3d[:, :, i], min_size)
        stack.append(slice_hull)

    AutoWrite.on = tmp_auto_write_on
    return np.dstack(stack).view(Image3D)


@transformation
def filter_median(im3d):
    itk_im = sitk.GetImageFromArray(im3d)
    median_filter = sitk.MedianImageFilter()
    itk_im = median_filter.Execute(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def gradient_magnitude(im3d):
    itk_im = sitk.GetImageFromArray(im3d)
    itk_im = sitk.GradientMagnitude(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def discrete_gaussian_filter(im3d, variance):
    itk_im = sitk.GetImageFromArray(im3d)
    gaussian_filter = sitk.DiscreteGaussianImageFilter()
    gaussian_filter.SetVariance(variance)
    itk_im = gaussian_filter.Execute(itk_im)
    return Image3D.from_array(sitk.GetArrayFromImage(itk_im),
                              log_in_history=False)


@transformation
def morphological_watershed(im3d, level):
    itk_im = sitk.GetImageFromArray(im3d)
    itk_im = sitk.MorphologicalWatershed(itk_im, level=level)
    return ColorImage3D.from_array(sitk.GetArrayFromImage(itk_im),
                                   log_in_history=False)


@transformation
def filter_cells_outside_mask(im3d, mask):
    inverse_mask = np.logical_not(mask)
    ids_to_filter = np.unique(im3d * inverse_mask)
    for i in ids_to_filter:
        im3d[im3d == i] = 0
    return im3d


def segment(stack):
    """Segment the stack into 3D regions representing cells."""
    mask = threshold_otsu(stack)
    mask = remove_small_objects_in_plane(mask, min_size=1000)
    mask = pseudo_convex_hull(mask)
    stack = identity(stack)
    stack = filter_median(stack)
    stack = gradient_magnitude(stack)
    stack = discrete_gaussian_filter(stack, 2.0)
    stack = morphological_watershed(stack, 0.664)
    stack = filter_cells_outside_mask(stack, mask)
    return stack
{'fraction_non_alphanumeric': 0.08441831163376733, 'fraction_numerical': 0.014279714405711887, 'mean_word_length': 3.7393034825870646, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1633492', 'n_tokens_mistral': 1689, 'n_tokens_neox': 1620, 'n_words': 285}