import streamlit as st
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageClassification
import torch

@st.cache_resource
def load_model():
    processor = AutoProcessor.from_pretrained("hamdan07/UltraSound-Lung")
    model = AutoModelForImageClassification.from_pretrained("hamdan07/UltraSound-Lung")
    return processor, model

processor, model = load_model()

st.set_page_config(page_title="Neonatal LUS AI", layout="centered")
st.title("🔬 AI-Based LUS Scoring for Neonatal Lung Ultrasound")

uploaded_files = st.file_uploader(
    "Upload one or more Lung Ultrasound Images (JPEG/PNG)", type=["jpg", "jpeg", "png"], accept_multiple_files=True
)

score_color_map = {
    "0": "🟢 (Normal A-lines)",
    "1": "🟡 (Moderate B-lines)",
    "2": "🟠 (Coalescent B-lines)",
    "3": "🔴 (Consolidation)"
}

if uploaded_files:
    for uploaded_file in uploaded_files:
        st.markdown("---")
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption=f"📷 {uploaded_file.name}", use_column_width=True)

        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class_id = logits.argmax(-1).item()
            predicted_score = model.config.id2label[predicted_class_id]

        st.markdown(f"**Predicted LUS Score:** `{predicted_score}` {score_color_map.get(predicted_score, '')}")
    st.warning("⚠️ These predictions are AI-based. Please verify clinically.")
else:
    st.info("Please upload at least one image to begin.")