Add spatial size as iput for infer class - UNETR configuration (#707)

* Add spatial size as iput for infer class - UNETR configuration

Signed-off-by: Andres <[email protected]>

* Change default network to train from main

Signed-off-by: Andres <[email protected]>

Author: diazandr3s

sample-apps/radiology/lib/configs/deepedit.py

@@ -63,13 +63,13 @@ def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **
             download_file(url, self.path[0])
 
         # Network
-        spatial_size = json.loads(self.conf.get("spatial_size", "[128, 128, 128]"))
+        self.spatial_size = json.loads(self.conf.get("spatial_size", "[128, 128, 128]"))
         if network == "unetr":
             self.network = UNETR(
                 spatial_dims=3,
                 in_channels=len(self.labels) + 1,
                 out_channels=len(self.labels),
-                img_size=spatial_size,
+                img_size=self.spatial_size,
                 feature_size=64,
                 hidden_size=1536,
                 mlp_dim=3072,
@@ -93,9 +93,15 @@ def init(self, name: str, model_dir: str, conf: Dict[str, str], planner: Any, **
 
     def infer(self) -> Union[InferTask, Dict[str, InferTask]]:
         return {
-            self.name: lib.infers.DeepEdit(path=self.path, network=self.network, labels=self.labels),
+            self.name: lib.infers.DeepEdit(
+                path=self.path, network=self.network, labels=self.labels, spatial_size=self.spatial_size
+            ),
             f"{self.name}_seg": lib.infers.DeepEdit(
-                path=self.path, network=self.network, labels=self.labels, type=InferType.SEGMENTATION
+                path=self.path,
+                network=self.network,
+                labels=self.labels,
+                spatial_size=self.spatial_size,
+                type=InferType.SEGMENTATION,
             ),
         }
 

sample-apps/radiology/main.py

@@ -219,6 +219,8 @@ def main():
     conf = {
         "models": "deepedit",
         "use_pretrained_model": "false",
+        # "network": "unetr",
+        # "spatial_size": "[128,128,128]",
     }
 
     app = MyApp(app_dir, studies, conf)