add figures of workflow and metrics, add invert transform

README.md

@@ -9,9 +9,11 @@ license: apache-2.0
 A pre-trained model for the endoscopic tool segmentation task.
 
 # Model Overview
-This model is trained using a flexible unet structure with an efficient-b0 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
+This model is trained using a flexible unet structure with an efficient-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
 The [pytorch model](https://drive.google.com/file/d/14r6WmzaZrgaWLGu0O9vSAzdeIGVFQ3cs/view?usp=sharing) and [torchscript model](https://drive.google.com/file/d/1i-e5xXHtmvmqitwUP8Q3JqvnmN3mlrEm/view?usp=sharing) are shared in google drive. Details can be found in large_files.yml file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
 
+![image](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_workflow.png)
+
 ## Data
 Datasets used in this work were provided by [Activ Surgical](https://www.activsurgical.com/).
 
@@ -36,6 +38,16 @@ This model achieves the following IoU score on the test dataset (our own split f
 
 Mean IoU = 0.87
 
+## Training Performance
+A graph showing the training loss over 100 epochs.
+
+![](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss.png) <br>
+
+## Validation Performance
+A graph showing the validation mean IoU over 100 epochs.
+
+![](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou.png) <br>
+
 ## commands example
 Execute training:
 
@@ -61,31 +73,6 @@ Export checkpoint to TorchScript file:
 python -m monai.bundle ckpt_export network_def --filepath models/model.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json
 ```
 
-Export checkpoint to onnx file, which has been tested on pytorch 1.12.0:
-
-```
-python scripts/export_to_onnx.py --model models/model.pt --outpath models/model.onnx
-```
-
-Export TorchScript file to a torchscript module targeting a TensorRT engine with float16 precision.
-
-```
-torchtrtc -p f16 models/model.ts models/model_trt.ts "[(1,3,736,480);(4,3,736,480);(8,3,736,480)]"
-```
-The last parameter is the dynamic input shape in which each parameter means "[(MIN_BATCH, MIN_CHANNEL, MIN_WIDTH, MIN_HEIGHT), (OPT_BATCH, .., ..., OPT_HEIGHT), (MAX_BATCH, .., ..., MAX_HEIGHT)]". Please notice if using docker, the TensorRT CUDA must match the environment CUDA and the Torch-TensorRT c++&python version must be installed. For more examples on how to use the Torch-TensorRT, you can go to this [link](https://pytorch.org/TensorRT/). The [github source code link](https://github.com/pytorch/TensorRT) here shows the detail about how to install it on your own environment.
-
-Export TensorRT float16 model from the onnx model:
-
-```
-polygraphy surgeon sanitize --fold-constants models/model.onnx -o models/new_model.onnx
-```
-
-```
-trtexec --onnx=models/new_model.onnx --saveEngine=models/model.trt --fp16 --minShapes=INPUT__0:1x3x736x480 --optShapes=INPUT__0:4x3x736x480 --maxShapes=INPUT__0:8x3x736x480 --shapes=INPUT__0:4x3x736x480
-```
-This command need TensorRT with correct CUDA installed in the environment. For the detail of installing TensorRT, please refer to [this link](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). In addition, there are padding operations in this FlexibleUNet structure that not support by TensorRT. Therefore, when tried to convert the onnx model to a TensorRT engine, an extra polygraphy command is needed to execute.
-
-
 # References
 [1] Tan, M. and Le, Q. V. Efficientnet: Rethinking model scaling for convolutional neural networks. ICML, 2019a. https://arxiv.org/pdf/1905.11946.pdf
 

configs/evaluate.json

@@ -2,6 +2,21 @@
     "validate#postprocessing": {
         "_target_": "Compose",
         "transforms": [
+            {
+                "_target_": "Invertd",
+                "keys": [
+                    "pred",
+                    "label"
+                ],
+                "transform": "@validate#preprocessing",
+                "orig_keys": "image",
+                "meta_key_postfix": "meta_dict",
+                "nearest_interp": [
+                    false,
+                    true
+                ],
+                "to_tensor": true
+            },
             {
                 "_target_": "AsDiscreted",
                 "keys": [
@@ -14,20 +29,14 @@
                 ],
                 "to_onehot": 2
             },
-            {
-                "_target_": "Lambdad",
-                "keys": [
-                    "pred"
-                ],
-                "func": "$lambda x : x[1:]"
-            },
             {
                 "_target_": "SaveImaged",
+                "_disabled_": true,
                 "keys": "pred",
                 "meta_keys": "pred_meta_dict",
                 "output_dir": "@output_dir",
                 "output_ext": ".png",
-                "scale": 255,
+                "resample": false,
                 "squeeze_end_dims": true
             }
         ]

configs/inference.json

@@ -12,9 +12,9 @@
         "_target_": "FlexibleUNet",
         "in_channels": 3,
         "out_channels": 2,
-        "backbone": "efficientnet-b0",
+        "backbone": "efficientnet-b2",
         "spatial_dims": 2,
-        "pretrained": true,
+        "pretrained": false,
         "is_pad": false
     },
     "network": "$@network_def.to(@device)",
@@ -27,12 +27,6 @@
                     "image"
                 ]
             },
-            {
-                "_target_": "ToTensord",
-                "keys": [
-                    "image"
-                ]
-            },
             {
                 "_target_": "AsChannelFirstd",
                 "keys": [
@@ -78,6 +72,15 @@
     "postprocessing": {
         "_target_": "Compose",
         "transforms": [
+            {
+                "_target_": "Invertd",
+                "keys": "pred",
+                "transform": "@preprocessing",
+                "orig_keys": "image",
+                "meta_key_postfix": "meta_dict",
+                "nearest_interp": false,
+                "to_tensor": true
+            },
             {
                 "_target_": "AsDiscreted",
                 "argmax": true,

configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20220324.json",
-    "version": "0.3.0",
+    "version": "0.3.1",
     "changelog": {
+        "0.3.1": "add figures of workflow and metrics, add invert transform",
         "0.3.0": "update dataset processing",
         "0.2.1": "update to use monai 1.0.1",
         "0.2.0": "update license files",

configs/train.json

@@ -11,7 +11,7 @@
     "dataset_dir": "/workspace/data/endoscopic_tool_dataset",
     "images": "$list(sorted(glob.glob(os.path.join(@dataset_dir,'train', '*', '*[!seg].jpg'))))",
     "labels": "$[x.replace('.jpg', '_seg.jpg') for x in @images]",
-    "val_images": "$list(sorted(glob.glob(os.path.join(@dataset_dir,'valid', '*' '*[!seg].jpg'))))",
+    "val_images": "$list(sorted(glob.glob(os.path.join(@dataset_dir,'val', '*', '*[!seg].jpg'))))",
     "val_labels": "$[x.replace('.jpg', '_seg.jpg') for x in @val_images]",
     "val_interval": 1,
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
@@ -19,8 +19,9 @@
         "_target_": "FlexibleUNet",
         "in_channels": 3,
         "out_channels": 2,
-        "backbone": "efficientnet-b0",
+        "backbone": "efficientnet-b2",
         "spatial_dims": 2,
+        "dropout": 0.5,
         "pretrained": true,
         "is_pad": false
     },
@@ -29,13 +30,20 @@
         "_target_": "DiceLoss",
         "include_background": false,
         "to_onehot_y": true,
-        "softmax": true
+        "softmax": true,
+        "jaccard": true
     },
     "optimizer": {
         "_target_": "torch.optim.Adam",
         "params": "$@network.parameters()",
         "lr": 0.0001
     },
+    "lr_scheduler": {
+        "_target_": "torch.optim.lr_scheduler.CosineAnnealingWarmRestarts",
+        "optimizer": "@optimizer",
+        "T_0": 100,
+        "T_mult": 1
+    },
     "train": {
         "deterministic_transforms": [
             {
@@ -45,13 +53,6 @@
                     "label"
                 ]
             },
-            {
-                "_target_": "ToTensord",
-                "keys": [
-                    "image",
-                    "label"
-                ]
-            },
             {
                 "_target_": "AsChannelFirstd",
                 "keys": "image"
@@ -150,6 +151,11 @@
                 "log_dir": "@output_dir",
                 "tag_name": "train_loss",
                 "output_transform": "$monai.handlers.from_engine(['loss'], first=True)"
+            },
+            {
+                "_target_": "LrScheduleHandler",
+                "lr_scheduler": "@lr_scheduler",
+                "print_lr": true
             }
         ],
         "key_metric": {
@@ -182,7 +188,7 @@
         },
         "trainer": {
             "_target_": "SupervisedTrainer",
-            "max_epochs": 60,
+            "max_epochs": 100,
             "device": "@device",
             "train_data_loader": "@train#dataloader",
             "network": "@network",

docs/README.md

@@ -2,9 +2,11 @@
 A pre-trained model for the endoscopic tool segmentation task.
 
 # Model Overview
-This model is trained using a flexible unet structure with an efficient-b0 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
+This model is trained using a flexible unet structure with an efficient-b2 [1] as the backbone and a UNet architecture [2] as the decoder. Datasets use private samples from [Activ Surgical](https://www.activsurgical.com/).
 The [pytorch model](https://drive.google.com/file/d/14r6WmzaZrgaWLGu0O9vSAzdeIGVFQ3cs/view?usp=sharing) and [torchscript model](https://drive.google.com/file/d/1i-e5xXHtmvmqitwUP8Q3JqvnmN3mlrEm/view?usp=sharing) are shared in google drive. Details can be found in large_files.yml file. Modify the "bundle_root" parameter specified in configs/train.json and configs/inference.json to reflect where models are downloaded. Expected directory path to place downloaded models is "models/" under "bundle_root".
 
+![image](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_workflow.png)
+
 ## Data
 Datasets used in this work were provided by [Activ Surgical](https://www.activsurgical.com/).
 
@@ -29,6 +31,16 @@ This model achieves the following IoU score on the test dataset (our own split f
 
 Mean IoU = 0.87
 
+## Training Performance
+A graph showing the training loss over 100 epochs.
+
+![](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_train_loss.png) <br>
+
+## Validation Performance
+A graph showing the validation mean IoU over 100 epochs.
+
+![](https://developer.download.nvidia.com/assets/Clara/Images/monai_endoscopic_tool_segmentation_val_iou.png) <br>
+
 ## commands example
 Execute training:
 
@@ -54,31 +66,6 @@ Export checkpoint to TorchScript file:
 python -m monai.bundle ckpt_export network_def --filepath models/model.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json
 ```
 
-Export checkpoint to onnx file, which has been tested on pytorch 1.12.0:
-
-```
-python scripts/export_to_onnx.py --model models/model.pt --outpath models/model.onnx
-```
-
-Export TorchScript file to a torchscript module targeting a TensorRT engine with float16 precision.
-
-```
-torchtrtc -p f16 models/model.ts models/model_trt.ts "[(1,3,736,480);(4,3,736,480);(8,3,736,480)]"
-```
-The last parameter is the dynamic input shape in which each parameter means "[(MIN_BATCH, MIN_CHANNEL, MIN_WIDTH, MIN_HEIGHT), (OPT_BATCH, .., ..., OPT_HEIGHT), (MAX_BATCH, .., ..., MAX_HEIGHT)]". Please notice if using docker, the TensorRT CUDA must match the environment CUDA and the Torch-TensorRT c++&python version must be installed. For more examples on how to use the Torch-TensorRT, you can go to this [link](https://pytorch.org/TensorRT/). The [github source code link](https://github.com/pytorch/TensorRT) here shows the detail about how to install it on your own environment.
-
-Export TensorRT float16 model from the onnx model:
-
-```
-polygraphy surgeon sanitize --fold-constants models/model.onnx -o models/new_model.onnx
-```
-
-```
-trtexec --onnx=models/new_model.onnx --saveEngine=models/model.trt --fp16 --minShapes=INPUT__0:1x3x736x480 --optShapes=INPUT__0:4x3x736x480 --maxShapes=INPUT__0:8x3x736x480 --shapes=INPUT__0:4x3x736x480
-```
-This command need TensorRT with correct CUDA installed in the environment. For the detail of installing TensorRT, please refer to [this link](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html). In addition, there are padding operations in this FlexibleUNet structure that not support by TensorRT. Therefore, when tried to convert the onnx model to a TensorRT engine, an extra polygraphy command is needed to execute.
-
-
 # References
 [1] Tan, M. and Le, Q. V. Efficientnet: Rethinking model scaling for convolutional neural networks. ICML, 2019a. https://arxiv.org/pdf/1905.11946.pdf
 

models/model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:85d0ef791f918a3e324a78014cd2db80b2e5d4e29165c936078355a40ea21bac
-size 30443509
+oid sha256:844e9e97c6c9e7ebab1dab660c42ceb923c085ab895cf74f989a3cd9c5a0b028
+size 46262677

models/model.ts

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d28be14c501cce79aab342e7082d21e24f3b4192472f0c9e4faf43cf451ae776
-size 30631441
+oid sha256:9fa9f747ea2cc8ddffd4839af0c8d8f1b62c63c1013a8d80af8baf412bb3e5f9
+size 46493609