Update app.py

app.py

@@ -1,27 +1,59 @@
 import streamlit as st
 from transformers import pipeline
-from PIL import Image
-import io
+from PIL import Image, ImageDraw
+import torch
 
-st.set_page_config(page_title="Knochenbrucherkennung", layout="centered")
+st.set_page_config(page_title="Multi-Model Fracture Detection", layout="wide")
 
 @st.cache_resource
-def load_model():
-    return pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray")
+def load_models():
+    models = {
+        "D3STRON (Object Detection)": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
+        "Heem2 (Classification)": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
+        "Akhileshav8 (Classification)": pipeline("image-classification", model="akhileshav8/image_classification_for_fracture"),
+        "Nandodeomkar (Classification)": pipeline("image-classification", model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388"),
+        "Anirban22 (Object Detection)": pipeline("object-detection", model="anirban22/detr-resnet-50-med_fracture")
+    }
+    return models
 
-def main():
-    st.title("🦴 Knochenbrucherkennung")
-    st.write("Laden Sie ein Röntgenbild hoch.")
+def draw_boxes(image, predictions):
+    draw = ImageDraw.Draw(image)
+    for pred in predictions:
+        box = pred['box']
+        label = f"{pred['label']} ({pred['score']:.2%})"
+        
+        draw.rectangle(
+            [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
+            outline="red",
+            width=3
+        )
+        
+        text_bbox = draw.textbbox((box['xmin'], box['ymin']), label)
+        draw.rectangle(text_bbox, fill="red")
+        draw.text((box['xmin'], box['ymin']), label, fill="white")
+    return image
 
-    pipe = load_model()
-    
-    uploaded_file = st.file_uploader(
-        "Röntgenbild auswählen", 
-        type=['png', 'jpg', 'jpeg']
-    )
+def process_classification(model, image, conf_threshold):
+    predictions = model(image)
+    results = []
+    for pred in predictions:
+        if pred['score'] >= conf_threshold:
+            results.append(f"{pred['label']}: {pred['score']:.2%}")
+    return results
+
+def process_detection(model, image, conf_threshold):
+    predictions = model(image)
+    return [pred for pred in predictions if pred['score'] >= conf_threshold]
 
+def main():
+    st.title("🦴 Multi-Model Fracture Detection")
+    
+    models = load_models()
+    
+    uploaded_file = st.file_uploader("Upload X-ray image", type=['png', 'jpg', 'jpeg'])
+    
     conf_threshold = st.slider(
-        "Konfidenzschwelle",
+        "Confidence threshold",
         min_value=0.0,
         max_value=1.0,
         value=0.3,
@@ -30,27 +62,38 @@ def main():
 
     if uploaded_file:
         image = Image.open(uploaded_file)
-        
-        # Redimensionner l'image
         max_size = (400, 400)
         image.thumbnail(max_size, Image.Resampling.LANCZOS)
         
-        st.image(image, caption="Hochgeladenes Bild")
-
-        with st.spinner("Analyse läuft..."):
-            predictions = pipe(image)
-            
-            st.subheader("Ergebnisse")
-            for pred in predictions:
-                if pred['score'] >= conf_threshold:
-                    label = "Bruch erkannt" if "fracture" in pred['label'].lower() else "Kein Bruch"
-                    st.write(f"• Diagnose: {label}")
-                    st.write(f"• Konfidenz: {pred['score']:.2%}")
-                    
-                    if "fracture" in pred['label'].lower() and pred['score'] >= conf_threshold:
-                        st.warning("⚠️ Möglicher Knochenbruch erkannt!")
-                    else:
-                        st.success("✅ Kein Bruch erkannt")
+        st.image(image, caption="Original Image", width=400)
+        
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.subheader("Classification Models")
+            for name, model in models.items():
+                if "Classification" in name:
+                    st.write(f"**{name}**")
+                    with st.spinner(f"Running {name}..."):
+                        results = process_classification(model, image, conf_threshold)
+                        for result in results:
+                            st.write(f"• {result}")
+        
+        with col2:
+            st.subheader("Object Detection Models")
+            for name, model in models.items():
+                if "Object Detection" in name:
+                    st.write(f"**{name}**")
+                    with st.spinner(f"Running {name}..."):
+                        detections = process_detection(model, image, conf_threshold)
+                        if detections:
+                            result_image = image.copy()
+                            result_image = draw_boxes(result_image, detections)
+                            st.image(result_image, caption=f"Results from {name}")
+                            for det in detections:
+                                st.write(f"• {det['label']}: {det['score']:.2%}")
+                        else:
+                            st.write("No detections above threshold")
 
 if __name__ == "__main__":
     main()