{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Vertebrae Classification
","metadata":{}},{"cell_type":"markdown","source":"![](\"https://www.sci-info-pages.com/wp-content/media/spinal-cord-segments.png\")
\n\nTask:
\n\n\nIn this notebook we will create a model that will classifies the cervical vertebraes from CT scans. Images are stored in DICOM format and vertebrae labels we have to extract from NIFTI segmentation masks. For each patient we have batch of CT scans from different depth levels. Single DICOM image may contain more than one vertebrae.\n
\n\nSolution:
\n\n\nFor vertebrae classification we will use
Monai framework. This tool help as easily process DICOM images and create pre-trained model for 2D image classification task. We will handle training and validation loop using
PyTorch Lightning.\n
Dataset
","metadata":{}},{"cell_type":"code","source":"class SpineSegmentationDataset(torch.utils.data.Dataset):\n def __init__(self, meta_df, segmentation_dir, num_class=7):\n self.meta_df = meta_df\n self.segmentation_dir = segmentation_dir\n self.num_class = num_class\n self.read_image = Compose([LoadImage(image_only=True, swap_ij=False), \n EnsureChannelFirst(), \n ScaleIntensity()])\n \n def __getitem__(self, idx):\n sample_meta = self.meta_df.loc[idx] \n image_path = os.path.join(TRAIN_DIR, sample_meta.StudyInstanceUID, str(sample_meta.slice) + \".dcm\")\n image = self.read_image(image_path)\n labels = sample_meta[LABEL_NAMES]\n slice_ratio = torch.tensor(sample_meta.slice_ratio).float()\n return image, torch.unsqueeze(slice_ratio, 0), torch.tensor(labels)\n \n def __len__(self):\n return self.meta_df.shape[0]\n \n def collate_fn(self, batch):\n return tuple(map(torch.stack, zip(*batch))) ","metadata":{"execution":{"iopub.status.busy":"2022-09-09T10:44:47.788142Z","iopub.execute_input":"2022-09-09T10:44:47.788597Z","iopub.status.idle":"2022-09-09T10:44:47.80122Z","shell.execute_reply.started":"2022-09-09T10:44:47.788555Z","shell.execute_reply":"2022-09-09T10:44:47.79903Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader \n\nclass SpineDataModule(pl.LightningDataModule):\n def __init__(self, train_full: pd.DataFrame, \n data_dir: str, \n valid_frac: float= 0.2, \n batch_size: int = 8,\n num_workers: int = 2):\n super().__init__()\n self.data_dir = data_dir\n self.valid_frac = valid_frac\n self.batch_size = batch_size\n self.num_workers = num_workers\n self.n_samples = train_full.shape[0]\n self.train_full = train_full\n \n def setup(self, stage: Optional[str] = None): \n \"\"\"Shuffle and split samples into training and validation partitions\"\"\"\n if not self.train_full.index.is_numeric():\n self.train_full = self.train_full.reset_index() \n \n valid_size = int(self.valid_frac * self.n_samples)\n perm = np.random.permutation(self.n_samples)\n\n train_meta = self.train_full.iloc[perm[valid_size:]]\n valid_meta = self.train_full.iloc[perm[:valid_size]]\n self.train_df = train_meta.reset_index(drop=True) \n self.valid_df = valid_meta.reset_index(drop=True)\n \n def train_dataloader(self):\n train_ds = SpineSegmentationDataset(self.train_df, self.data_dir)\n return DataLoader(train_ds, \n num_workers=self.num_workers,\n batch_size=self.batch_size, \n collate_fn=train_ds.collate_fn,\n shuffle=True)\n \n def val_dataloader(self):\n valid_ds = SpineSegmentationDataset(self.valid_df, self.data_dir)\n return DataLoader(valid_ds, \n num_workers=self.num_workers,\n batch_size=self.batch_size,\n collate_fn=valid_ds.collate_fn,\n shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2022-09-09T10:44:47.803263Z","iopub.execute_input":"2022-09-09T10:44:47.803629Z","iopub.status.idle":"2022-09-09T10:44:47.820715Z","shell.execute_reply.started":"2022-09-09T10:44:47.803597Z","shell.execute_reply":"2022-09-09T10:44:47.819236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Model training
","metadata":{}},{"cell_type":"code","source":"from torchmetrics import Accuracy\n\nclass SpineClassifier(pl.LightningModule):\n def __init__(self, in_channels: int=1, n_class: int=7, treshold: int=0.5):\n super().__init__()\n self.treshold = treshold\n self.n_class = n_class\n self.backbone = DenseNet121(spatial_dims=2, \n in_channels=in_channels,\n pretrained=True,\n out_channels=50) \n \n self.fc1 = nn.Linear(in_features=1, out_features=50)\n self.fc2 = nn.Linear(in_features=100, out_features=num_class) \n self.sigm = nn.Sigmoid()\n self.loss_fn = nn.BCELoss()\n self.metric = Accuracy(average=None, num_classes=n_class)\n \n def forward(self, x): \n images, ratio = x\n out1 = self.backbone(images)\n out2 = self.fc1(ratio)\n out = torch.cat((out1,out2),1)\n out = nn.functional.relu(out) # out\n out = self.fc2(out)\n \n out = self.sigm(out) \n return out\n \n def training_step(self, batch, batch_idx):\n images, ratio, labels = batch \n preds = self.forward((images, ratio))\n \n loss = self.loss_fn(preds, labels.float()) \n self.log(\"train_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n return loss\n \n def validation_step(self, batch, batch_idx):\n images, ratio, labels = batch \n preds = self.forward((images, ratio))\n \n loss = self.loss_fn(preds, labels.float()) \n self.metric.update(preds, labels.long())\n self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n return loss\n \n def predict_step(self, batch, batch_idx): \n inputs = (batch[0], batch[1]) \n scores = self.forward(inputs).detach().cpu().numpy()\n labels = np.array(scores > self.treshold, np.int8) \n output = {\"preds\": labels, \"scores\": scores}\n \n if len(batch) == 3:\n output[\"labels\"] = batch[2]\n \n return output\n \n def validation_epoch_end(self, outputs):\n accuracy = self.metric.compute()\n \n for i, CLS_NAME in enumerate(LABEL_NAMES):\n self.log(f\"val_acc_{CLS_NAME}\", accuracy[i], on_step=False, on_epoch=True)\n \n def configure_optimizers(self):\n return torch.optim.Adam(self.parameters(), lr=1e-5) ","metadata":{"execution":{"iopub.status.busy":"2022-09-09T10:44:47.822625Z","iopub.execute_input":"2022-09-09T10:44:47.823048Z","iopub.status.idle":"2022-09-09T10:44:47.844199Z","shell.execute_reply.started":"2022-09-09T10:44:47.823009Z","shell.execute_reply":"2022-09-09T10:44:47.842685Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pytorch_lightning.loggers import CSVLogger\n\ndevice = \"gpu\" if torch.cuda.is_available() else \"cpu\"\nlogger = CSVLogger(\"logs\", name=\"spine_classification\")\ndataset = SpineDataModule(train_segmentation, TRAIN_DIR)\nmodel = SpineClassifier()\n\ntrainer = pl.Trainer(accelerator=device, max_epochs=10, logger=logger)\ntrainer.fit(model, dataset)\ndel train_segmentation","metadata":{"execution":{"iopub.status.busy":"2022-09-09T10:44:47.846119Z","iopub.execute_input":"2022-09-09T10:44:47.846647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axs = plt.subplots(1, 2, figsize=(16, 4))\n\nmetrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\nmetrics.drop(\"step\", axis=1, inplace=True)\nmetrics = metrics.set_index(\"epoch\").stack()\nmetrics = metrics.unstack()\n\nmetrics[['train_loss', 'val_loss']].plot(ax=axs[0], xticks=metrics.index)\nmetrics[\"mean_accuracy\"] = metrics[[f\"val_acc_{cls_name}\" \n for cls_name in LABEL_NAMES]].mean(axis=1)\nmetrics[\"mean_accuracy\"].plot(ax=axs[1], xticks=metrics.index, legend=True)\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Evaluation
","metadata":{}},{"cell_type":"code","source":"pd.DataFrame(metrics.iloc[-1]).T","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"outputs = trainer.predict(model, dataset.val_dataloader())\n\npreds = []\nscores = []\nlabels = []\n\nfor output in outputs:\n preds.append(output[\"preds\"])\n scores.append(output[\"scores\"])\n labels.append(output[\"labels\"])\n \npreds = np.concatenate(preds, axis=0)\nscores = np.concatenate(scores, axis=0)\nlabels = np.concatenate(labels, axis=0)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.metrics import confusion_matrix\nimport seaborn as sns\n\n\nplt.figure(figsize=(16, 7))\n\nfor i in range(7):\n plt.subplot(2, 4, i + 1)\n sns.heatmap(confusion_matrix(preds[:, i], labels[:, i]), annot=True, fmt=\"d\")\n plt.title(LABEL_NAMES[i])\n plt.xlabel(\"Preds\")\n plt.ylabel(\"True\")\n \nplt.tight_layout()\n\ndel preds, scores, labels, outputs","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_preds(imgs, labels, preds, scores):\n n_rows = len(imgs) // 4\n fig, axs = plt.subplots(n_rows, 4, figsize=(18, 5 * n_rows))\n for i, ax in enumerate(axs.flat):\n cls = [str(int(c)) for c in labels[i]]\n prs = [str(int(p)) for p in preds[i]]\n scr = [str(round(p, 1)) for p in scores[i]]\n\n title = f\"Labels: {', '.join(cls)} \\nPreds: {', '.join(prs)} \\nScr: {', '.join(scr)}\"\n ax.set_title(title)\n ax.imshow(imgs[i], cmap=plt.cm.bone)\n \n \nimgs, ratio, labels = next(iter(dataset.val_dataloader()))\nscores = model((imgs, ratio))\nplot_preds(imgs.permute(0, 2, 3, 1), \n labels, \n np.array(scores > 0.5, np.int8), \n scores.detach().numpy())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(model, \"checkpoints\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}