File size: 2,282 Bytes
4c02d29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import streamlit as st
import numpy as np
from apps.data_preprocessing import preprocess_input
import os
import tempfile
import nibabel as nib
st.title("NIfTI File Uploader and Video Generator")
st.write("This tool accepts NIfTI `.nii` or `.nii.gz` files smaller than a specific size.")
uploaded_file = st.file_uploader("Upload a NIfTI file (.nii or .nii.gz)", type=["nii", "nii.gz"])
if uploaded_file is not None:
with tempfile.TemporaryDirectory() as temp_dir:
try:
# Save uploaded file to a temporary location
nifti_path = os.path.join(temp_dir, "uploaded_file.nii.gz")
with open(nifti_path, "wb") as f:
f.write(uploaded_file.read())
# Preprocess the uploaded NIfTI file
output_tensor, dataset = preprocess_input(nifti_path)
imgs = dataset[0]["CT"]["data"]
pred = output_tensor.argmax(0)
numpy_array = pred.numpy().astype(np.int16)
affine = np.eye(4)
# Save finalized mask
mask_path = os.path.join(temp_dir, "finalized_mask.nii")
nifti_image_mask = nib.Nifti1Image(numpy_array, affine)
nib.save(nifti_image_mask, mask_path)
# Save finalized image
nifti_array = imgs.squeeze(0).numpy()
image_path = os.path.join(temp_dir, "finalized_image.nii.gz")
nifti_image_img = nib.Nifti1Image(nifti_array, affine)
nib.save(nifti_image_img, image_path)
# Provide download buttons for the finalized files
with open(mask_path, "rb") as f:
st.download_button(
label="Download Finalized Mask",
data=f,
file_name="finalized_mask.nii",
mime="application/octet-stream"
)
with open(image_path, "rb") as f:
st.download_button(
label="Download Finalized Image",
data=f,
file_name="finalized_image.nii.gz",
mime="application/octet-stream"
)
except Exception as e:
st.error(f"An error occurred: {e}")
else:
st.info("Please upload a NIfTI file to generate the video.")
|