{"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":"\"Om Namah Shivaya!! 🙏🙏\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-03T01:59:00.605901Z","iopub.execute_input":"2024-12-03T01:59:00.606303Z","iopub.status.idle":"2024-12-03T01:59:00.640265Z","shell.execute_reply.started":"2024-12-03T01:59:00.606256Z","shell.execute_reply":"2024-12-03T01:59:00.639078Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Updates\n- Version 2: ","metadata":{}},{"cell_type":"code","source":"!pip install -q git+https://github.com/copick/copick-utils.git\n!pip install -q copick zarr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T01:59:00.642699Z","iopub.execute_input":"2024-12-03T01:59:00.643211Z","iopub.status.idle":"2024-12-03T01:59:56.097624Z","shell.execute_reply.started":"2024-12-03T01:59:00.643158Z","shell.execute_reply":"2024-12-03T01:59:56.096480Z"}},"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-amylase\",\n            \"is_particle\": true,\n            \"pdb_id\": \"1FA2\",\n            \"label\": 2,\n            \"color\": [153,  63,   0, 128],\n            \"radius\": 65,\n            \"map_threshold\": 0.035\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/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}\"\"\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T01:59:56.099394Z","iopub.execute_input":"2024-12-03T01:59:56.099815Z","iopub.status.idle":"2024-12-03T01:59:56.107728Z","shell.execute_reply.started":"2024-12-03T01:59:56.099771Z","shell.execute_reply":"2024-12-03T01:59:56.106664Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"copick_config_path = \"/kaggle/working/copick.config\"\noutput_png = \"/kaggle/working/train_png_normalized\"\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    \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        # 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        print(source_file, destination_file)\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":"2024-12-03T01:59:56.109059Z","iopub.execute_input":"2024-12-03T01:59:56.109397Z","iopub.status.idle":"2024-12-03T01:59:56.295413Z","shell.execute_reply.started":"2024-12-03T01:59:56.109363Z","shell.execute_reply":"2024-12-03T01:59:56.294218Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prepare Dataset","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nimport torch\nimport torchinfo\nimport zarr, copick\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T01:59:56.297940Z","iopub.execute_input":"2024-12-03T01:59:56.298309Z","iopub.status.idle":"2024-12-03T02:00:01.772857Z","shell.execute_reply.started":"2024-12-03T01:59:56.298275Z","shell.execute_reply":"2024-12-03T02:00:01.771754Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 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\"\nprecision = \"8bit\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:00:01.773876Z","iopub.execute_input":"2024-12-03T02:00:01.774582Z","iopub.status.idle":"2024-12-03T02:00:01.781743Z","shell.execute_reply.started":"2024-12-03T02:00:01.774535Z","shell.execute_reply":"2024-12-03T02:00:01.780562Z"}},"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":"%%writefile boundingboxes_from_picks.py\nimport numpy as np\nimport pandas as pd\nimport zarr\nimport copick\n\ndef from_picks(pick, \n               run, \n               seg_volume,\n               radius: float = 10.0, \n               label_value: int = 1,\n               voxel_spacing: float = 10):\n    \"\"\"\n    Generates bounding box information frame by frame for each pick.\n\n    Parameters:\n    -----------\n    pick : copick.models.CopickPicks\n        Copick object containing `points`, where each point has a `location` attribute with `x`, `y`, `z` coordinates.\n    run : copick.impl.filesystem.CopickRunFSSpec\n        Copick run contains specific information like `name`, voxel spacings and total number of picks.\n    seg_volume : numpy.ndarray\n        3D segmentation volume (numpy array) where the spheres are painted. Shape should be (Z, Y, X).\n    radius : float, optional\n        The radius of the spheres to be inserted in physical units (not voxel units). Default is 10.0.\n    label_value : int, optional\n        The integer value used to label the sphere regions in the segmentation volume. Default is 1.\n    voxel_spacing : float, optional\n        The spacing of voxels in the segmentation volume, used to scale the radius of the spheres. Default is 10.\n\n    Returns:\n    --------\n    pandas.DataFrame\n        A DataFrame containing bounding box information for each sphere, in the yolo format (x_center, y_center, h, w)\n    \"\"\"\n\n    # Adjust radius for voxel spacing\n    radius_voxel = radius / voxel_spacing\n    delta = int(np.ceil(radius_voxel))\n    \n    # Get volume dimensions, (184(z), 630(x), 630(y))\n    vol_shape_z, vol_shape_x, vol_shape_y = seg_volume.shape\n\n    # Adjust the pick's location for voxel spacing\n    cx_voxel = np.array([pnt.location.x / voxel_spacing for pnt in pick.points])\n    cy_voxel = np.array([pnt.location.y / voxel_spacing for pnt in pick.points])\n    cz_voxel = np.array([pnt.location.z / voxel_spacing for pnt in pick.points])\n\n    # Loop over frame by frame along z-axis\n    bboxes = []\n    for zframe in range(vol_shape_z):\n        # Get effective radius in each frame, for all points in a pick\n        rframe = np.sqrt(np.maximum(radius_voxel**2 - (cz_voxel-zframe)**2, 0.))\n\n        # Get cx, cy, r with raidus > 0\n        assert vol_shape_x == vol_shape_y\n        cx = np.clip(cx_voxel[rframe > 0], delta, vol_shape_x - delta) / vol_shape_x\n        cy = np.clip(cy_voxel[rframe > 0], delta, vol_shape_y - delta) / vol_shape_y\n        rframe = np.clip(rframe[rframe > 0], 0, delta) / vol_shape_x\n\n        # Convert into a dataframe\n        bboxes_frame = pd.DataFrame({\"exp_name\": run.name,\n                                     \"frame\": zframe, \n                                     \"label\": label_value,\n                                     \"x_center\": cx,\n                                     \"y_center\": cy,\n                                     \"height\": rframe * 2,\n                                     \"width\": rframe * 2})\n        bboxes.append(bboxes_frame)\n\n    # Concat all bboxes\n    bboxes = pd.concat(bboxes, axis=0).reset_index(drop=True)\n    return bboxes","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:00:01.783471Z","iopub.execute_input":"2024-12-03T02:00:01.783897Z","iopub.status.idle":"2024-12-03T02:00:01.798448Z","shell.execute_reply.started":"2024-12-03T02:00:01.783861Z","shell.execute_reply":"2024-12-03T02:00:01.797388Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import boundingboxes_from_picks\nfrom 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\ngenerate_bboxs = True\n\ntarget_objects = defaultdict(dict)\nfor 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# Generate masks for each pick\nif generate_masks:\n    for run in tqdm(root.runs):\n        tomo = run.get_voxel_spacing(voxel_size)\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 pickable_object.is_particle and 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                                                            voxel_spacing=voxel_size\n                                                            )\n\n        write.segmentation(run, target, copick_user_name, name=copick_segmentation_name)\n\n# Generate bbox for each pick\nbboxes_all = pd.DataFrame()\nif generate_bboxs:\n    for run in tqdm(root.runs):\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 pickable_object.is_particle and len(pick):\n                target_bboxs_df = boundingboxes_from_picks.from_picks(pick[0], \n                                                                      run,\n                                                                      target, \n                                                                      target_objects[pickable_object.name]['radius'],\n                                                                      target_objects[pickable_object.name]['label'],\n                                                                      voxel_spacing=voxel_size\n                                                                      )\n                \n                bboxes_all = bboxes_all._append(target_bboxs_df, ignore_index=True)\n\nbboxes_all.to_csv('train_bounding_boxes.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:00:01.799614Z","iopub.execute_input":"2024-12-03T02:00:01.799949Z","iopub.status.idle":"2024-12-03T02:00:22.839571Z","shell.execute_reply.started":"2024-12-03T02:00:01.799915Z","shell.execute_reply":"2024-12-03T02:00:22.838447Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3. Get tomograms and their segmentaion masks (from picks) arrays","metadata":{}},{"cell_type":"code","source":"from PIL import Image\n\ndef normalise_by_percentile(data, min=1, max=99):\n    min = np.percentile(data,min)\n    max = np.percentile(data,max)\n    data = np.clip(data,min,max)\n    data = (data-min)/(max-min)\n    return data\n\ndef write_tomogram(data, fpath):\n    D, H, W = data.shape\n    os.makedirs(fpath, exist_ok=True)\n    for i in range(D):\n        im = Image.fromarray(data[i])\n        im.save(fpath + f\"{i:03d}.png\")\n\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    \n    # Normalize tomogram persentile based \n    tomogram = normalise_by_percentile(tomogram)\n    if precision == '8bit':\n        tomogram = (tomogram * 255).astype(np.uint8)\n    else:\n        tomogram = (tomogram * 65535).astype(np.uint16)\n\n    write_tomogram(tomogram, f\"{output_png}_{precision}/{run.name}/images/\")\n    write_tomogram(segmentation, f\"{output_png}_{precision}/{run.name}/labels/\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:00:22.840869Z","iopub.execute_input":"2024-12-03T02:00:22.841576Z","iopub.status.idle":"2024-12-03T02:01:20.148374Z","shell.execute_reply.started":"2024-12-03T02:00:22.841538Z","shell.execute_reply":"2024-12-03T02:01:20.147320Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"precision = '16bit'\n\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    \n    # Normalize tomogram persentile based \n    tomogram = normalise_by_percentile(tomogram)\n    if precision == '8bit':\n        tomogram = (tomogram * 255).astype(np.uint8)\n    else:\n        tomogram = (tomogram * 65535).astype(np.uint16)\n\n    write_tomogram(tomogram, f\"{output_png}_{precision}/{run.name}/images/\")\n    write_tomogram(segmentation, f\"{output_png}_{precision}/{run.name}/labels/\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:01:20.149847Z","iopub.execute_input":"2024-12-03T02:01:20.150651Z","iopub.status.idle":"2024-12-03T02:02:48.509722Z","shell.execute_reply.started":"2024-12-03T02:01:20.150602Z","shell.execute_reply":"2024-12-03T02:02:48.508387Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -r /kaggle/working/overlay","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:02:48.511692Z","iopub.execute_input":"2024-12-03T02:02:48.512133Z","iopub.status.idle":"2024-12-03T02:02:49.586323Z","shell.execute_reply.started":"2024-12-03T02:02:48.512092Z","shell.execute_reply":"2024-12-03T02:02:49.584941Z"}},"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(tomogram[100],cmap='gray')\nplt.axis('off')\n\nplt.subplot(1, 2, 2)\nplt.title('Painted Segmentation from Picks')\nplt.imshow(segmentation[100], cmap='viridis')\nplt.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T02:02:49.588172Z","iopub.execute_input":"2024-12-03T02:02:49.588622Z","iopub.status.idle":"2024-12-03T02:02:50.099365Z","shell.execute_reply.started":"2024-12-03T02:02:49.588585Z","shell.execute_reply":"2024-12-03T02:02:50.098083Z"}},"outputs":[],"execution_count":null}]}