{"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":"code","source":"!pip install --upgrade \"numpy>=2.0,<3.0\"\n!pip install dask==2024.9.1\n!pip install --upgrade \"google-cloud-storage>=2.2.1,<3\"\n!pip install \"requests-toolbelt>=0.8.0,<1\"\n!pip install pydantic==1.10.0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-17T17:25:01.179809Z","iopub.execute_input":"2025-03-17T17:25:01.180231Z","iopub.status.idle":"2025-03-17T17:26:26.010826Z","shell.execute_reply.started":"2025-03-17T17:25:01.180198Z","shell.execute_reply":"2025-03-17T17:26:26.008938Z"}},"outputs":[],"execution_count":null},{"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-03-17T17:26:41.627854Z","iopub.execute_input":"2025-03-17T17:26:41.628327Z","iopub.status.idle":"2025-03-17T17:27:24.407317Z","shell.execute_reply.started":"2025-03-17T17:26:41.628285Z","shell.execute_reply":"2025-03-17T17:27:24.405462Z"}},"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-03-17T17:27:30.620724Z","iopub.execute_input":"2025-03-17T17:27:30.621185Z","iopub.status.idle":"2025-03-17T17:27:52.948232Z","shell.execute_reply.started":"2025-03-17T17:27:30.621144Z","shell.execute_reply":"2025-03-17T17:27:52.946239Z"}},"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    #将 config_blob 中的 JSON 字符串写入到 copick_config_path 指定的文件中。\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    #os.walk 返回一个三元组 (root, dirs, files)：root: 当前遍历的目录路径。dirs: 当前目录下的子目录列表。files: 当前目录下的文件列表。\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-03-17T17:28:32.403280Z","iopub.execute_input":"2025-03-17T17:28:32.403951Z","iopub.status.idle":"2025-03-17T17:28:32.682202Z","shell.execute_reply.started":"2025-03-17T17:28:32.403902Z","shell.execute_reply":"2025-03-17T17:28:32.680916Z"}},"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-03-17T17:28:40.808767Z","iopub.execute_input":"2025-03-17T17:28:40.809257Z","iopub.status.idle":"2025-03-17T17:28:46.842213Z","shell.execute_reply.started":"2025-03-17T17:28:40.809210Z","shell.execute_reply":"2025-03-17T17:28:46.841126Z"}},"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\"  #定义生成数据的用户 ID\ncopick_segmentation_name = \"paintedPicks\"  #定义输出分割掩膜的名称\nvoxel_size = 10  #体素大小（10 Å）\ntomo_type = \"denoised\" #指定要使用的断层扫描数据","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-17T17:28:59.146368Z","iopub.execute_input":"2025-03-17T17:28:59.147422Z","iopub.status.idle":"2025-03-17T17:28:59.155306Z","shell.execute_reply.started":"2025-03-17T17:28:59.147375Z","shell.execute_reply":"2025-03-17T17:28:59.153704Z"}},"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      #用于将生成的分割掩膜写入 Copick 项目\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:     #遍历配置文件中的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):  #遍历 Copick 项目中的每个 Run（即每个断层扫描数据集）即TS_ _等等这些文件夹\n        tomo = run.get_voxel_spacing(10)      #获取指定体素大小（10 Å）\n        tomo = tomo.get_tomogram(tomo_type).numpy() #选择denoised的断层扫描数据，并转换为 NumPy 数组。\n        target = np.zeros(tomo.shape, dtype=np.uint8)  #初始化一个全零数组，用于存储分割掩膜。\n        for pickable_object in root.pickable_objects:\n            #获取当前 Run 中指定粒子的标注坐标（通过 user_id=\"curation\" 指定标注来源）\n            pick = run.get_picks(object_name=pickable_object.name, user_id=\"curation\")\n            #根据不同粒子类型调整半径（例如，apo-ferritin 的半径缩小为 50%）。\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                #将粒子坐标转换为分割掩膜：pick[0]: 粒子坐标。target: 目标掩膜数组。radius * scale: 调整后的粒子半径。label: 粒子的标签值（用于区分不同粒子类型）。\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        #将生成的分割掩膜保存到 Copick 项目中：run: 当前 Run 对象。target: 分割掩膜数据。copick_user_name: 用户 ID。name=copick_segmentation_name: 掩膜名称（如 paintedPicks）\n        write.segmentation(run, target, copick_user_name, name=copick_segmentation_name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-17T17:29:02.677867Z","iopub.execute_input":"2025-03-17T17:29:02.678387Z","iopub.status.idle":"2025-03-17T17:29:22.234127Z","shell.execute_reply.started":"2025-03-17T17:29:02.678345Z","shell.execute_reply":"2025-03-17T17:29:22.232627Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Get tomograms and their segmentaion masks (from picks) arrays","metadata":{}},{"cell_type":"code","source":"#从 Copick 项目中提取断层扫描图像（tomogram）和对应的分割标签（segmentation mask），并将其整理成字典列表\ndata_dicts = []\nfor run in tqdm(root.runs):\n    #获取指定体素大小（voxel_size，如 10 Å）的数据，加载特定类型的断层扫描数据（如 denoised 去噪数据），将数据转换为 NumPy 数组，方便后续处理。\n    #tomogram: 3D 数组，表示断层扫描的灰度图像（形状为 (Z, Y, X)）\n    tomogram = run.get_voxel_spacing(voxel_size).get_tomogram(tomo_type).numpy()\n    #segmentation: 3D 数组，表示分割标签（形状与 tomogram 相同，每个像素值为整数标签）\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#np.unique(): 打印第一个 Run 的分割掩膜中包含哪些唯一标签值。 \nprint(np.unique(data_dicts[0]['label']))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-17T17:29:28.666790Z","iopub.execute_input":"2025-03-17T17:29:28.667250Z","iopub.status.idle":"2025-03-17T17:29:38.819689Z","shell.execute_reply.started":"2025-03-17T17:29:28.667205Z","shell.execute_reply":"2025-03-17T17:29:38.818281Z"}},"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-03-17T17:29:53.472738Z","iopub.execute_input":"2025-03-17T17:29:53.473183Z","iopub.status.idle":"2025-03-17T17:29:53.918552Z","shell.execute_reply.started":"2025-03-17T17:29:53.473148Z","shell.execute_reply":"2025-03-17T17:29:53.917113Z"}},"outputs":[],"execution_count":null}]}