import streamlit as st
from transformers import AutoModelForObjectDetection
import torch
from PIL import Image
import numpy as np
import cv2

st.set_page_config(page_title="Détection de nodules pulmonaires")
st.title("Détection de nodules pulmonaires sur images scanner")

@st.cache_resource
def load_model():
    model = AutoModelForObjectDetection.from_pretrained("monai-test/lung_nodule_ct_detection")
    model.eval()
    return model

def process_image(image):
    # Convertir en niveau de gris
    img_array = np.array(image.convert('L'))
    # Normaliser
    normalized = (img_array - img_array.min()) / (img_array.max() - img_array.min())
    # Redimensionner
    resized = cv2.resize(normalized, (512, 512))
    # Préparer pour PyTorch
    tensor = torch.FloatTensor(resized).unsqueeze(0).unsqueeze(0)
    return tensor

try:
    model = load_model()

    uploaded_file = st.file_uploader("Téléchargez une image scanner", type=["jpg", "jpeg", "png"])

    if uploaded_file:
        image = Image.open(uploaded_file)
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.image(image, caption="Image originale", use_container_width=True)
        
        with col2:
            with torch.no_grad():
                input_tensor = process_image(image)
                predictions = model(input_tensor)
                
            # Visualisation
            img_array = np.array(image)
            for pred in predictions:
                if pred['score'] > 0.5:
                    box = pred['box']
                    x1, y1, x2, y2 = map(int, [box['xmin'], box['ymin'], box['xmax'], box['ymax']])
                    cv2.rectangle(img_array, (x1, y1), (x2, y2), (255, 0, 0), 2)
                    text = f"Nodule: {pred['score']:.2f}"
                    cv2.putText(img_array, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            
            st.image(img_array, caption="Détections", use_container_width=True)
        
        # Résultats
        if len(predictions) > 0:
            st.warning(f"⚠️ {len(predictions)} nodules détectés")
            for i, pred in enumerate(predictions, 1):
                if pred['score'] > 0.5:
                    st.write(f"Nodule {i}: Confiance {pred['score']:.1%}")
        else:
            st.success("✅ Aucun nodule détecté")

except Exception as e:
    st.error(f"Erreur lors du chargement du modèle: {str(e)}")
    st.info("Veuillez vérifier que le modèle est correctement configuré sur Hugging Face.")