import streamlit as st
from ultralytics import YOLO
from PIL import Image
import torch

st.set_page_config(layout="centered")

@st.cache_resource
def load_model():
    model = YOLO('yolov8m.pt')  # Load base YOLOv8 model
    model.load('keremberke/yolov8m-chest-xray-classification.pt')  # Load weights
    return model

def main():
    st.title("Analyse Radiographie Thoracique")
    
    model = load_model()
    
    uploaded_file = st.file_uploader("Télécharger une radiographie", type=["jpg", "jpeg", "png"])
    
    if uploaded_file:
        image = Image.open(uploaded_file)
        resized_image = image.resize((640, 640))
        st.image(resized_image, width=400)
        
        if st.button("Analyser"):
            results = model.predict(source=resized_image)
            st.write(f"Résultat: {results[0].names[results[0].probs.argmax()]}")
            st.write(f"Confiance: {results[0].probs.max():.2%}")

if __name__ == "__main__":
    main()