trained model

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from PIL import Image
 from transformers import ViTForImageClassification
-from config import UNTRAINED, labels
+from config import UNTRAINED, labels, TRAINED
 from utils import predict
 
 
@@ -12,6 +12,13 @@ model_untrained = ViTForImageClassification.from_pretrained(
     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")
 
 col1, col2 = st.columns(2)
@@ -22,5 +29,15 @@ with col1:
     if file_name is not None:
         image = Image.open(file_name)
         col1.image(image, use_container_width=True)
-        label = predict
+        label = predict(model_untrained, image)
+        st.write(f"Predicted label: {label}")
+
+with col2:
+    st.markdown("## Fine-Tuned Model")
+    file_name = st.file_uploader("Upload a satellite image")
+
+    if file_name is not None:
+        image = Image.open(file_name)
+        col2.image(image, use_container_width=True)
+        label = predict(model_trained, image)
         st.write(f"Predicted label: {label}")

config.py

@@ -4,4 +4,6 @@ dataset_name = "jonathan-roberts1/Satellite-Images-of-Hurricane-Damage"
 ds = load_dataset(dataset_name)
 labels = ds['train'].features['label'].names
 
-UNTRAINED = 'google/vit-base-patch16-224-in21k'
+UNTRAINED = 'google/vit-base-patch16-224-in21k'
+
+TRAINED = '"till-onethousand/hurricane_model"'