base model

app.py

@@ -1,20 +1,26 @@
 import streamlit as st
-from transformers import pipeline
 from PIL import Image
+from transformers import ViTForImageClassification
+from config import UNTRAINED, labels
+from utils import predict
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 
-st.title("Hot Dog? Or Not?")
+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)},
+)
 
-file_name = st.file_uploader("Upload a hot dog candidate image")
+st.title("Detect Hurricane Damage")
 
-if file_name is not None:
-    col1, col2 = st.columns(2)
+col1, col2 = st.columns(2)
+with col1:
+    st.markdown("## Pre-Trained Model")
+    file_name = st.file_uploader("Upload a satellite image")
 
-    image = Image.open(file_name)
-    col1.image(image, use_column_width=True)
-    predictions = pipeline(image)
-
-    col2.header("Probabilities")
-    for p in predictions:
-        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+    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}")

config.py

@@ -0,0 +1,7 @@
+import datasets import load_dataset
+
+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'

utils.py

@@ -0,0 +1,15 @@
+import torch
+from transformers import ViTFeatureExtractor
+from config import UNTRAINED
+
+feature_extractor = ViTFeatureExtractor.from_pretrained(UNTRAINED)
+
+def predict(model, image):
+    inputs = feature_extractor(image, return_tensors="pt")
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    # model predicts one of the 1000 ImageNet classes
+    predicted_label = logits.argmax(-1).item()
+    return model.config.id2label[str(predicted_label)]