import streamlit as st
from transformers import ViTForImageClassification, ViTFeatureExtractor
from PIL import Image
import torch
import torch.nn.functional as F

# Load the feature extractor and model
model_name_or_path = 'google/vit-base-patch16-224-in21k'  # Replace with your actual model path
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
model = ViTForImageClassification.from_pretrained("./aryadytm-vit-vehicle-classifier")

def predict_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    probs = F.softmax(logits, dim=-1)
    predicted_label_id = probs.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_label_id]
    confidence = probs.max().item()

    return predicted_label, confidence

# Streamlit UI
st.markdown("## Vehicle Image Classification")
st.image('./assets/spotlight.png')
st.markdown("""
### Group 2
- Arya Adyatma - 2501985836
- Aldre Muhammad Keyzar - 2502006543
- Devin Eldrian Wijaya - 2501961363
- Rollando Marcellino Himmel Madison - 2502006575

This app lets you classify vehicle images using a pre-trained ViT model. You need to upload your own image.
- Kaggle dataset: https://www.kaggle.com/code/rydytm/vehicle-classification/edit.
- Colab Notebook: https://colab.research.google.com/drive/1El7RhY69KvE9Nj9vAxUPGjg42NwuNcPu?usp=sharing
""")
st.image('./assets/vit.png')
st.markdown("### Upload Your Image Here")

uploaded_file = st.file_uploader("Choose an image...", type=['png', 'jpg', 'jpeg'])

if uploaded_file is not None:
    image = (
        Image.open(uploaded_file)
        .convert("RGB")
        .resize((512, 512))
    )
    st.image(image, caption='Uploaded Image', use_column_width=True)
    st.write("")

    predicted_label, confidence = predict_image(image)

    st.write("### Prediction Result")
    st.write(f"Predicted label: **{predicted_label}**")
    st.write(f"Confidence: **{confidence:.2f}**")