Mradul's picture
Update app.py
449db7a
raw
history blame contribute delete
463 Bytes
import streamlit as st
from model import ENet
import numpy as np
import torch
import os
num_classes = 19
model = ENet(num_classes)
model_path = "path"
model.load_state_dict(model_path)
raw_image = st.file_uploader('Raw Input Image')
if raw_image is not None:
image = np.asarray(raw_image)
output = model(image)
h, w = output.shape
display = np.zeros((h, w, 3))
display[:,:,255] = np.array([0,0,0])
st.image(display, caption="Segmented Image")