diff --git a/.gitignore b/.gitignore index d0404310..d9550532 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,11 @@ 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 \ 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 +scripts/cooper/remove_h5key.py +scripts/cooper/analysis/calc_AZ_area.py \ No newline at end of file diff --git a/big_to_small_pixel_size.py b/big_to_small_pixel_size.py new file mode 100644 index 00000000..7f41b0ae --- /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_rescaled_tomograms" +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/AZ_segmentation_h5.py b/scripts/cooper/AZ_segmentation_h5.py new file mode 100644 index 00000000..00d2de43 --- /dev/null +++ b/scripts/cooper/AZ_segmentation_h5.py @@ -0,0 +1,189 @@ +import argparse +import h5py +import os +import json +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, compartment_seg): + 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 + + #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["/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) + 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") + + 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") + + + + +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: + + 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 + + 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 + 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( + "--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( + "--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." + ) + 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." + ) + parser.add_argument( + "--data_ext", "-d", default = ".h5", + help="Format extension of data to be segmented, default is .h5." + ) + parser.add_argument( + "--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" + ) + args = parser.parse_args() + + input_ = args.input_path + valid_files = get_dataset(args.json_path, input_) if args.json_path else None + + 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(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__": + main() \ No newline at end of file diff --git a/scripts/cooper/analysis/run_size_analysis.py b/scripts/cooper/analysis/run_size_analysis.py new file mode 100644 index 00000000..abad4405 --- /dev/null +++ b/scripts/cooper/analysis/run_size_analysis.py @@ -0,0 +1,174 @@ +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/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): + """ + 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: + compartment_mask = compartments == compartment_label + vesicle_mask = vesicles > 0 + + intersection = np.logical_and(compartment_mask, vesicle_mask) + 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 + + final_mask = compartments == best_compartment + + return final_mask + +# 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. --> 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) + + 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 + 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" + + 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. + 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 + ) + + + 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(): + 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"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 + 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" + + 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. + 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(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 + ) + + 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(): + 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"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() + +if __name__ == "__main__": + main() 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..69434844 --- /dev/null +++ b/scripts/cooper/analysis/run_spatial_distribution_analysis.py @@ -0,0 +1,187 @@ +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 +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_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): + """ + 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: + compartment_mask = compartments == compartment_label + vesicle_mask = vesicles > 0 + + intersection = np.logical_and(compartment_mask, vesicle_mask) + 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 + + final_mask = compartments == best_compartment + + return final_mask + +# 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_autoComp(separate_AZseg=False): + + 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 + 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" + + 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 + 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"][:] + + 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(): + 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) + +# 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 + 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" + + 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 + 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 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_autoComp(separate_AZseg=False) + #compute_per_vesicle_distance_to_AZ_manComp() + + + +if __name__ == "__main__": + main() diff --git a/scripts/cooper/compartment_segmentation_h5.py b/scripts/cooper/compartment_segmentation_h5.py new file mode 100644 index 00000000..573ac482 --- /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],boundary_threshold=0.2, postprocess_segments=False) + + 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/evaluate_AZ.py b/scripts/cooper/training/evaluate_AZ.py new file mode 100644 index 00000000..1647a921 --- /dev/null +++ b/scripts/cooper/training/evaluate_AZ.py @@ -0,0 +1,146 @@ +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 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)) + 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_v6"][:] + + if crop: + print("cropping the annotation and segmentation") + segmentation, gt = extract_gt_bounding_box(segmentation, gt) + + # Evaluate the match of ground truth and 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") + 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 + 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: + 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, precision_score=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, precision_score) + + + +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.") + 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, args.precision) + else: + evaluate_file(args.labels_path, segmentation_path, args.model_name, args.crop, args.precision) + + + +if __name__ == "__main__": + main() 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/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) diff --git a/scripts/cooper/training/postprocess_AZ.py b/scripts/cooper/training/postprocess_AZ.py new file mode 100644 index 00000000..91a417d4 --- /dev/null +++ b/scripts/cooper/training/postprocess_AZ.py @@ -0,0 +1,150 @@ +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 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("-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 + + 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__": + main() diff --git a/scripts/cooper/training/train_AZ.py b/scripts/cooper/training/train_AZ.py index 1468eaf7..1c6767e8 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_v1" +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" def _require_train_val_test_split(datasets): @@ -80,10 +80,15 @@ 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", - "06_hoi_wt_stem750_fm", - "12_chemical_fix_cryopreparation" + "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) @@ -93,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-v1" + model_name=f"3D-AZ-model-TEM_STEM_ChemFix_wichmann-v2" #checking for 2D training if training_2D: @@ -109,14 +114,16 @@ 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.8), 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), + 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, + BCE_loss=False, + sigmoid_layer=True, ) 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 9c8b1d11..27b5a14d 100644 --- a/scripts/cooper/vesicle_segmentation_h5.py +++ b/scripts/cooper/vesicle_segmentation_h5.py @@ -34,7 +34,21 @@ 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): + + 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) @@ -45,14 +59,19 @@ 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) - 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) @@ -68,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: @@ -84,7 +103,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 +116,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 +156,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,9 +171,95 @@ 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!") + +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/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), ) diff --git a/scripts/rizzoli/2D_vesicle_segmentation.py b/scripts/rizzoli/2D_vesicle_segmentation.py index 7974e3b3..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): +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,23 +70,41 @@ 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): + 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, 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) boundaries.append(boundaries_pred_slice) return processed_slices, foreground, boundaries - segmentation, foreground, boundaries = process_slices(input) + 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, mask) seg_output = _require_output_folders(output_path) file_name = Path(input_path).stem @@ -121,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) + 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.") @@ -152,6 +175,16 @@ 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." + ) + 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 @@ -159,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) + 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 1cae666e..5b5bbbd9 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 = 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) + + 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 86eedd17..ac2a28fa 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_2Dcooper_v1" 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,11 @@ 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" + "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,11 +88,13 @@ 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-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, @@ -97,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(5e4), + batch_size=8, + n_iterations=int(1.5e4), + n_samples_train=8000, + n_samples_val=50, ) 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/imod/to_imod.py b/synaptic_reconstruction/imod/to_imod.py index 307e645d..f97e4f01 100644 --- a/synaptic_reconstruction/imod/to_imod.py +++ b/synaptic_reconstruction/imod/to_imod.py @@ -130,6 +130,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, diff --git a/synaptic_reconstruction/inference/AZ.py b/synaptic_reconstruction/inference/AZ.py new file mode 100644 index 00000000..a1c9da83 --- /dev/null +++ b/synaptic_reconstruction/inference/AZ.py @@ -0,0 +1,89 @@ +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 +from synaptic_reconstruction.inference.postprocessing.postprocess_AZ import find_intersection_boundary + +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, + compartment: 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) + + #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/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") diff --git a/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py new file mode 100644 index 00000000..8ef2cd32 --- /dev/null +++ b/synaptic_reconstruction/inference/postprocessing/postprocess_AZ.py @@ -0,0 +1,35 @@ +import numpy as np +from skimage.segmentation import find_boundaries + +def find_intersection_boundary(segmented_AZ, segmented_compartment): + """ + 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, with multiple labels. + + Returns: + numpy.ndarray: 3D array with the cumulative intersection of all boundaries of segmented_compartment labels with segmented_AZ. + """ + # Step 0: Initialize an empty array to accumulate intersections + cumulative_intersection = np.zeros_like(segmented_AZ, dtype=bool) + + # 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 cumulative_intersection.astype(int) # Convert boolean array to int (1 for intersecting points, 0 elsewhere) 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( 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,