app.py

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(
        "object-detection",
        model="anirban22/detr-resnet-50-med_fracture",
        threshold=0.1  # Réduire le seuil de confiance
    )

model = load_model()

def enhance_image(image):
    # Convertir en niveaux de gris
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    # Améliorer le contraste
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    enhanced = clahe.apply(gray)
    # Reconvertir en RGB
    return cv2.cvtColor(enhanced, cv2.COLOR_GRAY2RGB)

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

if uploaded_file:
    image = Image.open(uploaded_file)
    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)
    
    # Améliorer l'image
    enhanced_image = enhance_image(image)
    
    # Obtenir les prédictions
    predictions = model(enhanced_image)
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.image(image, caption="Image originale", use_container_width=True)
    
    with col2:
        img_with_boxes = enhanced_image.copy()
        for pred in predictions:
            box = pred['box']
            score = pred['score']
            
            x1, y1, x2, y2 = [int(i) for i in [box['xmin'], box['ymin'], box['xmax'], box['ymax']]]
            # Rectangle rouge plus épais
            cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (255, 0, 0), 3)
            
            # Texte plus visible
            text = f"Fracture: {score:.2f}"
            cv2.putText(img_with_boxes, text, (x1, y1-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
        
        st.image(img_with_boxes, caption="Détection des fractures", use_container_width=True)
    
    st.subheader("Résultats")
    if predictions:
        for idx, pred in enumerate(predictions, 1):
            st.warning(f"⚠️ Fracture {idx} détectée (Confiance: {pred['score']*100:.1f}%)")
    else:
        st.warning("⚠️ Aucune fracture n'a été détectée avec certitude. Veuillez consulter un professionnel pour confirmation.")
        
else:
    st.info("Veuillez télécharger une image radiographique pour l'analyse.")