Spaces:
Build error
Build error
File size: 2,298 Bytes
6f774ac 199937f 6f774ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import streamlit as st
import torch
from PIL import Image
import torchvision.transforms as transforms
from train5 import deeplabv3_encoder_decoder
import numpy as np
# Function to load the model
def load_model(model_path):
model = deeplabv3_encoder_decoder()
try:
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()
return model
except Exception as e:
st.error(f"Error loading model: {e}")
return None
# Path to the model
model_path = 'model.pth'
# Load the trained model
model = load_model(model_path)
if model:
# Create a Streamlit app
st.title('Aerial Image Segmentation')
# Add a file uploader to the app
uploaded_file = st.file_uploader("Choose an image...", type="jpg")
if uploaded_file is not None:
image = Image.open(uploaded_file)
# Display the original image
st.image(image, caption='Uploaded Image.', use_column_width=True)
# Preprocess the image
data_transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor()]
)
image = data_transform(image)
image = image.unsqueeze(0) # add a batch dimension
# Pass the image through the model
with torch.no_grad():
output = model(image)
# Define the color map and class labels
color_map = {
0: np.array([255, 34, 133]), # Unlabeled
1: np.array([0, 252, 199]), # Early Blight
2: np.array([86, 0, 254]), # Late Blight
3: np.array([0, 0, 0]) # Leaf Minor
}
class_labels = {
0: 'Unlabeled',
1: 'Early Blight',
2: 'Late Blight',
3: 'Leaf Minor'
}
for k, v in class_labels.items():
st.sidebar.markdown(f'<div style="color:rgb{tuple(color_map[k])};">{v}</div>', unsafe_allow_html=True)
output = torch.argmax(output.squeeze(), dim=0).detach().cpu().numpy()
output_rgb = np.zeros((output.shape[0], output.shape[1], 3), dtype=np.uint8)
for k, v in color_map.items():
output_rgb[output == k] = v
st.image(output_rgb, caption='Segmented Image.', use_column_width=True)
|