Merge remote-tracking branch 'origin/HF_spatialtransformer' into HF_spatialtransformer

DeepDeformationMapRegistration/main.py

@@ -300,7 +300,6 @@ def main():
 
         LOGGER.info('Applying displacement map...')
         time_pred_img_start = time.time()
-        # pred_image = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]]).eval()
         pred_image = spatialtransformer_model.predict([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]])
         time_pred_img_end = time.time()
         LOGGER.info(f'\t... done ({time_pred_img_end - time_pred_img_start} s)')
@@ -313,6 +312,7 @@ def main():
             # disp_map = disp_map_or
             pred_image = zoom(pred_image, 1 / zoom_factors)
             pred_image = pad_crop_to_original_shape(pred_image, fixed_image_or.shape, crop_min)
+            pred_image = np.squeeze(pred_image, axis=-1)
             LOGGER.info('Done...')
 
         if args.original_resolution:

demo/src/gui.py

@@ -20,7 +20,7 @@ class WebUI:
         self.pred_images = []
 
         # @TODO: This should be dynamically set based on chosen volume size
-        self.nb_slider_items = 150
+        self.nb_slider_items = 128
 
         self.model_name = model_name
         self.cwd = cwd
@@ -60,8 +60,8 @@ class WebUI:
 
         self.fixed_images = load_ct_to_numpy(fixed_image_path)
         self.moving_images = load_ct_to_numpy(moving_image_path)
-        self.pred_images = np.ones_like(self.moving_images)
-        return self.pred_images
+        self.pred_images = load_ct_to_numpy(output_path + "pred_image.nii.gz")
+        return None
 
     def get_fixed_image(self, k):
         k = int(k) - 1
@@ -151,7 +151,7 @@ class WebUI:
                             for i in range(self.nb_slider_items):
                                 visibility = True if i == 1 else False
                                 t = gr.Image(
-                                    visible=visibility, elem_id="model-2d-fixed"
+                                    visible=visibility, elem_id="model-2d-fixed", label="fixed image", show_label=True,
                                 ).style(
                                     height=512,
                                     width=512,
@@ -162,7 +162,7 @@ class WebUI:
                             for i in range(self.nb_slider_items):
                                 visibility = True if i == 1 else False
                                 t = gr.Image(
-                                    visible=visibility, elem_id="model-2d-moving"
+                                    visible=visibility, elem_id="model-2d-moving", label="moving image", show_label=True,
                                 ).style(
                                     height=512,
                                     width=512,
@@ -173,7 +173,7 @@ class WebUI:
                             for i in range(self.nb_slider_items):
                                 visibility = True if i == 1 else False
                                 t = gr.Image(
-                                    visible=visibility, elem_id="model-2d-pred"
+                                    visible=visibility, elem_id="model-2d-pred", label="predicted fixed image", show_label=True,
                                 ).style(
                                     height=512,
                                     width=512,
@@ -183,7 +183,7 @@ class WebUI:
                             self.run_btn.click(
                                 fn=lambda x: self.process(x),
                                 inputs=file_output,
-                                outputs=t,
+                                outputs=None,
                             )
 
                         self.slider.input(

demo/src/utils.py

@@ -21,7 +21,6 @@ def load_ct_to_numpy(data_path):
     data = data / np.amax(data) * 255
     data = data.astype("uint8")
 
-    print(data.shape)
     return [data[..., i] for i in range(data.shape[-1])]
 
 
@@ -38,7 +37,6 @@ def load_pred_volume_to_numpy(data_path):
     data[data > 0] = 1
     data = data.astype("uint8")
 
-    print(data.shape)
     return [data[..., i] for i in range(data.shape[-1])]