import streamlit as st
from transformers import pipeline
from PIL import Image
import io

st.set_page_config(page_title="Knochenbrucherkennung", layout="centered")

@st.cache_resource
def load_model():
    return pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray")

def main():
    st.title("🦴 Knochenbrucherkennung")
    st.write("Laden Sie ein Röntgenbild hoch.")

    pipe = load_model()
    
    uploaded_file = st.file_uploader(
        "Röntgenbild auswählen", 
        type=['png', 'jpg', 'jpeg']
    )

    conf_threshold = st.slider(
        "Konfidenzschwelle",
        min_value=0.0,
        max_value=1.0,
        value=0.3,
        step=0.01
    )

    if uploaded_file:
        image = Image.open(uploaded_file)
        
        # Redimensionner l'image
        max_size = (400, 400)
        image.thumbnail(max_size, Image.Resampling.LANCZOS)
        
        st.image(image, caption="Hochgeladenes Bild")

        with st.spinner("Analyse läuft..."):
            predictions = pipe(image)
            
            st.subheader("Ergebnisse")
            for pred in predictions:
                if pred['score'] >= conf_threshold:
                    label = "Bruch erkannt" if "fracture" in pred['label'].lower() else "Kein Bruch"
                    st.write(f"• Diagnose: {label}")
                    st.write(f"• Konfidenz: {pred['score']:.2%}")
                    
                    if "fracture" in pred['label'].lower() and pred['score'] >= conf_threshold:
                        st.warning("⚠️ Möglicher Knochenbruch erkannt!")
                    else:
                        st.success("✅ Kein Bruch erkannt")

if __name__ == "__main__":
    main()