{"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":30787,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install copick git+https://github.com/copick/copick-utils.git git+https://github.com/copick/DeepFindET.git","metadata":{"execution":{"iopub.status.busy":"2025-03-03T14:58:31.547857Z","iopub.execute_input":"2025-03-03T14:58:31.548156Z","iopub.status.idle":"2025-03-03T15:01:01.745434Z","shell.execute_reply.started":"2025-03-03T14:58:31.548125Z","shell.execute_reply":"2025-03-03T15:01:01.744415Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Make a copick project\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}\"\"\"\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","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:01.747099Z","iopub.execute_input":"2025-03-03T15:01:01.747393Z","iopub.status.idle":"2025-03-03T15:01:01.754132Z","shell.execute_reply.started":"2025-03-03T15:01:01.747362Z","shell.execute_reply":"2025-03-03T15:01:01.753151Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Setup new overlay directory\nimport os\nimport shutil\n\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}\")\n","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:01.755484Z","iopub.execute_input":"2025-03-03T15:01:01.7565Z","iopub.status.idle":"2025-03-03T15:01:02.050099Z","shell.execute_reply.started":"2025-03-03T15:01:01.756453Z","shell.execute_reply":"2025-03-03T15:01:02.0493Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from deepfindET.entry_points import step1\nfrom deepfindET.utils import copick_tools\nimport matplotlib.pyplot as plt\nimport copick\n\n%matplotlib inline\n\n################## Input Parameters #################\n\n# Config File\nconfig = '/kaggle/working/copick.config'\n\n# Query Tomogram\nvoxel_size = 10 \ntomogram_algorithm = 'denoised'\n\n# Output Name for the Segmentation Targets\nout_name = 'remotetargets'\nout_user_id = 'deepfindET'\nout_session_id = '0'\n\n# Read Copick Directory\ncopickRoot = copick.from_file(config)\n\n","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:02.051807Z","iopub.execute_input":"2025-03-03T15:01:02.052066Z","iopub.status.idle":"2025-03-03T15:01:05.492745Z","shell.execute_reply.started":"2025-03-03T15:01:02.052039Z","shell.execute_reply":"2025-03-03T15:01:05.491847Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"[(obj.name, None, None, (obj.radius / voxel_size)) for obj in copickRoot.pickable_objects if obj.is_particle]","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:05.493836Z","iopub.execute_input":"2025-03-03T15:01:05.494665Z","iopub.status.idle":"2025-03-03T15:01:05.502062Z","shell.execute_reply.started":"2025-03-03T15:01:05.494633Z","shell.execute_reply":"2025-03-03T15:01:05.500934Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# Query Train Protein Coordiantes and any Associated Segmentations\ntrain_targets = {}\n\n# Define protein targets with their respective radii\n# We can Provide two forms of inputs, either \n# ('protein-name',radius) or ('protein-name', 'user-id', 'session-id', 'radius')\ntargets = [(obj.name, None, None, (obj.radius / voxel_size)) for obj in copickRoot.pickable_objects if obj.is_particle]\n\n# Set run_ids to None, indicating that targets will be generated for the entire CoPick project by default.\n# If specific Run-IDs were provided, this variable would contain a list of those IDs.\nrun_ids = None","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:05.503313Z","iopub.execute_input":"2025-03-03T15:01:05.50363Z","iopub.status.idle":"2025-03-03T15:01:05.516938Z","shell.execute_reply.started":"2025-03-03T15:01:05.503605Z","shell.execute_reply":"2025-03-03T15:01:05.516118Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generate train target information\nfor t in targets:\n    obj_name, user_id, session_id, radius = t\n    info = {\n        \"label\": copickRoot.get_object(obj_name).label,\n        \"user_id\": user_id,\n        \"session_id\": session_id,\n        \"radius\": radius,\n        \"is_particle_target\": True,\n    }\n    train_targets[obj_name] = info\n\n\n# Define segmentation target (e.g., membrane)\nseg_targets = [('membrane', None, None)]\n\n# Generate segmentation target information\nfor s in seg_targets:\n    obj_name, user_id, session_id = s\n    info = {\n        \"label\": copickRoot.get_object(obj_name).label,\n        \"user_id\": user_id,\n        \"session_id\": session_id,\n        \"radius\": None,       \n        \"is_particle_target\": False,                 \n    }\n    train_targets[obj_name] = info\n\n# Call the create_train_targets function from step1 to generate the training targets for the 3D U-Net model.\n# The function will use the parameters defined in the previous cells and the following inputs:\nstep1.create_train_targets(\n    config,              # The configuration file path specifying various settings and parameters for the project.\n    train_targets,       # A dictionary containing the target information for each protein or object to be segmented.\n    run_ids,             # The list of Run-IDs for which to generate targets. None means targets for the entire project.\n    voxel_size,          # The voxel size to be used in the tomogram data.\n    tomogram_algorithm,  # The reconstruction algorithm used for the tomograms, e.g., 'wbp' (weighted back projection).\n    out_name,            # The output name for the generated segmentation targets.\n    out_user_id,         # The user ID under which the output targets will be saved.\n    out_session_id,      # The session ID associated with the output, typically used for tracking purposes.\n)","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:05.517898Z","iopub.execute_input":"2025-03-03T15:01:05.518125Z","iopub.status.idle":"2025-03-03T15:01:18.137717Z","shell.execute_reply.started":"2025-03-03T15:01:05.518102Z","shell.execute_reply":"2025-03-03T15:01:18.136962Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Option 1: Query All RunIDs\n# Retrieve all available Run-IDs from the CoPick project. This generates a list of Run-IDs by iterating over all runs in copickRoot.\nrun_ids = [run.name for run in copickRoot.runs]\n\n# Option 2: Manually Specify Specific Run\n# Define a specific Run-ID manually. This is useful for extracting volumes for a specific run.\nrunID = 'TS_6_4'\n\n# Retrieve the specific run object from CoPick using the manually specified Run-ID.\ncopick_run = copickRoot.get_run(runID)\n\n# Extract the segmentation target associated with the specified run.\n# The function get_copick_segmentation retrieves the segmentation data (e.g., target volume) based on the run object,\n# segmentation name, user ID, and session ID.\ntrain_target = copick_tools.get_copick_segmentation(\n    copick_run,                 # The run object obtained from CoPick for the specific Run-ID.\n    segmentationName='remotetargets',  # The name of the segmentation target to retrieve.\n    userID='deepfindET',        # The user ID under which the segmentation data is saved.\n    sessionID='0'               # The session ID associated with the segmentation data.\n)\n\n# Retrieve the tomogram associated with the specified Run-ID from the CoPick project.\n# The function get_copick_tomogram extracts the tomogram data, using the voxel size, algorithm, and Run-ID.\ntrain_tomogram = copick_tools.get_copick_tomogram(\n    copickRoot,                 # The root object for the CoPick project, containing all runs and associated data.\n    voxelSize=voxel_size,       # The voxel size to be used for retrieving the tomogram.\n    tomoAlgorithm='wbp',        # The reconstruction algorithm used for the tomogram, e.g., 'wbp' (weighted back projection).\n    tomoID=runID                # The specific Run-ID for which the tomogram is being retrieved.\n)","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:18.138764Z","iopub.execute_input":"2025-03-03T15:01:18.139029Z","iopub.status.idle":"2025-03-03T15:01:18.195348Z","shell.execute_reply.started":"2025-03-03T15:01:18.139002Z","shell.execute_reply":"2025-03-03T15:01:18.194609Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot the images\nplt.figure(figsize=(15, 5))\n\n# Original Image\nplt.subplot(1, 2, 1)\nplt.title('Tomogram')\nplt.imshow(train_tomogram[90,],cmap='gray')\nplt.axis('off')\n\n# Original Image\nplt.subplot(1, 2, 2)\nplt.title('Train Target')\nplt.imshow(train_target[90,])\nplt.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:18.196541Z","iopub.execute_input":"2025-03-03T15:01:18.19714Z","iopub.status.idle":"2025-03-03T15:01:22.066994Z","shell.execute_reply.started":"2025-03-03T15:01:18.197101Z","shell.execute_reply":"2025-03-03T15:01:22.066066Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from deepfindET.entry_points import step2\n\n# Specify the directory where the training results will be saved.\ntraining_output_path = '/kaggle/working/train_results'\n\n# Set the model architecture to Residual U-Net ('res_unet'), \n# which combines U-Net with residual connections to improve training.\nmodel_name = 'res_unet'\n\n# Set to None to indicate that the model will betrained from scratch \n# without using pre-trained weights.\nmodel_pre_weights = None\n\n# Number of classes the model will predict. \n# Here, we are working with 8 different classes (6 proteins + membrane + background).\nn_class = 8\n\n# Input dimension size of the 3D volumes in voxels. Each input is a 72x72x72 voxel cube -- (72 Å)^3.\ndim_in = 72  # [voxels]","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:22.0692Z","iopub.execute_input":"2025-03-03T15:01:22.069529Z","iopub.status.idle":"2025-03-03T15:01:33.200292Z","shell.execute_reply.started":"2025-03-03T15:01:22.0695Z","shell.execute_reply":"2025-03-03T15:01:33.199365Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initiate the training of the DeepFindET 3D U-Net model.\nstep2.train_model(\n    config,                 # Configuration file with various settings for the project.\n    voxel_size,             # Voxel size used in the tomogram data.\n    tomogram_algorithm,     # Reconstruction algorithm used for the tomograms (e.g., 'wbp').\n    targets,                 # Target data for training the model.\n    training_output_path,   # Path where the training outputs will be saved.\n    model_name,             # Model architecture name ('res_unet').\n    model_pre_weights,      # Pre-trained weights (None for training from scratch).\n    n_class,                # Number of classes for segmentation (8 in this case).\n    path_valid=None,        # Path to validation data (None means internal splitting may be used).\n    dim_in=dim_in,          # Input dimension size in voxels.\n    n_sub_epoch=10,         # Number of epochs to train on tomograms prior to swapping to a new set of tomograms.\n    sample_size=3,          # Number of tomograms to extract per epoch.\n    batch_size=10,          # Batch size used during training.\n    epochs=70,              # Total number of training epochs.\n    steps_per_epoch=150,    # Number of steps per epoch.\n    n_valid=20,             # Number of validation samples.\n    model_filters=[48, 64, 128],  # Filters in the convolutional layers at each level of the U-Net.\n    model_dropout=0,        # Dropout rate (0 means no dropout applied).\n    target_name=\"remotetargets\",    # Name of the segmentation targets.\n    target_user_id=\"deepfindET\",  # User ID for the segmentation labels.\n    target_session_id=\"0\",    # Session ID associated with the labeling.\n    valid_tomo_ids=None,    # List of tomogram IDs for validation.\n    train_tomo_ids=None,     # List of tomogram IDs for training.\n    class_weights=(('apo-ferritin', 62400), ('beta-amylase', 4130), ('beta-galactosidase', 3080), ('ribosome', 1800), ('thyroglobulin', 10100), ('virus-like-particle', 8400))\n)\n# 10439:     # class_weights=(('membrane',1),('adp-mitochondrial',3000),('alkaline-phosphate',3000),('nucleosome',3000),('ribosome',750),('vault',500),('virus-like-capsid',750))\n","metadata":{"execution":{"iopub.status.busy":"2025-03-03T15:01:33.201527Z","iopub.execute_input":"2025-03-03T15:01:33.202617Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from deepfindET.utils import core\nimport h5py, os\n\n# Set the path to the training history file\nhistory_path = os.path.join(training_output_path, 'net_train_history.h5')\n\n# Convert the HDF5 file containing the training history into a dictionary format\n# This allows easy access to the training metrics like loss, accuracy, etc., stored during training\nhistory = core.convert_hdf5_to_dictionary(history_path)\n\n# Plot the training history to visualize the learning process\n# The plot_history function will generate curves for metrics such as training and validation loss, accuracy, etc.\ncore.plot_history(history, save_figure=False)\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}