add modality as input to config

Signed-off-by: Can-Zhao <canz@nvidia.com>

Author: Can-Zhao

generation/maisi/maisi_inference_tutorial.ipynb

@@ -209,7 +209,7 @@
     "        \"path\": \"models/autoencoder_epoch273.pt\",\n",
     "        \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials\"\n",
     "        \"/model_zoo/model_maisi_autoencoder_epoch273_alternative.pt\",\n",
-    "    },    \n",
+    "    },\n",
     "    {\n",
     "        \"path\": \"models/mask_generation_autoencoder.pt\",\n",
     "        \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\" \"/tutorials/mask_generation_autoencoder.pt\",\n",
@@ -218,7 +218,7 @@
     "        \"path\": \"models/mask_generation_diffusion_unet.pt\",\n",
     "        \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n",
     "        \"/tutorials/model_zoo/model_maisi_mask_generation_diffusion_unet_v2.pt\",\n",
-    "    },    \n",
+    "    },\n",
     "    {\n",
     "        \"path\": \"configs/all_anatomy_size_condtions.json\",\n",
     "        \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/all_anatomy_size_condtions.json\",\n",
@@ -258,14 +258,18 @@
     "        },\n",
     "        {\n",
     "            \"path\": \"models/controlnet_3d_rflow.pt\",\n",
-    "            \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/\"\n",
-    "            \"controlnet_rflow_epoch208.pt\",\n",
+    "            \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/\" \"controlnet_rflow_epoch208.pt\",\n",
     "        },\n",
     "        {\n",
     "            \"path\": \"configs/candidate_masks_flexible_size_and_spacing_4000.json\",\n",
     "            \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n",
     "            \"/tutorials/candidate_masks_flexible_size_and_spacing_4000.json\",\n",
-    "        },            \n",
+    "        },\n",
+    "        {\n",
+    "            \"path\": \"datasets/all_masks_flexible_size_and_spacing_4000.zip\",\n",
+    "            \"url\": \"https://developer.download.nvidia.com/assets/Clara/monai\"\n",
+    "            \"/tutorials/all_masks_flexible_size_and_spacing_4000.zip\",\n",
+    "        },\n",
     "    ]\n",
     "else:\n",
     "    raise ValueError(f\"maisi_version has to be chosen from ['maisi3d-ddpm', 'maisi3d-rflow'], yet got {maisi_version}.\")\n",
@@ -537,7 +541,7 @@
     "    mask_generation_num_inference_steps=args.mask_generation_num_inference_steps,\n",
     "    random_seed=args.random_seed,\n",
     "    autoencoder_sliding_window_infer_size=args.autoencoder_sliding_window_infer_size,\n",
-    "    autoencoder_sliding_window_infer_overlap=args.autoencoder_sliding_window_infer_overlap\n",
+    "    autoencoder_sliding_window_infer_overlap=args.autoencoder_sliding_window_infer_overlap,\n",
     ")"
    ]
   },

generation/maisi/scripts/augmentation.py

@@ -370,4 +370,4 @@ def augmentation(pt_nda, output_size, random_seed):
         print("augmenting body")
         pt_nda = augmentation_body(pt_nda, random_seed)
 
-    return pt_nda
+    return pt_nda

generation/maisi/scripts/diff_model_infer.py

