Update app.py

app.py

@@ -1,33 +1,73 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
+import torch
+import numpy as np
+import cv2
+
+st.set_page_config(page_title="Détection de fractures osseuses par rayons X")
+
+st.title("Détection de fractures osseuses par rayons X")
 
 @st.cache_resource
-def load_model():
-    return pipeline("image-classification", model="mrm8488/vit-base-patch16-224_finetuned-pneumothorax", top_k=2)
+def load_models():
+    processor = AutoImageProcessor.from_pretrained("Heem2/bone-fracture-detection-using-xray")
+    model = AutoModelForImageClassification.from_pretrained("Heem2/bone-fracture-detection-using-xray")
+    return processor, model
+
+processor, model = load_models()
 
-def main():
-    st.title("Détection de Pneumothorax")
+def generate_heatmap(image, model, processor):
+    # Préparer l'image
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Obtenir les activations
+    with torch.no_grad():
+        outputs = model(**inputs)
+        # Utiliser les dernières activations
+        features = model.classifier.weight.data
+    
+    # Créer la carte de chaleur
+    cam = torch.matmul(outputs.logits, features)
+    cam = cam.reshape(7, 7)  # Ajuster selon la taille de votre modèle
+    cam = cam.detach().numpy()
     
-    model = load_model()
+    # Normaliser
+    cam = (cam - cam.min()) / (cam.max() - cam.min())
     
-    uploaded_file = st.file_uploader("Télécharger une radiographie", type=["jpg", "jpeg", "png"])
+    # Redimensionner à la taille de l'image
+    cam = cv2.resize(cam, (image.size[0], image.size[1]))
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
     
-    if uploaded_file:
-        image = Image.open(uploaded_file).convert('RGB')
-        resized_image = image.resize((224, 224))
+    return heatmap
+
+uploaded_file = st.file_uploader("Téléchargez une image radiographique", type=["jpg", "jpeg", "png"])
+
+if uploaded_file:
+    # Afficher l'image originale
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Image originale", use_column_width=True)
+    
+    # Prédiction
+    pipe = pipeline("image-classification", model=model, feature_extractor=processor)
+    results = pipe(image)
+    
+    # Afficher les résultats
+    st.subheader("Résultats de l'analyse")
+    for result in results:
+        confidence = result['score'] * 100
+        label = "Fracture détectée" if result['label'] == "FRACTURE" else "Pas de fracture"
+        st.write(f"{label} (Confiance: {confidence:.2f}%)")
         
-        col1, col2 = st.columns(2)
-        with col1:
-            st.image(resized_image, width=300, caption="Image originale")
+        # Barre de confiance colorée
+        color = "red" if label == "Fracture détectée" else "green"
+        st.progress(result['score'])
         
-        if st.button("Analyser"):
-            with st.spinner("Analyse en cours..."):
-                results = model(resized_image)
-                with col2:
-                    st.write("Résultats:")
-                    for result in results:
-                        st.write(f"{result['label']}: {result['score']:.2%}")
-
-if __name__ == "__main__":
-    main()
+        # Générer et afficher la carte de chaleur si fracture détectée
+        if label == "Fracture détectée":
+            st.subheader("Localisation probable de la fracture")
+            heatmap = generate_heatmap(image, model, processor)
+            st.image(heatmap, caption="Carte de chaleur de la fracture", use_column_width=True)
+else:
+    st.write("Veuillez télécharger une image radiographique pour l'analyse.")