{"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":"nvidiaTeslaT4","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":"markdown","source":"## Adapted from the example code https://github.com/czimaginginstitute/2024_czii_mlchallenge_notebooks/blob/main/3d_unet_monai/train.ipynb\n\nRequesting computing resources: `Settings` -> `Accelerator` -> choose CPUs, GPUs, TPUs  \nWe can run 2 GPUs by choosing `GPU T4 x 2`.","metadata":{}},{"cell_type":"code","source":"!pip install git+https://github.com/copick/copick-utils.git matplotlib tqdm copick monai mlflow","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-09T02:10:09.448106Z","iopub.execute_input":"2024-11-09T02:10:09.448772Z","iopub.status.idle":"2024-11-09T02:10:51.833212Z","shell.execute_reply.started":"2024-11-09T02:10:09.448727Z","shell.execute_reply":"2024-11-09T02:10:51.830331Z"}},"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\": 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# 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":"2024-11-09T02:10:51.836026Z","iopub.execute_input":"2024-11-09T02:10:51.836458Z","iopub.status.idle":"2024-11-09T02:10:52.020172Z","shell.execute_reply.started":"2024-11-09T02:10:51.836407Z","shell.execute_reply":"2024-11-09T02:10:52.019248Z"}},"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 typing import Optional, Union, Tuple, List\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 pytorch_lightning as pl\nimport torch.distributed as dist\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.loggers import MLFlowLogger\nfrom pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T02:10:52.021291Z","iopub.execute_input":"2024-11-09T02:10:52.021586Z","iopub.status.idle":"2024-11-09T02:10:56.255949Z","shell.execute_reply.started":"2024-11-09T02:10:52.021553Z","shell.execute_reply":"2024-11-09T02:10:56.254843Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Generate the painted mask","metadata":{}},{"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\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T01:31:55.308047Z","iopub.execute_input":"2024-11-09T01:31:55.309596Z","iopub.status.idle":"2024-11-09T01:31:55.315486Z","shell.execute_reply.started":"2024-11-09T01:31:55.309539Z","shell.execute_reply":"2024-11-09T01:31:55.314456Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2.Generate multi-class segmentation masks from picks, and saved them to the copick overlay directory","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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T01:31:55.316852Z","iopub.execute_input":"2024-11-09T01:31:55.317219Z","iopub.status.idle":"2024-11-09T01:32:09.787666Z","shell.execute_reply.started":"2024-11-09T01:31:55.317178Z","shell.execute_reply":"2024-11-09T01:32:09.786639Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Define PyTorch lightning modules and the trainer","metadata":{}},{"cell_type":"code","source":"class Model(pl.LightningModule):\n    def __init__(\n        self, \n        spatial_dims: int = 3,\n        in_channels: int = 1,\n        out_channels: int = 8,\n        channels: Union[Tuple[int, ...], List[int]] = (48, 64, 80, 80),\n        strides: Union[Tuple[int, ...], List[int]] = (2, 2, 1),\n        num_res_units: int = 1,\n        lr: float=1e-3):\n    \n        super().__init__()\n        self.save_hyperparameters()\n        self.model = UNet(\n            spatial_dims=self.hparams.spatial_dims,\n            in_channels=self.hparams.in_channels,\n            out_channels=self.hparams.out_channels,\n            channels=self.hparams.channels,\n            strides=self.hparams.strides,\n            num_res_units=self.hparams.num_res_units,\n        )\n        self.loss_fn = TverskyLoss(include_background=True, to_onehot_y=True, softmax=True)  # softmax=True for multiclass\n        self.metric_fn = DiceMetric(include_background=False, reduction=\"mean\", ignore_empty=True)\n\n        self.train_loss = 0\n        self.val_metric = 0\n        self.num_train_batch = 0\n        self.num_val_batch = 0\n\n    def forward(self, x):\n        return self.model(x)\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch['image'], batch['label']\n        y_hat = self(x)\n        loss = self.loss_fn(y_hat, y)\n        self.train_loss += loss\n        self.num_train_batch += 1\n        torch.cuda.empty_cache()\n        return loss\n\n    def on_train_epoch_end(self):\n        loss_per_epoch = self.train_loss/self.num_train_batch\n        print(f\"Epoch {self.current_epoch} - Average Train Loss: {loss_per_epoch:.4f}\")\n        self.log('train_loss', loss_per_epoch, prog_bar=True)\n        self.train_loss = 0\n        self.num_train_batch = 0\n    \n    def validation_step(self, batch, batch_idx):\n        with torch.no_grad(): # This ensures that gradients are not stored in memory\n            x, y = batch['image'], batch['label']\n            y_hat = self(x)\n            metric_val_outputs = [AsDiscrete(argmax=True, to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y_hat)]\n            metric_val_labels = [AsDiscrete(to_onehot=self.hparams.out_channels)(i) for i in decollate_batch(y)]\n\n            # compute metric for current iteration\n            self.metric_fn(y_pred=metric_val_outputs, y=metric_val_labels)\n            metrics = self.metric_fn.aggregate(reduction=\"mean_batch\")\n            val_metric = torch.mean(metrics) # I used mean over all particle species as the metric. This can be explored.\n            self.val_metric += val_metric \n            self.num_val_batch += 1\n        torch.cuda.empty_cache()\n        return {'val_metric': val_metric}\n\n    def on_validation_epoch_end(self):\n        metric_per_epoch = self.val_metric/self.num_val_batch\n        print(f\"Epoch {self.current_epoch} - Average Val Metric: {metric_per_epoch:.4f}\")\n        self.log('val_metric', metric_per_epoch, prog_bar=True, sync_dist=False) # sync_dist=True for distributed training\n        self.val_metric = 0\n        self.num_val_batch = 0\n    \n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)\n\n    \n\nclass CopickDataModule(pl.LightningDataModule):\n    def __init__(\n        self, \n        copick_config_path: str, \n        train_batch_size: int,\n        val_batch_size: int,\n        num_random_samples_per_batch: int,\n        num_training_dataset: int=5 # the rest of the dataset is used for validation\n    ):\n    \n        super().__init__()\n        self.train_batch_size = train_batch_size\n        self.val_batch_size = val_batch_size\n        \n        self.data_dicts, self.nclasses = self.data_from_copick(copick_config_path)\n        self.train_files = self.data_dicts[:num_training_dataset]\n        self.val_files = self.data_dicts[num_training_dataset:]\n        print(f\"Number of training samples: {len(self.train_files)}\")\n        print(f\"Number of validation samples: {len(self.val_files)}\")\n\n        # Non-random transforms to be cached\n        self.non_random_transforms = Compose([\n            EnsureChannelFirstd(keys=[\"image\", \"label\"], channel_dim=\"no_channel\"),\n            NormalizeIntensityd(keys=\"image\"),\n            Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\")\n        ])\n\n        # Random transforms to be applied during training\n        self.random_transforms = Compose([\n            RandCropByLabelClassesd(\n                keys=[\"image\", \"label\"],\n                label_key=\"label\",\n                spatial_size=[96, 96, 96],\n                num_classes=self.nclasses,\n                num_samples=num_random_samples_per_batch \n            ),\n            RandRotate90d(keys=[\"image\", \"label\"], prob=0.5, spatial_axes=[0, 2]),\n            RandFlipd(keys=[\"image\", \"label\"], prob=0.5, spatial_axis=0),    \n        ])\n\n    def setup(self, stage: Optional[str] = None) -> None:\n        self.train_ds = CacheDataset(data=self.train_files, transform=self.non_random_transforms, cache_rate=1.0)\n        self.train_ds = Dataset(data=self.train_ds, transform=self.random_transforms)\n        self.val_ds = CacheDataset(data=self.val_files, transform=self.non_random_transforms, cache_rate=1.0)\n        self.val_ds = Dataset(data=self.val_ds, transform=self.random_transforms)\n\n    def train_dataloader(self) -> DataLoader:\n        return DataLoader(\n                self.train_ds,\n                batch_size=self.train_batch_size,\n                shuffle=True,\n                num_workers=1,\n                persistent_workers=False,\n                #pin_memory=torch.cuda.is_available(),\n            )\n\n    def val_dataloader(self) -> DataLoader:\n        return DataLoader(\n                self.val_ds,\n                batch_size=self.val_batch_size,\n                shuffle=False,  # Ensure the data order remains consistent\n                num_workers=1,\n                persistent_workers=False,\n                #pin_memory=torch.cuda.is_available(),\n            )\n\n    @staticmethod\n    def data_from_copick(copick_config_path, copick_user_name=\"copickUtils\", copick_segmentation_name=\"paintedPicks\", tomo_type=\"denoised\"):\n        root = copick.from_file(copick_config_path)\n        nclasses = len(root.pickable_objects) + 1\n        data_dicts = []\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        data_dicts = []\n        for run in tqdm(root.runs):\n            tomogram = run.get_voxel_spacing(10).get_tomogram(tomo_type).numpy()\n            segmentation = run.get_segmentations(user_id=copick_user_name, name=copick_segmentation_name, voxel_size=10, is_multilabel=True)[0].numpy()\n            # membrane_seg = run.get_segmentations(name=copick_segmentation_name, user_id=\"data-portal\")[0].numpy()\n            # segmentation[membrane_seg==1] = 1  \n            data_dicts.append({\"image\": tomogram, \"label\": segmentation})\n        \n        return data_dicts, nclasses","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T01:32:09.790121Z","iopub.execute_input":"2024-11-09T01:32:09.790707Z","iopub.status.idle":"2024-11-09T01:32:09.822745Z","shell.execute_reply.started":"2024-11-09T01:32:09.79066Z","shell.execute_reply":"2024-11-09T01:32:09.821874Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mlf_logger = MLFlowLogger('training-3D-UNet-model-for-the-cryoET-ML-Challenge')\n# Trainer callbacks\ncheckpoint_callback = ModelCheckpoint(monitor='val_metric', save_top_k=1, mode='max')\nlr_monitor = LearningRateMonitor(logging_interval='epoch')\n\n# Check if CUDA is available and then count the GPUs\nif torch.cuda.is_available():\n    num_gpus = torch.cuda.device_count()\n    print(f\"Number of GPUs available: {num_gpus}\")\nelse:\n    print(\"No GPU available. Running on CPU.\")\ndevices = list(range(num_gpus))\nprint(devices)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T01:37:43.561142Z","iopub.execute_input":"2024-11-09T01:37:43.561971Z","iopub.status.idle":"2024-11-09T01:37:43.572594Z","shell.execute_reply.started":"2024-11-09T01:37:43.561926Z","shell.execute_reply":"2024-11-09T01:37:43.571621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"channels = (48, 64, 80, 80)\nstrides_pattern = (2, 2, 1)       \nnum_res_units = 1\nlearning_rate = 1e-4\nnum_epochs = 10\n\nmodel = Model(channels=channels, strides=strides_pattern, num_res_units=num_res_units, lr=learning_rate)\ndatamodule = CopickDataModule(copick_config_path, 1, 1, 16)\n\n\n# Priotize performace over precision\ntorch.set_float32_matmul_precision('medium') # or torch.set_float32_matmul_precision('high')    \n\n# Trainer for distributed training with DDP\ntrainer = Trainer(\n    max_epochs=num_epochs,\n    logger=mlf_logger,\n    callbacks=[checkpoint_callback, lr_monitor],\n    strategy=\"ddp_notebook\", \n    accelerator=\"gpu\",\n    devices=devices,\n    num_nodes=1,\n    log_every_n_steps=1,\n    enable_progress_bar=True,\n)\n\ntrainer.fit(model, datamodule=datamodule)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-09T01:37:47.562428Z","iopub.execute_input":"2024-11-09T01:37:47.562854Z"}},"outputs":[],"execution_count":null}]}