@@ -136,8 +136,8 @@ def run_inference(
         np.ndarray: Generated synthetic image data.
     """
     include_body_region = unet.include_top_region_index_input
-    include_modality = (unet.num_class_embeds is not None)
-    
+    include_modality = unet.num_class_embeds is not None
+
     noise = torch.randn(
         (
             1,
@@ -178,18 +178,22 @@ def run_inference(
                 "timesteps": torch.Tensor((t,)).to(device),
                 "spacing_tensor": spacing_tensor,
             }
-            
+
             # Add extra arguments if include_body_region is True
             if include_body_region:
-                unet_inputs.update({
-                    "top_region_index_tensor": top_region_index_tensor,
-                    "bottom_region_index_tensor": bottom_region_index_tensor
-                }) 
+                unet_inputs.update(
+                    {
+                        "top_region_index_tensor": top_region_index_tensor,
+                        "bottom_region_index_tensor": bottom_region_index_tensor,
+                    }
+                )
 
             if include_modality:
-                unet_inputs.update({
-                    "class_labels": modality_tensor,
-                }) 
+                unet_inputs.update(
+                    {
+                        "class_labels": modality_tensor,
+                    }
+                )
             model_output = unet(**unet_inputs)
             if not isinstance(noise_scheduler, RFlowScheduler):
                 image, _ = noise_scheduler.step(model_output, t, image)  # type: ignore
@@ -241,9 +245,7 @@ def save_image(
 
 
 @torch.inference_mode()
-def diff_model_infer(
-    env_config_path: str, model_config_path: str, model_def_path: str, num_gpus: int
-) -> None:
+def diff_model_infer(env_config_path: str, model_config_path: str, model_def_path: str, num_gpus: int) -> None:
     """
     Main function to run the diffusion model inference.
 
@@ -301,7 +303,7 @@ def diff_model_infer(
         modality_tensor,
         output_size,
         divisor,
-        logger
+        logger,
     )
 
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
@@ -352,6 +354,4 @@ def diff_model_infer(
     )
 
     args = parser.parse_args()
-    diff_model_infer(
-        args.env_config, args.model_config, args.model_def, args.num_gpus
-    )
+    diff_model_infer(args.env_config, args.model_config, args.model_def, args.num_gpus)

generation/maisi/scripts/diff_model_train.py

@@ -226,8 +226,8 @@ def train_one_epoch(
         torch.Tensor: Training loss for the epoch.
     """
     include_body_region = unet.include_top_region_index_input
-    include_modality = (unet.num_class_embeds is not None)
-    
+    include_modality = unet.num_class_embeds is not None
+
     if local_rank == 0:
         current_lr = optimizer.param_groups[0]["lr"]
         logger.info(f"Epoch {epoch + 1}, lr {current_lr}.")
@@ -248,7 +248,7 @@ def train_one_epoch(
             bottom_region_index_tensor = train_data["bottom_region_index"].to(device)
         # We trained with only CT in this version
         if include_modality:
-            modality_tensor = torch.ones((len(images),),dtype=torch.long).to(device)
+            modality_tensor = torch.ones((len(images),), dtype=torch.long).to(device)
         spacing_tensor = train_data["spacing"].to(device)
 
         optimizer.zero_grad(set_to_none=True)
@@ -268,18 +268,22 @@ def train_one_epoch(
                 "x": noisy_latent,
                 "timesteps": timesteps,
                 "spacing_tensor": spacing_tensor,
-            }            
+            }
             # Add extra arguments if include_body_region is True
             if include_body_region:
-                unet_inputs.update({
-                    "top_region_index_tensor": top_region_index_tensor,
-                    "bottom_region_index_tensor": bottom_region_index_tensor
-                }) 
+                unet_inputs.update(
+                    {
+                        "top_region_index_tensor": top_region_index_tensor,
+                        "bottom_region_index_tensor": bottom_region_index_tensor,
+                    }
+                )
             if include_modality:
-                unet_inputs.update({
-                    "class_labels": modality_tensor,
-                }) 
-            model_output = unet(**unet_inputs)                
+                unet_inputs.update(
+                    {
+                        "class_labels": modality_tensor,
+                    }
+                )
+            model_output = unet(**unet_inputs)
 
             if noise_scheduler.prediction_type == DDPMPredictionType.EPSILON:
                 # predict noise
@@ -359,11 +363,7 @@ def save_checkpoint(
 
 
 def diff_model_train(
-    env_config_path: str,
-    model_config_path: str,
-    model_def_path: str,
-    num_gpus: int,
-    amp: bool = True
+    env_config_path: str, model_config_path: str, model_def_path: str, num_gpus: int, amp: bool = True
 ) -> None:
     """
     Main function to train a diffusion model.
@@ -424,8 +424,6 @@ def diff_model_train(
         include_body_region=include_body_region,
     )
 
-    
-
     scale_factor = calculate_scale_factor(train_loader, device, logger)
     optimizer = create_optimizer(unet, args.diffusion_unet_train["lr"])
 
@@ -455,7 +453,7 @@ def diff_model_train(
             device,
             logger,
             local_rank,
-            amp=amp
+            amp=amp,
         )
 
         loss_torch = loss_torch.tolist()
@@ -498,6 +496,4 @@ def diff_model_train(
     parser.add_argument("--no_amp", dest="amp", action="store_false", help="Disable automatic mixed precision training")
 
     args = parser.parse_args()
-    diff_model_train(
-        args.env_config, args.model_config, args.model_def, args.num_gpus, args.amp
-    )
+    diff_model_train(args.env_config, args.model_config, args.model_def, args.num_gpus, args.amp)

generation/maisi/scripts/infer_controlnet.py

@@ -111,8 +111,8 @@ def main():
     # define diffusion Model
     unet = define_instance(args, "diffusion_unet_def").to(device)
     include_body_region = unet.include_top_region_index_input
-    include_modality = (unet.num_class_embeds is not None)
-    
+    include_modality = unet.num_class_embeds is not None
+
     # load trained diffusion model
     if args.trained_diffusion_path is not None:
         if not os.path.exists(args.trained_diffusion_path):

generation/maisi/scripts/quality_check.py

@@ -113,7 +113,7 @@ def is_outlier(statistics, image_data, label_data, label_int_dict):
         high_thresh = max(stats["sigma_6_high"], stats["percentile_99_5"])  # or "sigma_12_high" depending on your needs
 
         if label_name == "bone":
-            high_thresh = 1000.
+            high_thresh = 1000.0
 
         # Retrieve the corresponding label integers
         labels = label_int_dict.get(label_name, [])
@@ -146,4 +146,4 @@ def is_outlier(statistics, image_data, label_data, label_int_dict):
             "high_thresh": high_thresh,
         }
 
-    return outlier_results
+    return outlier_results