{"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":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Overview\n\nIn this notebook, we’ll create 3D particle masks for tomograms, using the true particle sizes based on their center coordinates and radii. These accurate masks will then be visualized slice by slice.\n\nThis is helpful for showing the real particle density and can be used as a preprocessing step for training segmentation models.","metadata":{}},{"cell_type":"code","source":"!pip install zarr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:39:44.324391Z","iopub.execute_input":"2024-11-29T07:39:44.324834Z","iopub.status.idle":"2024-11-29T07:39:59.291982Z","shell.execute_reply.started":"2024-11-29T07:39:44.324800Z","shell.execute_reply":"2024-11-29T07:39:59.290936Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nimport json\nimport matplotlib.pyplot as plt\nimport zarr\nfrom tqdm import tqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:40:27.438725Z","iopub.execute_input":"2024-11-29T07:40:27.439160Z","iopub.status.idle":"2024-11-29T07:40:27.946957Z","shell.execute_reply.started":"2024-11-29T07:40:27.439126Z","shell.execute_reply":"2024-11-29T07:40:27.945948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Paths to data\nTRAIN_BASE = '/kaggle/input/czii-cryo-et-object-identification/train/static/ExperimentRuns'\nTEST_BASE = '/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns'\nOVERLAY_BASE = '/kaggle/input/czii-cryo-et-object-identification/train/overlay/ExperimentRuns'\n\n### Radii of particles of different types\n### (taken from public notebook provided by the team)\nPARTICLE_RADIUS = {\n    'ribosome': 150.0,\n    'virus-like-particle': 135.0,\n    'apo-ferritin': 60.0,\n    'beta-galactosidase': 90.0,\n    'thyroglobulin': 130.0,\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:40:27.993232Z","iopub.execute_input":"2024-11-29T07:40:27.994128Z","iopub.status.idle":"2024-11-29T07:40:28.000008Z","shell.execute_reply.started":"2024-11-29T07:40:27.994083Z","shell.execute_reply":"2024-11-29T07:40:27.998878Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Utils\n\ndef get_training_experiments():\n    return os.listdir(TRAIN_BASE)\n    \ndef get_test_experiments():\n    return os.listdir(TEST_BASE)\n\ndef get_experiment_path(experiment, type='denoised', test=False):\n    base = TEST_BASE if test else TRAIN_BASE\n    return base + '/' + experiment + '/VoxelSpacing10.000/' + type + '.zarr'\n\ndef open_experiment(experiment, test=False):\n    return zarr.open(get_experiment_path(experiment, test=test), mode='r')\n\ndef get_label_path(experiment, particle_type):\n    return OVERLAY_BASE + '/' + experiment + '/Picks/' + particle_type + '.json'\n\ndef open_labels_file(experiment, particle_type):\n    path = get_label_path(experiment, particle_type)\n    with open(path, 'r') as f:\n        data = json.load(f)\n    return data\n\ndef get_labels_with_radius(experiment, scaling_factor=10):\n    labels = {}\n    for particle_type in PARTICLE_RADIUS.keys():\n    \n        ### Load file with labels\n        raw_data = open_labels_file(experiment, particle_type)\n    \n        ### Save\n        labels[particle_type] = []\n        for item in raw_data['points']:\n            labels[particle_type].append(\n                (\n                    abs(item['location']['x'] // scaling_factor),\n                    abs(item['location']['y'] // scaling_factor),\n                    abs(item['location']['z'] // scaling_factor),\n                    abs(PARTICLE_RADIUS[particle_type] // scaling_factor)\n                )\n            )\n    return labels\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:40:28.569550Z","iopub.execute_input":"2024-11-29T07:40:28.569962Z","iopub.status.idle":"2024-11-29T07:40:28.581822Z","shell.execute_reply.started":"2024-11-29T07:40:28.569926Z","shell.execute_reply":"2024-11-29T07:40:28.580642Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"quality_to_scaling_factor = {\n    'high': 10,\n    'medium': 20,\n    'low': 40,\n}\nquality_to_dataset_index = {\n    'high': 0,\n    'medium': 1,\n    'low': 2,\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:40:29.144771Z","iopub.execute_input":"2024-11-29T07:40:29.145165Z","iopub.status.idle":"2024-11-29T07:40:29.150347Z","shell.execute_reply.started":"2024-11-29T07:40:29.145132Z","shell.execute_reply":"2024-11-29T07:40:29.149132Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Get sample and particles ","metadata":{}},{"cell_type":"code","source":"experiment = 'TS_5_4'\n\n### Get sample\nsample = open_experiment('TS_5_4', test=False)\n\n### Get particals coordinates for chosen quality\nquality = 'high'\nscaling_factor = quality_to_scaling_factor[quality] \nlabels = get_labels_with_radius(experiment, scaling_factor=scaling_factor)\n\ni = quality_to_dataset_index[quality]\ndataset = sample[i]\ndataset.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:40:33.908997Z","iopub.execute_input":"2024-11-29T07:40:33.909411Z","iopub.status.idle":"2024-11-29T07:40:33.924350Z","shell.execute_reply.started":"2024-11-29T07:40:33.909377Z","shell.execute_reply":"2024-11-29T07:40:33.923282Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# labels['thyroglobulin']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T08:20:07.232458Z","iopub.execute_input":"2024-11-29T08:20:07.232872Z","iopub.status.idle":"2024-11-29T08:20:07.237804Z","shell.execute_reply.started":"2024-11-29T08:20:07.232837Z","shell.execute_reply":"2024-11-29T08:20:07.236600Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create 3D mask \n\nTo create the 3D mask, we take each particle with its center coordinates and radius, then fill the mask with the corresponding class label for each particle.","metadata":{}},{"cell_type":"code","source":"PARTICLE_TO_CLASS = {\n    'ribosome': 1,\n    'virus-like-particle': 2,\n    'apo-ferritin': 3,\n    'beta-galactosidase': 4,\n    'thyroglobulin': 5,\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T07:46:32.847996Z","iopub.execute_input":"2024-11-29T07:46:32.848416Z","iopub.status.idle":"2024-11-29T07:46:32.853594Z","shell.execute_reply.started":"2024-11-29T07:46:32.848381Z","shell.execute_reply":"2024-11-29T07:46:32.852511Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### Create 3D mesh of coordinates \nz = np.arange(dataset.shape[0])\ny = np.arange(dataset.shape[1])\nx = np.arange(dataset.shape[2])\nZ, Y, X = np.meshgrid(z, y, x, indexing=\"ij\")\n\n### Init mask with zeros\nmask = np.zeros(dataset.shape, dtype=np.uint8)\n\nfor particle_type in PARTICLE_TO_CLASS.keys():\n    ### Add each circle to the mask\n    for cx, cy, cz, r in tqdm(labels[particle_type]):\n        \n        ### Calculate the distance to the circle's center\n        distance = np.sqrt((X - cx)**2 + (Y - cy)**2 + (Z - cz)**2)\n        \n        ### Fill the circle with class number\n        mask[distance <= r] = PARTICLE_TO_CLASS[particle_type]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T08:22:00.564772Z","iopub.execute_input":"2024-11-29T08:22:00.565210Z","iopub.status.idle":"2024-11-29T08:24:59.111171Z","shell.execute_reply.started":"2024-11-29T08:22:00.565173Z","shell.execute_reply":"2024-11-29T08:24:59.110069Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Visualization","metadata":{}},{"cell_type":"code","source":"for layer in [100, 108, 116, 120]:\n\n    fig, ax = plt.subplots(1, 3, figsize=(15, 6), constrained_layout=True)\n\n    ax[0].imshow(dataset[layer], cmap='gray', vmin=-0.00005, vmax=0.00005)\n\n    ax[1].imshow(dataset[layer], cmap='gray', vmin=-0.00005, vmax=0.00005)\n    ax[1].imshow(mask[layer], alpha=0.5) \n\n    mask_plot = ax[2].imshow(mask[layer]) \n    \n    cbar = fig.colorbar(mask_plot,  orientation='vertical', pad=0.1, aspect=5)\n\n    cbar.set_ticks(list(PARTICLE_TO_CLASS.values()))  \n    cbar.set_ticklabels(list(PARTICLE_TO_CLASS.keys())) \n\n    for j in range(3):\n        ax[j].set_xticks([])\n        ax[j].set_yticks([])\n\n    plt.suptitle('Layer {}: image & annotation'.format(layer), fontsize=16, y=0.95)\n    plt.show()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-29T08:20:40.734761Z","iopub.execute_input":"2024-11-29T08:20:40.735201Z","iopub.status.idle":"2024-11-29T08:20:48.295823Z","shell.execute_reply.started":"2024-11-29T08:20:40.735165Z","shell.execute_reply":"2024-11-29T08:20:48.294695Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Conclusion\nAs we can see, **the true particle density is quite high**. This isn’t obvious when visualizing slices with only the particles whose centers lie on that particular slice.\nBy using 3D masks, we get a clearer picture of the overall density and better understand the distribution of particles across the entire volume.","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}