Update app.py

app.py

@@ -5,58 +5,59 @@ 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")
+    return pipeline("object-detection", model="anirban22/detr-resnet-50-med_fracture")
 
 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
+    # Get predictions
+    predictions = model(image)
     
-    # Create columns for side by side display
+    # Create columns for 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
+        # Draw bounding boxes
+        img_with_boxes = image_array.copy()
+        for pred in predictions:
+            box = pred['box']
+            score = pred['score']
+            label = pred['label']
+            
+            # Draw rectangle
+            x1, y1, x2, y2 = [int(i) for i in [box['xmin'], box['ymin'], box['xmax'], box['ymax']]]
+            cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (255, 0, 0), 2)
             
-        # Blend images
-        output = cv2.addWeighted(image_array, 1, overlay, alpha, 0)
-        st.image(output, caption="Image analysée", use_container_width=True)
+            # Add label and score
+            text = f"{label}: {score:.2f}"
+            cv2.putText(img_with_boxes, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+        
+        st.image(img_with_boxes, caption="Fractures détectées", 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}%)")
+    # Display results
+    st.subheader("Résultats")
+    if predictions:
+        for pred in predictions:
+            st.warning(f"⚠️ {pred['label']} détectée (Confiance: {pred['score']*100:.1f}%)")
     else:
-        st.success(f"✅ Pas de fracture détectée (Confiance: {result['score']*100:.1f}%)")
+        st.success("✅ Aucune fracture détectée")
         
 else:
     st.info("Veuillez télécharger une image radiographique pour l'analyse.")