--- a +++ b/streamlit_app.py @@ -0,0 +1,131 @@ +import streamlit as st +from PIL import Image +import numpy as np +from zipfile import ZipFile +import os +import cv2 +from src.production import read_files, get_setup, make_masks, create_folder, make_legend + + +@st.cache +def cached_get_setup(): + return get_setup() + + +def main(): + models, transforms = cached_get_setup() + st.markdown( + f""" + <style> + .sidebar .sidebar-content {{ + background: url("https://i.ibb.co/BL3qFQW/background.png"); + background-repeat: repeat; + background-size: 100% auto; + }} + .reportview-container {{ + background: url("https://i.ibb.co/BL3qFQW/background.png"); + background-repeat: repeat; + background-size: 100% auto; + }} + .reportview-container .main .block-container{{ + max-width: 850px; + padding-top: 0rem; + padding-right: 0rem; + padding-left: 0rem; + padding-bottom: 0rem; + }} + </style> + """, + unsafe_allow_html=True, + ) + for folder in ['segmentations/', 'images/']: + create_folder(folder) + + st.title('Сегментация поражения легких коронавирусной пневмонией') + + st.subheader("Загрузка файлов") + filenames = st.file_uploader('Выберите или ператащите сюда снимки', type=['png', 'jpeg', 'jpg', '.nii', '.nii.gz'], + accept_multiple_files=True) + + multi_class = st.checkbox(label='Мульти-классовая сегментация', value=False) + show_legend = st.checkbox(label='Легенда на картинке', value=False) + + if st.button('Загрузить') and filenames: + paths, folder_name = read_files(filenames) + if not paths: + st.error('Неправильный формат или название файла') + else: + user_dir = "segmentations/" + folder_name + + # creating folders + create_folder(user_dir) + create_folder(os.path.join(user_dir, 'segmentations')) + create_folder(os.path.join(user_dir, 'annotations')) + + zip_obj = ZipFile(user_dir + 'segmentations.zip', 'w') + with st.expander("Информация о каждом фото"): + info = st.info('Делаем предсказания, пожалуйста, подождите') + for _paths in paths: + for img, annotation, original_path in make_masks(_paths, models, transforms, multi_class): + name = original_path.split('/')[-1].split('.')[0] + name = name.replace('\\', '/') + + # saving annotation + annotation_path = os.path.join(user_dir, 'annotations', name + '_annotation.txt') + with open(annotation_path, mode='w') as f: + f.write(annotation) + + info.empty() + + # name and annotation + st.markdown(f'<h3>{name}</h3>', unsafe_allow_html=True) + if not show_legend: + if len(annotation.split('\n')) == 3: + st.markdown('[red] - consolidation') + st.markdown('[green] - ground-glass') + else: + st.markdown('[yellow] - disease') + for line in annotation.split('\n'): + st.markdown(line) + + col1, col2 = st.columns(2) + + # original image + original = np.array(Image.open(original_path)) + col1.header("Оригинал") + col1.image(original, width=350) + + # refactoring image + if show_legend: + img = make_legend(img, annotation) + + # saving image + path = os.path.join(user_dir, 'segmentations', name + '_mask.png') + cv2.imwrite(path, img) + + # adding in zip + zip_obj.write(path) + zip_obj.write(annotation_path) + + # show segmentation + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = img / 255 # to [0;1] range + # print(img.shape, img.dtype, img) + col2.header("Сегментация") + col2.image(img, width=350) + + st.markdown('<br />', unsafe_allow_html=True) + + zip_obj.close() + + # download segmentation zip + with st.expander("Скачать сегментации"): + with open(os.path.join(user_dir, 'segmentations.zip'), 'rb') as file: + st.download_button( + label="Архив сегментаций", + data=file, + file_name="segmentations.zip") + + +if __name__ == '__main__': + main()