Diff of /streamlit_app.py [000000] .. [cbdc43]

Switch to side-by-side view

--- 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()