Refactor

Author: Qin Liu

.gitignore

@@ -1,7 +1,8 @@
 *ivs-demo
 
 # Models
-weights/
+*weights
+*data
 *saves
 *results
 *.pth
@@ -10,7 +11,7 @@ weights/
 # Data
 *videos
 *volumes
-workspace/
+*workspace
 *nii.gz
 
 # Byte-compiled / optimized / DLL files

README.md

@@ -41,11 +41,11 @@ pip3 install -r requirements.txt
 ```
 
 ## Getting Started
-First, download model weights and a medical volume for demo purposes. All downloaded data will be saved in a folder named ``saves``.
+First, download AbdomenCT-1K and model weights. AbdomenCT-1K will be saved in the ``data`` folder; model weight will be saved in the ``saves`` folder. AbdomenCT-1K can also be downloaded from the official [webset](https://github.com/JunMa11/AbdomenCT-1K). 
 ```
 python download.py
 ```
-Then run a demo:
+Unzip the AbdomenCT-1K in the ``data`` folder accordingly. Then run a demo:
 ```
 ./run_demo.sh
 ```

download.py

@@ -1,6 +1,7 @@
 import os
 import gdown
 
+# download models
 save_folder ='saves'
 os.makedirs(save_folder, exist_ok=True)
 
@@ -36,3 +37,12 @@
 weight_fusion_stcn_url = 'https://drive.google.com/uc?id=1hf3yHizbgcOQx6LLFCX_rvHR30p4ioVp'
 output = f'{save_folder}/fusion_stcn.pth'
 gdown.download(weight_fusion_stcn_url, output, quiet=False)
+
+# download data
+data_folder ='data'
+os.makedirs(data_folder, exist_ok=True)
+
+# AbdomenCT-1k
+data_url = 'https://drive.google.com/uc?id=1RBvKJ5eGRG5BUK4F3NUpLopuP-_cLjKC'
+output = f'{data_folder}/AbdomenCT-1K.zip'
+gdown.download(data_url, output, quiet=False)

maskprop/Med-STCN/abdomen1k-evaluation/evaluation_method.py

@@ -1,3 +1,4 @@
+from argparse import ArgumentParser
 import numpy as np
 import os
 from os import path
@@ -6,7 +7,6 @@
 import pandas as pd
 
 
-
 def evaluate(gt_folder, seg_folder, dataset_file, label):
 
     videos = []
@@ -33,10 +33,9 @@ def evaluate(gt_folder, seg_folder, dataset_file, label):
 
             # calculate J, M, and DSC across frames
             gt_frame_label = np.zeros_like(gt_frame)
-            gt_frame_label[gt_frame == label] = 1
+            gt_frame_label[gt_frame == int(label)] = 1
             if np.count_nonzero(gt_frame_label) > 10:
                 seg_frame_label = np.zeros_like(seg_frame)
-                # seg_frame_label[seg_frame == label] = 1
                 seg_frame_label[seg_frame > 0] = 1
                 iou_label = db_eval_iou(gt_frame_label, seg_frame_label)
                 iou_single_video_label.append(iou_label)
@@ -90,27 +89,13 @@ def evaluate(gt_folder, seg_folder, dataset_file, label):
 
 if __name__ == '__main__':
 
-    data_root_folder = '/playpen-raid2/qinliu/data/AbdomenCT-1K/Organ-12-Subset_frames/trainval'
-    # data_root_folder = '/playpen-raid2/qinliu/data/MSD/Task02_Heart_frames/trainval'
-
-    gt_folder = f'{data_root_folder}/Annotations/480p'
-    dataset_file = f'{data_root_folder}/ImageSets/val.txt'
-
-    # seg_root_folder = '/playpen-raid2/qinliu/projects/iSegFormer/maskprop/Med-STCN/results'
-    # seg_root_folder = '/playpen-raid2/qinliu/projects/iSegFormer/maskprop/Med-STCN/results/MSD'
-    # for label in range(1, 2):
-    #     seg_folder = f'{seg_root_folder}/stcn/label_{label}'
-    #     evaluate(gt_folder, seg_folder, dataset_file, label)
-
-    # for label in range(1, 2):
-    #     seg_folder = f'{seg_root_folder}/stcn_med_abdomen1k_Aug01_15.34.08_ft_s012_10k_no_cc/label_{label}'
-    #     evaluate(gt_folder, seg_folder, dataset_file, label)
-
-    # for label in range(1, 2):
-    #     seg_folder = f'{seg_root_folder}/stcn_med_abdomen1k_Aug01_22.03.33_ft_s012_10k_cc/label_{label}'
-    #     evaluate(gt_folder, seg_folder, dataset_file, label)
+    parser = ArgumentParser()
+    parser.add_argument('--data_folder')
+    parser.add_argument('--seg_folder')
+    parser.add_argument('--label')
+    args = parser.parse_args()
 
-    label=12
-    seg_root_folder = '/playpen-raid2/qinliu/projects/STM/test/ABD_STMval'
-    seg_folder = f'{seg_root_folder}/label_{label}'
-    evaluate(gt_folder, seg_folder, dataset_file, label)
+    gt_folder = f'{args.data_folder}/trainval/Annotations/480p'
+    dataset_file = f'{args.data_folder}/trainval/ImageSets/val.txt'
+    seg_folder = f'{args.seg_folder}/label_{args.label}'
+    evaluate(gt_folder, seg_folder, dataset_file, args.label)

maskprop/Med-STCN/eval_med_stcn.py

@@ -1,28 +1,25 @@
+from argparse import ArgumentParser
+import numpy as np
 import os
 from os import path
+from PIL import Image
 import time
-from argparse import ArgumentParser
-
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
-import numpy as np
-from PIL import Image
+from tqdm import tqdm
 
 from model.eval_network import STCN
 from dataset.med_test_dataset import MedTestDataset
 from util.tensor_util import unpad
 from inference_core import InferenceCore
 
-from progressbar import progressbar
-
-
 """
 Arguments loading
 """
 parser = ArgumentParser()
-parser.add_argument('--model', default='saves/stcn.pth')
-parser.add_argument('--med_path', default='/playpen-raid2/qinliu/data/AbdomenCT-1K/Organ-12-Subset_frames')
+parser.add_argument('--model')
+parser.add_argument('--med_path')
 parser.add_argument('--output')
 parser.add_argument('--split', help='val', default='val')
 parser.add_argument('--top', type=int, default=20)
@@ -36,7 +33,7 @@
 
 # Simple setup
 os.makedirs(out_path, exist_ok=True)
-palette = Image.open(path.expanduser(f'/playpen-raid2/qinliu/data/AbdomenCT-1K/Organ-12-Subset_frames/trainval/Annotations/480p/Organ12_0001/0000000.png')).getpalette()
+palette = Image.open(path.expanduser(f'{args.med_path}/trainval/Annotations/480p/Organ12_0001/0000000.png')).getpalette()
 
 torch.autograd.set_grad_enabled(False)
 
@@ -67,8 +64,8 @@ def evaluate_single_label(label):
     total_frames = 0
 
     # start evaluation. Only 1 round propagation.
-    for data in progressbar(test_loader, max_value=len(test_loader), redirect_stdout=True):
-
+    # for data in progressbar(test_loader, max_value=len(test_loader), redirect_stdout=True):
+    for data in tqdm(test_loader):
         with torch.cuda.amp.autocast(enabled=args.amp):
             # example shape
             # rgb: [1, T, 3, 480, 480]; msk: [N, T, 1, 480, 480]

maskprop/Med-STCN/util/hyper_para.py

@@ -13,12 +13,11 @@ def parse(self, unknown_arg_ok=False):
         parser.add_argument('--no_amp', action='store_true')
 
         # Data parameters
-        parser.add_argument('--static_root', help='Static training data root', default='../static')
-        parser.add_argument('--bl_root', help='Blender training data root', default='../BL30K')
+        parser.add_argument('--static_root', help='Static training data root', default='/work/data/static')
+        parser.add_argument('--bl_root', help='Blender training data root', default='/work/data/BL30K')
         parser.add_argument('--yv_root', help='YouTubeVOS data root', default='/work/data/YouTube')
         parser.add_argument('--davis_root', help='DAVIS data root', default='/work/data/DAVIS')
-        parser.add_argument('--abd1k_root', help='Abdomen1k data root', 
-                            default='/playpen-raid2/qinliu/data/AbdomenCT-1K/Organ-12-Subset_finetune')
+        parser.add_argument('--abd1k_root', help='Abdomen1k data root', default='/work/data/adb1k')
 
         parser.add_argument('--stage', help='Training stage (0-static images, 1-Blender dataset, 2-DAVIS+YouTubeVOS (300K), 3-DAVIS+YouTubeVOS (150K))', type=int, default=0)
         parser.add_argument('--num_workers', help='Number of datalaoder workers per process', type=int, default=8)

requirements.txt

@@ -0,0 +1 @@
+progressbar

run_eval_stcn.sh

@@ -0,0 +1,12 @@
+
+python3 ./maskprop/Med-STCN/eval_med_stcn.py \
+  --model ./saves/stcn_ft_with_cycle.pth \
+  --med_path ./data/AbdomenCT-1K/Organ-12-Subset_frames \
+  --output ./results \
+  --include_last
+
+# calcualte the segmentation metrics for an assigned label
+python ./maskprop/Med-STCN/abdomen1k-evaluation/evaluation_method.py \
+--data_folder ./data/AbdomenCT-1K/Organ-12-Subset_frames \
+--seg_folder ./results \
+--label 1 # Set label from 1 to 12 for 12 organs AbdomenCT-1K

run_train_stcn_with_cycle.sh

@@ -4,6 +4,7 @@
 torchrun ./maskprop/Med-STCN/train.py \
 --id retrain_s4_ft_from_med \
 --load_network ./saves/stcn.pth \
+--abd1k_root ./data/AbdomenCT-1K/Organ-12-Subset_finetune \
 --stage 4 \
 --batch_size 10 \
 --iterations 10000 \

run_train_stcn_without_cycle.sh

@@ -4,6 +4,7 @@
 torchrun ./maskprop/Med-STCN/train.py \
 --id retrain_s4_ft_from_med \
 --load_network ./saves/stcn.pth \
+--abd1k_root ./data/AbdomenCT-1K/Organ-12-Subset_finetune \
 --stage 4 \
 --batch_size 10 \
 --iterations 10000 \