import streamlit as st from streamlit_drawable_canvas import st_canvas from streamlit_image_coordinates import streamlit_image_coordinates from model.data_process.demo_data_process import process_ct_gt import numpy as np 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') ############################################# # init session_state if 'option' not in st.session_state: st.session_state.option = None if 'text_prompt' not in st.session_state: st.session_state.text_prompt = None if 'reset_demo_case' not in st.session_state: st.session_state.reset_demo_case = False 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 if 'points' not in st.session_state: st.session_state.points = [] if 'use_text_prompt' not in st.session_state: st.session_state.use_text_prompt = False if 'use_point_prompt' not in st.session_state: st.session_state.use_point_prompt = False if 'use_box_prompt' not in st.session_state: st.session_state.use_box_prompt = False if 'rectangle_3Dbox' not in st.session_state: st.session_state.rectangle_3Dbox = [0,0,0,0,0,0] if 'irregular_box' not in st.session_state: st.session_state.irregular_box = False if 'running' not in st.session_state: st.session_state.running = False if 'transparency' not in st.session_state: st.session_state.transparency = 0.25 case_list = [ 'model/asset/FLARE22_Tr_0002_0000.nii.gz', 'model/asset/FLARE22_Tr_0005_0000.nii.gz', 'model/asset/FLARE22_Tr_0034_0000.nii.gz', 'model/asset/FLARE22_Tr_0045_0000.nii.gz' ] ############################################# ############################################# # reset functions def clear_prompts(): st.session_state.points = [] st.session_state.rectangle_3Dbox = [0,0,0,0,0,0] def reset_demo_case(): st.session_state.data_item = None st.session_state.reset_demo_case = True clear_prompts() def clear_file(): st.session_state.option = None process_ct_gt.clear() reset_demo_case() clear_prompts() ############################################# st.image(Image.open('model/asset/overview back.png'), use_column_width=True) github_col, arxive_col = st.columns(2) with github_col: st.write('GitHub repo:https://github.com/BAAI-DCAI/SegVol') with arxive_col: st.write('Paper:https://arxiv.org/abs/2311.13385') # modify demo case here demo_type = st.radio( "Demo case source", ["Select", "Upload"], on_change=clear_file ) if demo_type=="Select": uploaded_file = st.selectbox( "Select a demo case", case_list, index=None, placeholder="Select a demo case...", on_change=reset_demo_case ) else: uploaded_file = st.file_uploader("Upload demo case(nii.gz)", type='nii.gz', on_change=reset_demo_case) st.session_state.option = uploaded_file if st.session_state.option is not None and \ st.session_state.reset_demo_case or (st.session_state.data_item is None and st.session_state.option is not None): 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) with prompt_col1: st.session_state.use_text_prompt = st.toggle('Sematic prompt') text_prompt_type = st.radio( "Sematic prompt type", ["Predefined", "Custom"], disabled=(not st.session_state.use_text_prompt) ) if text_prompt_type == "Predefined": pre_text = st.selectbox( "Predefined anatomical category:", ['liver', 'right kidney', 'spleen', 'pancreas', 'aorta', 'inferior vena cava', 'right adrenal gland', 'left adrenal gland', 'gallbladder', 'esophagus', 'stomach', 'duodenum', 'left kidney'], index=None, disabled=(not st.session_state.use_text_prompt) ) else: pre_text = st.text_input('Enter an Anatomical word or phrase:', None, max_chars=20, disabled=(not st.session_state.use_text_prompt)) if pre_text is None or len(pre_text) > 0: st.session_state.text_prompt = pre_text else: st.session_state.text_prompt = None with prompt_col2: spatial_prompt_on = st.toggle('Spatial prompt', on_change=clear_prompts) spatial_prompt = st.radio( "Spatial prompt type", ["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 st.session_state.use_box_prompt = False elif spatial_prompt == "Box prompt": st.session_state.use_box_prompt = True st.session_state.use_point_prompt = False else: st.session_state.use_point_prompt = False st.session_state.use_box_prompt = False if not spatial_prompt_on: st.session_state.use_point_prompt = False st.session_state.use_box_prompt = False if not st.session_state.use_text_prompt: st.session_state.text_prompt = None if st.session_state.option is None: st.write('please select demo case first') else: image_3D = st.session_state.data_item['z_image'][0].numpy() col_control1, col_control2 = st.columns(2) with col_control1: selected_index_z = st.slider('X-Y view', 0, image_3D.shape[0] - 1, 162, key='xy', disabled=st.session_state.running) with col_control2: selected_index_y = st.slider('X-Z view', 0, image_3D.shape[1] - 1, 162, key='xz', disabled=st.session_state.running) if st.session_state.use_box_prompt: top, bottom = st.select_slider( 'Top and bottom of box', options=range(0, 325), value=(0, 324), disabled=st.session_state.running ) st.session_state.rectangle_3Dbox[0] = top st.session_state.rectangle_3Dbox[3] = bottom col_image1, col_image2 = st.columns(2) if st.session_state.preds_3D is not None: st.session_state.transparency = st.slider('Mask opacity', 0.0, 1.0, 0.25, disabled=st.session_state.running) with col_image1: image_z_array = image_3D[selected_index_z] preds_z_array = None if st.session_state.preds_3D is not None: preds_z_array = st.session_state.preds_3D[selected_index_z] image_z = make_fig(image_z_array, preds_z_array, st.session_state.points, selected_index_z, 'xy') if st.session_state.use_point_prompt: value_xy = streamlit_image_coordinates(image_z, width=325) if value_xy is not None: point_ax_xy = (selected_index_z, value_xy['y'], value_xy['x']) if len(st.session_state.points) >= 3: st.warning('Max point num is 3', icon="⚠️") elif point_ax_xy not in st.session_state.points: st.session_state.points.append(point_ax_xy) print('point_ax_xy add rerun') st.rerun() elif st.session_state.use_box_prompt: canvas_result_xy = st_canvas( fill_color="rgba(255, 165, 0, 0.3)", # Fixed fill color with some opacity stroke_width=3, stroke_color='#2909F1', background_image=image_z, update_streamlit=True, height=325, width=325, drawing_mode='transform', point_display_radius=0, key="canvas_xy", initial_drawing=initial_rectangle, display_toolbar=True ) try: print(canvas_result_xy.json_data['objects'][0]['angle']) if canvas_result_xy.json_data['objects'][0]['angle'] != 0: st.warning('Rotating is undefined behavior', icon="⚠️") st.session_state.irregular_box = True else: st.session_state.irregular_box = False reflect_json_data_to_3D_box(canvas_result_xy.json_data, view='xy') except: print('exception') pass else: st.image(image_z, use_column_width=False) with col_image2: image_y_array = image_3D[:, selected_index_y, :] preds_y_array = None if st.session_state.preds_3D is not None: preds_y_array = st.session_state.preds_3D[:, selected_index_y, :] image_y = make_fig(image_y_array, preds_y_array, st.session_state.points, selected_index_y, 'xz') if st.session_state.use_point_prompt: value_yz = streamlit_image_coordinates(image_y, width=325) if value_yz is not None: point_ax_xz = (value_yz['y'], selected_index_y, value_yz['x']) if len(st.session_state.points) >= 3: st.warning('Max point num is 3', icon="⚠️") elif point_ax_xz not in st.session_state.points: st.session_state.points.append(point_ax_xz) print('point_ax_xz add rerun') st.rerun() elif st.session_state.use_box_prompt: if st.session_state.rectangle_3Dbox[1] <= selected_index_y and selected_index_y <= st.session_state.rectangle_3Dbox[4]: draw = ImageDraw.Draw(image_y) #rectangle xz view (upper-left and lower-right) rectangle_coords = [(st.session_state.rectangle_3Dbox[2], st.session_state.rectangle_3Dbox[0]), (st.session_state.rectangle_3Dbox[5], st.session_state.rectangle_3Dbox[3])] # Draw the rectangle on the image draw.rectangle(rectangle_coords, outline='#2909F1', width=3) st.image(image_y, use_column_width=False) else: st.image(image_y, use_column_width=False) col1, col2, col3 = st.columns(3) with col1: if st.button("Clear", use_container_width=True, 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'] 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) 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, disabled=( st.session_state.data_item is None or (st.session_state.text_prompt is None and len(st.session_state.points) == 0 and st.session_state.use_box_prompt is False) or st.session_state.irregular_box or st.session_state.running )): st.session_state.running = True st.rerun() if st.session_state.running: st.session_state.running = False with st.status("Running...", expanded=False) as status: run() st.rerun()