download support

app.py

@@ -9,6 +9,8 @@ import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw
 import monai.transforms as transforms
 from utils import show_points, make_fig, reflect_points_into_model, initial_rectangle, reflect_json_data_to_3D_box, reflect_box_into_model, run
+import nibabel as nib
+import tempfile
 
 print('script run')
 
@@ -24,6 +26,7 @@ if 'reset_demo_case' not in st.session_state:
 
 if 'preds_3D' not in st.session_state:
     st.session_state.preds_3D = None
+    st.session_state.preds_3D_ori = None
 
 if 'data_item' not in st.session_state:
     st.session_state.data_item = None
@@ -117,6 +120,7 @@ if  st.session_state.option is not None and \
     st.session_state.data_item = process_ct_gt(st.session_state.option)
     st.session_state.reset_demo_case = False
     st.session_state.preds_3D = None
+    st.session_state.preds_3D_ori = None
 
 prompt_col1, prompt_col2 = st.columns(2)
 
@@ -150,6 +154,7 @@ with prompt_col2:
         ["Point prompt", "Box prompt"],
         on_change=clear_prompts,
         disabled=(not spatial_prompt_on))
+    st.session_state.enforce_zoom = st.checkbox('Enforce zoom-out-zoom-in')
 
 if spatial_prompt == "Point prompt":
     st.session_state.use_point_prompt = True
@@ -284,8 +289,28 @@ with col1:
                  disabled=(st.session_state.option is None or (len(st.session_state.points)==0 and not st.session_state.use_box_prompt and st.session_state.preds_3D is None))):
         clear_prompts()
         st.session_state.preds_3D = None
+        st.session_state.preds_3D_ori = None
         st.rerun()
 
+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'])
+
+        with tempfile.NamedTemporaryFile(suffix=".nii.gz") as tmpfile:
+            nib.save(img_nii, tmpfile.name)
+            with open(tmpfile.name, "rb") as f:
+                bytes_data = f.read()
+                st.download_button(
+                    label="Download result(.nii.gz)",
+                    data=bytes_data,
+                    file_name="segvol_preds.nii.gz",
+                    mime="application/octet-stream",
+                    disabled=img_nii is None
+                )
+
 with col3:
     run_button_name = 'Run'if not st.session_state.running else 'Running'
     if st.button(run_button_name, type="primary", use_container_width=True,

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"]),
         ]
     )
@@ -70,13 +70,14 @@ def process_ct_gt(case_path, spatial_size=(32,256,256)):
     item = {}
     # generate ct_voxel_ndarray
     if type(case_path) is str:
-        ct_voxel_ndarray, _ = img_loader(case_path)
+        ct_voxel_ndarray, meta_tensor_dict = img_loader(case_path)
     else:
         bytes_data = case_path.read()
         with tempfile.NamedTemporaryFile(suffix='.nii.gz') as tmp:
             tmp.write(bytes_data)
             tmp.seek(0)
-            ct_voxel_ndarray, _ = img_loader(tmp.name)
+            ct_voxel_ndarray, meta_tensor_dict = img_loader(tmp.name)
+
     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
@@ -88,4 +89,5 @@ def process_ct_gt(case_path, spatial_size=(32,256,256)):
     
     item_z = z_transform(item)
     item['z_image'] = item_z['image']
+    item['meta'] = meta_tensor_dict
     return item

model/inference_cpu.py

@@ -167,7 +167,6 @@ def inference_case(_image, _image_zoom_out, _point_prompt, text_prompt, _box_pro
         transforms.Resize((325,325,325), mode='trilinear')
     ]
     )
-    logits = resize_transform(logits)[0]
-    print(logits.shape)
-    return (torch.sigmoid(logits) > 0.5).int().numpy()
+    logits_resize = resize_transform(logits)[0]
+    return (torch.sigmoid(logits_resize) > 0.5).int().numpy(), (torch.sigmoid(logits) > 0.5).int().numpy()
     

model/utils/monai_inferers_utils.py

@@ -192,7 +192,7 @@ def sliding_window_inference(
     slices = dense_patch_slices(image_size, roi_size, scan_interval)
     num_win = len(slices)  # number of windows per image
     total_slices = num_win * batch_size  # total number of windows
-    if total_slices > 10:
+    if total_slices > 10 and not st.session_state.enforce_zoom:
         return logits_global_single
 
     # Create window-level importance map

utils.py

@@ -61,7 +61,7 @@ def run():
     if st.session_state.use_box_prompt:
         box_prompt = reflect_box_into_model(st.session_state.rectangle_3Dbox)
     inference_case.clear()
-    st.session_state.preds_3D = inference_case(image, image_zoom_out, 
+    st.session_state.preds_3D, st.session_state.preds_3D_ori = inference_case(image, image_zoom_out, 
                                             text_prompt=text_prompt,
                                             _point_prompt=point_prompt,
                                             _box_prompt=box_prompt)