update search function to match monai 1.2

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.4.1",
+    "version": "0.4.2",
     "changelog": {
+        "0.4.2": "update search function to match monai 1.2",
         "0.4.1": "fix the wrong GPU index issue of multi-node",
         "0.4.0": "remove error dollar symbol in readme",
         "0.3.9": "add cpu ram requirement in readme",

scripts/prepare_datalist.py

@@ -11,11 +11,10 @@ def produce_sample_dict(line: str):
     return {"label": line, "image": line.replace("labelsTr", "imagesTr")}
 
 
-def produce_datalist(dataset_dir: str):
+def produce_datalist(dataset_dir: str, train_size: int = 196):
     """
     This function is used to split the dataset.
-    It will produce 200 samples for training, and the other samples are divided equally
-    into val and test sets.
+    It will produce "train_size" number of samples for training.
     """
 
     samples = sorted(glob.glob(os.path.join(dataset_dir, "labelsTr", "*"), recursive=True))
@@ -23,7 +22,7 @@ def produce_datalist(dataset_dir: str):
     datalist = []
     for line in samples:
         datalist.append(produce_sample_dict(line))
-    train_list, other_list = train_test_split(datalist, train_size=196)
+    train_list, other_list = train_test_split(datalist, train_size=train_size)
     val_list, test_list = train_test_split(other_list, train_size=0.66)
 
     return {"training": train_list, "validation": val_list, "testing": test_list}
@@ -37,7 +36,7 @@ def main(args):
     output_json = args.output
     # produce deterministic data splits
     monai.utils.set_determinism(seed=123)
-    datalist = produce_datalist(dataset_dir=data_file_base_dir)
+    datalist = produce_datalist(dataset_dir=data_file_base_dir, train_size=args.train_size)
     with open(output_json, "w") as f:
         json.dump(datalist, f, ensure_ascii=True, indent=4)
 
@@ -53,6 +52,7 @@ if __name__ == "__main__":
     parser.add_argument(
         "--output", type=str, default="dataset_0.json", help="relative path of output datalist json file."
     )
+    parser.add_argument("--train_size", type=int, default=196, help="number of training samples.")
     args = parser.parse_args()
 
     main(args)

scripts/search.py

@@ -28,7 +28,7 @@ from monai import transforms
 from monai.bundle import ConfigParser
 from monai.data import ThreadDataLoader, partition_dataset
 from monai.inferers import sliding_window_inference
-from monai.metrics import compute_meandice
+from monai.metrics import compute_dice
 from monai.utils import set_determinism
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.tensorboard import SummaryWriter
@@ -100,14 +100,12 @@ def run(config_file: Union[str, Sequence[str]]):
         train_files_w = partition_dataset(
             data=train_files_w, shuffle=True, num_partitions=world_size, even_divisible=True
         )[dist.get_rank()]
-    print("train_files_w:", len(train_files_w))
 
     train_files_a = train_files[len(train_files) // 2 :]
     if torch.cuda.device_count() > 1:
         train_files_a = partition_dataset(
             data=train_files_a, shuffle=True, num_partitions=world_size, even_divisible=True
         )[dist.get_rank()]
-    print("train_files_a:", len(train_files_a))
 
     # validation data
     files = []
@@ -125,7 +123,6 @@ def run(config_file: Union[str, Sequence[str]]):
         val_files = partition_dataset(data=val_files, shuffle=False, num_partitions=world_size, even_divisible=False)[
             dist.get_rank()
         ]
-    print("val_files:", len(val_files))
 
     # network architecture
     if torch.cuda.device_count() > 1:
@@ -421,7 +418,7 @@ def run(config_file: Union[str, Sequence[str]]):
                     val_labels = post_label(val_labels[0, ...])
                     val_labels = val_labels[None, ...]
 
-                    value = compute_meandice(y_pred=val_outputs, y=val_labels, include_background=False)
+                    value = compute_dice(y_pred=val_outputs, y=val_labels, include_background=False)
 
                     print(_index + 1, "/", len(val_loader), value)