Update app.py

app.py

@@ -1,73 +1,62 @@
 import streamlit as st
-from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
+from transformers import pipeline
 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.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_models():
-    processor = AutoImageProcessor.from_pretrained("Heem2/bone-fracture-detection-using-xray")
-    model = AutoModelForImageClassification.from_pretrained("Heem2/bone-fracture-detection-using-xray")
-    return processor, model
+def load_model():
+    return pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray")
 
-processor, model = load_models()
-
-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()
-    
-    # Normaliser
-    cam = (cam - cam.min()) / (cam.max() - cam.min())
-    
-    # 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)
-    
-    return heatmap
+model = load_model()
 
 uploaded_file = st.file_uploader("Téléchargez une image radiographique", type=["jpg", "jpeg", "png"])
 
 if uploaded_file:
-    # Afficher l'image originale
+    # Load and resize image
     image = Image.open(uploaded_file)
-    st.image(image, caption="Image originale", use_column_width=True)
+    # 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)
     
-    # Prédiction
-    pipe = pipeline("image-classification", model=model, feature_extractor=processor)
-    results = pipe(image)
+    # Convert to array for overlay
+    image_array = np.array(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}%)")
-        
-        # Barre de confiance colorée
-        color = "red" if label == "Fracture détectée" else "green"
-        st.progress(result['score'])
+    # 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}%)")
         
-        # 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.")
+    st.info("Veuillez télécharger une image radiographique pour l'analyse.")