Create main.py

main.py

@@ -0,0 +1,113 @@
+import streamlit as st
+from PIL import Image
+from transformers import pipeline
+import numpy as np
+import cv2
+import matplotlib.cm as cm
+import time
+import base64
+from io import BytesIO
+
+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 <b style='color: red; font-weight: 40px;'>road images</b>. 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)
+
+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>.
+</h3>""", unsafe_allow_html=True)
+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):
+    masked_image = image.copy()
+
+    colormap = cm.get_cmap('nipy_spectral')
+    
+    for i, result in enumerate(results):
+        mask = np.array(result['mask'])  
+        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)  
+        
+        color = colormap(i / len(results))[:3] 
+        color = tuple(int(c * 255) for c in color)  
+        
+        masked_image = np.where(mask, color, masked_image)
+
+        label_colors[color] = result['label']
+
+    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:
+    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)[:, :, :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)
+
+            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)
+
+            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) 
+
+
+pdf_url = "https://arxiv.org/pdf/2105.15203.pdf"
+
+st.markdown("""
+<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)
+
+st.markdown("Group Members:")
+for member in group_members:
+    st.markdown("- " + member)