crop foreground

app.py

@@ -296,8 +296,17 @@ with col2:
     img_nii = None
     if st.session_state.preds_3D_ori is not None and st.session_state.data_item is not None:
         meta_dict = st.session_state.data_item['meta']
-        pred_array = st.session_state.preds_3D_ori.transpose(2, 1, 0)
-        img_nii = nib.Nifti1Image(pred_array, affine=meta_dict['affine'])
+        foreground_start_coord = st.session_state.data_item['foreground_start_coord']
+        foreground_end_coord = st.session_state.data_item['foreground_end_coord']
+        original_shape = st.session_state.data_item['ori_shape']
+        pred_array = st.session_state.preds_3D_ori
+        original_array = np.zeros(original_shape)
+        original_array[foreground_start_coord[0]:foreground_end_coord[0], 
+                    foreground_start_coord[1]:foreground_end_coord[1], 
+                    foreground_start_coord[2]:foreground_end_coord[2]] = pred_array
+
+        original_array = original_array.transpose(2, 1, 0)
+        img_nii = nib.Nifti1Image(original_array, affine=meta_dict['affine'])
 
         with tempfile.NamedTemporaryFile(suffix=".nii.gz") as tmpfile:
             nib.save(img_nii, tmpfile.name)
@@ -323,9 +332,6 @@ with col3:
         st.session_state.running = True
         st.rerun()
 
-# if len(st.session_state.points) > 0:
-#     st.write(st.session_state.points)
-
 if st.session_state.running:
     st.session_state.running = False
     with st.status("Running...", expanded=False) as status:

model/data_process/demo_data_process.py

@@ -60,7 +60,7 @@ def process_ct_gt(case_path, spatial_size=(32,256,256)):
             DimTranspose(keys=["image"]),
             MinMaxNormalization(),
             transforms.SpatialPadd(keys=["image"], spatial_size=spatial_size, mode='constant'),
-            # transforms.CropForegroundd(keys=["image"], source_key="image"),
+            transforms.CropForegroundd(keys=["image"], source_key="image"),
             transforms.ToTensord(keys=["image"]),
         ]
     )
@@ -81,12 +81,14 @@ def process_ct_gt(case_path, spatial_size=(32,256,256)):
     ct_voxel_ndarray = np.array(ct_voxel_ndarray).squeeze()
     ct_voxel_ndarray = np.expand_dims(ct_voxel_ndarray, axis=0)
     item['image'] = ct_voxel_ndarray
+    ori_shape = np.swapaxes(ct_voxel_ndarray, -1, -3).shape[1:]
 
     # transform
     item = transform(item)
     item_zoom_out = zoom_out_transform(item)
     item['zoom_out_image'] = item_zoom_out['image']
-    
+    item['ori_shape'] = ori_shape
+
     item_z = z_transform(item)
     item['z_image'] = item_z['image']
     item['meta'] = meta_tensor_dict