{"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":"# 🦠 Cell Instance Segmentation with Lightning⚡Flash\n\n[Flash](https://lightning-flash.readthedocs.io/en/stable) makes complex AI recipes for over 15 tasks across 7 data domains accessible to all.\n\nIn a nutshell, Flash is the production grade research framework you always dreamed of but didn't have time to build.","metadata":{}},{"cell_type":"code","source":"! ls -l /kaggle/input/sartorius-cell-instance-segmentation\n! pip list | grep -E \"lightning|torch|icevision|Pillow\"","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:48:59.554506Z","iopub.execute_input":"2021-11-02T21:48:59.554891Z","iopub.status.idle":"2021-11-02T21:49:04.061662Z","shell.execute_reply.started":"2021-11-02T21:48:59.554795Z","shell.execute_reply":"2021-11-02T21:49:04.060678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ! pip install -q lightning-flash[image]\n! pip install -q 'https://github.com/PyTorchLightning/lightning-flash/archive/refs/heads/bugfix/icevision_memory_leak.zip#egg=lightning-flash[image]'\n# ! pip install -q icevision[all]\n! pip install -q 'https://github.com/Borda/icevision/archive/refs/heads/try/imageio.zip#egg=icevision[all]'\n! pip install -q pip install --upgrade pandas \"torchtext<0.10\" --force-reinstall\n! pip install -q -U Pillow\n! pip list | grep -E \"lightning|torch|icevision|Pillow\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-11-02T21:49:04.065016Z","iopub.execute_input":"2021-11-02T21:49:04.065325Z","iopub.status.idle":"2021-11-02T21:51:03.027330Z","shell.execute_reply.started":"2021-11-02T21:49:04.065287Z","shell.execute_reply":"2021-11-02T21:51:03.026437Z"},"_kg_hide-output":true,"_kg_hide-input":false,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data preparation\n\nMoving the data to be Coco complient in zero folder depth...","metadata":{}},{"cell_type":"code","source":"! mkdir -p /kaggle/working/dataset/annotations\n! mkdir -p /kaggle/working/dataset/images\n! cp /kaggle/input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/annotations/LIVECell/*.json /kaggle/working/dataset/annotations/\n! mkdir /kaggle/working/dataset/images/livecell_test_images/\n! mkdir /kaggle/working/dataset/images/livecell_train_val_images/\n! cp /kaggle/input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/images/livecell_test_images/*/*.tif /kaggle/working/dataset/images/livecell_test_images/\n! cp /kaggle/input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021/images/livecell_train_val_images/*/*.tif /kaggle/working/dataset/images/livecell_train_val_images/","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:51:03.029802Z","iopub.execute_input":"2021-11-02T21:51:03.030126Z","iopub.status.idle":"2021-11-02T21:51:52.298126Z","shell.execute_reply.started":"2021-11-02T21:51:03.030084Z","shell.execute_reply":"2021-11-02T21:51:52.297007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training with Flash Lightning\n\nSee the instance segm. docs: https://lightning-flash.readthedocs.io/en/stable/reference/instance_segmentation.html","metadata":{}},{"cell_type":"code","source":"%reload_ext autoreload\n%autoreload 2\n\nimport os\n\nPATH_PREDICT = \"/kaggle/input/sartorius-cell-instance-segmentation/test\"\nPATH_DATASET = \"/kaggle/input/sartorius-cell-instance-segmentation/LIVECell_dataset_2021\"\nLOCAL_DIR_DATASET = \"/kaggle/working/dataset\"\nLOCAL_DIR_ANNOTATIONS = os.path.join(LOCAL_DIR_DATASET, \"annotations\")\nLOCAL_DIR_IMAGES_TRAIN = os.path.join(LOCAL_DIR_DATASET, \"images\", \"livecell_train_val_images\")\nLOCAL_DIR_IMAGES_TEST = os.path.join(LOCAL_DIR_DATASET, \"images\", \"livecell_test_images\")","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:51:52.300249Z","iopub.execute_input":"2021-11-02T21:51:52.300568Z","iopub.status.idle":"2021-11-02T21:51:52.347282Z","shell.execute_reply.started":"2021-11-02T21:51:52.300524Z","shell.execute_reply":"2021-11-02T21:51:52.346557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nimport cv2\nfrom tqdm.auto import tqdm\n\nls_images = glob.glob(f\"{LOCAL_DIR_IMAGES_TRAIN}/*.tif\") + glob.glob(f\"{LOCAL_DIR_IMAGES_TEST}/*.tif\")\n\nfor img in tqdm(ls_images):\n  cv2.imwrite(img.replace(\".tif\", \".png\"), cv2.imread(img))","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:51:52.348532Z","iopub.execute_input":"2021-11-02T21:51:52.348817Z","iopub.status.idle":"2021-11-02T21:54:01.478347Z","shell.execute_reply.started":"2021-11-02T21:51:52.348777Z","shell.execute_reply":"2021-11-02T21:54:01.477265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Fixing the Coco annotations","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nNB_ANNOTTAIONS_THR = 900\nannots = glob.glob(os.path.join(LOCAL_DIR_ANNOTATIONS, \"*.json\"))\nddirs = dict(train=LOCAL_DIR_IMAGES_TRAIN, val=LOCAL_DIR_IMAGES_TRAIN, test=LOCAL_DIR_IMAGES_TEST)\n\nfor annot in tqdm(sorted(annots)):\n    with open(annot) as fp:\n        data = json.load(fp)\n    if isinstance(data['annotations'], dict):\n        data['annotations'] = list(data['annotations'].values())\n        \n    df_counts = pd.DataFrame(data['annotations']).groupby(['image_id']).size()\n    too_large = df_counts[df_counts >= NB_ANNOTTAIONS_THR].index.to_list()\n    df_counts.hist(bins=50, grid=True)\n    \n    data['annotations'] = [d for d in data['annotations'] if d['image_id'] not in too_large]\n        \n    for d in data['images']:\n        d['file_name'] = d['file_name'].replace(\".tif\", \".png\")\n\n    fname, _ = os.path.splitext(os.path.basename(annot))\n    img_dir = ddirs[fname.split(\"_\")[-1]]\n    miss_ = [d for d in data['images'] if not os.path.isfile(os.path.join(img_dir, d['file_name']))]\n    large_ = [d for d in data['images'] if d['id'] in too_large]\n    data['images'] = [d for d in data['images'] if os.path.isfile(os.path.join(img_dir, d['file_name'])) and d['id'] not in too_large]\n    print(f\"{len(data['images'])} (miss {len(miss_)}, large {len(large_)}) images and {len(data['annotations'])} annots in {os.path.basename(annot)}\")\n\n    with open(annot, 'w') as fp:\n        json.dump(data, fp)","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:54:01.480135Z","iopub.execute_input":"2021-11-02T21:54:01.480440Z","iopub.status.idle":"2021-11-02T21:57:10.379407Z","shell.execute_reply.started":"2021-11-02T21:54:01.480396Z","shell.execute_reply":"2021-11-02T21:57:10.378588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 0: import Flash","metadata":{}},{"cell_type":"code","source":"import flash\nfrom flash.core.utilities.imports import example_requires\nfrom flash.image import InstanceSegmentation, InstanceSegmentationData","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:57:10.380760Z","iopub.execute_input":"2021-11-02T21:57:10.381029Z","iopub.status.idle":"2021-11-02T21:57:21.202838Z","shell.execute_reply.started":"2021-11-02T21:57:10.380989Z","shell.execute_reply":"2021-11-02T21:57:21.201941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1. Create the DataModule","metadata":{}},{"cell_type":"code","source":"datamodule = InstanceSegmentationData.from_coco(\n    train_folder=LOCAL_DIR_IMAGES_TRAIN,\n    train_ann_file=os.path.join(LOCAL_DIR_ANNOTATIONS, \"livecell_coco_train.json\"),\n    val_folder=LOCAL_DIR_IMAGES_TRAIN,\n    val_ann_file=os.path.join(LOCAL_DIR_ANNOTATIONS, \"livecell_coco_val.json\"),\n    test_folder=LOCAL_DIR_IMAGES_TEST,\n    test_ann_file=os.path.join(LOCAL_DIR_ANNOTATIONS, \"livecell_coco_test.json\"),\n    predict_folder=PATH_PREDICT,\n#     train_transform: Optional[Dict[str, Callable]] = None,\n#     val_transform: Optional[Dict[str, Callable]] = None,\n#     test_transform: Optional[Dict[str, Callable]] = None,\n#     predict_transform: Optional[Dict[str, Callable]] = None,\n#     data_fetcher: Optional[BaseDataFetcher] = None,\n#     preprocess: Optional[Preprocess] = None,\n    batch_size=16,\n    num_workers=0,\n)","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:57:21.205380Z","iopub.execute_input":"2021-11-02T21:57:21.205690Z","iopub.status.idle":"2021-11-02T21:57:53.111516Z","shell.execute_reply.started":"2021-11-02T21:57:21.205648Z","shell.execute_reply":"2021-11-02T21:57:53.110574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Build the task","metadata":{}},{"cell_type":"code","source":"model = InstanceSegmentation(\n    head=\"mask_rcnn\",\n    backbone=\"resnet18_fpn\",\n    num_classes=datamodule.num_classes,\n)","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:57:53.113724Z","iopub.execute_input":"2021-11-02T21:57:53.114366Z","iopub.status.idle":"2021-11-02T21:57:53.985524Z","shell.execute_reply.started":"2021-11-02T21:57:53.114319Z","shell.execute_reply":"2021-11-02T21:57:53.984674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Create the trainer and finetune the model","metadata":{}},{"cell_type":"code","source":"import torch\nfrom pytorch_lightning.loggers import CSVLogger\n\nlogger = CSVLogger(save_dir='logs/')\ntrainer = flash.Trainer(\n    max_epochs=5,\n    gpus=torch.cuda.device_count(),\n    logger=logger,\n    progress_bar_refresh_rate=1,\n    precision=16,\n    benchmark=True,\n    accumulate_grad_batches=12,\n    auto_lr_find=True,\n)\n\n# ==============================\n\ntrainer.tune(\n    model, \n    datamodule=datamodule, \n    lr_find_kwargs=dict(min_lr=2e-5, max_lr=1e-2, num_training=25),\n    # scale_batch_size_kwargs=dict(max_trials=5),\n)\nprint(f\"Batch size: {datamodule.batch_size}\")\nprint(f\"Learning Rate: {model.learning_rate}\")\n\n# ==============================\n\ntrainer.finetune(model, datamodule=datamodule, strategy=\"freeze\")","metadata":{"execution":{"iopub.status.busy":"2021-11-02T21:57:53.986900Z","iopub.execute_input":"2021-11-02T21:57:53.987193Z","iopub.status.idle":"2021-11-02T22:36:20.029578Z","shell.execute_reply.started":"2021-11-02T21:57:53.987154Z","shell.execute_reply":"2021-11-02T22:36:20.028672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport seaborn as sns\nsns.set()\n\nmetrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\ndisplay(metrics.head())\nmetrics.set_index(\"step\", inplace=True)\ndel metrics[\"epoch\"]\nsns.relplot(data=metrics, kind=\"line\")","metadata":{"execution":{"iopub.status.busy":"2021-11-02T22:36:20.032818Z","iopub.execute_input":"2021-11-02T22:36:20.033141Z","iopub.status.idle":"2021-11-02T22:36:21.668281Z","shell.execute_reply.started":"2021-11-02T22:36:20.033095Z","shell.execute_reply":"2021-11-02T22:36:21.667350Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Detect objects in a few images!","metadata":{}},{"cell_type":"code","source":"import glob, os\nimport numpy as np\nimport matplotlib.pyplot as plt\n\npredict_imgs = sorted(glob.glob(os.path.join(PATH_PREDICT, \"*.png\")))\n\npredictions = model.predict(predict_imgs)\n\nfor p_img, pred in zip(predict_imgs, predictions):\n    print(p_img, pred.keys())\n    img = plt.imread(p_img)\n    mask = np.zeros(img.shape[:2])\n    for lb, m in enumerate(pred[\"masks\"]):\n        mask[m] = lb + 1\n    fig, axarr = plt.subplots(ncols=2)\n    axarr[0].imshow(img)\n    axarr[1].imshow(mask)","metadata":{"execution":{"iopub.status.busy":"2021-11-02T22:49:16.877211Z","iopub.execute_input":"2021-11-02T22:49:16.877525Z","iopub.status.idle":"2021-11-02T22:49:23.883524Z","shell.execute_reply.started":"2021-11-02T22:49:16.877482Z","shell.execute_reply":"2021-11-02T22:49:23.882594Z"},"trusted":true},"execution_count":null,"outputs":[]}]}