import streamlit as st
from PIL import Image
from transformers import pipeline
import numpy as np
import cv2
import matplotlib.cm as cm
import base64
from io import BytesIO
st.set_page_config(layout="wide")
with open("styles.css") as f:
st.markdown(''.format(f.read()), unsafe_allow_html=True)
st.markdown("Segformer Semantic Segmentation
", unsafe_allow_html=True)
st.markdown("""
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.
""", 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."]
st.markdown("""
âšī¸ You can get sample images of road scenes in this link.
""", unsafe_allow_html=True)
st.markdown("""
đ Read more about the paper here.
""", unsafe_allow_html=True)
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)
uploaded_file = st.file_uploader("", type=["jpg", "png"])
if uploaded_file is not None:
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-b5-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')
with col2:
st.image(image_with_masks_pil, caption='Segmented Image.', use_column_width=True)
html_segment = "Labels:
"
for color, label in label_colors.items():
html_segment += f"
"
buffered = BytesIO()
image_with_masks_pil.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode()
html_segment += f'
Download Segmented Image'
st.markdown(html_segment + "
Group 6 - Members:
"
for member in group_members:
html_members += "- " + member + "
"
st.markdown(html_members + "