File size: 1,180 Bytes
4b848f2 9506f43 7a1b8c6 9506f43 4b848f2 9506f43 4b848f2 7a1b8c6 9506f43 4b848f2 57822a3 9506f43 7a1b8c6 9506f43 |
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 |
import streamlit as st
from PIL import Image
from transformers import ViTForImageClassification
from config import UNTRAINED, labels, TRAINED
from utils import predict
model_untrained = ViTForImageClassification.from_pretrained(
UNTRAINED,
num_labels=len(labels),
id2label={str(i): c for i, c in enumerate(labels)},
label2id={c: str(i) for i, c in enumerate(labels)},
)
model_trained = ViTForImageClassification.from_pretrained(
TRAINED,
num_labels=len(labels),
id2label={str(i): c for i, c in enumerate(labels)},
label2id={c: str(i) for i, c in enumerate(labels)},
)
st.title("Detect Hurricane Damage")
file_name = st.file_uploader("Upload a satellite image")
if file_name is not None:
image = Image.open(file_name)
st.image(image, use_container_width=True)
col1, col2 = st.columns(2)
with col1:
st.markdown("## Pre-Trained Model")
if file_name is not None:
label = predict(model_untrained, image)
st.write(f"Predicted label: {label}")
with col2:
st.markdown("## Fine-Tuned Model")
if file_name is not None:
label = predict(model_trained, image)
st.write(f"Predicted label: {label}")
|