From e2c4f4b122b6c517615a2a24e7dc7ae567bc787a Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Wed, 23 Oct 2024 20:00:05 +0200 Subject: [PATCH 01/23] AZ segmentation --- scripts/cooper/AZ_segmentation_h5.py | 142 ++++++++++++++++++++++++ scripts/rizzoli/evaluation_2D.py | 4 +- synaptic_reconstruction/inference/AZ.py | 82 ++++++++++++++ 3 files changed, 226 insertions(+), 2 deletions(-) create mode 100644 scripts/cooper/AZ_segmentation_h5.py create mode 100644 synaptic_reconstruction/inference/AZ.py diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py new file mode 100644 index 00000000..07ff7180 --- /dev/null +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -0,0 +1,142 @@ +import argparse +import h5py +import os +from pathlib import Path + +from tqdm import tqdm +from elf.io import open_file + +from synaptic_reconstruction.inference.AZ import segment_AZ +from synaptic_reconstruction.inference.util import parse_tiling + +def _require_output_folders(output_folder): + #seg_output = os.path.join(output_folder, "segmentations") + seg_output = output_folder + os.makedirs(seg_output, exist_ok=True) + return seg_output + +def get_volume(input_path): + ''' + with h5py.File(input_path) as seg_file: + input_volume = seg_file["raw"][:] + ''' + with open_file(input_path, "r") as f: + + # Try to automatically derive the key with the raw data. + keys = list(f.keys()) + if len(keys) == 1: + key = keys[0] + elif "data" in keys: + key = "data" + elif "raw" in keys: + key = "raw" + + input_volume = f[key][:] + return input_volume + +def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, key_label): + tiling = parse_tiling(tile_shape, halo) + print(f"using tiling {tiling}") + input = get_volume(input_path) + + #check if we have a restricting mask for the segmentation + if mask_path is not None: + with open_file(mask_path, "r") as f: + mask = f[mask_key][:] + else: + mask = None + + foreground = segment_AZ(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, mask = mask) + + seg_output = _require_output_folders(output_path) + file_name = Path(input_path).stem + seg_path = os.path.join(seg_output, f"{file_name}.h5") + + #check + os.makedirs(Path(seg_path).parent, exist_ok=True) + + print(f"Saving results in {seg_path}") + with h5py.File(seg_path, "a") as f: + if "raw" in f: + print("raw image already saved") + else: + f.create_dataset("raw", data=input, compression="gzip") + + key=f"AZ/segment_from_{key_label}" + if key in f: + print("Skipping", input_path, "because", key, "exists") + else: + f.create_dataset(key, data=foreground, compression="gzip") + + if mask is not None: + if mask_key in f: + print("mask image already saved") + else: + f.create_dataset(mask_key, data = mask, compression = "gzip") + + + + +def segment_folder(args): + input_files = [] + for root, dirs, files in os.walk(args.input_path): + input_files.extend([ + os.path.join(root, name) for name in files if name.endswith(".h5") + ]) + print(input_files) + pbar = tqdm(input_files, desc="Run segmentation") + for input_path in pbar: + + filename = os.path.basename(input_path) + try: + mask_path = os.path.join(args.mask_path, filename) + except: + print(f"Mask file not found for {input_path}") + mask_path = None + + run_AZ_segmentation(input_path, args.output_path, args.model_path, mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) + +def main(): + parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") + parser.add_argument( + "--input_path", "-i", required=True, + help="The filepath to the mrc file or the directory containing the tomogram data." + ) + parser.add_argument( + "--output_path", "-o", required=True, + help="The filepath to directory where the segmentations will be saved." + ) + parser.add_argument( + "--model_path", "-m", required=True, help="The filepath to the vesicle model." + ) + parser.add_argument( + "--mask_path", help="The filepath to a h5 file with a mask that will be used to restrict the segmentation. Needs to be in combination with mask_key." + ) + parser.add_argument( + "--mask_key", help="Key name that holds the mask segmentation" + ) + parser.add_argument( + "--tile_shape", type=int, nargs=3, + help="The tile shape for prediction. Lower the tile shape if GPU memory is insufficient." + ) + parser.add_argument( + "--halo", type=int, nargs=3, + help="The halo for prediction. Increase the halo to minimize boundary artifacts." + ) + parser.add_argument( + "--key_label", "-k", default = "combined_vesicles", + help="Give the key name for saving the segmentation in h5." + ) + args = parser.parse_args() + + input_ = args.input_path + + if os.path.isdir(input_): + segment_folder(args) + else: + run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) + + print("Finished segmenting!") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/rizzoli/evaluation_2D.py b/scripts/rizzoli/evaluation_2D.py index 1cae666e..18fd4f13 100644 --- a/scripts/rizzoli/evaluation_2D.py +++ b/scripts/rizzoli/evaluation_2D.py @@ -58,8 +58,8 @@ def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key) #get the labels and vesicles with h5py.File(labels_path) as label_file: labels = label_file["labels"] - #vesicles = labels["vesicles"] - gt = labels[anno_key][:] + vesicles = labels["vesicles"] + gt = vesicles[anno_key][:] with h5py.File(vesicles_path) as seg_file: segmentation = seg_file["vesicles"] diff --git a/synaptic_reconstruction/inference/AZ.py b/synaptic_reconstruction/inference/AZ.py new file mode 100644 index 00000000..b93218f8 --- /dev/null +++ b/synaptic_reconstruction/inference/AZ.py @@ -0,0 +1,82 @@ +import time +from typing import Dict, List, Optional, Tuple, Union + +import elf.parallel as parallel +import numpy as np +import torch + +from synaptic_reconstruction.inference.util import get_prediction, _Scaler + + +def _run_segmentation( + foreground, verbose, min_size, + # blocking shapes for parallel computation + block_shape=(128, 256, 256), +): + + # get the segmentation via seeded watershed + t0 = time.time() + seg = parallel.label(foreground > 0.5, block_shape=block_shape, verbose=verbose) + if verbose: + print("Compute connected components in", time.time() - t0, "s") + + # size filter + t0 = time.time() + ids, sizes = parallel.unique(seg, return_counts=True, block_shape=block_shape, verbose=verbose) + filter_ids = ids[sizes < min_size] + seg[np.isin(seg, filter_ids)] = 0 + if verbose: + print("Size filter in", time.time() - t0, "s") + seg = np.where(seg > 0, 1, 0) + return seg + +def segment_AZ( + input_volume: np.ndarray, + model_path: Optional[str] = None, + model: Optional[torch.nn.Module] = None, + tiling: Optional[Dict[str, Dict[str, int]]] = None, + min_size: int = 500, + verbose: bool = True, + return_predictions: bool = False, + scale: Optional[List[float]] = None, + mask: Optional[np.ndarray] = None, +) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + """ + Segment mitochondria in an input volume. + + Args: + input_volume: The input volume to segment. + model_path: The path to the model checkpoint if `model` is not provided. + model: Pre-loaded model. Either `model_path` or `model` is required. + tiling: The tiling configuration for the prediction. + verbose: Whether to print timing information. + scale: The scale factor to use for rescaling the input volume before prediction. + mask: An optional mask that is used to restrict the segmentation. + + Returns: + The foreground mask as a numpy array. + """ + if verbose: + print("Segmenting AZ in volume of shape", input_volume.shape) + # Create the scaler to handle prediction with a different scaling factor. + scaler = _Scaler(scale, verbose) + input_volume = scaler.scale_input(input_volume) + + # Rescale the mask if it was given and run prediction. + if mask is not None: + mask = scaler.scale_input(mask, is_segmentation=True) + pred = get_prediction(input_volume, model_path=model_path, model=model, tiling=tiling, mask=mask, verbose=verbose) + + # Run segmentation and rescale the result if necessary. + foreground = pred[0] + #print(f"shape {foreground.shape}") + #foreground = pred[0, :, :, :] + print(f"shape {foreground.shape}") + + segmentation = _run_segmentation(foreground, verbose=verbose, min_size=min_size) + + if return_predictions: + pred = scaler.rescale_output(pred, is_segmentation=False) + return segmentation, pred + return segmentation + From a0f713f80ce3c75b0534cbdc2a79214c6dd52c6f Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Mon, 28 Oct 2024 13:43:56 +0100 Subject: [PATCH 02/23] updates --- .gitignore | 4 +++- scripts/cooper/training/train_AZ.py | 11 ++++++---- scripts/rizzoli/2D_vesicle_segmentation.py | 20 +++++++++++++------ scripts/rizzoli/train_2D_domain_adaptation.py | 16 ++++++++++----- 4 files changed, 35 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index d0404310..ca035777 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,6 @@ slurm/ scripts/cooper/evaluation_results/ scripts/cooper/training/copy_testset.py scripts/rizzoli/upsample_data.py -scripts/cooper/training/find_rec_testset.py \ No newline at end of file +scripts/cooper/training/find_rec_testset.py +scripts/rizzoli/combine_2D_slices.py +scripts/rizzoli/combine_2D_slices_raw.py \ No newline at end of file diff --git a/scripts/cooper/training/train_AZ.py b/scripts/cooper/training/train_AZ.py index 1468eaf7..9d7d2833 100644 --- a/scripts/cooper/training/train_AZ.py +++ b/scripts/cooper/training/train_AZ.py @@ -12,7 +12,7 @@ from synaptic_reconstruction.training import semisupervised_training TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/exported_imod_objects" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_v1" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_v2" def _require_train_val_test_split(datasets): @@ -80,8 +80,11 @@ def get_paths(split, datasets, testset=True): def train(key, ignore_label = None, training_2D = False, testset = True): + os.makedirs(OUTPUT_ROOT, exist_ok=True) + datasets = [ "01_hoi_maus_2020_incomplete", + "04_hoi_stem_examples", "06_hoi_wt_stem750_fm", "12_chemical_fix_cryopreparation" ] @@ -93,7 +96,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): print(len(val_paths), "tomograms for validation") patch_shape = [48, 256, 256] - model_name=f"3D-AZ-model-v1" + model_name=f"3D-AZ-model-v3" #checking for 2D training if training_2D: @@ -109,11 +112,11 @@ def train(key, ignore_label = None, training_2D = False, testset = True): val_paths=val_paths, label_key=f"/labels/{key}", patch_shape=patch_shape, batch_size=batch_size, - sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=1), + sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=1, p_reject = 0.95), n_samples_train=None, n_samples_val=25, check=check, save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_models", - n_iterations=int(5e3), + n_iterations=int(5e4), ignore_label= ignore_label, label_transform=torch_em.transform.label.labels_to_binary, out_channels = 1, diff --git a/scripts/rizzoli/2D_vesicle_segmentation.py b/scripts/rizzoli/2D_vesicle_segmentation.py index 7974e3b3..ddfdab70 100644 --- a/scripts/rizzoli/2D_vesicle_segmentation.py +++ b/scripts/rizzoli/2D_vesicle_segmentation.py @@ -57,7 +57,7 @@ def get_volume(input_path): input_volume = seg_file["raw"][:] return input_volume -def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, halo, include_boundary, key_label): +def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, halo, include_boundary, key_label, scale): tiling = get_2D_tiling() @@ -72,20 +72,24 @@ def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, ha device = "cuda" if torch.cuda.is_available() else "cpu" model = torch_em.util.load_model(checkpoint=model_path, device=device) - def process_slices(input_volume): + def process_slices(input_volume, scale): processed_slices = [] foreground = [] boundaries = [] for z in range(input_volume.shape[0]): slice_ = input_volume[z, :, :] - segmented_slice, prediction_slice = segment_vesicles(input_volume=slice_, model=model, verbose=False, tiling=tiling, return_predictions=True, exclude_boundary=not include_boundary) + segmented_slice, prediction_slice = segment_vesicles(input_volume=slice_, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary) processed_slices.append(segmented_slice) foreground_pred_slice, boundaries_pred_slice = prediction_slice[:2] foreground.append(foreground_pred_slice) boundaries.append(boundaries_pred_slice) return processed_slices, foreground, boundaries - segmentation, foreground, boundaries = process_slices(input) + if input.ndim == 2: + segmentation, prediction = segment_vesicles(input_volume=input, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary) + foreground, boundaries = prediction[:2] + else: + segmentation, foreground, boundaries = process_slices(input, scale) seg_output = _require_output_folders(output_path) file_name = Path(input_path).stem @@ -121,7 +125,7 @@ def segment_folder(args): print(input_files) pbar = tqdm(input_files, desc="Run segmentation") for input_path in pbar: - run_vesicle_segmentation(input_path, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label) + run_vesicle_segmentation(input_path, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale) def main(): parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") @@ -152,6 +156,10 @@ def main(): "--key_label", "-k", default = "combined_vesicles", help="Give the key name for saving the segmentation in h5." ) + parser.add_argument( + "--scale", "-s", type=float, nargs=2, + help="Scales the input data." + ) args = parser.parse_args() input_ = args.input_path @@ -159,7 +167,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_vesicle_segmentation(input_, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label) + run_vesicle_segmentation(input_, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale) print("Finished segmenting!") diff --git a/scripts/rizzoli/train_2D_domain_adaptation.py b/scripts/rizzoli/train_2D_domain_adaptation.py index 86eedd17..c8a94197 100644 --- a/scripts/rizzoli/train_2D_domain_adaptation.py +++ b/scripts/rizzoli/train_2D_domain_adaptation.py @@ -6,11 +6,13 @@ from sklearn.model_selection import train_test_split from synaptic_reconstruction.training.domain_adaptation import mean_teacher_adaptation -TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/rizzoli/extracted" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/2D_DA_training_rizzoli" +TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/2D_data" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/2D_DA_training_rizzoli_v4" def _require_train_val_test_split(datasets): train_ratio, val_ratio, test_ratio = 0.8, 0.1, 0.1 + if len(datasets) < 10: + train_ratio, val_ratio, test_ratio = 0.5, 0.25, 0.25 def _train_val_test_split(names): train, test = train_test_split(names, test_size=1 - train_ratio, shuffle=True) @@ -71,8 +73,12 @@ def get_paths(split, datasets, testset=True): return paths def vesicle_domain_adaptation(teacher_model, testset = True): + + os.makedirs(OUTPUT_ROOT, exist_ok=True) + datasets = [ - "upsampled_by2" + "maus_2020_tem2d_wt_unt_div14_exported_scaled_grouped", + "20241021_imig_2014_data_transfer_exported_grouped" ] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -83,7 +89,7 @@ def vesicle_domain_adaptation(teacher_model, testset = True): #adjustable parameters patch_shape = [1, 256, 256] #2D - model_name = "2D-vesicle-DA-rizzoli-v3" + model_name = "2D-vesicle-DA-rizzoli-v5" model_root = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" checkpoint_path = os.path.join(model_root, teacher_model) @@ -97,7 +103,7 @@ def vesicle_domain_adaptation(teacher_model, testset = True): save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_models", source_checkpoint=checkpoint_path, confidence_threshold=0.75, - n_iterations=int(5e4), + n_iterations=int(5e5), ) From ac1ac0082b0154e04729c8e55802b73a81ab3afd Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Mon, 28 Oct 2024 16:59:33 +0100 Subject: [PATCH 03/23] update 2D DA --- scripts/rizzoli/train_2D_domain_adaptation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/scripts/rizzoli/train_2D_domain_adaptation.py b/scripts/rizzoli/train_2D_domain_adaptation.py index c8a94197..3beb4871 100644 --- a/scripts/rizzoli/train_2D_domain_adaptation.py +++ b/scripts/rizzoli/train_2D_domain_adaptation.py @@ -7,7 +7,7 @@ from synaptic_reconstruction.training.domain_adaptation import mean_teacher_adaptation TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/2D_data" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/2D_DA_training_rizzoli_v4" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/2D_DA_training_2Dcooper_v1" def _require_train_val_test_split(datasets): train_ratio, val_ratio, test_ratio = 0.8, 0.1, 0.1 @@ -77,9 +77,8 @@ def vesicle_domain_adaptation(teacher_model, testset = True): os.makedirs(OUTPUT_ROOT, exist_ok=True) datasets = [ - "maus_2020_tem2d_wt_unt_div14_exported_scaled_grouped", "20241021_imig_2014_data_transfer_exported_grouped" -] +]#"maus_2020_tem2d_wt_unt_div14_exported_scaled_grouped", train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -89,7 +88,7 @@ def vesicle_domain_adaptation(teacher_model, testset = True): #adjustable parameters patch_shape = [1, 256, 256] #2D - model_name = "2D-vesicle-DA-rizzoli-v5" + model_name = "2D-vesicle-DA-2Dcooper-imig-v1" model_root = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" checkpoint_path = os.path.join(model_root, teacher_model) From 61c57faf1de13b2e29176a1f7896274dc7ae01ae Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Thu, 7 Nov 2024 10:52:06 +0100 Subject: [PATCH 04/23] small updates, compartment segmentation --- scripts/cooper/AZ_segmentation_h5.py | 8 +- scripts/cooper/compartment_segmentation_h5.py | 116 ++++++++++++++++++ scripts/cooper/training/evaluation.py | 18 ++- scripts/cooper/vesicle_segmentation_h5.py | 30 ++++- scripts/rizzoli/evaluation_2D.py | 30 ++++- scripts/rizzoli/train_2D_domain_adaptation.py | 11 +- synaptic_reconstruction/inference/vesicles.py | 4 +- 7 files changed, 194 insertions(+), 23 deletions(-) create mode 100644 scripts/cooper/compartment_segmentation_h5.py diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py index 07ff7180..4deadc87 100644 --- a/scripts/cooper/AZ_segmentation_h5.py +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -81,7 +81,7 @@ def segment_folder(args): input_files = [] for root, dirs, files in os.walk(args.input_path): input_files.extend([ - os.path.join(root, name) for name in files if name.endswith(".h5") + os.path.join(root, name) for name in files if name.endswith(args.data_ext) ]) print(input_files) pbar = tqdm(input_files, desc="Run segmentation") @@ -127,6 +127,10 @@ def main(): "--key_label", "-k", default = "combined_vesicles", help="Give the key name for saving the segmentation in h5." ) + parser.add_argument( + "--data_ext", "-d", default = ".h5", + help="Format extension of data to be segmented, default is .h5." + ) args = parser.parse_args() input_ = args.input_path @@ -134,7 +138,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) + run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.data_ext) print("Finished segmenting!") diff --git a/scripts/cooper/compartment_segmentation_h5.py b/scripts/cooper/compartment_segmentation_h5.py new file mode 100644 index 00000000..1d0020a0 --- /dev/null +++ b/scripts/cooper/compartment_segmentation_h5.py @@ -0,0 +1,116 @@ +import argparse +import h5py +import os +from pathlib import Path + +from tqdm import tqdm +from elf.io import open_file + +from synaptic_reconstruction.inference.compartments import segment_compartments +from synaptic_reconstruction.inference.util import parse_tiling + +def _require_output_folders(output_folder): + #seg_output = os.path.join(output_folder, "segmentations") + seg_output = output_folder + os.makedirs(seg_output, exist_ok=True) + return seg_output + +def get_volume(input_path): + + with open_file(input_path, "r") as f: + + # Try to automatically derive the key with the raw data. + keys = list(f.keys()) + if len(keys) == 1: + key = keys[0] + elif "data" in keys: + key = "data" + elif "raw" in keys: + key = "raw" + + input_volume = f[key][:] + return input_volume + +def run_compartment_segmentation(input_path, output_path, model_path, tile_shape, halo, key_label): + tiling = parse_tiling(tile_shape, halo) + print(f"using tiling {tiling}") + input = get_volume(input_path) + + segmentation, prediction = segment_compartments(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, scale=[0.25, 0.25, 0.25]) + + seg_output = _require_output_folders(output_path) + file_name = Path(input_path).stem + seg_path = os.path.join(seg_output, f"{file_name}.h5") + + #check + os.makedirs(Path(seg_path).parent, exist_ok=True) + + print(f"Saving results in {seg_path}") + with h5py.File(seg_path, "a") as f: + if "raw" in f: + print("raw image already saved") + else: + f.create_dataset("raw", data=input, compression="gzip") + + key=f"compartments/segment_from_{key_label}" + if key in f: + print("Skipping", input_path, "because", key, "exists") + else: + f.create_dataset(key, data=segmentation, compression="gzip") + f.create_dataset(f"compartment_pred_{key_label}/foreground", data = prediction, compression="gzip") + + + + +def segment_folder(args): + input_files = [] + for root, dirs, files in os.walk(args.input_path): + input_files.extend([ + os.path.join(root, name) for name in files if name.endswith(args.data_ext) + ]) + print(input_files) + pbar = tqdm(input_files, desc="Run segmentation") + for input_path in pbar: + run_compartment_segmentation(input_path, args.output_path, args.model_path, args.tile_shape, args.halo, args.key_label) + +def main(): + parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") + parser.add_argument( + "--input_path", "-i", required=True, + help="The filepath to the mrc file or the directory containing the tomogram data." + ) + parser.add_argument( + "--output_path", "-o", required=True, + help="The filepath to directory where the segmentations will be saved." + ) + parser.add_argument( + "--model_path", "-m", required=True, help="The filepath to the vesicle model." + ) + parser.add_argument( + "--tile_shape", type=int, nargs=3, + help="The tile shape for prediction. Lower the tile shape if GPU memory is insufficient." + ) + parser.add_argument( + "--halo", type=int, nargs=3, + help="The halo for prediction. Increase the halo to minimize boundary artifacts." + ) + parser.add_argument( + "--data_ext", "-d", default=".h5", help="The extension of the tomogram data. By default .h5." + ) + parser.add_argument( + "--key_label", "-k", default = "3Dmodel_v1", + help="Give the key name for saving the segmentation in h5." + ) + args = parser.parse_args() + + input_ = args.input_path + + if os.path.isdir(input_): + segment_folder(args) + else: + run_compartment_segmentation(input_, args.output_path, args.model_path, args.tile_shape, args.halo, args.key_label) + + print("Finished segmenting!") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/cooper/training/evaluation.py b/scripts/cooper/training/evaluation.py index d7aaf6e5..68fa863e 100644 --- a/scripts/cooper/training/evaluation.py +++ b/scripts/cooper/training/evaluation.py @@ -21,7 +21,7 @@ def summarize_eval(results): table = summary.to_markdown(index=False) print(table) -def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key): +def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key, mask_key = None): print(f"Evaluate labels {labels_path} and vesicles {vesicles_path}") ds_name = os.path.basename(os.path.dirname(labels_path)) @@ -33,11 +33,16 @@ def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key) #vesicles = labels["vesicles"] gt = labels[anno_key][:] + if mask_key is not None: + mask = labels[mask_key][:] + with h5py.File(vesicles_path) as seg_file: segmentation = seg_file["vesicles"] vesicles = segmentation[segment_key][:] - + if mask_key is not None: + gt[mask == 0] = 0 + vesicles[mask == 0] = 0 #evaluate the match of ground truth and vesicles scores = evaluate(gt, vesicles) @@ -65,7 +70,7 @@ def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key) summarize_eval(results) -def evaluate_folder(labels_path, vesicles_path, model_name, segment_key, anno_key): +def evaluate_folder(labels_path, vesicles_path, model_name, segment_key, anno_key, mask_key = None): print(f"Evaluating folder {vesicles_path}") print(f"Using labels stored in {labels_path}") @@ -75,7 +80,7 @@ def evaluate_folder(labels_path, vesicles_path, model_name, segment_key, anno_ke for vesicle_file in vesicles_files: if vesicle_file in label_files: - evaluate_file(os.path.join(labels_path, vesicle_file), os.path.join(vesicles_path, vesicle_file), model_name, segment_key, anno_key) + evaluate_file(os.path.join(labels_path, vesicle_file), os.path.join(vesicles_path, vesicle_file), model_name, segment_key, anno_key, mask_key) @@ -87,13 +92,14 @@ def main(): parser.add_argument("-n", "--model_name", required=True) parser.add_argument("-sk", "--segment_key", required=True) parser.add_argument("-ak", "--anno_key", required=True) + parser.add_argument("-m", "--mask_key") args = parser.parse_args() vesicles_path = args.vesicles_path if os.path.isdir(vesicles_path): - evaluate_folder(args.labels_path, vesicles_path, args.model_name, args.segment_key, args.anno_key) + evaluate_folder(args.labels_path, vesicles_path, args.model_name, args.segment_key, args.anno_key, args.mask_key) else: - evaluate_file(args.labels_path, vesicles_path, args.model_name, args.segment_key, args.anno_key) + evaluate_file(args.labels_path, vesicles_path, args.model_name, args.segment_key, args.anno_key, args.mask_key) diff --git a/scripts/cooper/vesicle_segmentation_h5.py b/scripts/cooper/vesicle_segmentation_h5.py index 9c8b1d11..1136f180 100644 --- a/scripts/cooper/vesicle_segmentation_h5.py +++ b/scripts/cooper/vesicle_segmentation_h5.py @@ -34,7 +34,7 @@ def get_volume(input_path): input_volume = f[key][:] return input_volume -def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, include_boundary, key_label): +def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, include_boundary, key_label, distance_threshold = None): tiling = parse_tiling(tile_shape, halo) print(f"using tiling {tiling}") input = get_volume(input_path) @@ -45,8 +45,17 @@ def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mas mask = f[mask_key][:] else: mask = None + if distance_threshold is not None: + segmentation, prediction = segment_vesicles( + input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, + exclude_boundary=not include_boundary, mask = mask, distance_threshold = distance_threshold + ) + else: + segmentation, prediction = segment_vesicles( + input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, + exclude_boundary=not include_boundary, mask = mask + ) - segmentation, prediction = segment_vesicles(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, exclude_boundary=not include_boundary, mask = mask) foreground, boundaries = prediction[:2] seg_output = _require_output_folders(output_path) @@ -84,7 +93,7 @@ def segment_folder(args): input_files = [] for root, dirs, files in os.walk(args.input_path): input_files.extend([ - os.path.join(root, name) for name in files if name.endswith(".h5") + os.path.join(root, name) for name in files if name.endswith(args.data_ext) ]) print(input_files) pbar = tqdm(input_files, desc="Run segmentation") @@ -97,7 +106,10 @@ def segment_folder(args): print(f"Mask file not found for {input_path}") mask_path = None - run_vesicle_segmentation(input_path, args.output_path, args.model_path, mask_path, args.mask_key, args.tile_shape, args.halo, args.include_boundary, args.key_label) + run_vesicle_segmentation( + input_path, args.output_path, args.model_path, mask_path, args.mask_key, + args.tile_shape, args.halo, args.include_boundary, args.key_label, args.distance_threshold + ) def main(): parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") @@ -134,6 +146,14 @@ def main(): "--key_label", "-k", default = "combined_vesicles", help="Give the key name for saving the segmentation in h5." ) + parser.add_argument( + "--distance_threshold", "-t", type=int, + help="Used for distance based segmentation." + ) + parser.add_argument( + "--data_ext", "-d", default = ".h5", + help="Format extension of data to be segmented, default is .h5." + ) args = parser.parse_args() input_ = args.input_path @@ -141,7 +161,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_vesicle_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.include_boundary, args.key_label) + run_vesicle_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.distance_threshold) print("Finished segmenting!") diff --git a/scripts/rizzoli/evaluation_2D.py b/scripts/rizzoli/evaluation_2D.py index 18fd4f13..9f918dff 100644 --- a/scripts/rizzoli/evaluation_2D.py +++ b/scripts/rizzoli/evaluation_2D.py @@ -6,8 +6,13 @@ import numpy as np from elf.evaluation import matching +from skimage.transform import rescale - +def transpose_tomo(tomogram): + data0 = np.swapaxes(tomogram, 0, -1) + data1 = np.fliplr(data0) + transposed_data = np.swapaxes(data1, 0, -1) + return transposed_data def evaluate(labels, vesicles): assert labels.shape == vesicles.shape @@ -54,21 +59,34 @@ def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key) ds_name = os.path.basename(os.path.dirname(labels_path)) tomo = os.path.basename(labels_path) - + use_mask = True #get the labels and vesicles with h5py.File(labels_path) as label_file: labels = label_file["labels"] - vesicles = labels["vesicles"] - gt = vesicles[anno_key][:] + #vesicles = labels["vesicles"] + gt = labels[anno_key][:] + gt = rescale(gt, scale=0.5, order=0, anti_aliasing=False, preserve_range=True).astype(gt.dtype) + gt = transpose_tomo(gt) + + if use_mask: + mask = labels["mask"][:] + mask = rescale(mask, scale=0.5, order=0, anti_aliasing=False, preserve_range=True).astype(mask.dtype) + mask = transpose_tomo(mask) with h5py.File(vesicles_path) as seg_file: segmentation = seg_file["vesicles"] vesicles = segmentation[segment_key][:] + if use_mask: + gt[mask == 0] = 0 + vesicles[mask == 0] = 0 - #evaluate the match of ground truth and vesicles - scores = evaluate_slices(gt, vesicles) + #evaluate the match of ground truth and vesicles + if len(vesicles.shape) == 3: + scores = evaluate_slices(gt, vesicles) + else: + scores = evaluate(gt,vesicles) #store results result_folder ="/user/muth9/u12095/synaptic-reconstruction/scripts/cooper/evaluation_results" os.makedirs(result_folder, exist_ok=True) diff --git a/scripts/rizzoli/train_2D_domain_adaptation.py b/scripts/rizzoli/train_2D_domain_adaptation.py index 3beb4871..ac2a28fa 100644 --- a/scripts/rizzoli/train_2D_domain_adaptation.py +++ b/scripts/rizzoli/train_2D_domain_adaptation.py @@ -78,7 +78,7 @@ def vesicle_domain_adaptation(teacher_model, testset = True): datasets = [ "20241021_imig_2014_data_transfer_exported_grouped" -]#"maus_2020_tem2d_wt_unt_div14_exported_scaled_grouped", +] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -88,11 +88,13 @@ def vesicle_domain_adaptation(teacher_model, testset = True): #adjustable parameters patch_shape = [1, 256, 256] #2D - model_name = "2D-vesicle-DA-2Dcooper-imig-v1" + model_name = "2D-vesicle-DA-2Dcooper-imig-v2" model_root = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" checkpoint_path = os.path.join(model_root, teacher_model) + patch_shape = [256, 256] if any("maus" in dataset for dataset in datasets) else [1, 256, 256] + mean_teacher_adaptation( name=model_name, unsupervised_train_paths=train_paths, @@ -102,7 +104,10 @@ def vesicle_domain_adaptation(teacher_model, testset = True): save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_models", source_checkpoint=checkpoint_path, confidence_threshold=0.75, - n_iterations=int(5e5), + batch_size=8, + n_iterations=int(1.5e4), + n_samples_train=8000, + n_samples_val=50, ) diff --git a/synaptic_reconstruction/inference/vesicles.py b/synaptic_reconstruction/inference/vesicles.py index 237d95af..4a56b0f9 100644 --- a/synaptic_reconstruction/inference/vesicles.py +++ b/synaptic_reconstruction/inference/vesicles.py @@ -49,6 +49,7 @@ def distance_based_vesicle_segmentation( # Get the segmentation via seeded watershed of components in the boundary distances. t0 = time.time() + print(f"using a distance thresholf of {distance_threshold} for distance based segmentation") seeds = parallel.label(bd_dist > distance_threshold, block_shape=block_shape, verbose=verbose) if verbose: print("Compute connected components in", time.time() - t0, "s") @@ -129,6 +130,7 @@ def segment_vesicles( min_size: int = 500, verbose: bool = True, distance_based_segmentation: bool = True, + distance_threshold: int = 8, return_predictions: bool = False, scale: Optional[List[float]] = None, exclude_boundary: bool = False, @@ -174,7 +176,7 @@ def segment_vesicles( if distance_based_segmentation: seg = distance_based_vesicle_segmentation( - foreground, boundaries, verbose=verbose, min_size=min_size, **kwargs + foreground, boundaries, verbose=verbose, min_size=min_size, distance_threshold = distance_threshold, **kwargs ) else: seg = simple_vesicle_segmentation( From 40e965ed72c87250b5c2ecf02e8ba1bb06b1e2e9 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 7 Nov 2024 17:24:28 +0100 Subject: [PATCH 05/23] Implement code for first analysis --- scripts/cooper/analysis/run_analysis_1.py | 66 +++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 scripts/cooper/analysis/run_analysis_1.py diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py new file mode 100644 index 00000000..55181893 --- /dev/null +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -0,0 +1,66 @@ +# This is the code for the first analysis for the cooper data. +# Here, we only compute the vesicle numbers and size distributions for the STEM tomograms +# in the 04 dataset. + +import os +from glob import glob + +import pandas as pd +import h5py +from tqdm import tqdm +from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres + +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/ground_truth/04Dataset_for_vesicle_eval" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04Dataset_for_vesicle_eval/model_segmentation" # noqa +RESULT_FOLDER = "./analysis_results/analysis_1" + + +# We compute the sizes for all vesicles in the compartment masks. +# We use the same logic in the size computation as for the vesicle extraction to IMOD, +# including the radius correction factor. +# The number of vesicles is automatically computed as the length of the size list. +def compute_sizes_for_all_tomorams(): + os.makedirs(RESULT_FOLDER, exist_ok=True) + + resolution = (0.8681,) * 3 + radius_factor = 1.3 + estimate_radius_2d = True + + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): + ds_name, fname = os.path.split(tomo) + ds_name = os.path.split(ds_name)[1] + fname = os.path.splitext(fname)[0] + output_path = os.path.join(RESULT_FOLDER, f"{ds_name}_{fname}.csv") + if os.path.exists(output_path): + continue + + # Load the vesicle segmentation from the predictions. + with h5py.File(tomo, "r") as f: + segmentation = f["/vesicles/segment_from_combined_vesicles"][:] + + input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") + assert os.path.exists(input_path), input_path + # Load the compartment mask from the tomogram + with h5py.File(input_path, "r") as f: + mask = f["labels/compartment"][:] + + segmentation[mask == 0] = 0 + _, sizes = convert_segmentation_to_spheres( + segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d + ) + + result = pd.DataFrame({ + "dataset": [ds_name] * len(sizes), + "tomogram": [fname] * len(sizes), + "sizes": sizes + }) + result.to_csv(output_path, index=False) + + +def main(): + compute_sizes_for_all_tomorams() + + +if __name__ == "__main__": + main() From 7be9ee8fd74d0667e61c84f0d3db104427527581 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Mon, 11 Nov 2024 20:21:47 +0100 Subject: [PATCH 06/23] 2D seg with mask --- .gitignore | 3 +- scripts/cooper/AZ_segmentation_h5.py | 2 +- scripts/rizzoli/2D_vesicle_segmentation.py | 37 ++++++++++++++++++---- scripts/rizzoli/evaluation_2D.py | 6 ++-- 4 files changed, 37 insertions(+), 11 deletions(-) diff --git a/.gitignore b/.gitignore index ca035777..4db569e9 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,5 @@ scripts/cooper/training/copy_testset.py scripts/rizzoli/upsample_data.py scripts/cooper/training/find_rec_testset.py scripts/rizzoli/combine_2D_slices.py -scripts/rizzoli/combine_2D_slices_raw.py \ No newline at end of file +scripts/rizzoli/combine_2D_slices_raw.py +scripts/cooper/remove_h5key.py \ No newline at end of file diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py index 4deadc87..2fb7045b 100644 --- a/scripts/cooper/AZ_segmentation_h5.py +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -138,7 +138,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.data_ext) + run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) print("Finished segmenting!") diff --git a/scripts/rizzoli/2D_vesicle_segmentation.py b/scripts/rizzoli/2D_vesicle_segmentation.py index ddfdab70..159be283 100644 --- a/scripts/rizzoli/2D_vesicle_segmentation.py +++ b/scripts/rizzoli/2D_vesicle_segmentation.py @@ -7,6 +7,7 @@ import torch import torch_em import numpy as np +from elf.io import open_file from synaptic_reconstruction.inference.vesicles import segment_vesicles from synaptic_reconstruction.inference.util import parse_tiling @@ -57,7 +58,7 @@ def get_volume(input_path): input_volume = seg_file["raw"][:] return input_volume -def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, halo, include_boundary, key_label, scale): +def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, halo, include_boundary, key_label, scale, mask_path, mask_key): tiling = get_2D_tiling() @@ -69,16 +70,29 @@ def run_vesicle_segmentation(input_path, output_path, model_path, tile_shape, ha tiling = parse_tiling(tile_shape, halo) input = get_volume(input_path) + #check if we have a restricting mask for the segmentation + if mask_path is not None: + with open_file(mask_path, "r") as f: + mask = f[mask_key][:] + else: + mask = None + device = "cuda" if torch.cuda.is_available() else "cpu" model = torch_em.util.load_model(checkpoint=model_path, device=device) - def process_slices(input_volume, scale): + def process_slices(input_volume, scale, mask): processed_slices = [] foreground = [] boundaries = [] for z in range(input_volume.shape[0]): slice_ = input_volume[z, :, :] - segmented_slice, prediction_slice = segment_vesicles(input_volume=slice_, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary) + #check if we have a restricting mask for the segmentation + if mask is not None: + mask_slice = mask[z, :, :] + segmented_slice, prediction_slice = segment_vesicles(input_volume=slice_, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary, mask = mask_slice) + else: + segmented_slice, prediction_slice = segment_vesicles(input_volume=slice_, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary) + processed_slices.append(segmented_slice) foreground_pred_slice, boundaries_pred_slice = prediction_slice[:2] foreground.append(foreground_pred_slice) @@ -86,10 +100,11 @@ def process_slices(input_volume, scale): return processed_slices, foreground, boundaries if input.ndim == 2: + #TODO: check if we have a restricting mask for the segmentation segmentation, prediction = segment_vesicles(input_volume=input, model=model, verbose=False, tiling=tiling, return_predictions=True, scale = scale, exclude_boundary=not include_boundary) foreground, boundaries = prediction[:2] else: - segmentation, foreground, boundaries = process_slices(input, scale) + segmentation, foreground, boundaries = process_slices(input, scale, mask) seg_output = _require_output_folders(output_path) file_name = Path(input_path).stem @@ -125,7 +140,11 @@ def segment_folder(args): print(input_files) pbar = tqdm(input_files, desc="Run segmentation") for input_path in pbar: - run_vesicle_segmentation(input_path, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale) + if args.mask_path is not None: + mask_path_for_file = os.path.join(args.mask_path, os.path.basename(input_path)) + else: + mask_path_for_file = None + run_vesicle_segmentation(input_path, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale, mask_path_for_file, args.mask_key) def main(): parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") @@ -160,6 +179,12 @@ def main(): "--scale", "-s", type=float, nargs=2, help="Scales the input data." ) + parser.add_argument( + "--mask_path", help="The filepath to a h5 file with a mask that will be used to restrict the segmentation. Needs to be in combination with mask_key." + ) + parser.add_argument( + "--mask_key", help="Key name that holds the mask segmentation" + ) args = parser.parse_args() input_ = args.input_path @@ -167,7 +192,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_vesicle_segmentation(input_, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale) + run_vesicle_segmentation(input_, args.output_path, args.model_path, args.tile_shape, args.halo, args.include_boundary, args.key_label, args.scale, args.mask_path, args.mask_key) print("Finished segmenting!") diff --git a/scripts/rizzoli/evaluation_2D.py b/scripts/rizzoli/evaluation_2D.py index 9f918dff..5b5bbbd9 100644 --- a/scripts/rizzoli/evaluation_2D.py +++ b/scripts/rizzoli/evaluation_2D.py @@ -59,14 +59,14 @@ def evaluate_file(labels_path, vesicles_path, model_name, segment_key, anno_key) ds_name = os.path.basename(os.path.dirname(labels_path)) tomo = os.path.basename(labels_path) - use_mask = True + use_mask = False #get the labels and vesicles with h5py.File(labels_path) as label_file: labels = label_file["labels"] #vesicles = labels["vesicles"] gt = labels[anno_key][:] - gt = rescale(gt, scale=0.5, order=0, anti_aliasing=False, preserve_range=True).astype(gt.dtype) - gt = transpose_tomo(gt) + #gt = rescale(gt, scale=0.5, order=0, anti_aliasing=False, preserve_range=True).astype(gt.dtype) + #gt = transpose_tomo(gt) if use_mask: mask = labels["mask"][:] From f85e4452bd3c32132a2cd94ff139ea47334e4566 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Mon, 11 Nov 2024 22:47:38 +0100 Subject: [PATCH 07/23] spatial distribution analysis --- scripts/cooper/analysis/run_analysis_1.py | 6 +- .../run_spatial_distribution_analysis.py | 75 +++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 scripts/cooper/analysis/run_spatial_distribution_analysis.py diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 55181893..523aba1b 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -10,8 +10,8 @@ from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/ground_truth/04Dataset_for_vesicle_eval" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/04Dataset_for_vesicle_eval/model_segmentation" # noqa +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/20241102_TOMO_DATA_Imig2014" # noqa RESULT_FOLDER = "./analysis_results/analysis_1" @@ -22,7 +22,7 @@ def compute_sizes_for_all_tomorams(): os.makedirs(RESULT_FOLDER, exist_ok=True) - resolution = (0.8681,) * 3 + resolution = (0.8681,) * 3 #change for each dataset radius_factor = 1.3 estimate_radius_2d = True diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py new file mode 100644 index 00000000..20024319 --- /dev/null +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -0,0 +1,75 @@ +import os +from glob import glob +import pandas as pd +import h5py +from tqdm import tqdm +from synaptic_reconstruction.distance_measurements import measure_segmentation_to_object_distances + +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/20241102_TOMO_DATA_Imig2014" # noqa +RESULT_FOLDER = "./analysis_results" + + +# We compute the distances for all vesicles in the compartment masks to the AZ. +# We use the same different resolution, depending on dataset. +# The closest distance is calculated, i.e., the closest point on the outer membrane of the vesicle to the AZ. +def compute_sizes_for_all_tomorams(): + os.makedirs(RESULT_FOLDER, exist_ok=True) + + resolution = (0.8681,) * 3 # Change for each dataset + + # Dictionary to hold the results for each dataset + dataset_results = {} + + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): + ds_name, fname = os.path.split(tomo) + ds_name = os.path.split(ds_name)[1] + fname = os.path.splitext(fname)[0] + + # Initialize a new dictionary entry for each dataset if not already present + if ds_name not in dataset_results: + dataset_results[ds_name] = {} + + # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name]: + continue + + # Load the vesicle segmentation from the predictions + with h5py.File(tomo, "r") as f: + segmentation = f["/vesicles/segment_from_combined_vesicles"][:] + segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] + + input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") + assert os.path.exists(input_path), input_path + + # Load the compartment mask from the tomogram + with h5py.File(input_path, "r") as f: + mask = f["labels/compartment"][:] + + segmentation[mask == 0] = 0 + distances, _, _, _ = measure_segmentation_to_object_distances( + segmentation, segmented_object=segmented_object, resolution=resolution + ) + + # Add distances to the dataset dictionary under the tomogram name + dataset_results[ds_name][fname] = distances + + # Save each dataset's results to a single CSV file + for ds_name, tomogram_data in dataset_results.items(): + # Create a DataFrame where each column is a tomogram's distances + result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + + # Define the output file path + output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) + + +def main(): + compute_sizes_for_all_tomorams() + + +if __name__ == "__main__": + main() From 8ef16bcafd4289b1793e1275d826a386a2cf8c87 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Tue, 12 Nov 2024 13:46:23 +0100 Subject: [PATCH 08/23] intersection between compartment boundary and AZ segmentaiton --- scripts/cooper/AZ_segmentation_h5.py | 35 ++++++++++++++++--- .../run_spatial_distribution_analysis.py | 2 +- synaptic_reconstruction/inference/AZ.py | 9 ++++- .../postprocessing/postprocess_AZ.py | 25 +++++++++++++ 4 files changed, 65 insertions(+), 6 deletions(-) create mode 100644 synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py index 2fb7045b..baa72254 100644 --- a/scripts/cooper/AZ_segmentation_h5.py +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -34,7 +34,7 @@ def get_volume(input_path): input_volume = f[key][:] return input_volume -def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, key_label): +def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, key_label, compartment_seg): tiling = parse_tiling(tile_shape, halo) print(f"using tiling {tiling}") input = get_volume(input_path) @@ -46,7 +46,14 @@ def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key else: mask = None - foreground = segment_AZ(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, mask = mask) + #check if intersection with compartment is necessary + if compartment_seg is None: + foreground = segment_AZ(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, mask = mask) + intersection = None + else: + with open_file(compartment_seg, "r") as f: + compartment = f["/compartments/segment_from_3Dmodel_v1"][:] + foreground, intersection = segment_AZ(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, mask = mask, compartment=compartment) seg_output = _require_output_folders(output_path) file_name = Path(input_path).stem @@ -73,6 +80,13 @@ def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key print("mask image already saved") else: f.create_dataset(mask_key, data = mask, compression = "gzip") + + if intersection is not None: + intersection_key = "AZ/compartment_AZ_intersection" + if intersection_key in f: + print("intersection already saved") + else: + f.create_dataset(intersection_key, data = intersection, compression = "gzip") @@ -93,8 +107,15 @@ def segment_folder(args): except: print(f"Mask file not found for {input_path}") mask_path = None + + if args.compartment_seg is not None: + try: + compartment_seg = os.path.join(args.compartment_seg, os.path.splitext(filename)[0] + '.h5') + except: + print(f"compartment file not found for {input_path}") + compartment_seg = None - run_AZ_segmentation(input_path, args.output_path, args.model_path, mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) + run_AZ_segmentation(input_path, args.output_path, args.model_path, mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, compartment_seg) def main(): parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") @@ -131,6 +152,12 @@ def main(): "--data_ext", "-d", default = ".h5", help="Format extension of data to be segmented, default is .h5." ) + parser.add_argument( + "--compartment_seg", "-c", + help="Path to compartment segmentation." + "If the compartment segmentation was executed before, this will add a key to output file that stores the intersection between compartment boundary and AZ." + "Maybe need to adjust the compartment key that the segmentation is stored under" + ) args = parser.parse_args() input_ = args.input_path @@ -138,7 +165,7 @@ def main(): if os.path.isdir(input_): segment_folder(args) else: - run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label) + run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.compartment_seg) print("Finished segmenting!") diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index 20024319..f8d07080 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -16,7 +16,7 @@ def compute_sizes_for_all_tomorams(): os.makedirs(RESULT_FOLDER, exist_ok=True) - resolution = (0.8681,) * 3 # Change for each dataset + resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset # Dictionary to hold the results for each dataset dataset_results = {} diff --git a/synaptic_reconstruction/inference/AZ.py b/synaptic_reconstruction/inference/AZ.py index b93218f8..a1c9da83 100644 --- a/synaptic_reconstruction/inference/AZ.py +++ b/synaptic_reconstruction/inference/AZ.py @@ -6,7 +6,7 @@ import torch from synaptic_reconstruction.inference.util import get_prediction, _Scaler - +from synaptic_reconstruction.inference.postprocessing.postprocess_AZ import find_intersection_boundary def _run_segmentation( foreground, verbose, min_size, @@ -40,6 +40,7 @@ def segment_AZ( return_predictions: bool = False, scale: Optional[List[float]] = None, mask: Optional[np.ndarray] = None, + compartment: Optional[np.ndarray] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ Segment mitochondria in an input volume. @@ -75,8 +76,14 @@ def segment_AZ( segmentation = _run_segmentation(foreground, verbose=verbose, min_size=min_size) + #returning prediciton and intersection not possible atm, but currently do not need prediction anyways if return_predictions: pred = scaler.rescale_output(pred, is_segmentation=False) return segmentation, pred + + if compartment is not None: + intersection = find_intersection_boundary(segmentation, compartment) + return segmentation, intersection + return segmentation diff --git a/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py new file mode 100644 index 00000000..54d9e2c3 --- /dev/null +++ b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py @@ -0,0 +1,25 @@ +import numpy as np +from scipy.ndimage import binary_erosion + +def find_intersection_boundary(segmented_AZ, segmented_compartment): + """ + Find the intersection of the boundary of segmented_compartment with segmented_AZ. + + Parameters: + segmented_AZ (numpy.ndarray): 3D array representing the active zone (AZ). + segmented_compartment (numpy.ndarray): 3D array representing the compartment. + + Returns: + numpy.ndarray: 3D array with the intersection of the boundary of segmented_compartment and segmented_AZ. + """ + # Step 0: Binarize the segmented_compartment + binarized_compartment = (segmented_compartment > 0).astype(int) + + # Step 1: Create a binary mask of the compartment boundary + eroded_compartment = binary_erosion(binarized_compartment) + boundary_compartment = binarized_compartment - eroded_compartment + + # Step 2: Find the intersection with the AZ + intersection = np.logical_and(boundary_compartment, segmented_AZ) + + return intersection.astype(int) # Convert boolean array to int (1 for intersecting points, 0 elsewhere) From 09f6c846cfc391096ee6c86fca49d4d5f56d8799 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Tue, 12 Nov 2024 17:20:26 +0100 Subject: [PATCH 09/23] Update compartment postprocessing --- .../inference/compartments.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/synaptic_reconstruction/inference/compartments.py b/synaptic_reconstruction/inference/compartments.py index a822d9f8..701c2222 100644 --- a/synaptic_reconstruction/inference/compartments.py +++ b/synaptic_reconstruction/inference/compartments.py @@ -77,6 +77,12 @@ def _segment_compartments_2d( mask = np.logical_or(binary_closing(mask, iterations=4), mask) segmentation[bb][mask] = prop.label + # import napari + # v = napari.Viewer() + # v.add_image(boundaries) + # v.add_labels(segmentation) + # napari.run() + return segmentation @@ -117,6 +123,7 @@ def _segment_compartments_3d( boundary_threshold=0.4, n_slices_exclude=0, min_z_extent=10, + postprocess_segments=False, ): distances = distance_transform_edt(prediction < boundary_threshold).astype("float32") seg_2d = np.zeros(prediction.shape, dtype="uint32") @@ -132,7 +139,8 @@ def _segment_compartments_3d( seg_2d[z] = seg_z seg = _merge_segmentation_3d(seg_2d, min_z_extent) - seg = _postprocess_seg_3d(seg) + if postprocess_segments: + seg = _postprocess_seg_3d(seg) # import napari # v = napari.Viewer() @@ -155,6 +163,9 @@ def segment_compartments( scale: Optional[List[float]] = None, mask: Optional[np.ndarray] = None, n_slices_exclude: int = 0, + boundary_threshold: float = 0.4, + min_z_extent: int = 10, + postprocess_segments: bool = False, **kwargs, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ @@ -194,9 +205,14 @@ def segment_compartments( # We may want to expose some of the parameters here. t0 = time.time() if input_volume.ndim == 2: - seg = _segment_compartments_2d(pred) + seg = _segment_compartments_2d(pred, boundary_threshold=boundary_threshold) else: - seg = _segment_compartments_3d(pred, n_slices_exclude=n_slices_exclude) + seg = _segment_compartments_3d( + pred, + boundary_threshold=boundary_threshold, + n_slices_exclude=n_slices_exclude, + postprocess_segments=postprocess_segments, + ) if verbose: print("Run segmentation in", time.time() - t0, "s") From f893d2300fab14a737b3d85adecb72af04644e22 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Wed, 13 Nov 2024 12:05:09 +0100 Subject: [PATCH 10/23] updating data analysis on smaller details --- .gitignore | 1 + scripts/cooper/AZ_segmentation_h5.py | 2 +- scripts/cooper/analysis/run_analysis_1.py | 106 +++++++++++++++++- .../run_spatial_distribution_analysis.py | 64 +++++++++-- scripts/cooper/compartment_segmentation_h5.py | 2 +- .../postprocessing/postprocess_AZ.py | 36 +++--- 6 files changed, 185 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index 4db569e9..0377c4a0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ models/*/ run_sbatch.sbatch slurm/ scripts/cooper/evaluation_results/ +analysis_results/ scripts/cooper/training/copy_testset.py scripts/rizzoli/upsample_data.py scripts/cooper/training/find_rec_testset.py diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py index baa72254..da694c19 100644 --- a/scripts/cooper/AZ_segmentation_h5.py +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -52,7 +52,7 @@ def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key intersection = None else: with open_file(compartment_seg, "r") as f: - compartment = f["/compartments/segment_from_3Dmodel_v1"][:] + compartment = f["/labels/compartment"][:] foreground, intersection = segment_AZ(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, mask = mask, compartment=compartment) seg_output = _require_output_folders(output_path) diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 523aba1b..b166a717 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -5,15 +5,56 @@ import os from glob import glob +import numpy as np import pandas as pd import h5py from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/20241102_TOMO_DATA_Imig2014" # noqa -RESULT_FOLDER = "./analysis_results/analysis_1" +DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" +def get_compartment_with_max_overlap(compartments, vesicles): + """ + Given 3D numpy arrays of compartments and vesicles, this function returns a binary mask + of the compartment with the most overlap with vesicles based on the number of overlapping voxels. + + Parameters: + compartments (numpy.ndarray): 3D array of compartment labels. + vesicles (numpy.ndarray): 3D array of vesicle labels or binary mask. + + Returns: + numpy.ndarray: Binary mask of the compartment with the most overlap with vesicles. + """ + + unique_compartments = np.unique(compartments) + if 0 in unique_compartments: + unique_compartments = unique_compartments[unique_compartments != 0] + + max_overlap_count = 0 + best_compartment = None + + # Iterate over each compartment and calculate the overlap with vesicles + for compartment_label in unique_compartments: + # Create a binary mask for the current compartment + compartment_mask = compartments == compartment_label + vesicle_mask = vesicles > 0 + + intersection = np.logical_and(compartment_mask, vesicle_mask) + + # Calculate the number of overlapping voxels + overlap_count = np.sum(intersection) + + # Track the compartment with the most overlap in terms of voxel count + if overlap_count > max_overlap_count: + max_overlap_count = overlap_count + best_compartment = compartment_label + + # Create the final mask for the compartment with the most overlap + final_mask = compartments == best_compartment + + return final_mask # We compute the sizes for all vesicles in the compartment masks. # We use the same logic in the size computation as for the vesicle extraction to IMOD, @@ -57,9 +98,66 @@ def compute_sizes_for_all_tomorams(): }) result.to_csv(output_path, index=False) +def compute_sizes_for_all_tomorams_manComp(): + os.makedirs(RESULT_FOLDER, exist_ok=True) + + resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset + radius_factor = 1.3 + estimate_radius_2d = True + + # Dictionary to hold the results for each dataset + dataset_results = {} + + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): + ds_name, fname = os.path.split(tomo) + ds_name = os.path.split(ds_name)[1] + fname = os.path.splitext(fname)[0] + # Initialize a new dictionary entry for each dataset if not already present + if ds_name not in dataset_results: + dataset_results[ds_name] = {} + + # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name]: + continue + + # Load the vesicle segmentation from the predictions. + with h5py.File(tomo, "r") as f: + segmentation = f["/vesicles/segment_from_combined_vesicles"][:] + + input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") + assert os.path.exists(input_path), input_path + # Load the compartment mask from the tomogram + with h5py.File(input_path, "r") as f: + compartments = f["/compartments/segment_from_3Dmodel_v2"][:] + mask = get_compartment_with_max_overlap(compartments, segmentation) + + #if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask + if np.sum(segmentation[mask == 0] > 0) > (0.5 * np.sum(segmentation > 0)): + print("using no mask") + else: + segmentation[mask == 0] = 0 + _, sizes = convert_segmentation_to_spheres( + segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d + ) + + # Add sizes to the dataset dictionary under the tomogram name + dataset_results[ds_name][fname] = sizes + + # Save each dataset's results to a single CSV file + for ds_name, tomogram_data in dataset_results.items(): + # Create a DataFrame where each column is a tomogram's sizes + result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + + # Define the output file path + output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) def main(): - compute_sizes_for_all_tomorams() + #compute_sizes_for_all_tomorams() + compute_sizes_for_all_tomorams_manComp() if __name__ == "__main__": diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index f8d07080..fca7eed2 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -4,12 +4,54 @@ import h5py from tqdm import tqdm from synaptic_reconstruction.distance_measurements import measure_segmentation_to_object_distances +import numpy as np -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/20241102_TOMO_DATA_Imig2014" # noqa -RESULT_FOLDER = "./analysis_results" +DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" +def get_compartment_with_max_overlap(compartments, vesicles): + """ + Given 3D numpy arrays of compartments and vesicles, this function returns a binary mask + of the compartment with the most overlap with vesicles based on the number of overlapping voxels. + + Parameters: + compartments (numpy.ndarray): 3D array of compartment labels. + vesicles (numpy.ndarray): 3D array of vesicle labels or binary mask. + + Returns: + numpy.ndarray: Binary mask of the compartment with the most overlap with vesicles. + """ + + unique_compartments = np.unique(compartments) + if 0 in unique_compartments: + unique_compartments = unique_compartments[unique_compartments != 0] + + max_overlap_count = 0 + best_compartment = None + + # Iterate over each compartment and calculate the overlap with vesicles + for compartment_label in unique_compartments: + # Create a binary mask for the current compartment + compartment_mask = compartments == compartment_label + vesicle_mask = vesicles > 0 + + intersection = np.logical_and(compartment_mask, vesicle_mask) + + # Calculate the number of overlapping voxels + overlap_count = np.sum(intersection) + + # Track the compartment with the most overlap in terms of voxel count + if overlap_count > max_overlap_count: + max_overlap_count = overlap_count + best_compartment = compartment_label + + # Create the final mask for the compartment with the most overlap + final_mask = compartments == best_compartment + + return final_mask + # We compute the distances for all vesicles in the compartment masks to the AZ. # We use the same different resolution, depending on dataset. # The closest distance is calculated, i.e., the closest point on the outer membrane of the vesicle to the AZ. @@ -38,16 +80,24 @@ def compute_sizes_for_all_tomorams(): # Load the vesicle segmentation from the predictions with h5py.File(tomo, "r") as f: segmentation = f["/vesicles/segment_from_combined_vesicles"][:] - segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] + segmented_object = f["/AZ/compartment_AZ_intersection"][:] + #if AZ intersect is small, compartment seg didn't align with AZ so we use the normal AZ and not intersect + if (segmented_object == 0).all() or np.sum(segmented_object == 1) < 2000: + segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") assert os.path.exists(input_path), input_path # Load the compartment mask from the tomogram with h5py.File(input_path, "r") as f: - mask = f["labels/compartment"][:] - - segmentation[mask == 0] = 0 + compartments = f["/compartments/segment_from_3Dmodel_v2"][:] + mask = get_compartment_with_max_overlap(compartments, segmentation) + + #if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask + if np.sum(segmentation[mask == 0] > 0) > (0.5 * np.sum(segmentation > 0)): + print("using no mask") + else: + segmentation[mask == 0] = 0 distances, _, _, _ = measure_segmentation_to_object_distances( segmentation, segmented_object=segmented_object, resolution=resolution ) diff --git a/scripts/cooper/compartment_segmentation_h5.py b/scripts/cooper/compartment_segmentation_h5.py index 1d0020a0..573ac482 100644 --- a/scripts/cooper/compartment_segmentation_h5.py +++ b/scripts/cooper/compartment_segmentation_h5.py @@ -36,7 +36,7 @@ def run_compartment_segmentation(input_path, output_path, model_path, tile_shape print(f"using tiling {tiling}") input = get_volume(input_path) - segmentation, prediction = segment_compartments(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, scale=[0.25, 0.25, 0.25]) + segmentation, prediction = segment_compartments(input_volume=input, model_path=model_path, verbose=False, tiling=tiling, return_predictions=True, scale=[0.25, 0.25, 0.25],boundary_threshold=0.2, postprocess_segments=False) seg_output = _require_output_folders(output_path) file_name = Path(input_path).stem diff --git a/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py index 54d9e2c3..8ef2cd32 100644 --- a/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py +++ b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py @@ -1,25 +1,35 @@ import numpy as np -from scipy.ndimage import binary_erosion +from skimage.segmentation import find_boundaries def find_intersection_boundary(segmented_AZ, segmented_compartment): """ - Find the intersection of the boundary of segmented_compartment with segmented_AZ. + Find the cumulative intersection of the boundary of each label in segmented_compartment with segmented_AZ. Parameters: segmented_AZ (numpy.ndarray): 3D array representing the active zone (AZ). - segmented_compartment (numpy.ndarray): 3D array representing the compartment. + segmented_compartment (numpy.ndarray): 3D array representing the compartment, with multiple labels. Returns: - numpy.ndarray: 3D array with the intersection of the boundary of segmented_compartment and segmented_AZ. + numpy.ndarray: 3D array with the cumulative intersection of all boundaries of segmented_compartment labels with segmented_AZ. """ - # Step 0: Binarize the segmented_compartment - binarized_compartment = (segmented_compartment > 0).astype(int) + # Step 0: Initialize an empty array to accumulate intersections + cumulative_intersection = np.zeros_like(segmented_AZ, dtype=bool) - # Step 1: Create a binary mask of the compartment boundary - eroded_compartment = binary_erosion(binarized_compartment) - boundary_compartment = binarized_compartment - eroded_compartment - - # Step 2: Find the intersection with the AZ - intersection = np.logical_and(boundary_compartment, segmented_AZ) + # Step 1: Loop through each unique label in segmented_compartment (excluding 0 if it represents background) + labels = np.unique(segmented_compartment) + labels = labels[labels != 0] # Exclude background label (0) if necessary + + for label in labels: + # Step 2: Create a binary mask for the current label + label_mask = (segmented_compartment == label) + + # Step 3: Find the boundary of the current label's compartment + boundary_compartment = find_boundaries(label_mask, mode='outer') + + # Step 4: Find the intersection with the AZ for this label's boundary + intersection = np.logical_and(boundary_compartment, segmented_AZ) + + # Step 5: Accumulate intersections for each label + cumulative_intersection = np.logical_or(cumulative_intersection, intersection) - return intersection.astype(int) # Convert boolean array to int (1 for intersecting points, 0 elsewhere) + return cumulative_intersection.astype(int) # Convert boolean array to int (1 for intersecting points, 0 elsewhere) From 08c56b9c08bb244869ad77d689efcc5ca294dfa8 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Wed, 13 Nov 2024 13:03:23 +0100 Subject: [PATCH 11/23] minor updates data analysis --- scripts/cooper/analysis/run_analysis_1.py | 44 ++++++++----- .../run_spatial_distribution_analysis.py | 62 +++++++++++++++++-- 2 files changed, 87 insertions(+), 19 deletions(-) diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index b166a717..459d490b 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -11,9 +11,9 @@ from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" def get_compartment_with_max_overlap(compartments, vesicles): """ @@ -60,20 +60,27 @@ def get_compartment_with_max_overlap(compartments, vesicles): # We use the same logic in the size computation as for the vesicle extraction to IMOD, # including the radius correction factor. # The number of vesicles is automatically computed as the length of the size list. -def compute_sizes_for_all_tomorams(): +def compute_sizes_for_all_tomorams_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) - resolution = (0.8681,) * 3 #change for each dataset + resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset radius_factor = 1.3 estimate_radius_2d = True + # Dictionary to hold the results for each dataset + dataset_results = {} + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) for tomo in tqdm(tomograms): ds_name, fname = os.path.split(tomo) ds_name = os.path.split(ds_name)[1] fname = os.path.splitext(fname)[0] - output_path = os.path.join(RESULT_FOLDER, f"{ds_name}_{fname}.csv") - if os.path.exists(output_path): + # Initialize a new dictionary entry for each dataset if not already present + if ds_name not in dataset_results: + dataset_results[ds_name] = {} + + # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name]: continue # Load the vesicle segmentation from the predictions. @@ -91,14 +98,21 @@ def compute_sizes_for_all_tomorams(): segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d ) - result = pd.DataFrame({ - "dataset": [ds_name] * len(sizes), - "tomogram": [fname] * len(sizes), - "sizes": sizes - }) - result.to_csv(output_path, index=False) + # Add sizes to the dataset dictionary under the tomogram name + dataset_results[ds_name][fname] = sizes -def compute_sizes_for_all_tomorams_manComp(): + # Save each dataset's results to a single CSV file + for ds_name, tomogram_data in dataset_results.items(): + # Create a DataFrame where each column is a tomogram's sizes + result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + + # Define the output file path + output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) + +def compute_sizes_for_all_tomorams_autoComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset @@ -156,8 +170,8 @@ def compute_sizes_for_all_tomorams_manComp(): result_df.to_csv(output_path, index=False) def main(): - #compute_sizes_for_all_tomorams() compute_sizes_for_all_tomorams_manComp() + #compute_sizes_for_all_tomorams_autoComp() if __name__ == "__main__": diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index fca7eed2..9a890a19 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -6,9 +6,9 @@ from synaptic_reconstruction.distance_measurements import measure_segmentation_to_object_distances import numpy as np -DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" def get_compartment_with_max_overlap(compartments, vesicles): @@ -116,9 +116,63 @@ def compute_sizes_for_all_tomorams(): # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) + +def compute_sizes_for_all_tomorams_manComp(): + os.makedirs(RESULT_FOLDER, exist_ok=True) + + resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset + + # Dictionary to hold the results for each dataset + dataset_results = {} + + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): + ds_name, fname = os.path.split(tomo) + ds_name = os.path.split(ds_name)[1] + fname = os.path.splitext(fname)[0] + # Initialize a new dictionary entry for each dataset if not already present + if ds_name not in dataset_results: + dataset_results[ds_name] = {} + + # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name]: + continue + + # Load the vesicle segmentation from the predictions + with h5py.File(tomo, "r") as f: + segmentation = f["/vesicles/segment_from_combined_vesicles"][:] + segmented_object = f["/AZ/compartment_AZ_intersection_manComp"][:] + + input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") + assert os.path.exists(input_path), input_path + + # Load the compartment mask from the tomogram + with h5py.File(input_path, "r") as f: + mask = f["/labels/compartment"][:] + + segmentation[mask == 0] = 0 + distances, _, _, _ = measure_segmentation_to_object_distances( + segmentation, segmented_object=segmented_object, resolution=resolution + ) + + # Add distances to the dataset dictionary under the tomogram name + dataset_results[ds_name][fname] = distances + + # Save each dataset's results to a single CSV file + for ds_name, tomogram_data in dataset_results.items(): + # Create a DataFrame where each column is a tomogram's distances + result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + + # Define the output file path + output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) + def main(): - compute_sizes_for_all_tomorams() + #compute_sizes_for_all_tomorams() + compute_sizes_for_all_tomorams_manComp() if __name__ == "__main__": From 49d1b7ce8834dd9e81c594f7a3e96f91f6b2208e Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Thu, 14 Nov 2024 16:52:48 +0100 Subject: [PATCH 12/23] calculation of AZ area --- scripts/cooper/analysis/calc_AZ_area.py | 227 ++++++++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 scripts/cooper/analysis/calc_AZ_area.py diff --git a/scripts/cooper/analysis/calc_AZ_area.py b/scripts/cooper/analysis/calc_AZ_area.py new file mode 100644 index 00000000..e9fcb522 --- /dev/null +++ b/scripts/cooper/analysis/calc_AZ_area.py @@ -0,0 +1,227 @@ +import h5py +import numpy as np +import os +import csv +from scipy.ndimage import binary_opening, median_filter,zoom, binary_closing +from skimage.measure import label, regionprops +from synaptic_reconstruction.morphology import compute_object_morphology +from skimage.morphology import ball +from scipy.spatial import ConvexHull +from skimage.draw import polygon + +def calculate_AZ_area_per_slice(AZ_slice, pixel_size_nm=1.554): + """ + Calculate the area of the AZ in a single 2D slice after applying error-reducing processing. + + Parameters: + - AZ_slice (numpy array): 2D array representing a single slice of the AZ segmentation. + - pixel_size_nm (float): Size of a pixel in nanometers. + + Returns: + - slice_area_nm2 (float): The area of the AZ in the slice in square nanometers. + """ + # Apply binary opening or median filter to reduce small segmentation errors + AZ_slice_filtered = binary_opening(AZ_slice, structure=np.ones((3, 3))).astype(int) + + # Calculate area in this slice + num_AZ_pixels = np.sum(AZ_slice_filtered == 1) + slice_area_nm2 = num_AZ_pixels * (pixel_size_nm ** 2) + + return slice_area_nm2 + +def calculate_total_AZ_area(tomo_path, pixel_size_nm=1.554): + """ + Calculate the total area of the AZ across all slices in a 3D tomogram file. + + Parameters: + - tomo_path (str): Path to the tomogram file (HDF5 format). + - pixel_size_nm (float): Size of a pixel in nanometers. + + Returns: + - total_AZ_area_nm2 (float): The total area of the AZ in square nanometers. + """ + with h5py.File(tomo_path, "r") as f: + AZ_intersect_seg = f["/AZ/compartment_AZ_intersection_manComp"][:] + + # Calculate the AZ area for each slice along the z-axis + total_AZ_area_nm2 = 0 + for z_slice in AZ_intersect_seg: + slice_area_nm2 = calculate_AZ_area_per_slice(z_slice, pixel_size_nm) + total_AZ_area_nm2 += slice_area_nm2 + + return total_AZ_area_nm2 + +def calculate_AZ_area_simple(tomo_path, pixel_size_nm=1.554): + """ + Calculate the volume of the AZ (active zone) in a 3D tomogram file. + + Parameters: + - tomo_path (str): Path to the tomogram file (HDF5 format). + - pixel_size_nm (float): Size of a pixel in nanometers (default is 1.554 nm). + + Returns: + - AZ_volume_nm3 (float): The volume of the AZ in cubic nanometers. + """ + # Open the file and read the AZ intersection segmentation data + with h5py.File(tomo_path, "r") as f: + AZ_intersect_seg = f["/AZ/compartment_AZ_intersection_manComp"][:] + + # Count voxels with label = 1 + num_AZ_voxels = np.sum(AZ_intersect_seg == 1) + + # Calculate the volume in cubic nanometers + AZ_area_nm2 = num_AZ_voxels * (pixel_size_nm ** 2) + + return AZ_area_nm2 + +def calculate_AZ_surface(tomo_path, pixel_size_nm=1.554): + with h5py.File(tomo_path, "r") as f: + AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] + + # Apply binary closing to smooth the segmented regions + struct_elem = ball(1) # Use a small 3D structuring element + AZ_seg_smoothed = binary_closing(AZ_seg > 0, structure=struct_elem, iterations=20) + + labeled_seg = label(AZ_seg_smoothed) + + regions = regionprops(labeled_seg) + if regions: + # Sort regions by area and get the label of the largest region + largest_region = max(regions, key=lambda r: r.area) + largest_label = largest_region.label + + largest_component_mask = (labeled_seg == largest_label) + AZ_seg_filtered = largest_component_mask.astype(np.uint8) + + else: + # If no regions found, return an empty array + AZ_seg_filtered = np.zeros_like(AZ_seg_interp, dtype=np.uint8) + + morphology_data = compute_object_morphology(AZ_seg_filtered, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) + surface_column = "surface [nm^2]" #if resolution is not None else "surface [pixel^2]" + surface_area = morphology_data[surface_column].iloc[0] + + return surface_area + +def calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm=1.554): + with h5py.File(tomo_path, "r") as f: + AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] + + # Apply binary closing to smooth the segmented regions + struct_elem = ball(1) # Use a small 3D structuring element + AZ_seg_smoothed = binary_closing(AZ_seg > 0, structure=struct_elem, iterations=10) + + labeled_seg = label(AZ_seg_smoothed) + + regions = regionprops(labeled_seg) + if regions: + # Sort regions by area and get the label of the largest region + largest_region = max(regions, key=lambda r: r.area) + largest_label = largest_region.label + + largest_component_mask = (labeled_seg == largest_label) + AZ_seg_filtered = largest_component_mask.astype(np.uint8) + AZ_seg = AZ_seg_filtered + # Extract coordinates of non-zero points + points = np.argwhere(AZ_seg > 0) # Get the coordinates of non-zero (foreground) pixels + + if points.shape[0] < 4: + # ConvexHull requires at least 4 points in 3D to form a valid hull + AZ_seg_filtered = np.zeros_like(AZ_seg, dtype=np.uint8) + else: + # Apply ConvexHull to the points + hull = ConvexHull(points) + + # Create a binary mask for the convex hull + convex_hull_mask = np.zeros_like(AZ_seg, dtype=bool) + + # Iterate over each simplex (facet) of the convex hull and fill in the polygon + for simplex in hull.simplices: + # For each face of the convex hull, extract the vertices and convert to a 2D polygon + polygon_coords = points[simplex] + rr, cc = polygon(polygon_coords[:, 0], polygon_coords[:, 1]) + convex_hull_mask[rr, cc] = True + + # Optional: Label the convex hull mask + labeled_seg = label(convex_hull_mask) + regions = regionprops(labeled_seg) + + if regions: + # Sort regions by area and get the label of the largest region + largest_region = max(regions, key=lambda r: r.area) + largest_label = largest_region.label + + largest_component_mask = (labeled_seg == largest_label) + AZ_seg_filtered = largest_component_mask.astype(np.uint8) + + else: + AZ_seg_filtered = np.zeros_like(AZ_seg, dtype=np.uint8) + + # Calculate surface area + morphology_data = compute_object_morphology(AZ_seg_filtered, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) + surface_column = "surface [nm^2]" + surface_area = morphology_data[surface_column].iloc[0] + + return surface_area + +def process_datasets(folder_path, output_csv="AZ_areas.csv", pixel_size_nm=1.554): + """ + Process all tomograms in multiple datasets within a folder and save results to a CSV. + + Parameters: + - folder_path (str): Path to the folder containing dataset folders with tomograms. + - output_csv (str): Filename for the output CSV file. + - pixel_size_nm (float): Size of a pixel in nanometers. + """ + results = [] + + # Iterate over each dataset folder + for dataset_name in os.listdir(folder_path): + dataset_path = os.path.join(folder_path, dataset_name) + + # Check if it's a directory (skip files in the main folder) + if not os.path.isdir(dataset_path): + continue + + # Iterate over each tomogram file in the dataset folder + for tomo_file in os.listdir(dataset_path): + tomo_path = os.path.join(dataset_path, tomo_file) + + # Check if the file is an HDF5 file (optional) + if tomo_file.endswith(".h5") or tomo_file.endswith(".hdf5"): + try: + # Calculate AZ area + #AZ_area = calculate_total_AZ_area(tomo_path, pixel_size_nm) + #AZ_area = calculate_AZ_area_simple(tomo_path, pixel_size_nm) + #AZ_surface_area = calculate_AZ_surface(tomo_path, pixel_size_nm) + AZ_surface_area = calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm) + # Append results to list + results.append({ + "Dataset": dataset_name, + "Tomogram": tomo_file, + "AZ_surface_area": AZ_surface_area + }) + except Exception as e: + print(f"Error processing {tomo_file} in {dataset_name}: {e}") + + # Write results to a CSV file + with open(output_csv, mode="w", newline="") as csvfile: + fieldnames = ["Dataset", "Tomogram", "AZ_surface_area"] + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + + writer.writeheader() + for result in results: + writer.writerow(result) + + print(f"Results saved to {output_csv}") + +def main(): + # Define the path to the folder containing dataset folders + folder_path = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" + output_csv = "./analysis_results/AZ_intersect_manualCompartment/AZ_surface_area.csv" + # Call the function to process datasets and save results + process_datasets(folder_path, output_csv = output_csv) + +# Call main +if __name__ == "__main__": + main() From 8a515d1704409f328d2fffd7cbc99e9d63986ba5 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Thu, 14 Nov 2024 18:01:29 +0100 Subject: [PATCH 13/23] corrected radius factor --- scripts/cooper/analysis/run_analysis_1.py | 53 ++++++++++++++--------- 1 file changed, 33 insertions(+), 20 deletions(-) diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 459d490b..4077ea3b 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -11,9 +11,9 @@ from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" def get_compartment_with_max_overlap(compartments, vesicles): """ @@ -64,7 +64,7 @@ def compute_sizes_for_all_tomorams_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - radius_factor = 1.3 + radius_factor = 0.7 estimate_radius_2d = True # Dictionary to hold the results for each dataset @@ -112,14 +112,21 @@ def compute_sizes_for_all_tomorams_manComp(): # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) +import os +import pandas as pd +import numpy as np +from glob import glob +import h5py +from tqdm import tqdm + def compute_sizes_for_all_tomorams_autoComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - radius_factor = 1.3 + radius_factor = 0.7 estimate_radius_2d = True - # Dictionary to hold the results for each dataset + # Dictionary to hold the results for each dataset and category (CTRL or DKO) dataset_results = {} tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) @@ -127,12 +134,16 @@ def compute_sizes_for_all_tomorams_autoComp(): ds_name, fname = os.path.split(tomo) ds_name = os.path.split(ds_name)[1] fname = os.path.splitext(fname)[0] - # Initialize a new dictionary entry for each dataset if not already present + + # Determine if the tomogram is 'CTRL' or 'DKO' + category = "CTRL" if "CTRL" in fname else "DKO" + + # Initialize a new dictionary entry for each dataset and category if not already present if ds_name not in dataset_results: - dataset_results[ds_name] = {} + dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} # Skip if this tomogram already exists in the dataset dictionary - if fname in dataset_results[ds_name]: + if fname in dataset_results[ds_name][category]: continue # Load the vesicle segmentation from the predictions. @@ -146,32 +157,34 @@ def compute_sizes_for_all_tomorams_autoComp(): compartments = f["/compartments/segment_from_3Dmodel_v2"][:] mask = get_compartment_with_max_overlap(compartments, segmentation) - #if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask + # if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask if np.sum(segmentation[mask == 0] > 0) > (0.5 * np.sum(segmentation > 0)): - print("using no mask") + print(f"using no mask for {tomo}") else: segmentation[mask == 0] = 0 _, sizes = convert_segmentation_to_spheres( segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d ) - # Add sizes to the dataset dictionary under the tomogram name - dataset_results[ds_name][fname] = sizes + # Add sizes to the dataset dictionary under the appropriate category + dataset_results[ds_name][category][fname] = sizes - # Save each dataset's results to a single CSV file - for ds_name, tomogram_data in dataset_results.items(): - # Create a DataFrame where each column is a tomogram's sizes - result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + # Save each dataset's results into separate CSV files for CTRL and DKO tomograms + for ds_name, categories in dataset_results.items(): + for category, tomogram_data in categories.items(): + # Sort tomograms by name within the category + sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names + result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() # Define the output file path - output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}.csv") + output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}.csv") # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) def main(): - compute_sizes_for_all_tomorams_manComp() - #compute_sizes_for_all_tomorams_autoComp() + #compute_sizes_for_all_tomorams_manComp() + compute_sizes_for_all_tomorams_autoComp() if __name__ == "__main__": From b1449d23ff090a8b8c8c2ebd5845ea578e7f9be1 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Tue, 19 Nov 2024 17:57:53 +0100 Subject: [PATCH 14/23] minor changes --- scripts/cooper/analysis/calc_AZ_area.py | 20 ++++- scripts/cooper/analysis/run_analysis_1.py | 53 ++++++----- .../run_spatial_distribution_analysis.py | 87 +++++++++++++++++-- 3 files changed, 122 insertions(+), 38 deletions(-) diff --git a/scripts/cooper/analysis/calc_AZ_area.py b/scripts/cooper/analysis/calc_AZ_area.py index e9fcb522..592b043e 100644 --- a/scripts/cooper/analysis/calc_AZ_area.py +++ b/scripts/cooper/analysis/calc_AZ_area.py @@ -76,7 +76,8 @@ def calculate_AZ_area_simple(tomo_path, pixel_size_nm=1.554): def calculate_AZ_surface(tomo_path, pixel_size_nm=1.554): with h5py.File(tomo_path, "r") as f: - AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] + #AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] + AZ_seg = f["/filtered_az"][:] # Apply binary closing to smooth the segmented regions struct_elem = ball(1) # Use a small 3D structuring element @@ -103,6 +104,16 @@ def calculate_AZ_surface(tomo_path, pixel_size_nm=1.554): return surface_area +def calculate_AZ_surface_simple(tomo_path, pixel_size_nm=1.554): + with h5py.File(tomo_path, "r") as f: + AZ_seg = f["/labels/AZ"][:] + + morphology_data = compute_object_morphology(AZ_seg, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) + surface_column = "surface [nm^2]" #if resolution is not None else "surface [pixel^2]" + surface_area = morphology_data[surface_column].iloc[0] + + return surface_area + def calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm=1.554): with h5py.File(tomo_path, "r") as f: AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] @@ -194,7 +205,8 @@ def process_datasets(folder_path, output_csv="AZ_areas.csv", pixel_size_nm=1.554 #AZ_area = calculate_total_AZ_area(tomo_path, pixel_size_nm) #AZ_area = calculate_AZ_area_simple(tomo_path, pixel_size_nm) #AZ_surface_area = calculate_AZ_surface(tomo_path, pixel_size_nm) - AZ_surface_area = calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm) + #AZ_surface_area = calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm) + AZ_surface_area = calculate_AZ_surface_simple(tomo_path, pixel_size_nm) # Append results to list results.append({ "Dataset": dataset_name, @@ -217,8 +229,8 @@ def process_datasets(folder_path, output_csv="AZ_areas.csv", pixel_size_nm=1.554 def main(): # Define the path to the folder containing dataset folders - folder_path = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" - output_csv = "./analysis_results/AZ_intersect_manualCompartment/AZ_surface_area.csv" + folder_path = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" + output_csv = "./analysis_results/manual_AZ_exported/AZ_surface_area.csv" # Call the function to process datasets and save results process_datasets(folder_path, output_csv = output_csv) diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 4077ea3b..3afde5dc 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -11,9 +11,9 @@ from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_autoCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa +RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" def get_compartment_with_max_overlap(compartments, vesicles): """ @@ -64,10 +64,10 @@ def compute_sizes_for_all_tomorams_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - radius_factor = 0.7 + radius_factor = 1 estimate_radius_2d = True - # Dictionary to hold the results for each dataset + # Dictionary to hold the results for each dataset and category (CTRL or DKO) dataset_results = {} tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) @@ -75,14 +75,18 @@ def compute_sizes_for_all_tomorams_manComp(): ds_name, fname = os.path.split(tomo) ds_name = os.path.split(ds_name)[1] fname = os.path.splitext(fname)[0] - # Initialize a new dictionary entry for each dataset if not already present + + # Determine if the tomogram is 'CTRL' or 'DKO' + category = "CTRL" if "CTRL" in fname else "DKO" + + # Initialize a new dictionary entry for each dataset and category if not already present if ds_name not in dataset_results: - dataset_results[ds_name] = {} + dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} # Skip if this tomogram already exists in the dataset dictionary - if fname in dataset_results[ds_name]: + if fname in dataset_results[ds_name][category]: continue - + # Load the vesicle segmentation from the predictions. with h5py.File(tomo, "r") as f: segmentation = f["/vesicles/segment_from_combined_vesicles"][:] @@ -98,32 +102,27 @@ def compute_sizes_for_all_tomorams_manComp(): segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d ) - # Add sizes to the dataset dictionary under the tomogram name - dataset_results[ds_name][fname] = sizes + # Add sizes to the dataset dictionary under the appropriate category + dataset_results[ds_name][category][fname] = sizes - # Save each dataset's results to a single CSV file - for ds_name, tomogram_data in dataset_results.items(): - # Create a DataFrame where each column is a tomogram's sizes - result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() + # Save each dataset's results into separate CSV files for CTRL and DKO tomograms + for ds_name, categories in dataset_results.items(): + for category, tomogram_data in categories.items(): + # Sort tomograms by name within the category + sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names + result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() # Define the output file path - output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}.csv") + output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}_rf1.csv") # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) -import os -import pandas as pd -import numpy as np -from glob import glob -import h5py -from tqdm import tqdm - def compute_sizes_for_all_tomorams_autoComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - radius_factor = 0.7 + radius_factor = 1 estimate_radius_2d = True # Dictionary to hold the results for each dataset and category (CTRL or DKO) @@ -177,14 +176,14 @@ def compute_sizes_for_all_tomorams_autoComp(): result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() # Define the output file path - output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}.csv") + output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}_rf1.csv") # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) def main(): - #compute_sizes_for_all_tomorams_manComp() - compute_sizes_for_all_tomorams_autoComp() + compute_sizes_for_all_tomorams_manComp() + #compute_sizes_for_all_tomorams_autoComp() if __name__ == "__main__": diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index 9a890a19..cdc4c0d8 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -6,9 +6,9 @@ from synaptic_reconstruction.distance_measurements import measure_segmentation_to_object_distances import numpy as np -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +RESULT_FOLDER = "./analysis_results/AZ_filtered_autoComp" def get_compartment_with_max_overlap(compartments, vesicles): @@ -55,7 +55,7 @@ def get_compartment_with_max_overlap(compartments, vesicles): # We compute the distances for all vesicles in the compartment masks to the AZ. # We use the same different resolution, depending on dataset. # The closest distance is calculated, i.e., the closest point on the outer membrane of the vesicle to the AZ. -def compute_sizes_for_all_tomorams(): +def compute_per_vesicle_distance_to_AZ(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset @@ -116,8 +116,80 @@ def compute_sizes_for_all_tomorams(): # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) +def compute_per_vesicle_distance_to_filteredAZ(): + filtered_AZ_path = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/az_seg_filtered" + os.makedirs(RESULT_FOLDER, exist_ok=True) + + resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset + + # Dictionary to hold the results for each dataset and category (CTRL or DKO) + dataset_results = {} + + tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): + ds_name, fname = os.path.split(tomo) + ds_name = os.path.split(ds_name)[1] + fname = os.path.splitext(fname)[0] + + # Determine if the tomogram is 'CTRL' or 'DKO' + category = "CTRL" if "CTRL" in fname else "DKO" + + # Initialize a new dictionary entry for each dataset and category if not already present + if ds_name not in dataset_results: + dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} + + # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name][category]: + continue + + #Load the AZ segmentations + AZ_path = os.path.join(filtered_AZ_path, ds_name, f"{fname}.h5") + with h5py.File(AZ_path, "r") as f: + segmented_object = f["/filtered_az"][:] + + # Load the vesicle segmentation from the predictions + with h5py.File(tomo, "r") as f: + segmentation = f["/vesicles/segment_from_combined_vesicles"][:] + + #if AZ intersect is small, compartment seg didn't align with AZ so we use the normal AZ and not intersect + if (segmented_object == 0).all() or np.sum(segmented_object == 1) < 2000: + segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] + + input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") + assert os.path.exists(input_path), input_path + + # Load the compartment mask from the tomogram + with h5py.File(input_path, "r") as f: + compartments = f["/compartments/segment_from_3Dmodel_v2"][:] + mask = get_compartment_with_max_overlap(compartments, segmentation) + + #if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask + if np.sum(segmentation[mask == 0] > 0) > (0.5 * np.sum(segmentation > 0)): + print("using no mask") + else: + segmentation[mask == 0] = 0 + distances, _, _, _ = measure_segmentation_to_object_distances( + segmentation, segmented_object=segmented_object, resolution=resolution + ) + + # Add distances to the dataset dictionary under the appropriate category + dataset_results[ds_name][category][fname] = distances + + # Save each dataset's results into separate CSV files for CTRL and DKO tomograms + for ds_name, categories in dataset_results.items(): + for category, tomogram_data in categories.items(): + # Sort tomograms by name within the category + sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names + result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() + + # Define the output file path + output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}_{category}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) + -def compute_sizes_for_all_tomorams_manComp(): +def compute_per_vesicle_distance_to_AZ_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset @@ -171,8 +243,9 @@ def compute_sizes_for_all_tomorams_manComp(): result_df.to_csv(output_path, index=False) def main(): - #compute_sizes_for_all_tomorams() - compute_sizes_for_all_tomorams_manComp() + #compute_per_vesicle_distance_to_AZ() + #compute_per_vesicle_distance_to_AZ_manComp() + compute_per_vesicle_distance_to_filteredAZ() if __name__ == "__main__": From db89b441b25770ec298bbe51eb8d82233b92a0f6 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Sat, 23 Nov 2024 14:43:53 +0100 Subject: [PATCH 15/23] evaluation of AZ seg --- scripts/cooper/analysis/run_analysis_1.py | 11 +- .../run_spatial_distribution_analysis.py | 2 +- scripts/cooper/training/evaluate_AZ.py | 107 ++++++++++++++++++ synaptic_reconstruction/imod/to_imod.py | 31 +++++ 4 files changed, 144 insertions(+), 7 deletions(-) create mode 100644 scripts/cooper/training/evaluate_AZ.py diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 3afde5dc..94d0a625 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -11,9 +11,9 @@ from tqdm import tqdm from synaptic_reconstruction.imod.to_imod import convert_segmentation_to_spheres -DATA_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" # noqa -PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg_manComp" # noqa -RESULT_FOLDER = "./analysis_results/AZ_intersect_manualCompartment" +DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa +RESULT_FOLDER = "./analysis_results/AZ_filtered_autoComp" def get_compartment_with_max_overlap(compartments, vesicles): """ @@ -182,9 +182,8 @@ def compute_sizes_for_all_tomorams_autoComp(): result_df.to_csv(output_path, index=False) def main(): - compute_sizes_for_all_tomorams_manComp() - #compute_sizes_for_all_tomorams_autoComp() - + #compute_sizes_for_all_tomorams_manComp() + compute_sizes_for_all_tomorams_autoComp() if __name__ == "__main__": main() diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index cdc4c0d8..edd8308e 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -188,7 +188,6 @@ def compute_per_vesicle_distance_to_filteredAZ(): # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) - def compute_per_vesicle_distance_to_AZ_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) @@ -248,5 +247,6 @@ def main(): compute_per_vesicle_distance_to_filteredAZ() + if __name__ == "__main__": main() diff --git a/scripts/cooper/training/evaluate_AZ.py b/scripts/cooper/training/evaluate_AZ.py new file mode 100644 index 00000000..fc32214e --- /dev/null +++ b/scripts/cooper/training/evaluate_AZ.py @@ -0,0 +1,107 @@ +import argparse +import os + +import h5py +import pandas as pd +import numpy as np + +from elf.evaluation.dice import dice_score + +def extract_gt_bounding_box(segmentation, gt, halo=[20, 320, 320]): + # Find the bounding box for the ground truth + bb = np.where(gt > 0) + bb = tuple(slice( + max(int(b.min() - ha), 0), # Ensure indices are not below 0 + min(int(b.max() + ha), sh) # Ensure indices do not exceed shape dimensions + ) for b, sh, ha in zip(bb, gt.shape, halo)) + + # Apply the bounding box to both segmentations + segmentation_cropped = segmentation[bb] + gt_cropped = gt[bb] + + return segmentation_cropped, gt_cropped + +def evaluate(labels, segmentation): + assert labels.shape == segmentation.shape + score = dice_score(segmentation, labels) + return score + +def evaluate_file(labels_path, segmentation_path, model_name, crop= False): + print(f"Evaluate labels {labels_path} and vesicles {segmentation_path}") + + ds_name = os.path.basename(os.path.dirname(labels_path)) + tomo = os.path.basename(labels_path) + + #get the labels and segmentation + with h5py.File(labels_path) as label_file: + gt = label_file["/labels/AZ"][:] + + with h5py.File(segmentation_path) as seg_file: + segmentation = seg_file["/AZ/segment_from_AZmodel_v3"][:] + + if crop: + print("cropping the annotation and segmentation") + segmentation, gt = extract_gt_bounding_box(segmentation, gt) + + # Evaluate the match of ground truth and segmentation + dice_score = evaluate(gt, segmentation) + + # Store results + result_folder = "/user/muth9/u12095/synaptic-reconstruction/scripts/cooper/evaluation_results" + os.makedirs(result_folder, exist_ok=True) + result_path = os.path.join(result_folder, f"evaluation_{model_name}.csv") + print("Evaluation results are saved to:", result_path) + + # Load existing results if the file exists + if os.path.exists(result_path): + results = pd.read_csv(result_path) + else: + results = None + + # Create a new DataFrame for the current evaluation + res = pd.DataFrame( + [[ds_name, tomo, dice_score]], columns=["dataset", "tomogram", "dice_score"] + ) + + # Combine with existing results or initialize with the new results + if results is None: + results = res + else: + results = pd.concat([results, res]) + + # Save the results to the CSV file + results.to_csv(result_path, index=False) + +def evaluate_folder(labels_path, segmentation_path, model_name, crop = False): + print(f"Evaluating folder {segmentation_path}") + print(f"Using labels stored in {labels_path}") + + label_files = os.listdir(labels_path) + vesicles_files = os.listdir(segmentation_path) + + for vesicle_file in vesicles_files: + if vesicle_file in label_files: + + evaluate_file(os.path.join(labels_path, vesicle_file), os.path.join(segmentation_path, vesicle_file), model_name, crop) + + + +def main(): + + parser = argparse.ArgumentParser() + parser.add_argument("-l", "--labels_path", required=True) + parser.add_argument("-v", "--segmentation_path", required=True) + parser.add_argument("-n", "--model_name", required=True) + parser.add_argument("--crop", action="store_true", help="Crop around the annotation.") + args = parser.parse_args() + + segmentation_path = args.segmentation_path + if os.path.isdir(segmentation_path): + evaluate_folder(args.labels_path, segmentation_path, args.model_name, args.crop) + else: + evaluate_file(args.labels_path, segmentation_path, args.model_name, args.crop) + + + +if __name__ == "__main__": + main() diff --git a/synaptic_reconstruction/imod/to_imod.py b/synaptic_reconstruction/imod/to_imod.py index 7a984690..6b217aa4 100644 --- a/synaptic_reconstruction/imod/to_imod.py +++ b/synaptic_reconstruction/imod/to_imod.py @@ -121,6 +121,37 @@ def coords_and_rads(prop): rads = [re[1] for re in res] return np.array(coords), np.array(rads) + def coords_and_rads(prop): + seg_id = prop.label + + bbox = prop.bbox + bb = np.s_[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]] + mask = segmentation[bb] == seg_id + + if estimate_radius_2d: + dists = np.array([distance_transform_edt(ma, sampling=resolution[1:]) for ma in mask]) + else: + dists = distance_transform_edt(mask, sampling=resolution) + + max_coord = np.unravel_index(np.argmax(dists), mask.shape) + radius = dists[max_coord] * radius_factor + + offset = np.array(bbox[:3]) + coord = np.array(max_coord) + offset + return coord, radius, seg_id + + with futures.ThreadPoolExecutor(num_workers) as tp: + res = list(tqdm( + tp.map(coords_and_rads, props), disable=not verbose, total=len(props), + desc="Compute coordinates and radii" + )) + + coords = [re[0] for re in res] + rads = [re[1] for re in res] + label_indxes = [re[2] for re in res] + return np.array(coords), np.array(rads), np.array(label_indxes) + + def write_points_to_imod( coordinates: np.ndarray, From aa5d78e2c3257a653a576c5adb006b686aedcb7f Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Sat, 23 Nov 2024 16:38:35 +0100 Subject: [PATCH 16/23] clean up --- scripts/cooper/analysis/calc_AZ_area.py | 239 ------------------ scripts/cooper/analysis/run_analysis_1.py | 35 +-- .../run_spatial_distribution_analysis.py | 149 +++-------- 3 files changed, 52 insertions(+), 371 deletions(-) delete mode 100644 scripts/cooper/analysis/calc_AZ_area.py diff --git a/scripts/cooper/analysis/calc_AZ_area.py b/scripts/cooper/analysis/calc_AZ_area.py deleted file mode 100644 index 592b043e..00000000 --- a/scripts/cooper/analysis/calc_AZ_area.py +++ /dev/null @@ -1,239 +0,0 @@ -import h5py -import numpy as np -import os -import csv -from scipy.ndimage import binary_opening, median_filter,zoom, binary_closing -from skimage.measure import label, regionprops -from synaptic_reconstruction.morphology import compute_object_morphology -from skimage.morphology import ball -from scipy.spatial import ConvexHull -from skimage.draw import polygon - -def calculate_AZ_area_per_slice(AZ_slice, pixel_size_nm=1.554): - """ - Calculate the area of the AZ in a single 2D slice after applying error-reducing processing. - - Parameters: - - AZ_slice (numpy array): 2D array representing a single slice of the AZ segmentation. - - pixel_size_nm (float): Size of a pixel in nanometers. - - Returns: - - slice_area_nm2 (float): The area of the AZ in the slice in square nanometers. - """ - # Apply binary opening or median filter to reduce small segmentation errors - AZ_slice_filtered = binary_opening(AZ_slice, structure=np.ones((3, 3))).astype(int) - - # Calculate area in this slice - num_AZ_pixels = np.sum(AZ_slice_filtered == 1) - slice_area_nm2 = num_AZ_pixels * (pixel_size_nm ** 2) - - return slice_area_nm2 - -def calculate_total_AZ_area(tomo_path, pixel_size_nm=1.554): - """ - Calculate the total area of the AZ across all slices in a 3D tomogram file. - - Parameters: - - tomo_path (str): Path to the tomogram file (HDF5 format). - - pixel_size_nm (float): Size of a pixel in nanometers. - - Returns: - - total_AZ_area_nm2 (float): The total area of the AZ in square nanometers. - """ - with h5py.File(tomo_path, "r") as f: - AZ_intersect_seg = f["/AZ/compartment_AZ_intersection_manComp"][:] - - # Calculate the AZ area for each slice along the z-axis - total_AZ_area_nm2 = 0 - for z_slice in AZ_intersect_seg: - slice_area_nm2 = calculate_AZ_area_per_slice(z_slice, pixel_size_nm) - total_AZ_area_nm2 += slice_area_nm2 - - return total_AZ_area_nm2 - -def calculate_AZ_area_simple(tomo_path, pixel_size_nm=1.554): - """ - Calculate the volume of the AZ (active zone) in a 3D tomogram file. - - Parameters: - - tomo_path (str): Path to the tomogram file (HDF5 format). - - pixel_size_nm (float): Size of a pixel in nanometers (default is 1.554 nm). - - Returns: - - AZ_volume_nm3 (float): The volume of the AZ in cubic nanometers. - """ - # Open the file and read the AZ intersection segmentation data - with h5py.File(tomo_path, "r") as f: - AZ_intersect_seg = f["/AZ/compartment_AZ_intersection_manComp"][:] - - # Count voxels with label = 1 - num_AZ_voxels = np.sum(AZ_intersect_seg == 1) - - # Calculate the volume in cubic nanometers - AZ_area_nm2 = num_AZ_voxels * (pixel_size_nm ** 2) - - return AZ_area_nm2 - -def calculate_AZ_surface(tomo_path, pixel_size_nm=1.554): - with h5py.File(tomo_path, "r") as f: - #AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] - AZ_seg = f["/filtered_az"][:] - - # Apply binary closing to smooth the segmented regions - struct_elem = ball(1) # Use a small 3D structuring element - AZ_seg_smoothed = binary_closing(AZ_seg > 0, structure=struct_elem, iterations=20) - - labeled_seg = label(AZ_seg_smoothed) - - regions = regionprops(labeled_seg) - if regions: - # Sort regions by area and get the label of the largest region - largest_region = max(regions, key=lambda r: r.area) - largest_label = largest_region.label - - largest_component_mask = (labeled_seg == largest_label) - AZ_seg_filtered = largest_component_mask.astype(np.uint8) - - else: - # If no regions found, return an empty array - AZ_seg_filtered = np.zeros_like(AZ_seg_interp, dtype=np.uint8) - - morphology_data = compute_object_morphology(AZ_seg_filtered, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) - surface_column = "surface [nm^2]" #if resolution is not None else "surface [pixel^2]" - surface_area = morphology_data[surface_column].iloc[0] - - return surface_area - -def calculate_AZ_surface_simple(tomo_path, pixel_size_nm=1.554): - with h5py.File(tomo_path, "r") as f: - AZ_seg = f["/labels/AZ"][:] - - morphology_data = compute_object_morphology(AZ_seg, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) - surface_column = "surface [nm^2]" #if resolution is not None else "surface [pixel^2]" - surface_area = morphology_data[surface_column].iloc[0] - - return surface_area - -def calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm=1.554): - with h5py.File(tomo_path, "r") as f: - AZ_seg = f["/AZ/segment_from_AZmodel_v3"][:] - - # Apply binary closing to smooth the segmented regions - struct_elem = ball(1) # Use a small 3D structuring element - AZ_seg_smoothed = binary_closing(AZ_seg > 0, structure=struct_elem, iterations=10) - - labeled_seg = label(AZ_seg_smoothed) - - regions = regionprops(labeled_seg) - if regions: - # Sort regions by area and get the label of the largest region - largest_region = max(regions, key=lambda r: r.area) - largest_label = largest_region.label - - largest_component_mask = (labeled_seg == largest_label) - AZ_seg_filtered = largest_component_mask.astype(np.uint8) - AZ_seg = AZ_seg_filtered - # Extract coordinates of non-zero points - points = np.argwhere(AZ_seg > 0) # Get the coordinates of non-zero (foreground) pixels - - if points.shape[0] < 4: - # ConvexHull requires at least 4 points in 3D to form a valid hull - AZ_seg_filtered = np.zeros_like(AZ_seg, dtype=np.uint8) - else: - # Apply ConvexHull to the points - hull = ConvexHull(points) - - # Create a binary mask for the convex hull - convex_hull_mask = np.zeros_like(AZ_seg, dtype=bool) - - # Iterate over each simplex (facet) of the convex hull and fill in the polygon - for simplex in hull.simplices: - # For each face of the convex hull, extract the vertices and convert to a 2D polygon - polygon_coords = points[simplex] - rr, cc = polygon(polygon_coords[:, 0], polygon_coords[:, 1]) - convex_hull_mask[rr, cc] = True - - # Optional: Label the convex hull mask - labeled_seg = label(convex_hull_mask) - regions = regionprops(labeled_seg) - - if regions: - # Sort regions by area and get the label of the largest region - largest_region = max(regions, key=lambda r: r.area) - largest_label = largest_region.label - - largest_component_mask = (labeled_seg == largest_label) - AZ_seg_filtered = largest_component_mask.astype(np.uint8) - - else: - AZ_seg_filtered = np.zeros_like(AZ_seg, dtype=np.uint8) - - # Calculate surface area - morphology_data = compute_object_morphology(AZ_seg_filtered, "AZ Structure", resolution=(pixel_size_nm, pixel_size_nm, pixel_size_nm)) - surface_column = "surface [nm^2]" - surface_area = morphology_data[surface_column].iloc[0] - - return surface_area - -def process_datasets(folder_path, output_csv="AZ_areas.csv", pixel_size_nm=1.554): - """ - Process all tomograms in multiple datasets within a folder and save results to a CSV. - - Parameters: - - folder_path (str): Path to the folder containing dataset folders with tomograms. - - output_csv (str): Filename for the output CSV file. - - pixel_size_nm (float): Size of a pixel in nanometers. - """ - results = [] - - # Iterate over each dataset folder - for dataset_name in os.listdir(folder_path): - dataset_path = os.path.join(folder_path, dataset_name) - - # Check if it's a directory (skip files in the main folder) - if not os.path.isdir(dataset_path): - continue - - # Iterate over each tomogram file in the dataset folder - for tomo_file in os.listdir(dataset_path): - tomo_path = os.path.join(dataset_path, tomo_file) - - # Check if the file is an HDF5 file (optional) - if tomo_file.endswith(".h5") or tomo_file.endswith(".hdf5"): - try: - # Calculate AZ area - #AZ_area = calculate_total_AZ_area(tomo_path, pixel_size_nm) - #AZ_area = calculate_AZ_area_simple(tomo_path, pixel_size_nm) - #AZ_surface_area = calculate_AZ_surface(tomo_path, pixel_size_nm) - #AZ_surface_area = calculate_AZ_surface_convexHull(tomo_path, pixel_size_nm) - AZ_surface_area = calculate_AZ_surface_simple(tomo_path, pixel_size_nm) - # Append results to list - results.append({ - "Dataset": dataset_name, - "Tomogram": tomo_file, - "AZ_surface_area": AZ_surface_area - }) - except Exception as e: - print(f"Error processing {tomo_file} in {dataset_name}: {e}") - - # Write results to a CSV file - with open(output_csv, mode="w", newline="") as csvfile: - fieldnames = ["Dataset", "Tomogram", "AZ_surface_area"] - writer = csv.DictWriter(csvfile, fieldnames=fieldnames) - - writer.writeheader() - for result in results: - writer.writerow(result) - - print(f"Results saved to {output_csv}") - -def main(): - # Define the path to the folder containing dataset folders - folder_path = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/exported/" - output_csv = "./analysis_results/manual_AZ_exported/AZ_surface_area.csv" - # Call the function to process datasets and save results - process_datasets(folder_path, output_csv = output_csv) - -# Call main -if __name__ == "__main__": - main() diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_analysis_1.py index 94d0a625..abad4405 100644 --- a/scripts/cooper/analysis/run_analysis_1.py +++ b/scripts/cooper/analysis/run_analysis_1.py @@ -1,7 +1,3 @@ -# This is the code for the first analysis for the cooper data. -# Here, we only compute the vesicle numbers and size distributions for the STEM tomograms -# in the 04 dataset. - import os from glob import glob @@ -37,13 +33,10 @@ def get_compartment_with_max_overlap(compartments, vesicles): # Iterate over each compartment and calculate the overlap with vesicles for compartment_label in unique_compartments: - # Create a binary mask for the current compartment compartment_mask = compartments == compartment_label vesicle_mask = vesicles > 0 intersection = np.logical_and(compartment_mask, vesicle_mask) - - # Calculate the number of overlapping voxels overlap_count = np.sum(intersection) # Track the compartment with the most overlap in terms of voxel count @@ -51,14 +44,13 @@ def get_compartment_with_max_overlap(compartments, vesicles): max_overlap_count = overlap_count best_compartment = compartment_label - # Create the final mask for the compartment with the most overlap final_mask = compartments == best_compartment return final_mask -# We compute the sizes for all vesicles in the compartment masks. +# We compute the sizes for all vesicles in the MANUALLY ANNOTATED compartment masks. # We use the same logic in the size computation as for the vesicle extraction to IMOD, -# including the radius correction factor. +# including the radius correction factor. --> not needed here # The number of vesicles is automatically computed as the length of the size list. def compute_sizes_for_all_tomorams_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) @@ -66,8 +58,6 @@ def compute_sizes_for_all_tomorams_manComp(): resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset radius_factor = 1 estimate_radius_2d = True - - # Dictionary to hold the results for each dataset and category (CTRL or DKO) dataset_results = {} tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) @@ -79,11 +69,9 @@ def compute_sizes_for_all_tomorams_manComp(): # Determine if the tomogram is 'CTRL' or 'DKO' category = "CTRL" if "CTRL" in fname else "DKO" - # Initialize a new dictionary entry for each dataset and category if not already present if ds_name not in dataset_results: dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} - - # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name][category]: continue @@ -93,6 +81,7 @@ def compute_sizes_for_all_tomorams_manComp(): input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") assert os.path.exists(input_path), input_path + # Load the compartment mask from the tomogram with h5py.File(input_path, "r") as f: mask = f["labels/compartment"][:] @@ -102,30 +91,30 @@ def compute_sizes_for_all_tomorams_manComp(): segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d ) - # Add sizes to the dataset dictionary under the appropriate category + dataset_results[ds_name][category][fname] = sizes # Save each dataset's results into separate CSV files for CTRL and DKO tomograms for ds_name, categories in dataset_results.items(): for category, tomogram_data in categories.items(): - # Sort tomograms by name within the category sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() - # Define the output file path output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}_rf1.csv") # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) +# We compute the sizes for all vesicles in the AUTOMATIC SEGMENTED compartment masks. +# We use the same logic in the size computation as for the vesicle extraction to IMOD, +# including the radius correction factor. --> not needed here +# The number of vesicles is automatically computed as the length of the size list. def compute_sizes_for_all_tomorams_autoComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset radius_factor = 1 estimate_radius_2d = True - - # Dictionary to hold the results for each dataset and category (CTRL or DKO) dataset_results = {} tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) @@ -137,11 +126,9 @@ def compute_sizes_for_all_tomorams_autoComp(): # Determine if the tomogram is 'CTRL' or 'DKO' category = "CTRL" if "CTRL" in fname else "DKO" - # Initialize a new dictionary entry for each dataset and category if not already present if ds_name not in dataset_results: dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} - # Skip if this tomogram already exists in the dataset dictionary if fname in dataset_results[ds_name][category]: continue @@ -151,6 +138,7 @@ def compute_sizes_for_all_tomorams_autoComp(): input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") assert os.path.exists(input_path), input_path + # Load the compartment mask from the tomogram with h5py.File(input_path, "r") as f: compartments = f["/compartments/segment_from_3Dmodel_v2"][:] @@ -165,17 +153,14 @@ def compute_sizes_for_all_tomorams_autoComp(): segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d ) - # Add sizes to the dataset dictionary under the appropriate category dataset_results[ds_name][category][fname] = sizes # Save each dataset's results into separate CSV files for CTRL and DKO tomograms for ds_name, categories in dataset_results.items(): for category, tomogram_data in categories.items(): - # Sort tomograms by name within the category sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() - # Define the output file path output_path = os.path.join(RESULT_FOLDER, f"size_analysis_for_{ds_name}_{category}_rf1.csv") # Save the DataFrame to CSV diff --git a/scripts/cooper/analysis/run_spatial_distribution_analysis.py b/scripts/cooper/analysis/run_spatial_distribution_analysis.py index edd8308e..69434844 100644 --- a/scripts/cooper/analysis/run_spatial_distribution_analysis.py +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -9,7 +9,8 @@ DATA_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa PREDICTION_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/segmentation/for_spatial_distribution_analysis/final_Imig2014_seg/" # noqa RESULT_FOLDER = "./analysis_results/AZ_filtered_autoComp" - +AZ_PATH = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/az_seg_filtered" + def get_compartment_with_max_overlap(compartments, vesicles): """ @@ -33,13 +34,10 @@ def get_compartment_with_max_overlap(compartments, vesicles): # Iterate over each compartment and calculate the overlap with vesicles for compartment_label in unique_compartments: - # Create a binary mask for the current compartment compartment_mask = compartments == compartment_label vesicle_mask = vesicles > 0 intersection = np.logical_and(compartment_mask, vesicle_mask) - - # Calculate the number of overlapping voxels overlap_count = np.sum(intersection) # Track the compartment with the most overlap in terms of voxel count @@ -47,85 +45,20 @@ def get_compartment_with_max_overlap(compartments, vesicles): max_overlap_count = overlap_count best_compartment = compartment_label - # Create the final mask for the compartment with the most overlap final_mask = compartments == best_compartment return final_mask -# We compute the distances for all vesicles in the compartment masks to the AZ. -# We use the same different resolution, depending on dataset. +# We compute the distances for all vesicles in the AUTOMATIC SEGMENTED compartment masks to the AZ. +# We use different resolution, depending on dataset. # The closest distance is calculated, i.e., the closest point on the outer membrane of the vesicle to the AZ. -def compute_per_vesicle_distance_to_AZ(): - os.makedirs(RESULT_FOLDER, exist_ok=True) - - resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - - # Dictionary to hold the results for each dataset - dataset_results = {} - - tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) - for tomo in tqdm(tomograms): - ds_name, fname = os.path.split(tomo) - ds_name = os.path.split(ds_name)[1] - fname = os.path.splitext(fname)[0] - - # Initialize a new dictionary entry for each dataset if not already present - if ds_name not in dataset_results: - dataset_results[ds_name] = {} - - # Skip if this tomogram already exists in the dataset dictionary - if fname in dataset_results[ds_name]: - continue - - # Load the vesicle segmentation from the predictions - with h5py.File(tomo, "r") as f: - segmentation = f["/vesicles/segment_from_combined_vesicles"][:] - segmented_object = f["/AZ/compartment_AZ_intersection"][:] - #if AZ intersect is small, compartment seg didn't align with AZ so we use the normal AZ and not intersect - if (segmented_object == 0).all() or np.sum(segmented_object == 1) < 2000: - segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] - - input_path = os.path.join(DATA_ROOT, ds_name, f"{fname}.h5") - assert os.path.exists(input_path), input_path - - # Load the compartment mask from the tomogram - with h5py.File(input_path, "r") as f: - compartments = f["/compartments/segment_from_3Dmodel_v2"][:] - mask = get_compartment_with_max_overlap(compartments, segmentation) - - #if more than half of the vesicles (approximation, its checking pixel and not label) would get filtered by mask it means the compartment seg didn't work and thus we won't use the mask - if np.sum(segmentation[mask == 0] > 0) > (0.5 * np.sum(segmentation > 0)): - print("using no mask") - else: - segmentation[mask == 0] = 0 - distances, _, _, _ = measure_segmentation_to_object_distances( - segmentation, segmented_object=segmented_object, resolution=resolution - ) +def compute_per_vesicle_distance_to_AZ_autoComp(separate_AZseg=False): - # Add distances to the dataset dictionary under the tomogram name - dataset_results[ds_name][fname] = distances - - # Save each dataset's results to a single CSV file - for ds_name, tomogram_data in dataset_results.items(): - # Create a DataFrame where each column is a tomogram's distances - result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() - - # Define the output file path - output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}.csv") - - # Save the DataFrame to CSV - result_df.to_csv(output_path, index=False) - -def compute_per_vesicle_distance_to_filteredAZ(): - filtered_AZ_path = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/20241102_TOMO_DATA_Imig2014/az_seg_filtered" os.makedirs(RESULT_FOLDER, exist_ok=True) - resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - - # Dictionary to hold the results for each dataset and category (CTRL or DKO) dataset_results = {} - tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): ds_name, fname = os.path.split(tomo) ds_name = os.path.split(ds_name)[1] @@ -134,23 +67,27 @@ def compute_per_vesicle_distance_to_filteredAZ(): # Determine if the tomogram is 'CTRL' or 'DKO' category = "CTRL" if "CTRL" in fname else "DKO" - # Initialize a new dictionary entry for each dataset and category if not already present if ds_name not in dataset_results: dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} - - # Skip if this tomogram already exists in the dataset dictionary + if fname in dataset_results[ds_name][category]: continue - #Load the AZ segmentations - AZ_path = os.path.join(filtered_AZ_path, ds_name, f"{fname}.h5") - with h5py.File(AZ_path, "r") as f: - segmented_object = f["/filtered_az"][:] # Load the vesicle segmentation from the predictions with h5py.File(tomo, "r") as f: segmentation = f["/vesicles/segment_from_combined_vesicles"][:] - + + #Check if AZ seg is stored in a different tomo or same + if separate_AZseg: + print(f"using AZ segmentation from {AZ_PATH}") + #Load the AZ segmentations + AZ_path = os.path.join(AZ_PATH, ds_name, f"{fname}.h5") + with h5py.File(AZ_path, "r") as f_AZ: + segmented_object = f_AZ["/thin_az"][:] + else: + segmented_object = f["/AZ/compartment_AZ_intersection"][:] + #if AZ intersect is small, compartment seg didn't align with AZ so we use the normal AZ and not intersect if (segmented_object == 0).all() or np.sum(segmented_object == 1) < 2000: segmented_object = f["/AZ/segment_from_AZmodel_v3"][:] @@ -168,6 +105,7 @@ def compute_per_vesicle_distance_to_filteredAZ(): print("using no mask") else: segmentation[mask == 0] = 0 + distances, _, _, _ = measure_segmentation_to_object_distances( segmentation, segmented_object=segmented_object, resolution=resolution ) @@ -178,36 +116,35 @@ def compute_per_vesicle_distance_to_filteredAZ(): # Save each dataset's results into separate CSV files for CTRL and DKO tomograms for ds_name, categories in dataset_results.items(): for category, tomogram_data in categories.items(): - # Sort tomograms by name within the category sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() - - # Define the output file path output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}_{category}.csv") # Save the DataFrame to CSV result_df.to_csv(output_path, index=False) +# We compute the distances for all vesicles in the MANUALLY ANNOTATED compartment masks to the AZ. +# We use different resolution, depending on dataset. +# The closest distance is calculated, i.e., the closest point on the outer membrane of the vesicle to the AZ. def compute_per_vesicle_distance_to_AZ_manComp(): os.makedirs(RESULT_FOLDER, exist_ok=True) resolution = (1.554,) * 3 # Change for each dataset #1.554 for Munc and snap #0.8681 for 04 dataset - - # Dictionary to hold the results for each dataset dataset_results = {} - tomograms = sorted(glob(os.path.join(PREDICTION_ROOT, "**/*.h5"), recursive=True)) + for tomo in tqdm(tomograms): ds_name, fname = os.path.split(tomo) ds_name = os.path.split(ds_name)[1] fname = os.path.splitext(fname)[0] - # Initialize a new dictionary entry for each dataset if not already present - if ds_name not in dataset_results: - dataset_results[ds_name] = {} + # Determine if the tomogram is 'CTRL' or 'DKO' + category = "CTRL" if "CTRL" in fname else "DKO" - # Skip if this tomogram already exists in the dataset dictionary - if fname in dataset_results[ds_name]: + if ds_name not in dataset_results: + dataset_results[ds_name] = {'CTRL': {}, 'DKO': {}} + + if fname in dataset_results[ds_name][category]: continue # Load the vesicle segmentation from the predictions @@ -223,28 +160,26 @@ def compute_per_vesicle_distance_to_AZ_manComp(): mask = f["/labels/compartment"][:] segmentation[mask == 0] = 0 + distances, _, _, _ = measure_segmentation_to_object_distances( segmentation, segmented_object=segmented_object, resolution=resolution ) - # Add distances to the dataset dictionary under the tomogram name - dataset_results[ds_name][fname] = distances - - # Save each dataset's results to a single CSV file - for ds_name, tomogram_data in dataset_results.items(): - # Create a DataFrame where each column is a tomogram's distances - result_df = pd.DataFrame.from_dict(tomogram_data, orient='index').transpose() - - # Define the output file path - output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}.csv") - - # Save the DataFrame to CSV - result_df.to_csv(output_path, index=False) + # Add distances to the dataset dictionary under the appropriate category + dataset_results[ds_name][category][fname] = distances + # Save each dataset's results into separate CSV files for CTRL and DKO tomograms + for ds_name, categories in dataset_results.items(): + for category, tomogram_data in categories.items(): + sorted_data = dict(sorted(tomogram_data.items())) # Sort by tomogram names + result_df = pd.DataFrame.from_dict(sorted_data, orient='index').transpose() + output_path = os.path.join(RESULT_FOLDER, f"spatial_distribution_analysis_for_{ds_name}_{category}.csv") + + # Save the DataFrame to CSV + result_df.to_csv(output_path, index=False) def main(): - #compute_per_vesicle_distance_to_AZ() + compute_per_vesicle_distance_to_AZ_autoComp(separate_AZseg=False) #compute_per_vesicle_distance_to_AZ_manComp() - compute_per_vesicle_distance_to_filteredAZ() From 20e429b3b7b5e000a3aeb48c603326ceb2be1c5c Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Sat, 23 Nov 2024 16:40:26 +0100 Subject: [PATCH 17/23] clean up --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 0377c4a0..d9550532 100644 --- a/.gitignore +++ b/.gitignore @@ -14,4 +14,5 @@ scripts/rizzoli/upsample_data.py scripts/cooper/training/find_rec_testset.py scripts/rizzoli/combine_2D_slices.py scripts/rizzoli/combine_2D_slices_raw.py -scripts/cooper/remove_h5key.py \ No newline at end of file +scripts/cooper/remove_h5key.py +scripts/cooper/analysis/calc_AZ_area.py \ No newline at end of file From 19f618e388a510884c82475622b66074237c46ba Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Sat, 23 Nov 2024 16:42:00 +0100 Subject: [PATCH 18/23] clean up --- .../cooper/analysis/{run_analysis_1.py => run_size_analysis.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename scripts/cooper/analysis/{run_analysis_1.py => run_size_analysis.py} (100%) diff --git a/scripts/cooper/analysis/run_analysis_1.py b/scripts/cooper/analysis/run_size_analysis.py similarity index 100% rename from scripts/cooper/analysis/run_analysis_1.py rename to scripts/cooper/analysis/run_size_analysis.py From 622da1e618953d385e6a93285e7e167a049bddc5 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Wed, 27 Nov 2024 12:00:34 +0100 Subject: [PATCH 19/23] update AZ evaluation --- scripts/cooper/training/evaluate_AZ.py | 63 ++++++++++--- scripts/cooper/training/postprocess_AZ.py | 107 ++++++++++++++++++++++ 2 files changed, 158 insertions(+), 12 deletions(-) create mode 100644 scripts/cooper/training/postprocess_AZ.py diff --git a/scripts/cooper/training/evaluate_AZ.py b/scripts/cooper/training/evaluate_AZ.py index fc32214e..dbf8d675 100644 --- a/scripts/cooper/training/evaluate_AZ.py +++ b/scripts/cooper/training/evaluate_AZ.py @@ -26,7 +26,37 @@ def evaluate(labels, segmentation): score = dice_score(segmentation, labels) return score -def evaluate_file(labels_path, segmentation_path, model_name, crop= False): +def compute_precision(ground_truth, segmentation): + """ + Computes the Precision score for 3D arrays representing the ground truth and segmentation. + + Parameters: + - ground_truth (np.ndarray): 3D binary array where 1 represents the ground truth region. + - segmentation (np.ndarray): 3D binary array where 1 represents the predicted segmentation region. + + Returns: + - precision (float): The precision score, or 0 if the segmentation is empty. + """ + assert ground_truth.shape == segmentation.shape + # Ensure inputs are binary arrays + ground_truth = (ground_truth > 0).astype(np.int32) + segmentation = (segmentation > 0).astype(np.int32) + + # Compute intersection: overlap between segmentation and ground truth + intersection = np.sum(segmentation * ground_truth) + + # Compute total predicted (segmentation region) + total_predicted = np.sum(segmentation) + + # Handle case where there are no predictions + if total_predicted == 0: + return 0.0 # Precision is undefined; returning 0 + + # Calculate precision + precision = intersection / total_predicted + return precision + +def evaluate_file(labels_path, segmentation_path, model_name, crop= False, precision_score=False): print(f"Evaluate labels {labels_path} and vesicles {segmentation_path}") ds_name = os.path.basename(os.path.dirname(labels_path)) @@ -34,22 +64,25 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False): #get the labels and segmentation with h5py.File(labels_path) as label_file: - gt = label_file["/labels/AZ"][:] + gt = label_file["/labels/thin_az"][:] with h5py.File(segmentation_path) as seg_file: - segmentation = seg_file["/AZ/segment_from_AZmodel_v3"][:] + segmentation = seg_file["/AZ/thin_az"][:] if crop: print("cropping the annotation and segmentation") segmentation, gt = extract_gt_bounding_box(segmentation, gt) # Evaluate the match of ground truth and segmentation - dice_score = evaluate(gt, segmentation) + if precision_score: + precision = compute_precision(gt, segmentation) + else: + dice_score = evaluate(gt, segmentation) # Store results result_folder = "/user/muth9/u12095/synaptic-reconstruction/scripts/cooper/evaluation_results" os.makedirs(result_folder, exist_ok=True) - result_path = os.path.join(result_folder, f"evaluation_{model_name}.csv") + result_path = os.path.join(result_folder, f"evaluation_{model_name}_dice_thinpred_thinanno.csv") print("Evaluation results are saved to:", result_path) # Load existing results if the file exists @@ -59,9 +92,14 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False): results = None # Create a new DataFrame for the current evaluation - res = pd.DataFrame( - [[ds_name, tomo, dice_score]], columns=["dataset", "tomogram", "dice_score"] - ) + if precision_score: + res = pd.DataFrame( + [[ds_name, tomo, precision]], columns=["dataset", "tomogram", "precision"] + ) + else: + res = pd.DataFrame( + [[ds_name, tomo, dice_score]], columns=["dataset", "tomogram", "dice_score"] + ) # Combine with existing results or initialize with the new results if results is None: @@ -72,7 +110,7 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False): # Save the results to the CSV file results.to_csv(result_path, index=False) -def evaluate_folder(labels_path, segmentation_path, model_name, crop = False): +def evaluate_folder(labels_path, segmentation_path, model_name, crop = False, precision_score=False): print(f"Evaluating folder {segmentation_path}") print(f"Using labels stored in {labels_path}") @@ -82,7 +120,7 @@ def evaluate_folder(labels_path, segmentation_path, model_name, crop = False): for vesicle_file in vesicles_files: if vesicle_file in label_files: - evaluate_file(os.path.join(labels_path, vesicle_file), os.path.join(segmentation_path, vesicle_file), model_name, crop) + evaluate_file(os.path.join(labels_path, vesicle_file), os.path.join(segmentation_path, vesicle_file), model_name, crop, precision_score) @@ -93,13 +131,14 @@ def main(): parser.add_argument("-v", "--segmentation_path", required=True) parser.add_argument("-n", "--model_name", required=True) parser.add_argument("--crop", action="store_true", help="Crop around the annotation.") + parser.add_argument("--precision", action="store_true", help="Calculate precision score.") args = parser.parse_args() segmentation_path = args.segmentation_path if os.path.isdir(segmentation_path): - evaluate_folder(args.labels_path, segmentation_path, args.model_name, args.crop) + evaluate_folder(args.labels_path, segmentation_path, args.model_name, args.crop, args.precision) else: - evaluate_file(args.labels_path, segmentation_path, args.model_name, args.crop) + evaluate_file(args.labels_path, segmentation_path, args.model_name, args.crop, args.precision) diff --git a/scripts/cooper/training/postprocess_AZ.py b/scripts/cooper/training/postprocess_AZ.py new file mode 100644 index 00000000..e2b849e0 --- /dev/null +++ b/scripts/cooper/training/postprocess_AZ.py @@ -0,0 +1,107 @@ +import os +from glob import glob +import argparse + +import h5py +import numpy as np +from tqdm import tqdm +from scipy.ndimage import binary_closing +from skimage.measure import label +from synaptic_reconstruction.ground_truth.shape_refinement import edge_filter +from synaptic_reconstruction.morphology import skeletonize_object + + + +def filter_az(path, output_path): + """Filter the active zone (AZ) data from the HDF5 file.""" + ds, fname = os.path.split(path) + dataset_name = os.path.basename(ds) + out_file_path = os.path.join(output_path, "postprocessed_AZ", dataset_name, fname) + + os.makedirs(os.path.dirname(out_file_path), exist_ok=True) + + if os.path.exists(out_file_path): + return + + with h5py.File(path, "r") as f: + raw = f["raw"][:] + az = f["AZ/segment_from_AZmodel_v3"][:] + + hmap = edge_filter(raw, sigma=1.0, method="sato", per_slice=True, n_threads=8) + + # Filter the active zone by combining a bunch of things: + # 1. Find a mask with high values in the ridge filter. + threshold_hmap = 0.5 + az_filtered = hmap > threshold_hmap + # 2. Intersect it with the active zone predictions. + az_filtered = np.logical_and(az_filtered, az) + + # Postprocessing of the filtered active zone: + # 1. Apply connected components and only keep the largest component. + az_filtered = label(az_filtered) + ids, sizes = np.unique(az_filtered, return_counts=True) + ids, sizes = ids[1:], sizes[1:] + az_filtered = (az_filtered == ids[np.argmax(sizes)]).astype("uint8") + # 2. Apply binary closing. + az_filtered = np.logical_or(az_filtered, binary_closing(az_filtered, iterations=4)).astype("uint8") + + # Save the result. + with h5py.File(out_file_path, "a") as f: + f.create_dataset("AZ/filtered_az", data=az_filtered, compression="gzip") + + +def process_az(path, view=False): + """Skeletonize the filtered AZ data to obtain a 1D representation.""" + key = "AZ/thin_az" + with h5py.File(path, "r") as f: + if key in f and not view: + return + az_seg = f["AZ/filtered_az"][:] + + az_thin = skeletonize_object(az_seg) + + if view: + import napari + ds, fname = os.path.split(path) + raw_path = os.path.join(ROOT, ds, fname) + with h5py.File(raw_path, "r") as f: + raw = f["raw"][:] + v = napari.Viewer() + v.add_image(raw) + v.add_labels(az_seg) + v.add_labels(az_thin) + napari.run() + else: + with h5py.File(path, "a") as f: + f.create_dataset(key, data=az_thin, compression="gzip") + + +def filter_all_azs(input_path, output_path): + """Apply filtering to all AZ data in the specified directory.""" + files = sorted(glob(os.path.join(input_path, "**/*.h5"), recursive=True)) + for ff in tqdm(files, desc="Filtering AZ segmentations"): + filter_az(ff, output_path) + + +def process_all_azs(output_path): + """Apply skeletonization to all filtered AZ data.""" + files = sorted(glob(os.path.join(output_path, "postprocessed_AZ", "**/*.h5"), recursive=True)) + for ff in tqdm(files, desc="Thinning AZ segmentations"): + process_az(ff, view=False) + + +def main(): + parser = argparse.ArgumentParser(description="Filter and process AZ data.") + parser.add_argument("input_path", type=str, help="Path to the root directory containing datasets.") + parser.add_argument("output_path", type=str, help="Path to the root directory for saving processed data.") + args = parser.parse_args() + + input_path = args.input_path + output_path = args.output_path + + filter_all_azs(input_path, output_path) + process_all_azs(output_path) + + +if __name__ == "__main__": + main() From 686b018fc4a2af32e8b7a9e8747e9708ce9d2d66 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Thu, 28 Nov 2024 21:28:44 +0100 Subject: [PATCH 20/23] erosion dilation filtering of AZ --- scripts/cooper/training/filter_AZ.py | 67 ++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 scripts/cooper/training/filter_AZ.py diff --git a/scripts/cooper/training/filter_AZ.py b/scripts/cooper/training/filter_AZ.py new file mode 100644 index 00000000..78b8ba72 --- /dev/null +++ b/scripts/cooper/training/filter_AZ.py @@ -0,0 +1,67 @@ +import os +import h5py +import numpy as np +from scipy.ndimage import binary_erosion, binary_dilation, label + +def process_labels(label_file_path, erosion_structure=None, dilation_structure=None): + """ + Process the labels: perform erosion, find the largest connected component, + and perform dilation on it. + + Args: + label_file_path (str): Path to the HDF5 file containing the label data. + erosion_structure (ndarray, optional): Structuring element for erosion. + dilation_structure (ndarray, optional): Structuring element for dilation. + + Returns: + None: The processed data is saved back into the HDF5 file under a new key. + """ + with h5py.File(label_file_path, "r+") as label_file: + # Read the ground truth data + gt = label_file["/labels/filtered_az"][:] + + # Perform binary erosion + eroded = binary_erosion(gt, structure=erosion_structure) + + # Label connected components + labeled_array, num_features = label(eroded) + + # Identify the largest connected component + if num_features > 0: + largest_component_label = np.argmax(np.bincount(labeled_array.flat, weights=eroded.flat)[1:]) + 1 + largest_component = (labeled_array == largest_component_label) + else: + largest_component = np.zeros_like(gt, dtype=bool) + + # Perform binary dilation on the largest connected component + dilated = binary_dilation(largest_component, structure=dilation_structure) + + # Save the result back into the HDF5 file + if "labels/erosion_filtered_az" in label_file: + del label_file["labels/erosion_filtered_az"] # Remove if it already exists + label_file.create_dataset("labels/erosion_filtered_az", data=dilated.astype(np.uint8), compression="gzip") + +def process_folder(folder_path, erosion_structure=None, dilation_structure=None): + """ + Process all HDF5 files in a folder. + + Args: + folder_path (str): Path to the folder containing HDF5 files. + erosion_structure (ndarray, optional): Structuring element for erosion. + dilation_structure (ndarray, optional): Structuring element for dilation. + + Returns: + None + """ + for file_name in os.listdir(folder_path): + if file_name.endswith(".h5") or file_name.endswith(".hdf5"): + label_file_path = os.path.join(folder_path, file_name) + print(f"Processing {label_file_path}...") + process_labels(label_file_path, erosion_structure, dilation_structure) + +# Example usage +if __name__ == "__main__": + folder_path = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_v2/postprocessed_AZ/12_chemical_fix_cryopreparation" # Replace with the path to your folder + erosion_structure = np.ones((3, 3, 3)) # Example structuring element + dilation_structure = np.ones((3, 3, 3)) # Example structuring element + process_folder(folder_path, erosion_structure, dilation_structure) From 6b54e4aaeb265f119220f8b20410377b19b182a5 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Mon, 31 Mar 2025 14:16:08 +0200 Subject: [PATCH 21/23] stuff for revision --- scripts/cooper/AZ_segmentation_h5.py | 38 +++++++++++++------ scripts/cooper/training/evaluate_AZ.py | 6 +-- scripts/cooper/training/train_AZ.py | 15 ++++---- scripts/cooper/training/train_vesicles.py | 21 ++++------ scripts/cooper/vesicle_segmentation_h5.py | 22 ++++++++--- .../cryo/vesicles/train_domain_adaptation.py | 19 +++++----- 6 files changed, 71 insertions(+), 50 deletions(-) diff --git a/scripts/cooper/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py index da694c19..00d2de43 100644 --- a/scripts/cooper/AZ_segmentation_h5.py +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -1,6 +1,7 @@ import argparse import h5py import os +import json from pathlib import Path from tqdm import tqdm @@ -91,13 +92,11 @@ def run_AZ_segmentation(input_path, output_path, model_path, mask_path, mask_key -def segment_folder(args): - input_files = [] - for root, dirs, files in os.walk(args.input_path): - input_files.extend([ - os.path.join(root, name) for name in files if name.endswith(args.data_ext) - ]) +def segment_folder(args, valid_files): + input_files = [os.path.join(root, name) for root, _, files in os.walk(args.input_path) for name in files if name.endswith(args.data_ext)] + input_files = [f for f in input_files if f in valid_files] if valid_files else input_files print(input_files) + pbar = tqdm(input_files, desc="Run segmentation") for input_path in pbar: @@ -114,9 +113,17 @@ def segment_folder(args): except: print(f"compartment file not found for {input_path}") compartment_seg = None + else: + compartment_seg = None run_AZ_segmentation(input_path, args.output_path, args.model_path, mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, compartment_seg) +def get_dataset(json_file, input_path, sets=["test"]): + with open(json_file, 'r') as f: + data = json.load(f) + return {os.path.join(input_path, f) for dataset in sets for f in data.get(dataset, [])} + + def main(): parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") parser.add_argument( @@ -130,6 +137,9 @@ def main(): parser.add_argument( "--model_path", "-m", required=True, help="The filepath to the vesicle model." ) + parser.add_argument( + "--json_path", "-j", help="The filepath to the json file that stores the train, val, and test split." + ) parser.add_argument( "--mask_path", help="The filepath to a h5 file with a mask that will be used to restrict the segmentation. Needs to be in combination with mask_key." ) @@ -153,7 +163,7 @@ def main(): help="Format extension of data to be segmented, default is .h5." ) parser.add_argument( - "--compartment_seg", "-c", + "--compartment_seg", "-c", default = None, help="Path to compartment segmentation." "If the compartment segmentation was executed before, this will add a key to output file that stores the intersection between compartment boundary and AZ." "Maybe need to adjust the compartment key that the segmentation is stored under" @@ -161,12 +171,18 @@ def main(): args = parser.parse_args() input_ = args.input_path + valid_files = get_dataset(args.json_path, input_) if args.json_path else None - if os.path.isdir(input_): - segment_folder(args) + if valid_files: + if len(valid_files) == 1: + run_AZ_segmentation(next(iter(valid_files)), args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.compartment_seg) + else: + segment_folder(args, valid_files) + elif os.path.isdir(args.input_path): + segment_folder(args, valid_files) else: - run_AZ_segmentation(input_, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.compartment_seg) - + run_AZ_segmentation(args.input_path, args.output_path, args.model_path, args.mask_path, args.mask_key, args.tile_shape, args.halo, args.key_label, args.compartment_seg) + print("Finished segmenting!") if __name__ == "__main__": diff --git a/scripts/cooper/training/evaluate_AZ.py b/scripts/cooper/training/evaluate_AZ.py index dbf8d675..4ff7be6d 100644 --- a/scripts/cooper/training/evaluate_AZ.py +++ b/scripts/cooper/training/evaluate_AZ.py @@ -64,10 +64,10 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False, preci #get the labels and segmentation with h5py.File(labels_path) as label_file: - gt = label_file["/labels/thin_az"][:] + gt = label_file["/labels/az"][:] with h5py.File(segmentation_path) as seg_file: - segmentation = seg_file["/AZ/thin_az"][:] + segmentation = seg_file["/AZ/segment_from_AZmodel_v5"][:] if crop: print("cropping the annotation and segmentation") @@ -82,7 +82,7 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False, preci # Store results result_folder = "/user/muth9/u12095/synaptic-reconstruction/scripts/cooper/evaluation_results" os.makedirs(result_folder, exist_ok=True) - result_path = os.path.join(result_folder, f"evaluation_{model_name}_dice_thinpred_thinanno.csv") + result_path = os.path.join(result_folder, f"evaluation_{model_name}.csv") print("Evaluation results are saved to:", result_path) # Load existing results if the file exists diff --git a/scripts/cooper/training/train_AZ.py b/scripts/cooper/training/train_AZ.py index 9d7d2833..3f310fac 100644 --- a/scripts/cooper/training/train_AZ.py +++ b/scripts/cooper/training/train_AZ.py @@ -11,8 +11,8 @@ from synaptic_reconstruction.training import supervised_training from synaptic_reconstruction.training import semisupervised_training -TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/exported_imod_objects" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_v2" +TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/exported_imod_objects/postprocessed_AZ" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_after1stRevision" def _require_train_val_test_split(datasets): @@ -83,10 +83,11 @@ def train(key, ignore_label = None, training_2D = False, testset = True): os.makedirs(OUTPUT_ROOT, exist_ok=True) datasets = [ - "01_hoi_maus_2020_incomplete", + "01data_withoutInvertedFiles", "04_hoi_stem_examples", "06_hoi_wt_stem750_fm", - "12_chemical_fix_cryopreparation" + "12_chemical_fix_cryopreparation", + "wichmann_withAZ" ] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -96,7 +97,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): print(len(val_paths), "tomograms for validation") patch_shape = [48, 256, 256] - model_name=f"3D-AZ-model-v3" + model_name=f"3D-AZ-model-v6" #checking for 2D training if training_2D: @@ -115,8 +116,8 @@ def train(key, ignore_label = None, training_2D = False, testset = True): sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=1, p_reject = 0.95), n_samples_train=None, n_samples_val=25, check=check, - save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_models", - n_iterations=int(5e4), + save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/for_revison/models", + n_iterations=int(1e5), ignore_label= ignore_label, label_transform=torch_em.transform.label.labels_to_binary, out_channels = 1, diff --git a/scripts/cooper/training/train_vesicles.py b/scripts/cooper/training/train_vesicles.py index 7f609438..0a85e6bb 100644 --- a/scripts/cooper/training/train_vesicles.py +++ b/scripts/cooper/training/train_vesicles.py @@ -8,8 +8,8 @@ from synaptic_reconstruction.training import supervised_training from synaptic_reconstruction.training import semisupervised_training -TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/vesicles_processed_v2" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_v2" +TRAIN_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/cryo-et/cryoVesNet/" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/for_revison/training_CryoVesNet" def _require_train_val_test_split(datasets): @@ -76,14 +76,7 @@ def get_paths(split, datasets, testset=True): def train(key, ignore_label = None, training_2D = False, testset = True): datasets = [ - "01_hoi_maus_2020_incomplete", - "02_hcc_nanogold", - "03_hog_cs1sy7", - "05_stem750_sv_training", - "07_hoi_s1sy7_tem250_ihgp", - "10_tem_single_release", - "11_tem_multiple_release", - "12_chemical_fix_cryopreparation" + "exported" ] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -93,7 +86,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): print(len(val_paths), "tomograms for validation") patch_shape = [48, 256, 256] - model_name=f"3D-vesicles-model-new_postprocessing_{key}" + model_name=f"3D-vesicles-model-for_revison_{key}" #checking for 2D training if training_2D: @@ -107,11 +100,13 @@ def train(key, ignore_label = None, training_2D = False, testset = True): name=model_name, train_paths=train_paths, val_paths=val_paths, - label_key=f"/labels/vesicles/{key}", + label_key=f"/labels/vesicles",#/{key}", + raw_key="raw", patch_shape=patch_shape, batch_size=batch_size, + n_iterations=int(5e4), n_samples_train=None, n_samples_val=25, check=check, - save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2", + save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/for_revison/models", ignore_label= ignore_label, ) diff --git a/scripts/cooper/vesicle_segmentation_h5.py b/scripts/cooper/vesicle_segmentation_h5.py index 1136f180..2e5bbbca 100644 --- a/scripts/cooper/vesicle_segmentation_h5.py +++ b/scripts/cooper/vesicle_segmentation_h5.py @@ -35,6 +35,20 @@ def get_volume(input_path): return input_volume def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mask_key,tile_shape, halo, include_boundary, key_label, distance_threshold = None): + + seg_output = _require_output_folders(output_path) + file_name = Path(input_path).stem + seg_path = os.path.join(seg_output, f"{file_name}.h5") + + try: + with h5py.File(seg_path, "r") as f: + key=f"vesicles/segment_from_{key_label}" + if key in f: + print("Skipping", input_path, "because", key, "exists") + return + except: + print("need to create new seg file") + tiling = parse_tiling(tile_shape, halo) print(f"using tiling {tiling}") input = get_volume(input_path) @@ -58,10 +72,6 @@ def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mas foreground, boundaries = prediction[:2] - seg_output = _require_output_folders(output_path) - file_name = Path(input_path).stem - seg_path = os.path.join(seg_output, f"{file_name}.h5") - #check os.makedirs(Path(seg_path).parent, exist_ok=True) @@ -77,8 +87,8 @@ def run_vesicle_segmentation(input_path, output_path, model_path, mask_path, mas print("Skipping", input_path, "because", key, "exists") else: f.create_dataset(key, data=segmentation, compression="gzip") - f.create_dataset(f"prediction_{key_label}/foreground", data = foreground, compression="gzip") - f.create_dataset(f"prediction_{key_label}/boundaries", data = boundaries, compression="gzip") + #f.create_dataset(f"prediction_{key_label}/foreground", data = foreground, compression="gzip") + #f.create_dataset(f"prediction_{key_label}/boundaries", data = boundaries, compression="gzip") if mask is not None: if mask_key in f: diff --git a/scripts/cryo/vesicles/train_domain_adaptation.py b/scripts/cryo/vesicles/train_domain_adaptation.py index 9270ba41..dca33da0 100644 --- a/scripts/cryo/vesicles/train_domain_adaptation.py +++ b/scripts/cryo/vesicles/train_domain_adaptation.py @@ -6,8 +6,8 @@ from sklearn.model_selection import train_test_split from synaptic_reconstruction.training.domain_adaptation import mean_teacher_adaptation -TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/fernandez-busnadiego/from_arsen" -OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_training" +TRAIN_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/cryo-et/cryoVesNet/" +OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_training/CryoVesNet" def _require_train_val_test_split(datasets): train_ratio, val_ratio, test_ratio = 0.8, 0.1, 0.1 @@ -24,7 +24,7 @@ def _train_val_test_split(names): if os.path.exists(split_path): continue - file_paths = sorted(glob(os.path.join(TRAIN_ROOT, ds, "*.mrc"))) + file_paths = sorted(glob(os.path.join(TRAIN_ROOT, ds, "*.h5"))) file_names = [os.path.basename(path) for path in file_paths] train, val, test = _train_val_test_split(file_names) @@ -45,7 +45,7 @@ def _train_val_split(names): if os.path.exists(split_path): continue - file_paths = sorted(glob(os.path.join(TRAIN_ROOT, ds, "*.mrc"))) + file_paths = sorted(glob(os.path.join(TRAIN_ROOT, ds, "*.h5"))) file_names = [os.path.basename(path) for path in file_paths] train, val = _train_val_split(file_names) @@ -72,8 +72,7 @@ def get_paths(split, datasets, testset=True): def vesicle_domain_adaptation(teacher_model, testset = True): datasets = [ - "tomos_deconv_18924", - "old_data" + "exported" ] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -84,21 +83,21 @@ def vesicle_domain_adaptation(teacher_model, testset = True): #adjustable parameters patch_shape = [48, 256, 256] - model_name = "vesicle-DA-cryo-v2" + model_name = "vesicle-DA-CryoVesNet-v1" - model_root = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" + model_root = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" # cryo DA model: "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_models/checkpoints" #cooper model: "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/models_v2/checkpoints/" checkpoint_path = os.path.join(model_root, teacher_model) mean_teacher_adaptation( name=model_name, unsupervised_train_paths=train_paths, unsupervised_val_paths=val_paths, - raw_key="data", + raw_key="raw", patch_shape=patch_shape, save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/DA_models", source_checkpoint=checkpoint_path, confidence_threshold=0.75, - n_iterations=int(5e4), + n_iterations=int(5e3), ) From 7d675ab2c32d452064c99967f546b621692ad762 Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Thu, 22 May 2025 10:36:27 +0200 Subject: [PATCH 22/23] everything after 1st revision relating to training, inference, postprocessing az --- big_to_small_pixel_size.py | 87 ++++++++++++++++++ scripts/cooper/training/evaluate_AZ.py | 2 +- scripts/cooper/training/postprocess_AZ.py | 51 ++++++++++- scripts/cooper/training/train_AZ.py | 19 ++-- scripts/cooper/vesicle_segmentation_h5.py | 88 ++++++++++++++++++- small_to_big_pixel_size.py | 69 +++++++++++++++ .../training/supervised_training.py | 13 ++- 7 files changed, 314 insertions(+), 15 deletions(-) create mode 100644 big_to_small_pixel_size.py create mode 100644 small_to_big_pixel_size.py diff --git a/big_to_small_pixel_size.py b/big_to_small_pixel_size.py new file mode 100644 index 00000000..6f4be4ac --- /dev/null +++ b/big_to_small_pixel_size.py @@ -0,0 +1,87 @@ +import os +import numpy as np +import h5py +from glob import glob +from scipy.ndimage import zoom +from scipy.ndimage import label +from skimage.morphology import closing, ball + +# Input and output folders +input_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ" +output_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ_right_rescaled" +os.makedirs(output_folder, exist_ok=True) + +# Define scaling factors +old_pixel_size = np.array([1.75, 1.75, 1.75]) +new_pixel_size = np.array([1.55, 1.55, 1.55]) +scaling_factors = old_pixel_size / new_pixel_size + +# Utility function to process segmentation +def rescale_and_fix_segmentation(segmentation, scaling_factors): + """ + Rescale the segmentation and ensure labels are preserved. + Args: + segmentation (numpy.ndarray): The input segmentation array with integer labels. + scaling_factors (list or array): Scaling factors for each axis. + Returns: + numpy.ndarray: Rescaled and hole-free segmentation with preserved labels. + """ + # Rescale segmentation using nearest-neighbor interpolation + rescaled_segmentation = zoom(segmentation, scaling_factors, order=0) + + # Initialize an array to hold the processed segmentation + processed_segmentation = np.zeros_like(rescaled_segmentation) + + # Ensure no holes for each label + unique_labels = np.unique(rescaled_segmentation) + for label_id in unique_labels: + if label_id == 0: # Skip the background + continue + + # Extract binary mask for the current label + label_mask = rescaled_segmentation == label_id + + # Apply morphological closing to fill holes + closed_mask = closing(label_mask, ball(1)) + + # Add the processed label back to the output segmentation + processed_segmentation[closed_mask] = label_id + + return processed_segmentation.astype(segmentation.dtype) + + +# Get all .h5 files in the specified input folder +h5_files = glob(os.path.join(input_folder, "*.h5")) +existing_files = {os.path.basename(f) for f in glob(os.path.join(output_folder, "*.h5"))} + +for h5_file in h5_files: + print(f"Processing {h5_file}...") + + if os.path.basename(h5_file) in existing_files: + print(f"Skipping {h5_file} as it already exists in the output folder.") + continue + + # Read data from the .h5 file + with h5py.File(h5_file, "r") as f: + raw = f["raw"][:] # Assuming the dataset is named "raw" + az = f["labels/az"][:] + + print(f"Original shape - raw: {raw.shape}; az: {az.shape}") + + # Process raw data (tomogram) with linear interpolation + print("Rescaling raw data...") + rescaled_raw = zoom(raw, scaling_factors, order=1) + + # Process az segmentation + print("Rescaling and fixing az segmentation...") + rescaled_az = rescale_and_fix_segmentation(az, scaling_factors) + + # Save the processed data to a new .h5 file + output_path = os.path.join(output_folder, os.path.basename(h5_file)) + with h5py.File(output_path, "w") as f: + f.create_dataset("raw", data=rescaled_raw, compression="gzip") + f.create_dataset("labels/az", data=rescaled_az, compression="gzip") + + print(f"Saved rescaled data to {output_path}") + +print("Processing complete. Rescaled files are saved in:", output_folder) diff --git a/scripts/cooper/training/evaluate_AZ.py b/scripts/cooper/training/evaluate_AZ.py index 4ff7be6d..1647a921 100644 --- a/scripts/cooper/training/evaluate_AZ.py +++ b/scripts/cooper/training/evaluate_AZ.py @@ -67,7 +67,7 @@ def evaluate_file(labels_path, segmentation_path, model_name, crop= False, preci gt = label_file["/labels/az"][:] with h5py.File(segmentation_path) as seg_file: - segmentation = seg_file["/AZ/segment_from_AZmodel_v5"][:] + segmentation = seg_file["/AZ/segment_from_AZmodel_v6"][:] if crop: print("cropping the annotation and segmentation") diff --git a/scripts/cooper/training/postprocess_AZ.py b/scripts/cooper/training/postprocess_AZ.py index e2b849e0..91a417d4 100644 --- a/scripts/cooper/training/postprocess_AZ.py +++ b/scripts/cooper/training/postprocess_AZ.py @@ -89,18 +89,61 @@ def process_all_azs(output_path): for ff in tqdm(files, desc="Thinning AZ segmentations"): process_az(ff, view=False) +def subtract_SVseg_from_AZseg(input_path, SV_path, output_path): + """ + Modifies AZ segmentation by setting regions where SV segmentation is not zero to zero. + + Parameters: + input_path (str): Path to the folder containing AZ segmentation H5 files. + SV_path (str): Path to the folder containing corresponding SV segmentation H5 files. + output_path (str): Path to save modified AZ segmentation H5 files. + """ + + # Ensure output directory exists + os.makedirs(output_path, exist_ok=True) + + # Iterate over all AZ segmentation files + for file_name in os.listdir(input_path): + if file_name.endswith(".h5"): + az_file = os.path.join(input_path, file_name) + sv_file = os.path.join(SV_path, file_name) + + # Check if corresponding SV segmentation file exists + if os.path.exists(sv_file): + with h5py.File(sv_file, "r") as f_sv, h5py.File(az_file, "r+") as f_az: + SV = f_sv["/vesicles/segment_from_combined_vesicles"][:] + AZ = f_az["/labels/az"][:] + raw = f_az["raw"][:] + + # Modify AZ segmentation where SV segmentation is not zero + AZ[SV != 0] = 0 + + # Save modified AZ segmentation + output_file = os.path.join(output_path, file_name) + with h5py.File(output_file, "a") as f: + f.create_dataset("raw", data=raw, compression="gzip") + f.create_dataset("labels/az", data=AZ, compression="gzip") + + print(f"Processed: {file_name}") + + else: + print(f"Skipping {file_name}, corresponding SV file not found.") + def main(): parser = argparse.ArgumentParser(description="Filter and process AZ data.") - parser.add_argument("input_path", type=str, help="Path to the root directory containing datasets.") - parser.add_argument("output_path", type=str, help="Path to the root directory for saving processed data.") + parser.add_argument("-i", "--input_path", required=True, type=str, help="Path to the root directory containing datasets.") + parser.add_argument("-o", "--output_path", required=True, type=str, help="Path to the root directory for saving processed data.") + parser.add_argument("-sv", "--SV_path", type=str, help="Path to the root directory that contains the SV segmentations.") args = parser.parse_args() input_path = args.input_path output_path = args.output_path - filter_all_azs(input_path, output_path) - process_all_azs(output_path) + subtract_SVseg_from_AZseg(input_path, args.SV_path, output_path) + + #filter_all_azs(input_path, output_path) + #process_all_azs(output_path) if __name__ == "__main__": diff --git a/scripts/cooper/training/train_AZ.py b/scripts/cooper/training/train_AZ.py index 3f310fac..53c91cfa 100644 --- a/scripts/cooper/training/train_AZ.py +++ b/scripts/cooper/training/train_AZ.py @@ -11,7 +11,7 @@ from synaptic_reconstruction.training import supervised_training from synaptic_reconstruction.training import semisupervised_training -TRAIN_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/exported_imod_objects/postprocessed_AZ" +TRAIN_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ" OUTPUT_ROOT = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/training_AZ_after1stRevision" @@ -83,11 +83,12 @@ def train(key, ignore_label = None, training_2D = False, testset = True): os.makedirs(OUTPUT_ROOT, exist_ok=True) datasets = [ - "01data_withoutInvertedFiles", - "04_hoi_stem_examples", - "06_hoi_wt_stem750_fm", - "12_chemical_fix_cryopreparation", - "wichmann_withAZ" + "01data_withoutInvertedFiles_minusSVseg_corrected", + "04_hoi_stem_examples_fidi_and_sarah_corrected_rescaled_tomograms", + "04_hoi_stem_examples_minusSVseg_cropped_corrected_rescaled_tomograms", + "06_hoi_wt_stem750_fm_minusSVseg_cropped_corrected_rescaled_tomograms", + "12_chemical_fix_cryopreparation_minusSVseg_corrected", + "wichmann_withAZ_rescaled_tomograms" ] train_paths = get_paths("train", datasets=datasets, testset=testset) val_paths = get_paths("val", datasets=datasets, testset=testset) @@ -97,7 +98,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): print(len(val_paths), "tomograms for validation") patch_shape = [48, 256, 256] - model_name=f"3D-AZ-model-v6" + model_name=f"3D-AZ-model-TEM_STEM_ChemFix_wichmann-v1" #checking for 2D training if training_2D: @@ -113,7 +114,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): val_paths=val_paths, label_key=f"/labels/{key}", patch_shape=patch_shape, batch_size=batch_size, - sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=1, p_reject = 0.95), + sampler = torch_em.data.sampler.MinInstanceSampler(min_num_instances=1, p_reject = 0.8), n_samples_train=None, n_samples_val=25, check=check, save_root="/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/for_revison/models", @@ -121,6 +122,8 @@ def train(key, ignore_label = None, training_2D = False, testset = True): ignore_label= ignore_label, label_transform=torch_em.transform.label.labels_to_binary, out_channels = 1, + BCE_loss=False, + sigmoid_layer=True, ) diff --git a/scripts/cooper/vesicle_segmentation_h5.py b/scripts/cooper/vesicle_segmentation_h5.py index 2e5bbbca..27b5a14d 100644 --- a/scripts/cooper/vesicle_segmentation_h5.py +++ b/scripts/cooper/vesicle_segmentation_h5.py @@ -175,5 +175,91 @@ def main(): print("Finished segmenting!") +def segment_folder2(args, valid_files): + input_files = [os.path.join(root, name) for root, _, files in os.walk(args.input_path) for name in files if name.endswith(args.data_ext)] + input_files = [f for f in input_files if f in valid_files] if valid_files else input_files + print(input_files) + for root, dirs, files in os.walk(args.input_path): + input_files.extend([ + os.path.join(root, name) for name in files if name.endswith(args.data_ext) + ]) + print(input_files) + pbar = tqdm(input_files, desc="Run segmentation") + for input_path in pbar: + + filename = os.path.basename(input_path) + try: + mask_path = os.path.join(args.mask_path, filename) + except: + print(f"Mask file not found for {input_path}") + mask_path = None + + run_vesicle_segmentation( + input_path, args.output_path, args.model_path, mask_path, args.mask_key, + args.tile_shape, args.halo, args.include_boundary, args.key_label, args.distance_threshold + ) + +def get_dataset(testfolder_path, input_path): + print("Getting valid test set") + # Get base filenames without '_processed' and extension + test_files = [ + os.path.splitext(f)[0].replace('_processed', '') + for f in os.listdir(testfolder_path) + if f.endswith('.h5') + ] + + # Build a set for faster lookup + test_files_set = set(test_files) + + # List files in input_path that match names in test_files_set + valid_files = [ + os.path.join(input_path, f) + for f in os.listdir(input_path) + if os.path.splitext(f)[0] in test_files_set and f.endswith('.h5') + ] + + return valid_files + +def main2(): + parser = argparse.ArgumentParser(description="Segment vesicles in EM tomograms.") + parser.add_argument("--input_path", "-i", required=True, help="Path to MRC file or directory.") + parser.add_argument("--output_path", "-o", required=True, help="Output directory.") + parser.add_argument("--model_path", "-m", required=True, help="Model file path.") + parser.add_argument("--mask_path", help="Optional mask file.") + parser.add_argument("--mask_key", help="Key in mask file.") + parser.add_argument("--tile_shape", type=int, nargs=3, help="Tile shape for prediction.") + parser.add_argument("--halo", type=int, nargs=3, help="Halo size.") + parser.add_argument("--include_boundary", action="store_true", help="Include edge vesicles.") + parser.add_argument("--key_label", "-k", default="combined_vesicles", help="Output H5 key.") + parser.add_argument("--distance_threshold", "-t", type=int, help="Distance threshold.") + parser.add_argument("--data_ext", "-d", default=".h5", help="File extension.") + parser.add_argument("--testfolder_path", help="Folder path to testset.") + + args = parser.parse_args() + + input_ = args.input_path + valid_files = get_dataset(args.testfolder_path, input_) if args.testfolder_path else None + + if valid_files: + if len(valid_files) == 1: + run_vesicle_segmentation( + input_, args.output_path, args.model_path, args.mask_path, + args.mask_key, args.tile_shape, args.halo, + args.include_boundary, args.key_label, args.distance_threshold + ) + else: + segment_folder2(args, valid_files) + elif os.path.isdir(args.input_path): + segment_folder2(args, valid_files) + else: + run_vesicle_segmentation( + input_, args.output_path, args.model_path, args.mask_path, + args.mask_key, args.tile_shape, args.halo, + args.include_boundary, args.key_label, args.distance_threshold + ) + + print("Finished segmenting!") + if __name__ == "__main__": - main() \ No newline at end of file + #main() + main2() #needed this to do the vesicle segmentation for only the data that is in the testset for paper but the uncropped version of the only GT annotated data \ No newline at end of file diff --git a/small_to_big_pixel_size.py b/small_to_big_pixel_size.py new file mode 100644 index 00000000..b5db4abb --- /dev/null +++ b/small_to_big_pixel_size.py @@ -0,0 +1,69 @@ +import os +from glob import glob +import h5py +import numpy as np +from skimage.transform import rescale + + +def rescale_h5_files(input_folder, output_folder, current_resolution, target_resolution): + # Compute the scale factor: target / current + scale_factor = current_resolution / target_resolution + + # Ensure output folder exists + os.makedirs(output_folder, exist_ok=True) + + # Get all .h5 files in the input folder + h5_files = glob(os.path.join(input_folder, "*.h5")) + existing_files = {os.path.basename(f) for f in glob(os.path.join(output_folder, "*.h5"))} + + for h5_file in h5_files: + filename = os.path.basename(h5_file) + output_path = os.path.join(output_folder, filename) + + if filename in existing_files: + print(f"Skipping {filename}, already exists in output.") + continue + + print(f"Processing {filename}...") + + with h5py.File(h5_file, "r") as f: + raw = f["raw"][:] + az = f["labels/az"][:] + + # Rescale in 3D + raw_rescaled = rescale( + raw, + scale=(scale_factor, scale_factor, scale_factor), + order=3, + preserve_range=True, + anti_aliasing=True + ).astype(raw.dtype) + + az_rescaled = rescale( + az, + scale=(scale_factor, scale_factor, scale_factor), + order=0, # Nearest neighbor for label maps + preserve_range=True, + anti_aliasing=False + ).astype(az.dtype) + + # Save to new .h5 file + with h5py.File(output_path, "w") as f_out: + f_out.create_dataset("raw", data=raw_rescaled, compression="gzip") + f_out.create_dataset("labels/az", data=az_rescaled, compression="gzip") + + print(f"Saved rescaled data to {output_path}") + + +def main(): + input_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ" + output_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ_rescaled_tomograms" + + current_resolution = 0.8681 # nm + target_resolution = 1.55 # nm + + rescale_h5_files(input_folder, output_folder, current_resolution, target_resolution) + + +if __name__ == "__main__": + main() diff --git a/synaptic_reconstruction/training/supervised_training.py b/synaptic_reconstruction/training/supervised_training.py index 72a32f01..9612481b 100644 --- a/synaptic_reconstruction/training/supervised_training.py +++ b/synaptic_reconstruction/training/supervised_training.py @@ -184,6 +184,8 @@ def supervised_training( label_transform: Optional[callable] = None, out_channels: int = 2, mask_channel: bool = False, + BCE_loss:bool = False, + sigmoid_layer:bool = True, **loader_kwargs, ): """Run supervised segmentation training. @@ -243,7 +245,10 @@ def supervised_training( if is_2d: model = get_2d_model(out_channels=out_channels) else: - model = get_3d_model(out_channels=out_channels) + if sigmoid_layer: + model = get_3d_model(out_channels=out_channels) + else: + model = get_3d_model(out_channels=out_channels, final_activation=None) loss, metric = None, None # No ignore label -> we can use default loss. @@ -268,6 +273,12 @@ def supervised_training( metric = loss else: raise ValueError + + #check for BCE loss (for AZ training) + if BCE_loss: + loss=torch.nn.BCEWithLogitsLoss() + #loss=torch_em.loss.dice.BCEDiceLossWithLogits() + metric = loss trainer = torch_em.default_segmentation_trainer( name=name, From f052b98e1ca7352cdc7e16a58413fddc8f2cf9fb Mon Sep 17 00:00:00 2001 From: SarahMuth Date: Sun, 25 May 2025 19:07:04 +0200 Subject: [PATCH 23/23] minor things --- big_to_small_pixel_size.py | 2 +- scripts/cooper/training/train_AZ.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/big_to_small_pixel_size.py b/big_to_small_pixel_size.py index 6f4be4ac..7f41b0ae 100644 --- a/big_to_small_pixel_size.py +++ b/big_to_small_pixel_size.py @@ -8,7 +8,7 @@ # Input and output folders input_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ" -output_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ_right_rescaled" +output_folder = "/mnt/lustre-emmy-hdd/usr/u12095/synaptic_reconstruction/AZ_data_after1stRevision/recorrected_length_of_AZ/wichmann_withAZ_rescaled_tomograms" os.makedirs(output_folder, exist_ok=True) # Define scaling factors diff --git a/scripts/cooper/training/train_AZ.py b/scripts/cooper/training/train_AZ.py index 53c91cfa..1c6767e8 100644 --- a/scripts/cooper/training/train_AZ.py +++ b/scripts/cooper/training/train_AZ.py @@ -98,7 +98,7 @@ def train(key, ignore_label = None, training_2D = False, testset = True): print(len(val_paths), "tomograms for validation") patch_shape = [48, 256, 256] - model_name=f"3D-AZ-model-TEM_STEM_ChemFix_wichmann-v1" + model_name=f"3D-AZ-model-TEM_STEM_ChemFix_wichmann-v2" #checking for 2D training if training_2D: