{"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":"**Note:** This is training notebook only. Inference ain't included in . \nAnybody who wants to use this notebook for inference purposes is most welcome.","metadata":{}},{"cell_type":"code","source":"!pip install -q git+https://github.com/copick/copick-utils.git matplotlib tqdm copick \n!pip install -q \"monai-weekly[mlflow]\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:54:19.741472Z","iopub.execute_input":"2025-01-07T14:54:19.741872Z","iopub.status.idle":"2025-01-07T14:55:29.881651Z","shell.execute_reply.started":"2025-01-07T14:54:19.741824Z","shell.execute_reply":"2025-01-07T14:55:29.880311Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q zarr\n!pip install -q copick","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:55:29.884205Z","iopub.execute_input":"2025-01-07T14:55:29.884560Z","iopub.status.idle":"2025-01-07T14:55:51.270818Z","shell.execute_reply.started":"2025-01-07T14:55:29.884526Z","shell.execute_reply":"2025-01-07T14:55:51.269379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Make a copick project\nimport os\nimport shutil\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\": 2,\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\": 3,\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\": 4,\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\": 5,\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\": 6,\n            \"color\": [100, 100, 100, 128]\n        },\n        {\n            \"name\": \"background\",\n            \"is_particle\": false,\n            \"label\": 7,\n            \"color\": [10, 150, 200, 128]\n        }\n    ],\n\n    \"overlay_root\": \"/kaggle/working/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\ncopick_config_path = \"/kaggle/working/copick.config\"\noutput_overlay = \"/kaggle/working/overlay\"\n\nwith open(copick_config_path, \"w\") as f:\n    f.write(config_blob)\n    \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/overlay'\n\n# Walk through the source directory\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    # Copy and rename each file\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        \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)\n        print(f\"Copied {source_file} to {destination_file}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:55:51.273414Z","iopub.execute_input":"2025-01-07T14:55:51.273940Z","iopub.status.idle":"2025-01-07T14:55:51.478726Z","shell.execute_reply.started":"2025-01-07T14:55:51.273882Z","shell.execute_reply":"2025-01-07T14:55:51.477593Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom pathlib import Path\nimport torch\nimport torchinfo\nimport zarr, copick\nfrom tqdm import tqdm\nfrom glob import glob\nimport json","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:55:51.480106Z","iopub.execute_input":"2025-01-07T14:55:51.480439Z","iopub.status.idle":"2025-01-07T14:55:56.740948Z","shell.execute_reply.started":"2025-01-07T14:55:51.480405Z","shell.execute_reply":"2025-01-07T14:55:56.739862Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prepare the dataset\n## 1. Get copick root","metadata":{}},{"cell_type":"code","source":"root = copick.from_file(copick_config_path)\n\ncopick_user_name = \"copickUtils\"\ncopick_segmentation_name = \"paintedPicks\"\nvoxel_size = 10\ntomo_type = \"denoised\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:55:56.742489Z","iopub.execute_input":"2025-01-07T14:55:56.743710Z","iopub.status.idle":"2025-01-07T14:55:56.750409Z","shell.execute_reply.started":"2025-01-07T14:55:56.743657Z","shell.execute_reply":"2025-01-07T14:55:56.749174Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Generate multi-class segmentation masks from picks, and saved them to the copick overlay directory (one-time)","metadata":{}},{"cell_type":"code","source":"from copick_utils.segmentation import segmentation_from_picks\nimport copick_utils.writers.write as write\nfrom collections import defaultdict\n\n# Just do this once\ngenerate_masks = True\n\nif generate_masks:\n    target_objects = defaultdict(dict)\n    for object in root.pickable_objects:\n        if object.is_particle:\n            target_objects[object.name]['label'] = object.label\n            target_objects[object.name]['radius'] = object.radius\n\n\n    for run in tqdm(root.runs):\n        tomo = run.get_voxel_spacing(10)\n        tomo = tomo.get_tomogram(tomo_type).numpy()\n        target = np.zeros(tomo.shape, dtype=np.uint8)\n        for pickable_object in root.pickable_objects:\n            pick = run.get_picks(object_name=pickable_object.name, user_id=\"curation\")\n            if len(pick):  \n                if pick[0].pickable_object_name == 'apo-ferritin': \n                    scale = 0.5\n                elif pick[0].pickable_object_name == 'beta-galactosidase': \n                    scale = 0.5\n                elif pick[0].pickable_object_name == 'ribosome':\n                    scale = 1/3\n                elif pick[0].pickable_object_name == 'thyroglobulin':\n                    scale = 1/3\n                elif pick[0].pickable_object_name == 'virus-like-particle':\n                    scale = 1/3\n                print(target_objects[pickable_object.name], target_objects[pickable_object.name]['radius'] * scale)\n                target = segmentation_from_picks.from_picks(pick[0], \n                                                            target, \n                                                            target_objects[pickable_object.name]['radius'] * scale,\n                                                            target_objects[pickable_object.name]['label']\n                                                            )\n            #print(target_objects)\n            if len(pick) != 0: \n                print(' ')\n        write.segmentation(run, target, copick_user_name, name=copick_segmentation_name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:55:56.771325Z","iopub.execute_input":"2025-01-07T14:55:56.771667Z","iopub.status.idle":"2025-01-07T14:56:14.596714Z","shell.execute_reply.started":"2025-01-07T14:55:56.771636Z","shell.execute_reply":"2025-01-07T14:56:14.595653Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Get tomograms and their segmentaion masks (from picks) arrays","metadata":{}},{"cell_type":"code","source":"data_dicts = []\nfor run in tqdm(root.runs):\n    tomogram = run.get_voxel_spacing(voxel_size).get_tomogram(tomo_type).numpy()\n    segmentation = run.get_segmentations(name=copick_segmentation_name, user_id=copick_user_name, voxel_size=voxel_size, is_multilabel=True)[0].numpy()\n    data_dicts.append({\"image\": tomogram, \"label\": segmentation})\n    \nprint(np.unique(data_dicts[0]['label']))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:56:14.598206Z","iopub.execute_input":"2025-01-07T14:56:14.598478Z","iopub.status.idle":"2025-01-07T14:56:23.073922Z","shell.execute_reply.started":"2025-01-07T14:56:14.598451Z","shell.execute_reply":"2025-01-07T14:56:23.072726Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Visualize the tomogram and painted segmentation from ground-truth picks","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Plot the images\nplt.figure(figsize=(15, 5))\n\nplt.subplot(1, 2, 1)\nplt.title('Tomogram')\nplt.imshow(data_dicts[1]['image'][110],cmap='gray')\nplt.axis('off')\n\nplt.subplot(1, 2, 1)\nplt.title('Painted Segmentation from Picks')\nplt.imshow(data_dicts[1]['label'][110], cmap='viridis', alpha=0.5)\nplt.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:56:24.875149Z","iopub.execute_input":"2025-01-07T14:56:24.875510Z","iopub.status.idle":"2025-01-07T14:56:25.316493Z","shell.execute_reply.started":"2025-01-07T14:56:24.875460Z","shell.execute_reply":"2025-01-07T14:56:25.315314Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Plot the images\nplt.figure(figsize=(15, 5))\n\nplt.subplot(1, 2, 1)\nplt.title('Tomogram')\nplt.imshow(data_dicts[1]['image'][110],cmap='gray')\nplt.axis('off')\n\nplt.subplot(1, 2, 1)\nplt.title('Painted Segmentation from Picks')\nplt.imshow(data_dicts[1]['label'][110], cmap='viridis', alpha=0.5)\nplt.axis('off')\n\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-07T14:56:25.317910Z","iopub.execute_input":"2025-01-07T14:56:25.318238Z","iopub.status.idle":"2025-01-07T14:56:25.680227Z","shell.execute_reply.started":"2025-01-07T14:56:25.318206Z","shell.execute_reply":"2025-01-07T14:56:25.679067Z"}},"outputs":[],"execution_count":null}]}