{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"}],"dockerImageVersionId":30804,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Convert to tensor data\n\n## Intro\nThis code uses simlar methods to https://www.kaggle.com/code/yoshio13/how-to-convert-json-to-segmentation-images/notebook but saves the processed dataset as tensor files instead of numpy files. Also the segmentation masks are created in a one-hot implementation that preserves overlap in the full radius. The final output is compressed due to kaggle storage limits.\n\n## Overview\n\nThis script processes CryoET data, converting tomograms and segmentations into PyTorch tensors and organizing them into a structured format. Here's a breakdown of what the code does:\n1. It iterates through each experimental run in the dataset.\n2. For each run, it loads the tomogram data and converts it to a PyTorch tensor.\n3. It then processes each pickable object, getting its segmentation data and converting it to a binary mask.\n4. All segmentations for a run are combined into a single tensor.\n5. The tomogram and combined segmentation are saved as PyTorch tensor files.\n6. The files for each run are compressed into a zip file to save space.\n7. The original uncompressed folder is deleted after zipping.\n\nThis process helps to standardize the data format and reduce storage requirements while preserving all necessary information for further analysis or machine learning tasks.","metadata":{}},{"cell_type":"markdown","source":"## Train Dataset","metadata":{}},{"cell_type":"code","source":"!pip install -q zarr\n!pip install -q copick\n!pip install -q git+https://github.com/copick/copick-utils.git","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:21:16.538239Z","iopub.execute_input":"2025-02-04T18:21:16.538640Z","iopub.status.idle":"2025-02-04T18:22:34.323287Z","shell.execute_reply.started":"2025-02-04T18:21:16.538606Z","shell.execute_reply":"2025-02-04T18:22:34.321674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\n# Update the overlay\n# Define source and destination directories\nsource_dir = '/kaggle/input/czii-cryo-et-object-identification/train/overlay'\ndestination_dir = '/kaggle/working/train_overlay'\n\n# Create the destination directory if it doesn't exist\nos.makedirs(destination_dir, exist_ok=True)\n\nfor root, dirs, files in os.walk(source_dir):\n    # Create corresponding subdirectories in the destination\n    relative_path = os.path.relpath(root, source_dir)\n    target_dir = os.path.join(destination_dir, relative_path)\n    os.makedirs(target_dir, exist_ok=True)\n\n    for file in files:\n        if file.startswith(\"curation_0_\"):\n            new_filename = file\n        else:\n            new_filename = f\"curation_0_{file}\"\n\n        # Define full paths for the source and destination files\n        source_file = os.path.join(root, file)\n        destination_file = os.path.join(target_dir, new_filename)\n        \n        # Copy the file with the new name\n        shutil.copy2(source_file, destination_file)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:22:34.326294Z","iopub.execute_input":"2025-02-04T18:22:34.327335Z","iopub.status.idle":"2025-02-04T18:22:34.540219Z","shell.execute_reply.started":"2025-02-04T18:22:34.327279Z","shell.execute_reply":"2025-02-04T18:22:34.539089Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Next we create a Copick project using the CryoET ML Challenge dataset characteristics and define our data paths and settings.","metadata":{}},{"cell_type":"code","source":"# Import required libraries\nimport copick\nimport numpy as np\n\n# Define the configuration for the CryoET ML Challenge project\nconfig_blob = \"\"\"{\n    \"name\": \"czii_cryoet_mlchallenge_2024\",\n    \"description\": \"2024 CZII CryoET ML Challenge training data.\",\n    \"version\": \"1.0.0\",\n\n    \"pickable_objects\": [\n        {\n            \"name\": \"apo-ferritin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"4V1W\",\n            \"label\": 1,\n            \"color\": [  0, 117, 220, 128],\n            \"radius\": 60,\n            \"map_threshold\": 0.0418\n        },\n        {\n            \"name\": \"beta-galactosidase\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6X1Q\",\n            \"label\": 3,\n            \"color\": [ 76,   0,  92, 128],\n            \"radius\": 90,\n            \"map_threshold\": 0.0578\n        },\n        {\n            \"name\": \"ribosome\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6EK0\",\n            \"label\": 4,\n            \"color\": [  0,  92,  49, 128],\n            \"radius\": 150,\n            \"map_threshold\": 0.0374\n        },\n        {\n            \"name\": \"thyroglobulin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6SCJ\",\n            \"label\": 5,\n            \"color\": [ 43, 206,  72, 128],\n            \"radius\": 130,\n            \"map_threshold\": 0.0278\n        },\n        {\n            \"name\": \"virus-like-particle\",\n            \"is_particle\": true,\n            \"label\": 6,\n            \"color\": [255, 204, 153, 128],\n            \"radius\": 135,\n            \"map_threshold\": 0.201\n        },\n        {\n            \"name\": \"membrane\",\n            \"is_particle\": false,\n            \"label\": 8,\n            \"color\": [100, 100, 100, 128]\n        },\n        {\n            \"name\": \"background\",\n            \"is_particle\": false,\n            \"label\": 9,\n            \"color\": [10, 150, 200, 128]\n        }\n    ],\n\n    \"overlay_root\": \"/kaggle/working/train_overlay\",\n\n    \"overlay_fs_args\": {\n        \"auto_mkdir\": true\n    },\n\n    \"static_root\": \"/kaggle/input/czii-cryo-et-object-identification/train/static\"\n}\"\"\"\n\n# Set the path for the Copick configuration file\ncopick_config_path = \"/kaggle/working/copick.config\"\n\n# Set the output directory for overlay data\noutput_overlay = \"/kaggle/working/train_overlay\"\n\n# Write the configuration to a file\nwith open(copick_config_path, \"w\") as f:\n    f.write(config_blob)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:22:34.541985Z","iopub.execute_input":"2025-02-04T18:22:34.542427Z","iopub.status.idle":"2025-02-04T18:22:37.010396Z","shell.execute_reply.started":"2025-02-04T18:22:34.542384Z","shell.execute_reply":"2025-02-04T18:22:37.009234Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the Copick project configuration from the specified file path\nroot = copick.from_file(copick_config_path)\n\n# Set the username for Copick operations\ncopick_user_name = \"copickUtils\"\n\n# Set the name for the segmentation to be created\ncopick_segmentation_name = \"paintedPicks\"\n\n# Set the voxel size (in Angstroms) for the tomogram\nvoxel_size = 10\n\n# Specify the type of tomogram to be used (denoised version)\ntomo_type = \"denoised\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:22:37.012471Z","iopub.execute_input":"2025-02-04T18:22:37.013242Z","iopub.status.idle":"2025-02-04T18:22:37.020026Z","shell.execute_reply.started":"2025-02-04T18:22:37.013204Z","shell.execute_reply":"2025-02-04T18:22:37.018871Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Next we create a dictionary of particle objects, storing their labels and radii. The use of defaultdict allows for easy creation of nexted dictionaries without explicitly checking if a key exists before assigning values.","metadata":{}},{"cell_type":"code","source":"from collections import defaultdict\n\n# Create a defaultdict that will automatically create an empty dictionary as the default value\nobject_dict = defaultdict(dict)\n\n# Iterate through each pickable object in root.pickable_objects\nfor object in root.pickable_objects:\n    # Check if the object is a particle\n    if object.is_particle:\n        # If it is a particle, add its label and radius to the object_dict\n        # Using the object's name as the key\n        object_dict[object.name]['label'] = object.label\n        object_dict[object.name]['radius'] = object.radius","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:22:37.021482Z","iopub.execute_input":"2025-02-04T18:22:37.021998Z","iopub.status.idle":"2025-02-04T18:22:37.129210Z","shell.execute_reply.started":"2025-02-04T18:22:37.021942Z","shell.execute_reply":"2025-02-04T18:22:37.128181Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Next we process the tomogram data and create 3D segmentation labels for each pickable object (particle type) in each run. Then we write these segmentations to files using the copick_utils library.","metadata":{}},{"cell_type":"code","source":"from copick_utils.segmentation import segmentation_from_picks\nimport copick_utils.writers.write as write\n\n# Create a dictionary to store segmentations\nsegmentation_dict = defaultdict(dict)\n\n# Iterate through each run in the dataset\nfor run in root.runs:\n    # Get the tomogram for the current run with voxel spacing of 10\n    tomo = run.get_voxel_spacing(10).get_tomogram(tomo_type).numpy()\n\n    # Iterate through each pickable object\n    for pickable_object in root.pickable_objects:\n        # Create an empty segmentation label array\n        seg_label_zero = np.zeros(tomo.shape, dtype=np.uint8)\n        \n        # Get the picks for the current object\n        pick = run.get_picks(object_name=pickable_object.name, user_id=\"curation\")\n        \n        # If picks exist for this object\n        if len(pick): \n            # Generate segmentation from picks\n            seg_label = segmentation_from_picks.from_picks(pick[0], \n                                                        seg_label_zero, \n                                                        object_dict[pickable_object.name]['radius'],\n                                                        object_dict[pickable_object.name]['label']\n                                                        )\n        # Write the segmentation to file\n        write.segmentation(run, seg_label, copick_user_name, name=pickable_object.name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:22:37.130730Z","iopub.execute_input":"2025-02-04T18:22:37.131231Z","iopub.status.idle":"2025-02-04T18:23:03.473385Z","shell.execute_reply.started":"2025-02-04T18:22:37.131197Z","shell.execute_reply":"2025-02-04T18:23:03.472474Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Next we process the CryoET data, converting the tomograms and segmentations into PyTorch tensors and organizing them into a structured format. \n\nDetailed code breakdown:\n1. It iterates through each experimental run in the dataset.\n2. For each run, it loads the tomogram data and converts it to a PyTorch tensor.\n3. It processes each pickable object (particle type) getting its segmentation data and converting it to a binary mask.\n4. All segmentations for a run are combined into a single one-hot segmentation tensor.\n5. The tomogram and segmentation data are saved as PyTorch tensor files.\n6. The files for each run are compressed into a zip file to save space.\n7. The original uncompressed folder is deleted after zipping.","metadata":{}},{"cell_type":"code","source":"import torch\nimport zipfile\nfrom tqdm import tqdm\n\n# Define the output directory for processed tensor data\noutput_root = \"/kaggle/working/CZII-CryoETObjectIdentification_TensorData\"\n\n# Iterate through each experimental run\nfor run in tqdm(root.runs):\n    # Load the tomogram data and convert to PyTorch tensor\n    tomogram = torch.from_numpy(run.get_voxel_spacing(voxel_size).get_tomogram(tomo_type).numpy())\n    \n    # Initialize a list to store segmentations for each pickable object\n    all_segmentations = []\n\n    # Create output folder for this run\n    output_folder = os.path.join(output_root, run.name)\n    if not os.path.exists(output_folder):\n        os.makedirs(output_folder)\n\n    # Process each pickable object\n    for pickable_object in root.pickable_objects:\n        if pickable_object.is_particle==True:\n            # Get segmentation data for the current object\n            segmentation = run.get_segmentations(name=pickable_object.name, user_id=copick_user_name, voxel_size=voxel_size, is_multilabel=True)[0].numpy()\n            # Convert segmentation to binary mask and to PyTorch tensor\n            segmentation = torch.from_numpy((segmentation > 0).astype(int))\n            # Append the segmentation tensor to the list\n            all_segmentations.append(segmentation)\n            print(pickable_object.name, len(all_segmentations), torch.sum(segmentation))\n\n    # Combine all segmentations into one tensor\n    combined_segmentation = torch.stack(all_segmentations, dim=0)\n    print(combined_segmentation.shape)\n    print(tomogram.shape)\n\n    # Save tomogram and combined segmentation as PyTorch tensors\n    torch.save(tomogram, os.path.join(output_folder, run.name+\"_0_image.pt\"))\n    torch.save(combined_segmentation, os.path.join(output_folder, run.name+\"_0_label.pt\"))\n    \n    errors = []\n    \n    # Before we zip the file check that all of the picks are segmented properly\n    # (Commented out for faster processing)\n    \"\"\"\n    # Check if all picks are properly segmented\n    for obj in root.pickable_objects:\n        picks = run.get_picks(object_name=pickable_object.name, user_id=\"curation\")\n        for pick in picks:\n            location = pick['location']\n            x, y, z = int(location['x'] / voxel_size), int(location['y'] / voxel_size), int(location['z'] / voxel_size)\n            segmentation_value = combined_segmentation[obj.label-1, z, y, x].item()\n            print(segmentation_value)\n            if segmentation_value != 1:\n                errors.append({\n                    'run': run_folder,\n                    'object': obj.name,\n                    'pick_location': location,\n                    'segmentation_label': seg_label,\n                    'expected_label': obj.label\n                })\n\n    # Report any segmentation errors\n    if errors:\n        error_message = f\"Found {len(errors)} mismatches:\\n\"\n        # Limit to first 10 errors to avoid overwhelming output\n        for error in errors[:10]:  # Limit to first 10 errors to avoid overwhelming output\n            error_message += f\"Run: {error['run']}, Object: {error['object']}, \"\n            error_message += f\"Pick location: {error['pick_location']}, \"\n            error_message += f\"Segmentation label: {error['segmentation_label']}, \"\n            error_message += f\"Expected label: {error['expected_label']}\\n\"\n        if len(errors) > 10:\n            error_message += f\"... and {len(errors) - 10} more errors\"\n        raise ValueError(error_message)\n    else:\n        print(\"All picks match their corresponding segmentation labels.\")\n    \"\"\"\n    \n    # Create a zip file for the run\n    zip_filename = os.path.join(output_root, run.name + \".zip\")\n    with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:\n        for rt, dirs, files in os.walk(output_folder):\n            for file in files:\n                zipf.write(os.path.join(rt, file), \n                           os.path.relpath(os.path.join(rt, file), \n                                           os.path.join(output_folder, '..')))\n    \n    # Delete the original folder to save space\n    shutil.rmtree(output_folder)\n\nprint(\"All run folders have been zipped and original folders deleted.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:44:27.120479Z","iopub.execute_input":"2025-02-04T18:44:27.120943Z","iopub.status.idle":"2025-02-04T18:50:11.243673Z","shell.execute_reply.started":"2025-02-04T18:44:27.120909Z","shell.execute_reply":"2025-02-04T18:50:11.242417Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create a zip archive of the output_root directory\nshutil.make_archive(output_root, 'zip', output_root)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:51:55.365998Z","iopub.execute_input":"2025-02-04T18:51:55.366417Z","iopub.status.idle":"2025-02-04T18:53:04.462322Z","shell.execute_reply.started":"2025-02-04T18:51:55.366387Z","shell.execute_reply":"2025-02-04T18:53:04.461141Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Visualize Segmentation\n\nimport matplotlib.pyplot as plt\nimport numpy as np\n\ndef plot_tomogram_with_segmentation(tomogram, labels):\n    fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n    \n    # Define middle slice index for each dimension\n    slices = [tomogram.shape[i] // 2 for i in range(3)]\n    \n    # Define a colormap for the labels\n    cmap = plt.get_cmap('tab10')\n    \n    for i, ax in enumerate(axes):\n        # Plot the tomogram slice for each plane\n        if i == 0:\n            tomo_slice = tomogram[slices[0], :, :]\n            label_slice = labels[:, slices[0], :, :]\n            ax.set_title('YZ Plane')\n        elif i == 1:\n            tomo_slice = tomogram[:, slices[1], :]\n            label_slice = labels[:, :, slices[1], :]\n            ax.set_title('XZ Plane')\n        else:\n            tomo_slice = tomogram[:, :, slices[2]]\n            label_slice = labels[:, :, :, slices[2]]\n            ax.set_title('XY Plane')\n\n        # Plot the tomogram slice in grayscale\n        ax.imshow(tomo_slice, cmap='gray')\n        \n        # Overlay segmentation for each label\n        for j in range(labels.shape[0]):\n            # Create a mask for the current label\n            mask = label_slice[j] > 0\n            # Plot the mask with a unique color and some transparency\n            ax.imshow(np.ma.masked_where(~mask, mask), \n                      cmap=plt.cm.colors.ListedColormap([cmap(j)]), \n                      alpha=0.3)\n        ax.axis('off')\n    \n    plt.tight_layout()\n    plt.show()\n\n# Convert PyTorch tensors to numpy arrays for matplotlib compatibility\ntomogram_np = tomogram.numpy()\nlabels_np = combined_segmentation.numpy()\n\n# Call the function to create and display the plot\nplot_tomogram_with_segmentation(tomogram_np, labels_np)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-04T18:53:04.464024Z","iopub.execute_input":"2025-02-04T18:53:04.464343Z","iopub.status.idle":"2025-02-04T18:53:05.677008Z","shell.execute_reply.started":"2025-02-04T18:53:04.464314Z","shell.execute_reply":"2025-02-04T18:53:05.675920Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Test Dataset\n\nExports test dataset as torch files in same format.","metadata":{}},{"cell_type":"code","source":"import os\nimport shutil\n# Update the overlay\n# Define source and destination directories\nsource_dir = '/kaggle/input/czii-cryo-et-object-identification/test/static'\ndestination_dir = '/kaggle/working/test_static'\n\n# Create the destination directory if it doesn't exist\nos.makedirs(destination_dir, exist_ok=True)\n\nfor root, dirs, files in os.walk(source_dir):\n    # Create corresponding subdirectories in the destination\n    relative_path = os.path.relpath(root, source_dir)\n    target_dir = os.path.join(destination_dir, relative_path)\n    os.makedirs(target_dir, exist_ok=True)\n\n    for file in files:\n        if file.startswith(\"curation_0_\"):\n            new_filename = file\n        else:\n            new_filename = f\"curation_0_{file}\"\n\n        # Define full paths for the source and destination files\n        source_file = os.path.join(root, file)\n        destination_file = os.path.join(target_dir, new_filename)\n        \n        # Copy the file with the new name\n        shutil.copy2(source_file, destination_file)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:06:57.705128Z","iopub.execute_input":"2025-01-07T14:06:57.705600Z","iopub.status.idle":"2025-01-07T14:07:10.626477Z","shell.execute_reply.started":"2025-01-07T14:06:57.705561Z","shell.execute_reply":"2025-01-07T14:07:10.625107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Make a copick project\nimport copick\nimport numpy as np\n\nconfig_blob = \"\"\"{\n    \"name\": \"czii_cryoet_mlchallenge_2024\",\n    \"description\": \"2024 CZII CryoET ML Challenge training data.\",\n    \"version\": \"1.0.0\",\n\n    \"pickable_objects\": [\n        {\n            \"name\": \"apo-ferritin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"4V1W\",\n            \"label\": 1,\n            \"color\": [  0, 117, 220, 128],\n            \"radius\": 60,\n            \"map_threshold\": 0.0418\n        },\n        {\n            \"name\": \"beta-galactosidase\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6X1Q\",\n            \"label\": 3,\n            \"color\": [ 76,   0,  92, 128],\n            \"radius\": 90,\n            \"map_threshold\": 0.0578\n        },\n        {\n            \"name\": \"ribosome\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6EK0\",\n            \"label\": 4,\n            \"color\": [  0,  92,  49, 128],\n            \"radius\": 150,\n            \"map_threshold\": 0.0374\n        },\n        {\n            \"name\": \"thyroglobulin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6SCJ\",\n            \"label\": 5,\n            \"color\": [ 43, 206,  72, 128],\n            \"radius\": 130,\n            \"map_threshold\": 0.0278\n        },\n        {\n            \"name\": \"virus-like-particle\",\n            \"is_particle\": true,\n            \"label\": 6,\n            \"color\": [255, 204, 153, 128],\n            \"radius\": 135,\n            \"map_threshold\": 0.201\n        },\n        {\n            \"name\": \"membrane\",\n            \"is_particle\": false,\n            \"label\": 8,\n            \"color\": [100, 100, 100, 128]\n        },\n        {\n            \"name\": \"background\",\n            \"is_particle\": false,\n            \"label\": 9,\n            \"color\": [10, 150, 200, 128]\n        }\n    ],\n\n    \"overlay_root\": \"/kaggle/working/test_static\",\n\n    \"overlay_fs_args\": {\n        \"auto_mkdir\": true\n    },\n\n    \"static_root\": \"/kaggle/input/czii-cryo-et-object-identification/test/static\"\n}\"\"\"\n\ntest_copick_config_path = \"/kaggle/working/test_copick.config\"\ntest_output_overlay = \"/kaggle/working/test_static\"\n\nwith open(test_copick_config_path, \"w\") as f:\n    f.write(config_blob)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:07:10.632300Z","iopub.execute_input":"2025-01-07T14:07:10.632870Z","iopub.status.idle":"2025-01-07T14:07:10.643268Z","shell.execute_reply.started":"2025-01-07T14:07:10.632817Z","shell.execute_reply":"2025-01-07T14:07:10.641892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_root = copick.from_file(test_copick_config_path)\nvoxel_size = 10\ntomo_type = \"denoised\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:07:10.645311Z","iopub.execute_input":"2025-01-07T14:07:10.645879Z","iopub.status.idle":"2025-01-07T14:07:10.666444Z","shell.execute_reply.started":"2025-01-07T14:07:10.645826Z","shell.execute_reply":"2025-01-07T14:07:10.665240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport zipfile\nfrom tqdm import tqdm\n\noutput_root = \"/kaggle/working/TEST_CZII-CryoETObjectIdentification_TensorData\"\nfor run in tqdm(test_root.runs):\n    tomogram = torch.from_numpy(run.get_voxel_spacing(voxel_size).get_tomogram(tomo_type).numpy())\n    output_folder = os.path.join(output_root, run.name)\n\n    if not os.path.exists(output_folder):\n        os.makedirs(output_folder)\n    \n    torch.save(tomogram, os.path.join(output_folder, run.name+\"_0_image.pt\"))\n    \n    # Create a zip file for the run\n    zip_filename = os.path.join(output_root, run.name + \".zip\")\n    with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:\n        for rt, dirs, files in os.walk(output_folder):\n            for file in files:\n                zipf.write(os.path.join(rt, file), \n                           os.path.relpath(os.path.join(rt, file), \n                                           os.path.join(output_folder, '..')))\n    \n    # Delete the original folder\n    shutil.rmtree(output_folder)\n\nprint(\"All run folders have been zipped and original folders deleted.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:07:10.668184Z","iopub.execute_input":"2025-01-07T14:07:10.668584Z","iopub.status.idle":"2025-01-07T14:08:08.112728Z","shell.execute_reply.started":"2025-01-07T14:07:10.668550Z","shell.execute_reply":"2025-01-07T14:08:08.110946Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"shutil.make_archive(output_root, 'zip', output_root) ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:08:08.114660Z","iopub.execute_input":"2025-01-07T14:08:08.115191Z","iopub.status.idle":"2025-01-07T14:08:37.820357Z","shell.execute_reply.started":"2025-01-07T14:08:08.115137Z","shell.execute_reply":"2025-01-07T14:08:37.818873Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\n\ndef plot_tomogram_with_segmentation(tomogram):\n    fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n    \n    # Define slice indices for each dimension\n    slices = [tomogram.shape[i] // 2 for i in range(3)]\n    \n    # Define a colormap for the labels\n    cmap = plt.get_cmap('tab10')\n    \n    for i, ax in enumerate(axes):\n        # Plot the tomogram slice\n        if i == 0:\n            tomo_slice = tomogram[slices[0], :, :]\n            ax.set_title('YZ Plane')\n        elif i == 1:\n            tomo_slice = tomogram[:, slices[1], :]\n            ax.set_title('XZ Plane')\n        else:\n            tomo_slice = tomogram[:, :, slices[2]]\n            ax.set_title('XY Plane')\n        \n        ax.imshow(tomo_slice, cmap='gray')    \n        ax.axis('off')\n    \n    plt.tight_layout()\n    plt.show()\n\n# Assuming tomogram and labels are your PyTorch tensors\n# Convert them to numpy arrays for matplotlib\ntomogram_np = tomogram.numpy()\n\nplot_tomogram_with_segmentation(tomogram_np)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:08:37.822380Z","iopub.execute_input":"2025-01-07T14:08:37.822744Z","iopub.status.idle":"2025-01-07T14:08:38.502458Z","shell.execute_reply.started":"2025-01-07T14:08:37.822712Z","shell.execute_reply":"2025-01-07T14:08:38.500820Z"}},"outputs":[],"execution_count":null}]}