Update app.py

app.py

@@ -4,7 +4,6 @@ from transformers import pipeline
 import numpy as np
 import cv2
 import matplotlib.cm as cm
-import time
 import base64
 from io import BytesIO
 
@@ -13,7 +12,6 @@ st.set_page_config(layout="wide")
 with open("styles.css") as f:
     st.markdown('<style>{}</style>'.format(f.read()), unsafe_allow_html=True)
 
-
 st.markdown("<h1 class='title'>Segformer Semantic Segmentation</h1>", unsafe_allow_html=True)
 st.markdown("""
 <div class='text-center'>
@@ -25,10 +23,6 @@ road scene and the app will return the image with semantic segmentation applied.
 
 group_members = ["Ang Ngo Ching, Josh Darren W.", "Bautista, Ryan Matthew M.", "Lacuesta, Angelo Giuseppe M.", "Reyes, Kenwin Hans", "Ting, Sidney Mitchell O."]
 
-
-# model_versions = ["b1", "b2", "b3", "b4", "b5"]
-# selected_model_version = st.selectbox("Select a model version:", model_versions)
-
 st.markdown("""
             <h3 class='text-center' style='margin-top: 0.5rem;'>
 ℹ️ You can get sample images of road scenes in this <a href='https://drive.google.com/drive/folders/1202EMeXAHnN18NuhJKWWme34vg0V-svY?fbclid=IwAR3kyjGS895nOBKi9aGT_P4gLX9jvSNrV5b5y3GH49t2Pvg2sZSRA58LLxs' target='_blank'>link</a>.
@@ -39,14 +33,6 @@ st.markdown("""
 📜 Read more about the paper <a href='https://arxiv.org/pdf/2105.15203.pdf' target='_blank'>here</a>.
 </h3>""", unsafe_allow_html=True)
 
-st.markdown("Group Members:")
-for member in group_members:
-    st.markdown("- " + member)
-    
-semantic_segmentation = pipeline("image-segmentation", f"nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
-
-new_file_uploaded = False
-uploaded_file = st.file_uploader("", type=["jpg", "png"])
 label_colors = {}
 
 def draw_masks_fromDict(image, results):
@@ -68,45 +54,47 @@ def draw_masks_fromDict(image, results):
     masked_image = masked_image.astype(np.uint8)
     return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)
 
-col1, col2 = st.columns(2)
-
-if "uploaded_file" not in st.session_state:
-    st.session_state.uploaded_file = None
+uploaded_file = st.file_uploader("", type=["jpg", "png"])
 
 if uploaded_file is not None:
-    st.session_state.uploaded_file = uploaded_file
-
-if st.session_state.uploaded_file is not None:
-    image = Image.open(st.session_state.uploaded_file)
+    image = Image.open(uploaded_file)
+    
     col1, col2 = st.columns(2)
 
     with col1:
         st.image(image, caption='Uploaded Image.', use_column_width=True)
 
+    semantic_segmentation = pipeline("image-segmentation", f"nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
+    segmentation_results = semantic_segmentation(image)
+    image_with_masks = draw_masks_fromDict(np.array(image)[:, :, :3], segmentation_results)
+    image_with_masks_pil = Image.fromarray(image_with_masks, 'RGB')
 
-    while True:
-        with st.spinner('Processing...'):  
-            segmentation_results = semantic_segmentation(image)
-            image_with_masks = draw_masks_fromDict(np.array(image)[:, :, :3], segmentation_results)
-            image_with_masks_pil = Image.fromarray(image_with_masks, 'RGB')
+    with col2:
+        st.image(image_with_masks_pil, caption='Segmented Image.', use_column_width=True)
 
-        with col2:
-            st.image(image_with_masks_pil, caption='Segmented Image.', use_column_width=True)
-
-            st.markdown("**Labels:**")
-            for color, label in label_colors.items():
-                st.markdown(f"<div style='display: flex; align-items: center; margin-bottom: 0.5rem;'><span style='display: inline-block; width: 20px; height: 20px; background-color: rgb{color}; margin-right: 1rem; border-radius: 10px;'></span><p style='margin: 0;'>{label}</p></div>", unsafe_allow_html=True)
+        html_segment = "<div class='container'><h3>Labels:</h3>"
+        
+        for color, label in label_colors.items():
+            html_segment += f"<div style='display: flex; align-items: center; margin-bottom: 0.5rem;'><span style='display: inline-block; width: 20px; height: 20px; background-color: rgb{color}; margin-right: 1rem; border-radius: 10px;'></span><p style='margin: 0;'>{label}</p></div>"
+        
+        buffered = BytesIO()
+        image_with_masks_pil.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        html_segment += f'<a href="data:file/png;base64,{img_str}" download="segmented_{uploaded_file.name}">Download Segmented Image</a>'
+        st.markdown(html_segment + "</div>", unsafe_allow_html=True)
 
-            buffered = BytesIO()
-            image_with_masks_pil.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            href = f'<a href="data:file/png;base64,{img_str}" download="segmented_{st.session_state.uploaded_file.name}">Download Segmented Image</a>'
-            st.markdown(href, unsafe_allow_html=True)
 
-        new_file_uploaded = False 
+pdf_url = "https://arxiv.org/pdf/2105.15203.pdf"
 
-        label_colors = {}
+st.markdown("""
+<hr><h3 style='text-align: center; margin-top: 2rem;'>
+Read more about the paper below👇
+</h5>
+""", unsafe_allow_html=True)
+st.markdown(f'<iframe class="pdf" src={pdf_url}></iframe>', unsafe_allow_html=True)
 
-        while not new_file_uploaded: 
-            time.sleep(1) 
+html_members = "<hr><div style='display: flex; justify-content: center;'><h3>Group 6 - Members:</h3><ul>"
+for member in group_members:
+    html_members += "<li>" + member + "</li>"
 
+st.markdown(html_members + "</ul></div>", unsafe_allow_html=True)