{"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":"gpu","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"}],"dockerImageVersionId":30805,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Install pkgs","metadata":{}},{"cell_type":"code","source":"!pip install git+https://github.com/copick/copick-utils.git matplotlib tqdm copick \n!pip install -q \"monai-weekly[mlflow]\"","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:18:05.998573Z","iopub.execute_input":"2024-11-09T16:18:05.999048Z","iopub.status.idle":"2024-11-09T16:19:29.656992Z","shell.execute_reply.started":"2024-11-09T16:18:05.998988Z","shell.execute_reply":"2024-11-09T16:19:29.655369Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install zarr","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:19:29.660409Z","iopub.execute_input":"2024-11-09T16:19:29.661008Z","iopub.status.idle":"2024-11-09T16:19:45.251438Z","shell.execute_reply.started":"2024-11-09T16:19:29.66095Z","shell.execute_reply":"2024-11-09T16:19:45.2499Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install copick","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:19:45.253312Z","iopub.execute_input":"2024-11-09T16:19:45.253779Z","iopub.status.idle":"2024-11-09T16:20:01.362623Z","shell.execute_reply.started":"2024-11-09T16:19:45.253732Z","shell.execute_reply":"2024-11-09T16:20:01.360173Z"},"trusted":true},"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":{"execution":{"iopub.status.busy":"2024-11-09T16:20:01.367011Z","iopub.execute_input":"2024-11-09T16:20:01.367641Z","iopub.status.idle":"2024-11-09T16:20:01.609167Z","shell.execute_reply.started":"2024-11-09T16:20:01.367575Z","shell.execute_reply":"2024-11-09T16:20:01.60755Z"},"trusted":true},"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 monai.data import DataLoader, Dataset, CacheDataset, decollate_batch\nfrom monai.transforms import (\n    Compose, \n    EnsureChannelFirstd, \n    Orientationd,  \n    AsDiscrete,  \n    RandFlipd, \n    RandRotate90d, \n    NormalizeIntensityd,\n    RandCropByLabelClassesd,\n)\nfrom monai.networks.nets import UNet\nfrom monai.losses import DiceLoss, FocalLoss, TverskyLoss\nfrom monai.metrics import DiceMetric, ConfusionMatrixMetric\nimport mlflow\nimport mlflow.pytorch","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:20:01.611079Z","iopub.execute_input":"2024-11-09T16:20:01.611682Z","iopub.status.idle":"2024-11-09T16:21:12.448673Z","shell.execute_reply.started":"2024-11-09T16:20:01.611625Z","shell.execute_reply":"2024-11-09T16:21:12.447458Z"},"trusted":true},"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":{"execution":{"iopub.status.busy":"2024-11-09T16:21:12.450231Z","iopub.execute_input":"2024-11-09T16:21:12.4515Z","iopub.status.idle":"2024-11-09T16:21:12.463855Z","shell.execute_reply.started":"2024-11-09T16:21:12.45145Z","shell.execute_reply":"2024-11-09T16:21:12.462085Z"},"trusted":true},"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                target = segmentation_from_picks.from_picks(pick[0], \n                                                            target, \n                                                            target_objects[pickable_object.name]['radius'] * 0.8,\n                                                            target_objects[pickable_object.name]['label']\n                                                            )\n        write.segmentation(run, target, copick_user_name, name=copick_segmentation_name)","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:21:12.465291Z","iopub.execute_input":"2024-11-09T16:21:12.465759Z","iopub.status.idle":"2024-11-09T16:21:40.603322Z","shell.execute_reply.started":"2024-11-09T16:21:12.465713Z","shell.execute_reply":"2024-11-09T16:21:40.6021Z"},"trusted":true},"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":{"execution":{"iopub.status.busy":"2024-11-09T16:21:40.604969Z","iopub.execute_input":"2024-11-09T16:21:40.60552Z","iopub.status.idle":"2024-11-09T16:21:49.459852Z","shell.execute_reply.started":"2024-11-09T16:21:40.605463Z","shell.execute_reply":"2024-11-09T16:21:49.458434Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_dicts[0]['label'].shape","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:21:49.461527Z","iopub.execute_input":"2024-11-09T16:21:49.462011Z","iopub.status.idle":"2024-11-09T16:21:49.473774Z","shell.execute_reply.started":"2024-11-09T16:21:49.461965Z","shell.execute_reply":"2024-11-09T16:21:49.472282Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_dicts[0]['image'].shape","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:21:49.477443Z","iopub.execute_input":"2024-11-09T16:21:49.478494Z","iopub.status.idle":"2024-11-09T16:21:49.490877Z","shell.execute_reply.started":"2024-11-09T16:21:49.478439Z","shell.execute_reply":"2024-11-09T16:21:49.489083Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"image_names = ['TS_5_4', 'TS_69_2', 'TS_6_4', 'TS_73_6', 'TS_86_3', 'TS_99_9', 'TS_6_6']\n\nfor i in range(7):\n    \n    with open(f'train_image_{image_names[i]}.npy', 'wb') as f:\n        np.save(f, data_dicts[i]['image'])\n        \n    with open(f'train_label_{image_names[i]}.npy', 'wb') as f:\n        np.save(f, data_dicts[i]['label'])","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:21:49.492786Z","iopub.execute_input":"2024-11-09T16:21:49.493406Z","iopub.status.idle":"2024-11-09T16:21:53.979191Z","shell.execute_reply.started":"2024-11-09T16:21:49.493319Z","shell.execute_reply":"2024-11-09T16:21:53.977708Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls -lh","metadata":{"execution":{"iopub.status.busy":"2024-11-09T16:21:53.980853Z","iopub.execute_input":"2024-11-09T16:21:53.981271Z","iopub.status.idle":"2024-11-09T16:21:55.185045Z","shell.execute_reply.started":"2024-11-09T16:21:53.981229Z","shell.execute_reply":"2024-11-09T16:21:55.183409Z"},"trusted":true},"outputs":[],"execution_count":null}]}