update test cases

Can-Zhao · Can-Zhao · commit 210a063d705a · 2025-03-19T05:08:54.000Z
Signed-off-by: Can-Zhao &lt;canz@nvidia.com&gt;
diff --git a/generation/maisi/configs/config_infer.json b/generation/maisi/configs/config_infer.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_16g_256x256x128.json b/generation/maisi/configs/config_infer_16g_256x256x128.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_16g_256x256x256.json b/generation/maisi/configs/config_infer_16g_256x256x256.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_16g_512x512x128.json b/generation/maisi/configs/config_infer_16g_512x512x128.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_24g_256x256x256.json b/generation/maisi/configs/config_infer_24g_256x256x256.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_24g_512x512x128.json b/generation/maisi/configs/config_infer_24g_512x512x128.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_24g_512x512x512.json b/generation/maisi/configs/config_infer_24g_512x512x512.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_32g_512x512x512.json b/generation/maisi/configs/config_infer_32g_512x512x512.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_80g_512x512x512.json b/generation/maisi/configs/config_infer_80g_512x512x512.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
-    "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "body_region": ["chest"],
+    "anatomy_list": ["lung tumor"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/configs/config_infer_80g_512x512x768.json b/generation/maisi/configs/config_infer_80g_512x512x768.json
@@ -1,7 +1,7 @@
 {
     "num_output_samples": 1,
     "body_region": ["abdomen"],
-    "anatomy_list": ["liver","hepatic tumor"],
+    "anatomy_list": ["liver"],
     "controllable_anatomy_size": [],
     "num_inference_steps": 30,
     "mask_generation_num_inference_steps": 1000,
diff --git a/generation/maisi/maisi_inference_tutorial.ipynb b/generation/maisi/maisi_inference_tutorial.ipynb
diff --git a/generation/maisi/scripts/find_masks.py b/generation/maisi/scripts/find_masks.py
@@ -102,24 +102,26 @@ def find_masks(
         db = json.load(f)
 
     # select candidate_masks
-    candidate_masks = []
+    candidate_masks = []    
     for _item in db:
         if not set(anatomy_list).issubset(_item["label_list"]):
             continue
 
-        # extract region indice (top_index and bottom_index) for candidate mask
-        top_index = [index for index, element in enumerate(_item["top_region_index"]) if element != 0]
-        top_index = top_index[0]
-        bottom_index = [index for index, element in enumerate(_item["bottom_region_index"]) if element != 0]
-        bottom_index = bottom_index[0]
-
         # whether to keep this mask, default to be True.
         keep_mask = True
 
-        # if candiate mask does not contain all the body_region, skip it
-        for _idx in body_region:
-            if _idx > bottom_index or _idx < top_index:
-                keep_mask = False
+        # extract region indice (top_index and bottom_index) for candidate mask
+        include_body_region = ("top_region_index" in _item.keys())
+        if include_body_region:
+            top_index = [index for index, element in enumerate(_item["top_region_index"]) if element != 0]
+            top_index = top_index[0]
+            bottom_index = [index for index, element in enumerate(_item["bottom_region_index"]) if element != 0]
+            bottom_index = bottom_index[0]
+
+            # if candiate mask does not contain all the body_region, skip it
+            for _idx in body_region:
+                if _idx > bottom_index or _idx < top_index:
+                    keep_mask = False
 
         for tumor_label in [23, 24, 26, 27, 128]:
             # we skip those mask with tumors if users do not provide tumor label in anatomy_list
@@ -138,9 +140,10 @@ def find_masks(
                 "pseudo_label": os.path.join(mask_foldername, _item["pseudo_label_filename"]),
                 "spacing": _item["spacing"],
                 "dim": _item["dim"],
-                "top_region_index": _item["top_region_index"],
-                "bottom_region_index": _item["bottom_region_index"],
             }
+            if include_body_region:
+                candidate["top_region_index"] = _item["top_region_index"]
+                candidate["bottom_region_index"] = _item["bottom_region_index"]
 
             # Conditionally add the label to the candidate dictionary
             if "label_filename" in _item:

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`	`1`	`{`
`2`	`2`	`"num_output_samples": 1,`
`3`	`3`	`"body_region": ["abdomen"],`
`4`		`- "anatomy_list": ["liver","hepatic tumor"],`
	`4`	`+ "anatomy_list": ["liver"],`
`5`	`5`	`"controllable_anatomy_size": [],`
`6`	`6`	`"num_inference_steps": 30,`
`7`	`7`	`"mask_generation_num_inference_steps": 1000,`