{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"\n",
"import torch\n",
"\n",
"from rtnls_inference import (\n",
" HeatmapRegressionEnsemble,\n",
" SegmentationEnsemble,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Segmentation of preprocessed images\n",
"\n",
"Here we segment images preprocessed using 0_preprocess.ipynb\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"ds_path = Path(\"../samples/fundus\")\n",
"\n",
"# input folders. these are the folders where we stored the preprocessed images\n",
"rgb_path = ds_path / \"rgb\"\n",
"ce_path = ds_path / \"ce\"\n",
"\n",
"# these are the output folders for:\n",
"av_path = ds_path / \"av\" # artery-vein segmentations\n",
"vessels_path = ds_path / \"vessels\" # vessels segmentation\n",
"discs_path = ds_path / \"discs\" # optic disc segmentations\n",
"overlays_path = ds_path / \"overlays\" # optional overlay visualizations\n",
"\n",
"device = torch.device(\"cuda:0\") # device to use for inference"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"rgb_paths = sorted(list(rgb_path.glob(\"*.png\")))\n",
"ce_paths = sorted(list(ce_path.glob(\"*.png\")))"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"data = [{\n",
" 'id': Path(rgb_path).stem,\n",
" 'image': rgb_path,\n",
" 'contrast_enhanced': ce_path\n",
"} for rgb_path, ce_path in zip(rgb_paths, ce_paths)]"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[{'id': 'CHASEDB1_08L',\n",
" 'image': PosixPath('../samples/fundus/rgb/CHASEDB1_08L.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/CHASEDB1_08L.png')},\n",
" {'id': 'CHASEDB1_12R',\n",
" 'image': PosixPath('../samples/fundus/rgb/CHASEDB1_12R.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/CHASEDB1_12R.png')},\n",
" {'id': 'DRIVE_22',\n",
" 'image': PosixPath('../samples/fundus/rgb/DRIVE_22.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/DRIVE_22.png')},\n",
" {'id': 'DRIVE_40',\n",
" 'image': PosixPath('../samples/fundus/rgb/DRIVE_40.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/DRIVE_40.png')},\n",
" {'id': 'HRF_04_g',\n",
" 'image': PosixPath('../samples/fundus/rgb/HRF_04_g.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/HRF_04_g.png')},\n",
" {'id': 'HRF_07_dr',\n",
" 'image': PosixPath('../samples/fundus/rgb/HRF_07_dr.png'),\n",
" 'contrast_enhanced': PosixPath('../samples/fundus/ce/HRF_07_dr.png')}]"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"data"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Artery-vein segmentation\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 1/1 [00:07<00:00, 7.53s/it]\n"
]
}
],
"source": [
"av_ensemble = SegmentationEnsemble.from_huggingface(\"Eyened/vascx:artery_vein/av_july24.pt\").to(device)\n",
"av_ensemble.predict_preprocessed(data, dest_path=av_path, num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Vessel segmentation (optional)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 1/1 [00:06<00:00, 6.81s/it]\n"
]
}
],
"source": [
"vessels_ensemble = SegmentationEnsemble.from_huggingface(\"Eyened/vascx:vessels/vessels_july24.pt\").to(device)\n",
"vessels_ensemble.predict_preprocessed(data, dest_path=vessels_path, num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Disc segmentation\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 1/1 [00:02<00:00, 2.34s/it]\n"
]
}
],
"source": [
"disc_ensemble = SegmentationEnsemble.from_huggingface(\"Eyened/vascx:disc/disc_july24.pt\").to(device)\n",
"disc_ensemble.predict_preprocessed(data, dest_path=discs_path, num_workers=2)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Fovea detection\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 1/1 [00:02<00:00, 2.04s/it]\n"
]
}
],
"source": [
"fovea_ensemble = HeatmapRegressionEnsemble.from_huggingface(\"Eyened/vascx:fovea/fovea_july24.pt\").to(device)\n",
"# note: this model does not use contrast enhanced images\n",
"df = fovea_ensemble.predict_preprocessed(data, num_workers=2)\n",
"df.columns = [\"mean_x\", \"mean_y\"]\n",
"df.to_csv(ds_path / \"fovea.csv\")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" mean_x | \n",
" mean_y | \n",
"
\n",
" \n",
" \n",
" \n",
" | CHASEDB1_08L | \n",
" 953.400024 | \n",
" 553.799988 | \n",
"
\n",
" \n",
" | CHASEDB1_12R | \n",
" 101.800003 | \n",
" 634.200012 | \n",
"
\n",
" \n",
" | DRIVE_22 | \n",
" 487.399994 | \n",
" 522.599976 | \n",
"
\n",
" \n",
" | DRIVE_40 | \n",
" 503.399994 | \n",
" 521.000000 | \n",
"
\n",
" \n",
" | HRF_04_g | \n",
" 551.000000 | \n",
" 535.400024 | \n",
"
\n",
" \n",
" | HRF_07_dr | \n",
" 545.000000 | \n",
" 534.599976 | \n",
"
\n",
" \n",
"
\n",
"