Update app.py

app.py

@@ -4,10 +4,37 @@ from transformers import pipeline
 import numpy as np
 import cv2
 import matplotlib.cm as cm
+import time
+import base64
+from io import BytesIO
 
-semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
+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'>
+This app uses the Segformer deep learning model to perform semantic segmentation on road images. The Transformer-based model is 
+trained on the CityScapes dataset which contains images of urban road scenes. Upload a 
+road scene and the app will return the image with semantic segmentation applied.
+</div>
+""", unsafe_allow_html=True)
+
+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)
+
+
+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"])
 
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"])
 
 def draw_masks_fromDict(image, results):
     masked_image = image.copy()
@@ -26,16 +53,53 @@ 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
 
 if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image.', use_column_width=True)
-    st.write("")
+    st.session_state.uploaded_file = uploaded_file
+
+if st.session_state.uploaded_file is not None:
+    image = Image.open(st.session_state.uploaded_file)
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.image(image, caption='Uploaded Image.', use_column_width=True)
+
+
+    while True:
+        with st.spinner('Processing...'):  
+            segmentation_results = semantic_segmentation(image)
+            image_with_masks = draw_masks_fromDict(np.array(image), 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)
+
+        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 
+
+        while not new_file_uploaded: 
+            time.sleep(1) 
+
+
 
-    segmentation_results = semantic_segmentation(image)
-    st.json(segmentation_results)
+pdf_url = "https://arxiv.org/pdf/2105.15203.pdf"
 
-    image_with_masks = draw_masks_fromDict(np.array(image), segmentation_results)
+st.markdown("""
+<h3 class='text-center'>
+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)
 
-    image_with_masks_pil = Image.fromarray(image_with_masks, 'RGB')
-    st.image(image_with_masks_pil, caption='Segmented Image', use_column_width=True)
+st.markdown("Group Members:")
+for member in group_members:
+    st.markdown("- " + member)