def __init__(self,
ds_factor=(1, 2, 2),
crop_factor=None,
crop_slice=None,
order=None,
**super_kwargs):
"""
:param ds_factor: If factor is 2, then downscaled to half-resolution
:param crop_factor: If factor is 2, the central crop of half-size is taken
:param order: downscaling order
"""
super(DownsampleAndCrop3D, self).__init__(**super_kwargs)
self.order = order
self.ds_factor = ds_factor
if crop_factor is not None:
assert isinstance(crop_factor, (tuple, list))
# assert crop_slice is None
if crop_slice is not None:
# assert crop_slice is not None
assert isinstance(crop_slice, str)
crop_slice = parse_data_slice(crop_slice)
self.crop_factor = crop_factor
self.crop_slice = crop_slice
def get_kwargs_for_each_run(self):
# TODO: add option to load GT and affs directly in run_agglomeration? (More memory efficient)
nb_iterations = self.get("postproc_config/nb_iterations", 1)
kwargs_collected = []
postproc_config = self.get("postproc_config", ensure_exists=True)
iterated_options = self.get("postproc_config/iterated_options", {})
# Assert to have a sample:
if "sample" not in iterated_options:
assert "sample" in postproc_config, "At least one sample-dataset should be given"
iterated_options["sample"] = postproc_config["sample"]
# Initialize default iterated options:
iterated_options.setdefault("noise_factor",
[postproc_config.get("noise_factor", 0.)])
iterated_options.setdefault("edge_prob",
[postproc_config.get("edge_prob", 0.)])
iterated_options.setdefault("preset",
[postproc_config.get("preset", None)])
iterated_options.setdefault(
"local_attraction",
[postproc_config.get("local_attraction", False)])
iterated_options.setdefault(
"crop_slice", [postproc_config.get("crop_slice", ":,:,:,:")])
iterated_options.setdefault(
"sub_crop_slice", postproc_config.get("sub_crop_slice",
[":,:,:,:"]))
# Make sure to have lists:
for iter_key in iterated_options:
if isinstance(iterated_options[iter_key], dict):
for dict_key in iterated_options[iter_key]:
iterated_options[iter_key][dict_key] = iterated_options[iter_key][dict_key] \
if isinstance(iterated_options[iter_key][dict_key], list) \
else [iterated_options[iter_key][dict_key]]
else:
iterated_options[iter_key] = iterated_options[iter_key] if isinstance(iterated_options[iter_key], list) \
else [iterated_options[iter_key]]
for _ in range(nb_iterations):
collected_data = {"affs": {}, "GT": {}, "affs_mask": {}}
for sample in iterated_options['sample']:
print("Loading...")
# Create new dict entry if needed:
for dt_type in collected_data:
collected_data[dt_type][sample] = {} if sample not in collected_data[dt_type] else \
collected_data[dt_type][sample]
# Check if we have a dictionary with single values for each sample:
all_crops = iterated_options['crop_slice'][sample] if isinstance(iterated_options['crop_slice'], dict) \
else iterated_options['crop_slice']
all_subcrops = iterated_options['sub_crop_slice'][sample] if isinstance(iterated_options['sub_crop_slice'], dict) \
else iterated_options['sub_crop_slice']
# ----------------------------------------------------------------------
# Load data (and possibly add noise or select long range edges):
# ----------------------------------------------------------------------
for crop in all_crops:
# Create new dict entry if needed:
for dt_type in collected_data:
collected_data[dt_type][sample][crop] = {} if crop not in collected_data[dt_type][sample] else \
collected_data[dt_type][sample][crop]
for sub_crop in all_subcrops:
# Create new dict entry if needed:
for dt_type in collected_data:
collected_data[dt_type][sample][crop][sub_crop] = {} \
if sub_crop not in collected_data[dt_type][sample][crop] else \
collected_data[dt_type][sample][crop][sub_crop]
noise_seed = np.random.randint(-100000, 100000)
GT_vol_config = deepcopy(self.get('volume_config/GT'))
affs_vol_config = deepcopy(
self.get('volume_config/affinities'))
# FIXME: if I pass multiple crops, they get ignored and I get an error below when I create the runs...
if "crop_slice" in GT_vol_config:
gt_crop_slc = GT_vol_config.pop("crop_slice")
else:
gt_crop_slc = segm_utils.parse_data_slice(crop)[1:]
GT = segm_utils.readHDF5_from_volume_config(
sample,
crop_slice=gt_crop_slc,
run_connected_components=False**GT_vol_config,
)
# print(GT.shape)
# print("NUmber of clusters before: ", np.unique(GT).shape)
# GT = vigra.analysis.labelVolumeWithBackground(GT.astype('uint32'))
# print("NUmber of clusters after: ", np.unique(GT).shape)
# raise ValueError
if self.get("volume_config/ignore_glia", False):
print("Ignoring glia during evaluation")
inner_path_glia = self.get(
"volume_config/glia_specs/inner_path",
ensure_exists=True)
glia_label = self.get(
"volume_config/glia_specs/glia_label",
ensure_exists=True)
ignore_label = self.get(
"volume_config/glia_specs/ignore_label",
ensure_exists=True)
GT_vol_config['inner_path'] = inner_path_glia
various_masks = segm_utils.readHDF5_from_volume_config(
sample,
crop_slice=gt_crop_slc,
run_connected_components=False**GT_vol_config,
)
GT[various_masks == glia_label] = ignore_label
# Optionally, affinity paths are deduced dynamically:
if self.get("affinities_dir_path") is not None:
affs_vol_config['path'] = \
os.path.join(self.get("affinities_dir_path"), "predictions_sample_{}.h5".format(sample))
if "crop_slice" in affs_vol_config:
affs_crop_slc = affs_vol_config.pop("crop_slice")
else:
affs_crop_slc = segm_utils.parse_data_slice(crop)
affinities = segm_utils.readHDF5_from_volume_config(
sample,
crop_slice=affs_crop_slc,
run_connected_components=False**affs_vol_config,
)
assert GT.shape == affinities.shape[
1:], "Loaded GT and affinities do not match: {} - {}".format(
GT.shape, affinities.shape[1:])
sub_crop_slc = segm_utils.parse_data_slice(sub_crop)
affinities = affinities[sub_crop_slc]
GT = GT[sub_crop_slc[1:]]
GT = vigra.analysis.labelVolumeWithBackground(
GT.astype('uint32'))
collected_data["GT"][sample][crop][sub_crop] = GT
collected_data["affs"][sample][crop][sub_crop] = {}
collected_data["affs_mask"][sample][crop][
sub_crop] = {}
for long_range_prob in iterated_options['edge_prob']:
for noise in iterated_options['noise_factor']:
if noise != 0.:
raise DeprecationWarning(
"Opensimplex noise no longer in nifty")
noise_mod = postproc_config.get(
"noise_mod", 'split-biased')
collected_data["affs"][sample][crop][sub_crop][noise] = \
postproc_utils.add_opensimplex_noise_to_affs(
affinities, noise,
mod=noise_mod,
target_affs='all',
seed=noise_seed
)
else:
collected_data["affs"][sample][crop][
sub_crop][noise] = affinities
# Fix already long-range edges that will be in the graph:
if long_range_prob < 1.0 and long_range_prob > 0.0:
collected_data["affs_mask"][sample][crop][
sub_crop][
long_range_prob] = np.random.random(
affinities.shape
) < long_range_prob
# ----------------------------------------------------------------------
# Create iterators:
# ----------------------------------------------------------------------
print("Creating pool instances...")
for crop in all_crops:
for sub_crop in all_subcrops:
assert collected_data["affs"][sample][crop][
sub_crop] is not None
for local_attr in iterated_options['local_attraction']:
for preset in iterated_options['preset']:
if local_attr and preset in [
'greedyFixation', 'GAEC'
]:
continue
for edge_prob in iterated_options['edge_prob']:
for noise in iterated_options[
'noise_factor']:
# Build a new run dictionary with the iterated options:
kwargs_collected.append({
'affinities':
collected_data['affs'][sample]
[crop][sub_crop][noise],
'GT':
collected_data['GT'][sample][crop]
[sub_crop],
'sample':
sample,
'noise_factor':
noise,
'edge_prob':
edge_prob,
'preset':
preset,
'local_attraction':
local_attr,
'crop_slice':
crop,
'sub_crop_slice':
sub_crop,
'mask_used_edges':
collected_data['affs_mask'][sample]
[crop][sub_crop].get(
edge_prob, None)
})
return kwargs_collected
self.merge_rule.merge_edges(aliveEdge, deadEdge)
if self.debug:
print("Merge edges", aliveEdge, deadEdge)
def mergeNodes(self, aliveNode, deadNode):
if self.debug:
print("Merge nodes", aliveNode, deadNode)
def contractEdgeDone(self, contractedEdge):
if self.debug:
print("############")
for sample in ["A+", "B+", "C+"]:
# for sample in ["C"]:
crop_slice = parse_data_slice(
config["volume_config"]["GT"]["crop_slice"][sample])
# Load glia and GT:
glia_prediction_path = os.path.join(
project_dir, glia_mask_exp, "predictions_sample_{}.h5".format(sample))
print("Loading glia for sample ", sample)
glia_mask = readHDF5(glia_prediction_path, "glia_mask")[0]
glia_mask = glia_mask[crop_slice]
invalid_glia_mask = np.logical_or(glia_mask < 0., glia_mask > 1.)
glia_mask[invalid_glia_mask] = 0
print("Average glia mask (should be close to zero): ", glia_mask.mean())
# This is already cropped
gt_segm = readHDF5_from_volume_config(sample,
**config["volume_config"]["GT"])
import segmfriends.vis as vis
from matplotlib import pyplot as plt
sample = "C"
# slices = {'A': ":,:,:,:", 'B': ":, :, 90:, 580: 1900", 'C': ":, :, 70:1450, 95:1425"}
slices = {
'A': ":,:,:,:",
'B': ":,:,:,:",
'C': ":,70:-6,50:-50,50:-50",
"0": ":,:,:,:",
"1": ":,:,:,:",
"2": ":,:,:,:",
"3": ":,:,:,:",
}
parsed_slice = parse_data_slice(slices[sample])
data_path = os.path.join(get_abailoni_hci_home_path(),
"datasets/new_cremi/sample{}.h5".format(sample))
# data_path = os.path.join(get_trendytukan_drive_path(), "datasets/new_cremi/fib25/sample{}.h5".format(sample))
with h5py.File(data_path, 'r') as f:
print([atr for atr in f['volumes/labels']])
# glia = f['volumes/labels/glia'][:]
raw = f['volumes/raw'][parsed_slice[1:]]
GT = f['volumes/labels/neuron_ids_fixed'][parsed_slice[1:]]
# Load affs:
from segmfriends.utils.various import writeHDF5, readHDF5
affs_path = os.path.join(
get_trendytukan_drive_path(),
"projects/pixel_embeddings/{}/predictions_sample_{}.h5".format(
"v4_addSparseAffs_eff", sample))
def forward(self, predictions, all_targets):
predictions = [predictions] if not isinstance(predictions, (list, tuple)) else predictions
all_targets = [all_targets] if not isinstance(all_targets, (list, tuple)) else all_targets
loss = 0
# # ----------------------------
# # Predict glia mask:
# # ----------------------------
if self.train_glia_mask:
assert not self.target_has_various_masks, "To be implemented"
frg_kwargs = self.model.models[-1].foreground_prediction_kwargs
if frg_kwargs is None:
# Legacy:
nb_glia_preds = 1
nb_glia_targets = [0]
else:
nb_glia_preds = len(frg_kwargs)
nb_glia_targets = [frg_kwargs[dpth]["nb_target"] for dpth in frg_kwargs]
all_glia_preds = predictions[-nb_glia_preds:]
predictions = predictions[:-nb_glia_preds]
loss_glia = 0
for counter, glia_pred, nb_tar in zip(range(len(all_glia_preds)), all_glia_preds, nb_glia_targets):
glia_target = all_targets[nb_tar][:,[-1]]
all_targets[nb_tar] = all_targets[nb_tar][:, :-1]
assert self.target_has_label_segm
gt_segm = all_targets[nb_tar][:,[0]]
glia_target = auto_crop_tensor_to_shape(glia_target, glia_pred.shape)
gt_segm = auto_crop_tensor_to_shape(gt_segm, glia_pred.shape)
# TODO: generalize ignore label:
valid_mask = (gt_segm != 0).float()
glia_pred = glia_pred * valid_mask
glia_target = glia_target * valid_mask
with warnings.catch_warnings(record=True) as w:
loss_glia_new = data_parallel(self.loss, (glia_pred, glia_target), self.devices).mean()
loss_glia = loss_glia + loss_glia_new
log_image("glia_target_d{}".format(counter), glia_target)
log_image("glia_pred_d{}".format(counter), glia_pred)
loss = loss + loss_glia
log_scalar("loss_glia", loss_glia)
for counter, nb_pred in enumerate(self.predictions_specs):
assert len(predictions) > nb_pred
pred = predictions[nb_pred]
# TODO: add precrop_pred?
# if self.precrop_pred is not None:
# from segmfriends.utils.various import parse_data_slice
# crop_slc = (slice(None), slice(None)) + parse_data_slice(self.precrop_pred)
# predictions = predictions[crop_slc]
pred_specs = self.predictions_specs[nb_pred]
target = all_targets[pred_specs.get("target", 0)]
target_dws_fact = pred_specs.get("target_dws_fact", None)
if target_dws_fact is not None:
assert isinstance(target_dws_fact, list) and len(target_dws_fact) == 3
target = target[(slice(None), slice(None)) + tuple(slice(None,None,dws) for dws in target_dws_fact)]
target = auto_crop_tensor_to_shape(target, pred.shape,
ignore_channel_and_batch_dims=True)
if self.target_has_label_segm:
if self.target_has_various_masks:
target = target[:, 2:]
else:
target = target[:,1:]
assert target.shape[1] % 2 == 0, "Target should include both affinities and masks"
# Get ignore-mask and affinities:
nb_channels = int(target.shape[1] / 2)
affs_channels = pred_specs.get("affs_channels", None)
if affs_channels is not None:
if isinstance(affs_channels, str):
affs_slice = parse_data_slice(affs_channels)[0]
elif isinstance(affs_channels, list):
# TODO: make as a tuple???
affs_slice = affs_channels
else:
raise ValueError("The passed affinities channels are not compatible")
else:
affs_slice = slice(None)
gt_affs = target[:,:nb_channels][:, affs_slice]
assert gt_affs.shape[1] == pred.shape[1], "Prediction has a wrong number of offset channels"
valid_pixels = target[:,nb_channels:][:, affs_slice]
# Invert affinities for Dice loss: (1 boundary, 0 otherwise)
gt_affs = 1. - gt_affs
pred = pred*valid_pixels
gt_affs = gt_affs*valid_pixels
with warnings.catch_warnings(record=True) as w:
loss_new = data_parallel(self.loss, (pred, gt_affs), self.devices).mean()
loss = loss + loss_new
log_scalar("loss_sparse_d{}".format(counter), loss_new)
# TODO: use Callback from Roman to run it every N iterations
gc.collect()
return loss
import matplotlib
matplotlib.rcParams.update({'font.size': 25})
sample = "C"
# slices = {'A': ":,:,:,:", 'B': ":, :, 90:, 580: 1900", 'C': ":, :, 70:1450, 95:1425"}
slices = {
'A': ":,:,:,:",
'B': ":,:,:,:",
'C': ":,70:-6,50:-50,50:-50",
"0": ":,:,:,:",
"1": ":,:,:,:",
"2": ":,:,:,:",
"3": ":,:,:,:",
}
parsed_slice = parse_data_slice(slices[sample])
conf_folder_path = os.path.join(
get_trendytukan_drive_path(),
"projects/pixel_embeddings/v4_addSparseAffs_avgDirectVar/scores")
prefix_math = [
"C__MEAN___",
# "C__MutexWatershed___",
"C__MWS__stride10___"
]
labels = {
"C__MEAN___": "Gasp Average",
"C__MWS__stride10___": "Mutex Watershed"
}
affs[affs_valid_mask==0] = 1
# Combine left and right affinities:
segment_mask = np.logical_and(affs[0], affs[1])
# This functions erode binary out_mask (segments 1, boundary 0)
eroded_segment_mask = segment_mask.copy()
for z in range(eroded_segment_mask.shape[0]):
eroded_segment_mask[z] = vigra.filters.multiBinaryErosion(segment_mask[z], radius=2.)
boundary_mask = np.logical_not(eroded_segment_mask)
BOUNDARY_LABEL = 2
DEFECTED_LABEL = 3
out_mask = glia.copy()
out_mask[boundary_mask] = BOUNDARY_LABEL
# Mask defected slices:
for slc in defected_slices[sample]:
out_mask[parse_data_slice(slc)] = DEFECTED_LABEL
# Copy GT from previous (to avoid weird connected components problems):
for z in copy_from_previous[sample]:
GT[z] = GT[z-1]
print("Now writing...")
from segmfriends.utils.various import writeHDF5
writeHDF5(out_mask, data_path, 'volumes/labels/various_masks')
writeHDF5(GT, data_path, 'volumes/labels/neuron_ids_fixed')
active_nodes[tuple(self.starting_coordinate)] = 2
return out_segm, active_nodes
import os
from vaeAffs.utils.path_utils import get_abailoni_hci_home_path
import h5py
sample = "A"
data_path = os.path.join(
get_abailoni_hci_home_path(),
"../ialgpu1_local_home/datasets/cremi/SOA_affinities/sample{}_train.h5".
format(sample))
crop_slice = ":,20:21,300:400,300:400"
from segmfriends.utils.various import parse_data_slice
crop_slice = parse_data_slice(crop_slice)
with h5py.File(data_path, 'r') as f:
affs = f['predictions']['full_affs'][crop_slice]
raw = f['raw'][crop_slice[1:]]
offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 0, 0], [0, -3, 0],
[0, 0, -3], [-3, 0, 0], [0, -9, 0], [0, 0, -9], [-4, 0, 0],
[0, -27, 0], [0, 0, -27]]
# Fake duplicate affinities:
duplicate_affs = np.empty_like(affs)
for i, off in enumerate(offsets):
duplicate_affs[i] = np.roll(affs[i], off)
affs = np.concatenate([affs, duplicate_affs], axis=0)
offsets = offsets + [[-off[0], -off[1], -off[2]] for off in offsets]
def forward(self, all_predictions, target):
mdl_kwargs = self.model_kwargs
ptch_kwargs = mdl_kwargs["patchNet_kwargs"]
nb_inputs = mdl_kwargs.get("number_multiscale_inputs")
# print([(pred.shape[-3], pred.shape[-2], pred.shape[-1]) for pred in all_predictions])
# print([(targ.shape[-3], targ.shape[-2], targ.shape[-1]) for targ in target])
# Plot some patches with the raw:
if self.model.return_input:
raw_inputs = all_predictions[-nb_inputs:]
all_predictions = all_predictions[:-nb_inputs]
loss = 0
# # ----------------------------
# # Predict glia mask:
# # ----------------------------
if self.train_glia_mask:
assert self.glia_label is not None
frg_kwargs = self.model.foreground_prediction_kwargs
if frg_kwargs is None:
# Legacy:
nb_glia_preds = 1
nb_glia_targets = [0]
else:
nb_glia_preds = len(frg_kwargs)
nb_glia_targets = [
frg_kwargs[dpth]["nb_target"] for dpth in frg_kwargs
]
all_glia_preds = all_predictions[-nb_glia_preds:]
all_predictions = all_predictions[:-nb_glia_preds]
loss_glia = 0
for counter, glia_pred, nb_tar in zip(range(len(all_glia_preds)),
all_glia_preds,
nb_glia_targets):
glia_target = (target[nb_tar][:,
[1]] == self.glia_label).float()
valid_mask = (target[nb_tar][:, [0]] !=
self.ignore_label).float()
glia_target = auto_crop_tensor_to_shape(
glia_target, glia_pred.shape)
valid_mask = auto_crop_tensor_to_shape(valid_mask,
glia_pred.shape)
glia_pred = glia_pred * valid_mask
glia_target = glia_target * valid_mask
with warnings.catch_warnings(record=True) as w:
loss_glia_cur = data_parallel(self.loss,
(glia_pred, glia_target),
self.devices).mean()
loss_glia = loss_glia + loss_glia_cur
log_image("glia_target_d{}".format(counter), glia_target)
log_image("glia_pred_d{}".format(counter), glia_pred)
loss = loss + loss_glia
log_scalar("loss_glia", loss_glia)
else:
glia_pred = all_predictions.pop(-1)
if self.train_sparse_loss:
loss = loss + self.sparse_multilevelDiceLoss(
all_predictions, target)
# Delete affinities from targets:
target = [tar[:, :2].int() for tar in target]
# IoU loss:
if self.add_IoU_loss:
assert self.boundary_label is None, "Not implemented"
assert self.indx_trained_patchNets is None
loss = loss + self.IoU_loss(all_predictions, target)
if self.indx_trained_patchNets is None:
nb_preds = len(all_predictions)
assert len(ptch_kwargs) == nb_preds
indx_trained_patchNets = zip(range(nb_preds), range(nb_preds))
else:
indx_trained_patchNets = self.indx_trained_patchNets
# ----------------------------
# Loss on patches:
# ----------------------------
for nb_patch_net, nb_pr in indx_trained_patchNets:
# ----------------------------
# Initializations:
# ----------------------------
pred = all_predictions[nb_pr]
kwargs = ptch_kwargs[nb_patch_net]
if isinstance(target, (list, tuple)):
assert "nb_target" in kwargs, "Multiple targets passed. Target should be specified"
gt_segm = target[kwargs["nb_target"]]
else:
gt_segm = target
# Collect options from config:
patch_shape_input = kwargs.get("patch_size")
assert all(i % 2 == 1
for i in patch_shape_input), "Patch should be odd"
patch_dws_fact = kwargs.get("patch_dws_fact", [1, 1, 1])
stride = tuple(kwargs.get("patch_stride", [1, 1, 1]))
pred_dws_fact = kwargs.get("pred_dws_fact", [1, 1, 1])
# print(nb_patch_net, patch_dws_fact, pred_dws_fact)
precrop_pred = kwargs.get("precrop_pred", None)
limit_nb_patches = kwargs.get("limit_nb_patches", None)
from segmfriends.utils.various import parse_data_slice
if precrop_pred is not None:
precrop_pred_slice = (
slice(None), slice(None)) + parse_data_slice(precrop_pred)
pred = pred[precrop_pred_slice]
central_shape = tuple(kwargs.get("central_shape", [1, 3, 3]))
max_random_crop = tuple(kwargs.get("max_random_crop", [0, 5, 5]))
if self.fix_bug_multiscale_patches:
real_patch_shape = tuple(
pt * fc - fc + 1
for pt, fc in zip(patch_shape_input, patch_dws_fact))
else:
real_patch_shape = tuple(
pt * fc
for pt, fc in zip(patch_shape_input, patch_dws_fact))
full_target_shape = gt_segm.shape[-3:]
assert all([
i <= j for i, j in zip(real_patch_shape, full_target_shape)
]), "Real-sized patch is too large!"
# ----------------------------
# Deduce crop size of the prediction and select target patches accordingly:
# ----------------------------
# print(pred.shape, full_target_shape, pred_dws_fact, real_patch_shape)
crop_slice_targets, crop_slice_prediction = get_slicing_crops(
pred.shape[2:], full_target_shape, pred_dws_fact,
real_patch_shape)
# print(crop_slice_prediction, crop_slice_targets, nb_patch_net)
gt_segm = gt_segm[crop_slice_targets]
pred = pred[crop_slice_prediction]
full_target_shape = gt_segm.shape[-3:]
# # ----------------------------
# # Plot some random patches with associated raw patch:
# # ----------------------------
if self.model.return_input and nb_patch_net < 5:
# raw = raw_inputs[kwargs["nb_target"]][crop_slice_targets]
# FIXME: raw is not correct for deeper ones
raw = raw_inputs[0][crop_slice_targets]
raw_to_plot, gt_labels_to_plot, gt_masks_to_plot, pred_emb_to_plot = [], [], [], []
for n in range(40):
# Select a random pixel and define sliding-window crop slices:
selected_coord = [
np.random.randint(shp) for shp in pred.shape[2:]
]
# selected_coord[0] = 4 # For plots, get always 4
full_patch_slice = (slice(None), slice(0, 1)) + tuple(
slice(selected_coord[i], selected_coord[i] +
real_patch_shape[i])
for i in range(len(selected_coord)))
emb_slice = (slice(None), slice(0, 1)) + tuple(
slice(
selected_coord[i] +
int(real_patch_shape[i] / 2), selected_coord[i] +
int(real_patch_shape[i] / 2) + 1)
for i in range(len(selected_coord)))
pred_center_coord = [
int(selected_coord[i] / pred_dws_fact[i])
for i in range(len(selected_coord))
]
emb_slice_pred = (slice(None), slice(None)) + tuple(
slice(pred_center_coord[i], pred_center_coord[i] + 1)
for i in range(len(selected_coord)))
# Collect data for current sliding window:
center_label = gt_segm[emb_slice]
center_label_repeated = center_label.repeat(
1, 1, *real_patch_shape)
gt_patch_labels = gt_segm[full_patch_slice]
gt_masks_to_plot.append(
gt_patch_labels != center_label_repeated)
gt_labels_to_plot.append(gt_patch_labels)
# ignore_mask_patch = (gt_patch_labels == 0)
pred_emb_to_plot.append(pred[emb_slice_pred])
raw_to_plot.append(raw[full_patch_slice])
# Highlight center pixel:
raw_to_plot = torch.cat(raw_to_plot, dim=0)
center_pixel_coord = (slice(None), 0) + tuple(
int(shp / 2) for shp in real_patch_shape)
raw_to_plot[center_pixel_coord] = raw_to_plot.min() - 1.
gt_labels_to_plot = torch.cat(gt_labels_to_plot, dim=0)
gt_masks_to_plot = torch.cat(gt_masks_to_plot, dim=0)
pred_emb_to_plot = torch.cat(pred_emb_to_plot, dim=0)
# Decode embeddings:
ptch_num = kwargs["patchNet_number"]
pred_patch_to_plot = data_parallel(
self.model.patch_models[ptch_num],
pred_emb_to_plot[:, :, 0, 0, 0], self.devices)
# Downscale and rescale targets:
down_sc_slice = (slice(None), slice(None)) + tuple(
slice(int(dws_fact / 2), None, dws_fact)
for dws_fact in patch_dws_fact)
gt_masks_to_plot = torch.nn.functional.interpolate(
gt_masks_to_plot[down_sc_slice].float(),
scale_factor=tuple(patch_dws_fact))
pred_patch_to_plot = torch.nn.functional.interpolate(
pred_patch_to_plot, scale_factor=tuple(patch_dws_fact))
gt_masks_to_plot = 1. - gt_masks_to_plot
if patch_dws_fact[1] <= 6:
pred_patch_to_plot = 1. - pred_patch_to_plot
log_image("raw_patch_l{}".format(nb_patch_net), raw_to_plot)
log_image("gt_label_patch_l{}".format(nb_patch_net),
gt_labels_to_plot)
log_image("gt_mask_patch_l{}".format(nb_patch_net),
gt_masks_to_plot)
log_image("pred_patch_l{}".format(nb_patch_net),
pred_patch_to_plot)
# # ----------------------------
# # Patch-Loss:
# # ----------------------------
if kwargs.get("skip_standard_patch_loss", False):
continue
# If multiple strides were given, process all of them:
all_strides = stride if isinstance(stride[0], list) else [stride]
if limit_nb_patches is not None:
all_limit_nb_patches = limit_nb_patches if isinstance(
limit_nb_patches[0], list) else [limit_nb_patches]
else:
all_limit_nb_patches = [None for _ in all_strides]
for nb_stride, stride, limit_nb_patches in zip(
range(len(all_strides)), all_strides,
all_limit_nb_patches):
# ----------------------------
# Get some random prediction embeddings:
# ----------------------------
pred_strides = get_prediction_strides(pred_dws_fact, stride)
pred_patches, crop_slice_pred, nb_patches = extract_patches_torch_new(
pred, (1, 1, 1),
stride=pred_strides,
max_random_crop=max_random_crop)
# Try to get some raw patches:
# TODO: the factor is simply the level in the UNet
# get_slicing_crops(pred.shape[2:], full_target_shape, [1,1,1], real_patch_shape)
# ----------------------------
# Collect gt_segm patches and corresponding center labels:
# ----------------------------
crop_slice_targets = tuple(
slice(sl.start, None) for sl in crop_slice_pred)
gt_patches, _, _ = extract_patches_torch_new(
gt_segm,
real_patch_shape,
stride=stride,
crop_slice=crop_slice_targets,
limit_patches_to=nb_patches)
gt_patches = gt_patches[:, [0]]
# Make sure to crop some additional border and get the centers correctly:
# TODO: this can be now easily done by cropping the gt_patches...
crop_slice_center_labels = (slice(None), slice(None)) + tuple(
slice(slc.start + int(sh / 2), slc.stop) for slc, sh in
zip(crop_slice_targets[2:], real_patch_shape))
target_at_patch_center, _, _ = extract_patches_torch_new(
gt_segm, (1, 1, 1),
stride=stride,
crop_slice=crop_slice_center_labels,
limit_patches_to=nb_patches)
# Get GT and other masks separately:
label_at_patch_center = target_at_patch_center[:, [0]]
mask_at_patch_center = target_at_patch_center[:, [1]]
# ----------------------------
# Ignore patches on the boundary or involving ignore-label:
# ----------------------------
# Ignore pixels involving ignore-labels:
ignore_masks = (gt_patches == self.ignore_label)
valid_patches = (label_at_patch_center != self.ignore_label)
assert self.boundary_label is not None, "Old boundary method is deprecated"
# # Exclude a patch from training if the central region contains more than one gt label
# # (i.e. it is really close to a boundary):
# central_crop = (slice(None), slice(None)) + convert_central_shape_to_crop_slice(gt_patches.shape[-3:], central_shape)
# mean_central_crop_labels = gt_patches[central_crop].mean(dim=-1, keepdim=True) \
# .mean(dim=-2, keepdim=True) \
# .mean(dim=-3, keepdim=True)
#
# valid_patches = valid_patches & (mean_central_crop_labels == center_labels)
# is_on_boundary_mask = None
patch_is_on_boundary = (
mask_at_patch_center == self.boundary_label).repeat(
1, 1, *real_patch_shape)
# Ignore patches that represent a glia:
if not self.train_patches_on_glia:
assert self.glia_label is not None
# print("Glia: ", (mask_at_patch_center != self.glia_label).min())
valid_patches = valid_patches & (mask_at_patch_center !=
self.glia_label)
# Delete redundant patches from batch:
valid_batch_indices = np.argwhere(
valid_patches[:, 0, 0, 0, 0].cpu().detach().numpy())[:, 0]
if limit_nb_patches is not None:
limit = limit_nb_patches[0]
if limit_nb_patches[1] == 'number':
if valid_batch_indices.shape[0] > limit:
valid_batch_indices = np.random.choice(
valid_batch_indices, limit, replace=False)
elif limit_nb_patches[1] == 'factor':
assert limit <= 1. and limit >= 0.
valid_batch_indices = np.random.choice(
valid_batch_indices,
int(limit * valid_batch_indices.shape[0]),
replace=False)
if valid_batch_indices.shape[0] == 0:
print(
"ZERO valid patches at level {}!".format(nb_patch_net))
# Avoid problems if all patches are invalid and torch complains that autograd cannot be performed:
loss += pred_patches.sum() * 0.
continue
# ----------------------------
# Compute the actual (inverted) MeMasks targets: (0 is me, 1 are the others)
# best targets for Dice loss (usually more me than others)
# ----------------------------
center_labels_repeated = label_at_patch_center.repeat(
1, 1, *real_patch_shape)
me_masks = gt_patches != center_labels_repeated
if patch_is_on_boundary is not None:
# If on boundary, we make (inverted) me_masks completely 1 (split from everything)
me_masks = me_masks | patch_is_on_boundary
# Downscale MeMasks using MaxPooling (preserve narrow processes):
# moreover, during the maxPool, better shrink me mask than expanding (avoid merge predictions)
if all(fctr == 1 for fctr in patch_dws_fact):
maxpool = Identity()
else:
maxpool = nn.MaxPool3d(kernel_size=patch_dws_fact,
stride=patch_dws_fact,
padding=0)
# Downscaling patch:
down_sc_slice = (slice(None), slice(None)) + tuple(
slice(int(dws_fact / 2), None, dws_fact)
for dws_fact in patch_dws_fact)
# Final targets:
patch_targets = me_masks[valid_batch_indices].float(
)[down_sc_slice]
patch_ignore_masks = ignore_masks[valid_batch_indices][
down_sc_slice].byte()
# Invert MeMasks:
# best targets for Dice loss are: meMask == 0; others == 1
# FIXME: generalize
if patch_dws_fact[1] > 6:
patch_targets = 1. - patch_targets
assert valid_batch_indices.max() < pred_patches.shape[
0], "Something went wrong, more target patches were collected than predicted: {} targets vs {} pred...".format(
valid_batch_indices.max(), pred_patches.shape[0])
pred_embed = pred_patches[valid_batch_indices]
pred_embed = pred_embed[:, :, 0, 0, 0]
# ----------------------------
# Expand embeddings to patches using PatchNet models:
# ----------------------------
if "model_number" in kwargs:
# FIXME: update this crap
# In this case we are training a stacked model:
mdl_num = kwargs["model_number"]
ptch_num = kwargs["patchNet_number"]
expanded_patches = data_parallel(
self.model.models[mdl_num].patch_models[ptch_num],
pred_embed, self.devices)
else:
expanded_patches = data_parallel(
self.model.patch_models[nb_patch_net], pred_embed,
self.devices)
# print(expanded_patches.shape)
assert expanded_patches.shape[
1] == 1, "PatchNet should output only a one-channel mask!"
# Some logs:
if nb_stride == 0:
log_image("ptc_trg_l{}".format(nb_patch_net),
patch_targets)
log_image("ptc_pred_l{}".format(nb_patch_net),
expanded_patches)
# log_image("ptc_ign_l{}".format(nb_patch_net), patch_ignore_masks)
log_scalar("avg_targets_l{}".format(nb_patch_net),
patch_targets.float().mean())
# Train only checkerboard pattern:
if self.apply_checkerboard:
checkerboard = np.zeros(patch_shape_input)
# Verticals:
center_coord = [int(sh / 2) for sh in patch_shape_input]
checkerboard[:, center_coord[1], :] = 1
checkerboard[:, :, center_coord[2]] = 1
# Two diagonals:
indices = np.indices(patch_shape_input)
checkerboard[indices[1] == indices[2]] = 1
checkerboard[indices[1] == (patch_shape_input[2] -
indices[2] - 1)] = 1
# Reduce z-context:
z_mask = np.zeros_like(checkerboard)
z_mask[center_coord[0]] = 1
for z in range(patch_shape_input[0]):
offs = abs(center_coord[0] - z)
if offs != 0:
z_mask[z, offs:-offs, offs:-offs] = 1
checkerboard[np.logical_not(z_mask)] = 0
# Expand channels and wrap:
checkerboard = torch.from_numpy(checkerboard).cuda(
patch_ignore_masks.get_device()).float()
checkerboard = checkerboard.unsqueeze(0).unsqueeze(0)
checkerboard = checkerboard.repeat(
*patch_ignore_masks.shape[:2], 1, 1, 1)
# ----------------------------
# Apply ignore mask and compute loss:
# ----------------------------
patch_valid_masks = 1. - patch_ignore_masks.float()
if self.apply_checkerboard:
patch_valid_masks = patch_valid_masks * checkerboard
expanded_patches = expanded_patches * patch_valid_masks
patch_targets = patch_targets * patch_valid_masks
with warnings.catch_warnings(record=True) as w:
loss_unet = data_parallel(
self.loss, (expanded_patches, patch_targets.float()),
self.devices).mean()
loss = loss + loss_unet
if nb_stride == 0:
log_scalar("loss_l{}".format(nb_patch_net), loss_unet)
log_scalar("nb_patches_l{}".format(nb_patch_net),
expanded_patches.shape[0])
# print("Loss done, memory {}", torch.cuda.memory_allocated(0)/1000000)
# TODO: use Callback from Roman to run it every N iterations
gc.collect()
return loss
from segmfriends.utils.various import parse_data_slice
import os
import h5py
from scipy.ndimage import zoom
crp_slices = {"B": ":,:-1,:", "C": ":,:,:-1"}
for sample in ["A", "B", "C", "0", "1", "2"]:
print("Sample", sample)
data_path = os.path.join(get_abailoni_hci_home_path(),
"datasets/new_cremi/sample{}.h5".format(sample))
if sample in crp_slices:
crp_slc = parse_data_slice(crp_slices[sample])
else:
crp_slc = slice(None)
with h5py.File(data_path, 'r+') as f:
print([atr for atr in f['volumes/labels']])
# glia = f['volumes/labels/glia'][:]
# raw = f['volumes/raw'][crp_slc]
# GT = f['volumes/labels/neuron_ids_fixed'][crp_slc]
# various_masks = f['volumes/labels/various_masks'][crp_slc]
various_masks = f['volumes/labels/various_masks_noDefects'][crp_slc]
print("Now writing...")
# f['volumes/raw_2x'] = zoom(raw, (1, 0.5, 0.5), order=3)
# f['volumes/labels/neuron_ids_fixed_2x'] = zoom(GT, (1, 0.5, 0.5), order=0)
if 'volumes/labels/various_masks_noDefects_2x' in f:
def forward(self, all_predictions, target):
mdl = self.model
nb_inputs = mdl.number_multiscale_inputs
# Plot some patches with the raw:
if self.model.return_input:
raw_inputs = all_predictions[-nb_inputs:]
all_predictions = all_predictions[:-nb_inputs]
loss = 0
if self.train_sparse_loss:
raise NotImplementedError
loss = loss + self.sparse_multilevelDiceLoss(all_predictions, target)
# Delete affinities from targets:
target = [tar[:, :2].int() for tar in target]
# ----------------------------
# Loss on patches:
# ----------------------------
for mask_dec_indx in range(len(all_predictions)):
# ----------------------------
# Initializations:
# ----------------------------
mask_dec = self.model.mask_decoders[mask_dec_indx]
pred = all_predictions[mask_dec_indx]
gt_segm = target[mask_dec.target_index]
# Collect options from config:
mask_shape = mask_dec.mask_shape
mask_dws_fact = mask_dec.mask_dws_fact
sample_strides = mask_dec.sample_strides
pred_dws_fact = mask_dec.pred_dws_fact
crop_slice_prediction = mask_dec.crop_slice_prediction
limit_nb_decoded_masks_to = mask_dec.limit_nb_decoded_masks_to
if crop_slice_prediction is not None:
precrop_pred_slice = (slice(None), slice(None)) + parse_data_slice(crop_slice_prediction)
pred = pred[precrop_pred_slice]
max_random_crop = mask_dec.max_random_crop
real_shape_mask = tuple(pt * fc for pt, fc in zip(mask_shape, mask_dws_fact))
full_target_shape = gt_segm.shape[-3:]
assert all([i <= j for i, j in zip(real_shape_mask, full_target_shape)]), "Real-sized patch is too large!"
# ----------------------------
# Deduce crop size of the prediction and select target patches accordingly:
# ----------------------------
# TODO: explain better what is going on here
crop_slice_targets, crop_slice_prediction = get_slicing_crops(pred.shape[2:], full_target_shape,
pred_dws_fact, real_shape_mask)
gt_segm = gt_segm[crop_slice_targets]
pred = pred[crop_slice_prediction]
full_target_shape = gt_segm.shape[-3:]
# # # ----------------------------
# # # Plot some random patches with associated raw patch:
# # # ----------------------------
# if self.model.return_input and mask_dec_indx<5:
# # raw = raw_inputs[kwargs["nb_target"]][crop_slice_targets]
# # FIXME: raw is not correct for deeper ones
# raw = raw_inputs[0][crop_slice_targets]
# raw_to_plot, gt_labels_to_plot, gt_masks_to_plot, pred_emb_to_plot = [], [], [], []
# for n in range(40):
# # Select a random pixel and define sliding-window crop slices:
# selected_coord = [np.random.randint(shp) for shp in pred.shape[2:]]
# # selected_coord[0] = 4 # For plots, get always 4
# full_patch_slice = (slice(None), slice(0,1)) + tuple(
# slice(selected_coord[i], selected_coord[i] + real_shape_mask[i]) for i in range(len(selected_coord)))
# emb_slice = (slice(None), slice(0,1)) + tuple(slice(selected_coord[i] + int(real_shape_mask[i] / 2),
# selected_coord[i] + int(
# real_shape_mask[i] / 2) + 1) for i in
# range(len(selected_coord)))
# pred_center_coord = [int(selected_coord[i] / pred_dws_fact[i]) for i in range(len(selected_coord))]
# emb_slice_pred = (slice(None), slice(None)) + tuple(
# slice(pred_center_coord[i], pred_center_coord[i] + 1)
# for i in range(len(selected_coord)))
#
# # Collect data for current sliding window:
# center_label = gt_segm[emb_slice]
# center_label_repeated = center_label.repeat(1, 1, *real_shape_mask)
# gt_patch_labels = gt_segm[full_patch_slice]
# gt_masks_to_plot.append(gt_patch_labels != center_label_repeated)
# gt_labels_to_plot.append(gt_patch_labels)
# # ignore_mask_patch = (gt_patch_labels == 0)
# pred_emb_to_plot.append(pred[emb_slice_pred])
#
# raw_to_plot.append(raw[full_patch_slice])
#
# # Highlight center pixel:
# raw_to_plot = torch.cat(raw_to_plot, dim=0)
# center_pixel_coord = (slice(None), 0) + tuple(int(shp / 2) for shp in real_shape_mask)
# raw_to_plot[center_pixel_coord] = raw_to_plot.min() - 1.
#
# gt_labels_to_plot = torch.cat(gt_labels_to_plot, dim=0)
# gt_masks_to_plot = torch.cat(gt_masks_to_plot, dim=0)
# pred_emb_to_plot = torch.cat(pred_emb_to_plot, dim=0)
#
# # Decode embeddings:
# ptch_num = kwargs["patchNet_number"]
# pred_patch_to_plot = data_parallel(self.model.patch_models[ptch_num], pred_emb_to_plot[:, :, 0, 0, 0], self.devices)
#
# # Downscale and rescale targets:
# down_sc_slice = (slice(None), slice(None)) + tuple(
# slice(int(dws_fact / 2), None, dws_fact) for dws_fact in mask_dws_fact)
# gt_masks_to_plot = torch.nn.functional.interpolate(gt_masks_to_plot[down_sc_slice].float(), scale_factor=tuple(mask_dws_fact))
# pred_patch_to_plot = torch.nn.functional.interpolate(pred_patch_to_plot,
# scale_factor=tuple(mask_dws_fact))
#
# gt_masks_to_plot = 1. - gt_masks_to_plot
# if mask_dws_fact[1] <= 6:
# pred_patch_to_plot = 1. - pred_patch_to_plot
#
# log_image("raw_patch_l{}".format(mask_dec_indx), raw_to_plot)
# log_image("gt_label_patch_l{}".format(mask_dec_indx), gt_labels_to_plot)
# log_image("gt_mask_patch_l{}".format(mask_dec_indx), gt_masks_to_plot)
# log_image("pred_patch_l{}".format(mask_dec_indx), pred_patch_to_plot)
# # ----------------------------
# # Patch-Loss:
# # ----------------------------
# If multiple strides were given, process all of them:
sample_strides = sample_strides if isinstance(sample_strides[0], list) else [sample_strides]
if limit_nb_decoded_masks_to is not None:
limit_nb_decoded_masks_to = limit_nb_decoded_masks_to if isinstance(limit_nb_decoded_masks_to[0],
list) else [
limit_nb_decoded_masks_to]
else:
limit_nb_decoded_masks_to = [None for _ in sample_strides]
for nb_stride, smpl_stride, max_nb_masks in zip(range(len(sample_strides)), sample_strides,
limit_nb_decoded_masks_to):
# ----------------------------
# Get some random prediction embeddings:
# ----------------------------
prediction_strides = get_prediction_strides(pred_dws_fact, smpl_stride)
selected_embeddings, crop_slice_pred, nb_selected_masks = extract_patches_torch(pred, (1, 1, 1),
stride=prediction_strides,
max_random_crop=max_random_crop)
# ----------------------------
# Collect gt_segm patches and corresponding center labels:
# ----------------------------
crop_slice_targets = tuple(slice(sl.start, None) for sl in crop_slice_pred)
gt_patches, _, _ = extract_patches_torch(gt_segm, real_shape_mask, stride=smpl_stride,
apply_specific_crop_slice=crop_slice_targets,
limit_patches_nb_to=nb_selected_masks)
gt_patches = gt_patches[:, [0]]
# Make sure to crop some additional border and get the centers correctly:
# TODO: this can be now easily done by cropping the gt_patches...
crop_slice_center_labels = (slice(None), slice(None)) + tuple(
slice(slc.start + int(sh / 2), slc.stop) for slc, sh in
zip(crop_slice_targets[2:], real_shape_mask))
target_at_patch_center, _, _ = extract_patches_torch(gt_segm, (1, 1, 1), stride=smpl_stride,
apply_specific_crop_slice=crop_slice_center_labels,
limit_patches_nb_to=nb_selected_masks)
# Get GT and other masks separately:
label_at_patch_center = target_at_patch_center[:, [0]]
mask_at_patch_center = target_at_patch_center[:, [1]]
# ----------------------------
# Ignore patches on the boundary or involving ignore-label:
# ----------------------------
# Ignore pixels involving ignore-labels:
ignore_masks = (gt_patches == self.ignore_label)
valid_patches = (label_at_patch_center != self.ignore_label)
patch_is_on_boundary = None
if self.boundary_label is not None:
patch_is_on_boundary = (mask_at_patch_center == self.boundary_label).repeat(1, 1, *real_shape_mask)
# Delete non-valid patches from batch:
valid_batch_indices = np.argwhere(valid_patches[:, 0, 0, 0, 0].cpu().detach().numpy())[:, 0]
if max_nb_masks is not None:
limit = max_nb_masks[0]
if max_nb_masks[1] == 'number':
if valid_batch_indices.shape[0] > limit:
valid_batch_indices = np.random.choice(valid_batch_indices, limit, replace=False)
elif max_nb_masks[1] == 'factor':
assert limit <= 1. and limit >= 0.
valid_batch_indices = np.random.choice(valid_batch_indices,
int(limit * valid_batch_indices.shape[0]), replace=False)
if valid_batch_indices.shape[0] == 0:
# Avoid problems if all patches are invalid and
# torch complaining that autograd cannot be performed:
loss += selected_embeddings.sum() * 0.
print("ZERO valid patches at level {}".format(mask_dec_indx))
continue
# ----------------------------
# Compute the actual (inverted) MeMasks targets: (0 is me, 1 are the others)
# best targets for Dice loss (usually more me than others)
# ----------------------------
center_labels_repeated = label_at_patch_center.repeat(1, 1, *real_shape_mask)
target_me_masks = gt_patches != center_labels_repeated
if patch_is_on_boundary is not None:
# If on boundary, we make (inverted) me_masks completely 1 (split from everything)
target_me_masks = target_me_masks | patch_is_on_boundary
# Downscaling patches:
down_sc_slice = (slice(None), slice(None)) + tuple(
slice(int(dws_fact / 2), None, dws_fact) for dws_fact in mask_dws_fact)
# Final targets:
target_me_masks = target_me_masks[valid_batch_indices].float()[down_sc_slice]
ignore_masks = ignore_masks[valid_batch_indices][down_sc_slice].byte()
# Invert MeMasks:
# best targets for Dice loss are: meMask == 0; others == 1
# TODO: generalize
if mask_dws_fact[1] > 6:
target_me_masks = 1. - target_me_masks
assert valid_batch_indices.max() < selected_embeddings.shape[
0], "Something went wrong, more target patches were collected than those predicted: {} targets vs {} pred...".format(
valid_batch_indices.max(), selected_embeddings.shape[0])
selected_embeddings = selected_embeddings[valid_batch_indices]
selected_embeddings = selected_embeddings[:, :, 0, 0, 0]
# ----------------------------
# Decode the actual predicted using the decoder models:
# ----------------------------
decoded_masks = data_parallel(mask_dec, selected_embeddings, self.devices)
# print(expanded_patches.shape)
assert decoded_masks.shape[1] == 1, "MaskDecoder should output only single-channel masks!"
# Some logs:
if nb_stride == 0:
log_image("ptc_trg_l{}".format(mask_dec_indx), target_me_masks)
log_image("ptc_pred_l{}".format(mask_dec_indx), decoded_masks)
# log_image("ptc_ign_l{}".format(nb_patch_net), patch_ignore_masks)
log_scalar("avg_targets_l{}".format(mask_dec_indx), target_me_masks.float().mean())
# ----------------------------
# Apply ignore mask and compute loss:
# ----------------------------
valid_masks = 1. - ignore_masks.float()
decoded_masks = decoded_masks * valid_masks
target_me_masks = target_me_masks * valid_masks
with warnings.catch_warnings(record=True) as w:
reconstruction_loss = data_parallel(self.loss, (decoded_masks, target_me_masks.float()),
self.devices).mean()
loss = loss + reconstruction_loss
if nb_stride == 0:
log_scalar("loss_l{}".format(mask_dec_indx), reconstruction_loss)
log_scalar("nb_patches_l{}".format(mask_dec_indx), decoded_masks.shape[0])
gc.collect()
return loss