| import streamlit as st | |
| from PIL import Image | |
| from transformers import ViTForImageClassification | |
| from config import UNTRAINED, labels | |
| 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)}, | |
| ) | |
| st.title("Detect Hurricane Damage") | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.markdown("## Pre-Trained Model") | |
| file_name = st.file_uploader("Upload a satellite image") | |
| if file_name is not None: | |
| image = Image.open(file_name) | |
| col1.image(image, use_container_width=True) | |
| label = predict | |
| st.write(f"Predicted label: {label}") | |