import streamlit as st
from transformers import pipeline
from PIL import Image
import numpy as np
import cv2

st.set_page_config(page_title="Détection de fractures osseuses")

st.title("Détection de fractures osseuses par rayons X")

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

model = load_model()

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

if uploaded_file:
    # Load and resize image
    image = Image.open(uploaded_file)
    # Resize to max 800px width while maintaining aspect ratio
    if image.size[0] > 800:
        ratio = 800.0 / image.size[0]
        size = (800, int(image.size[1] * ratio))
        image = image.resize(size, Image.Resampling.LANCZOS)
    
    # Convert to array for overlay
    image_array = np.array(image)
    
    # Make prediction
    result = model(image)[0]  # Get only top prediction
    
    # Create columns for side by side display
    col1, col2 = st.columns(2)
    
    with col1:
        st.image(image, caption="Image originale", use_container_width=True)
    
    with col2:
        # Create colored overlay based on prediction
        overlay = np.zeros_like(image_array)
        if result['label'] == "FRACTURE":
            overlay[..., 0] = 255  # Red tint for fracture
            alpha = 0.3
        else:
            overlay[..., 1] = 255  # Green tint for normal
            alpha = 0.2
            
        # Blend images
        output = cv2.addWeighted(image_array, 1, overlay, alpha, 0)
        st.image(output, caption="Image analysée", use_container_width=True)
    
    # Display result
    st.subheader("Résultat")
    if result['label'] == "FRACTURE":
        st.error(f"⚠️ Fracture détectée (Confiance: {result['score']*100:.1f}%)")
    else:
        st.success(f"✅ Pas de fracture détectée (Confiance: {result['score']*100:.1f}%)")
        
else:
    st.info("Veuillez télécharger une image radiographique pour l'analyse.")