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")