Merge branch 'maisi' of https://github.com/Can-Zhao/tutorials into maisi

Author: Can-Zhao

generation/maisi/scripts/diff_model_infer.py

@@ -149,7 +149,7 @@ def run_inference(
     if isinstance(noise_scheduler, RFlowScheduler):
         noise_scheduler.set_timesteps(
             num_inference_steps=args.diffusion_unet_inference["num_inference_steps"],
-            input_img_size_numel=torch.prod(torch.tensor(noise.shape[-3:]))
+            input_img_size_numel=torch.prod(torch.tensor(noise.shape[-3:])),
         )
     else:
         noise_scheduler.set_timesteps(num_inference_steps=args.diffusion_unet_inference["num_inference_steps"])
@@ -161,9 +161,9 @@ def run_inference(
     all_timesteps = noise_scheduler.timesteps
     all_next_timesteps = torch.cat((all_timesteps[1:], torch.tensor([0], dtype=all_timesteps.dtype)))
     progress_bar = tqdm(
-            zip(all_timesteps, all_next_timesteps),
-            total=min(len(all_timesteps), len(all_next_timesteps)),
-        )
+        zip(all_timesteps, all_next_timesteps),
+        total=min(len(all_timesteps), len(all_next_timesteps)),
+    )
     with torch.amp.autocast("cuda", enabled=True):
         for t, next_t in progress_bar:
             model_output = unet(
@@ -178,7 +178,6 @@ def run_inference(
             else:
                 image, _ = noise_scheduler.step(model_output, t, image, next_t)  # type: ignore
 
-        
         inferer = SlidingWindowInferer(
             roi_size=[80,80,80],
             sw_batch_size=1,
@@ -224,7 +223,9 @@ 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, include_body_region: bool = False ) -> None:
+def diff_model_infer(
+    env_config_path: str, model_config_path: str, model_def_path: str, num_gpus: int, include_body_region: bool = False
+) -> None:
     """
     Main function to run the diffusion model inference.
 
@@ -331,7 +332,14 @@ def diff_model_infer(env_config_path: str, model_config_path: str, model_def_pat
         default=1,
         help="Number of GPUs to use for distributed inference",
     )
-    parser.add_argument("--include_body_region", dest="include_body_region", action="store_true", help="Whether to include body region in data")
+    parser.add_argument(
+        "--include_body_region",
+        dest="include_body_region",
+        action="store_true",
+        help="Whether to include body region in data",
+    )
 
     args = parser.parse_args()
-    diff_model_infer(args.env_config, args.model_config, args.model_def, args.num_gpus, include_body_region=args.include_body_region)
+    diff_model_infer(
+        args.env_config, args.model_config, args.model_def, args.num_gpus, include_body_region=args.include_body_region
+    )

generation/maisi/scripts/diff_model_train.py

@@ -51,12 +51,12 @@ def load_filenames(data_list_path: str) -> list:
 
 
 def prepare_data(
-    train_files: list, 
-    device: torch.device, 
-    cache_rate: float, 
-    num_workers: int = 2, 
-    batch_size: int = 1, 
-    include_body_region: bool = False
+    train_files: list,
+    device: torch.device,
+    cache_rate: float,
+    num_workers: int = 2,
+    batch_size: int = 1,
+    include_body_region: bool = False,
 ) -> DataLoader:
     """
     Prepare training data.
@@ -78,11 +78,11 @@ def _load_data_from_file(file_path, key):
             return torch.FloatTensor(json.load(f)[key])
 
     train_transforms_list = [
-            monai.transforms.LoadImaged(keys=["image"]),
-            monai.transforms.EnsureChannelFirstd(keys=["image"]),            
-            monai.transforms.Lambdad(keys="spacing", func=lambda x: _load_data_from_file(x, "spacing")),            
-            monai.transforms.Lambdad(keys="spacing", func=lambda x: x * 1e2),
-        ]
+        monai.transforms.LoadImaged(keys=["image"]),
+        monai.transforms.EnsureChannelFirstd(keys=["image"]),
+        monai.transforms.Lambdad(keys="spacing", func=lambda x: _load_data_from_file(x, "spacing")),
+        monai.transforms.Lambdad(keys="spacing", func=lambda x: x * 1e2),
+    ]
     if include_body_region:
         train_transforms_list += [
             monai.transforms.Lambdad(
@@ -202,7 +202,7 @@ def train_one_epoch(
     logger: logging.Logger,
     local_rank: int,
     amp: bool = True,
-    include_body_region: bool = False
+    include_body_region: bool = False,
 ) -> torch.Tensor:
     """
     Train the model for one epoch.
@@ -284,9 +284,10 @@ def train_one_epoch(
                 # predict velocity
                 loss = loss_pt(model_output.float(), (images - noise).float())
             else:
-                raise ValueError("noise scheduler prediction type has to be chosen from ",
-                     f"[{DDPMPredictionType.EPSILON},{DDPMPredictionType.SAMPLE},{DDPMPredictionType.V_PREDICTION}]"
-                    )
+                raise ValueError(
+                    "noise scheduler prediction type has to be chosen from ",
+                    f"[{DDPMPredictionType.EPSILON},{DDPMPredictionType.SAMPLE},{DDPMPredictionType.V_PREDICTION}]",
+                )
 
         if amp:
             scaler.scale(loss).backward()
@@ -349,7 +350,12 @@ 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, include_body_region: bool = False
+    env_config_path: str,
+    model_config_path: str,
+    model_def_path: str,
+    num_gpus: int,
+    amp: bool = True,
+    include_body_region: bool = False,
 ) -> None:
     """
     Main function to train a diffusion model.
@@ -400,9 +406,11 @@ def diff_model_train(
         )[local_rank]
 
     train_loader = prepare_data(
-        train_files, device, args.diffusion_unet_train["cache_rate"], 
+        train_files,
+        device,
+        args.diffusion_unet_train["cache_rate"],
         batch_size=args.diffusion_unet_train["batch_size"],
-        include_body_region = include_body_region
+        include_body_region=include_body_region,
     )
 
     unet = load_unet(args, device, logger)
@@ -438,7 +446,7 @@ def diff_model_train(
             logger,
             local_rank,
             amp=amp,
-            include_body_region=include_body_region
+            include_body_region=include_body_region,
         )
 
         loss_torch = loss_torch.tolist()
@@ -479,7 +487,14 @@ def diff_model_train(
     )
     parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs to use for training")
     parser.add_argument("--no_amp", dest="amp", action="store_false", help="Disable automatic mixed precision training")
-    parser.add_argument("--include_body_region", dest="include_body_region", action="store_true", help="Whether to include body region in data")
+    parser.add_argument(
+        "--include_body_region",
+        dest="include_body_region",
+        action="store_true",
+        help="Whether to include body region in data",
+    )
 
     args = parser.parse_args()
-    diff_model_train(args.env_config, args.model_config, args.model_def, args.num_gpus, args.amp, args.include_body_region)
+    diff_model_train(
+        args.env_config, args.model_config, args.model_def, args.num_gpus, args.amp, args.include_body_region
+    )

generation/maisi/scripts/utils.py

@@ -712,14 +712,14 @@ def dynamic_infer(inferer, model, images):
         # Extract the spatial dimensions from the images tensor (H, W, D)
         spatial_dims = images.shape[2:]
         orig_roi = inferer.roi_size
-        
+
         # Check that roi has the same number of dimensions as spatial_dims
         if len(orig_roi) != len(spatial_dims):
             raise ValueError(f"ROI length ({len(orig_roi)}) does not match spatial dimensions ({len(spatial_dims)}).")
-        
+
         # Iterate and adjust each ROI dimension
         adjusted_roi = [min(roi_dim, img_dim) for roi_dim, img_dim in zip(orig_roi, spatial_dims)]
         inferer.roi_size = adjusted_roi
         output = inferer(network=model, inputs=images)
         inferer.roi_size = orig_roi
-        return output
+        return output