import streamlit as st
from PIL import Image
import torch
import json
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

extractor = AutoFeatureExtractor.from_pretrained("Amite5h/convnext-tiny-finetuned-eurosat")

model = AutoModelForImageClassification.from_pretrained("Amite5h/convnext-tiny-finetuned-eurosat")



st.title("EuroSAT Detection")

file_name = st.file_uploader("Upload a geospatial image")

if file_name is not None:
    col1, col2 = st.columns(2)

    image = Image.open(file_name)
    if image.mode != "RGB":
        image = image.convert("RGB")
    col1.image(image, use_column_width=True)
    # Convert grayscale image to RGB format
    image_tensor = extractor(images=image, return_tensors="pt")["pixel_values"]
    predictions = model(image_tensor)

    predicted_probabilities = torch.softmax(predictions.logits, dim=1)[0]
    predicted_labels = model.config.id2label
    
    # Create a dictionary of labels and probabilities
    label_probabilities = {
        predicted_labels[i]: predicted_probabilities[i].item() for i in range(len(predicted_labels))
    }
    
    # Convert the output to JSON string
    json_output = json.dumps(label_probabilities)
        
    #predicted_class = torch.argmax(predictions.logits, dim=1)
    col2.header("Probabilities")
    col2.subheader(json_output)

    # col2.header("Probabilities")
    # for p in predictions:
    #     col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")