Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 from transformers import pipeline
 from PIL import Image, ImageDraw
-import torch
+import numpy as np
 
+# Chargement des modèles
 def load_models():
     return {
         "KnochenAuge": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
@@ -11,116 +12,188 @@ def load_models():
             model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
     }
 
-def draw_boxes(image, predictions, conf_threshold=0.6):
-    draw = ImageDraw.Draw(image)
-    fractures_found = False
+def translate_label(label):
+    translations = {
+        "fracture": "Knochenbruch",
+        "no fracture": "Kein Knochenbruch",
+        "normal": "Normal",
+        "abnormal": "Auffällig",
+        "F1": "Knochenbruch",
+        "NF": "Kein Knochenbruch"
+    }
+    return translations.get(label.lower(), label)
+
+def create_heatmap_overlay(image, box, score):
+    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
+    draw = ImageDraw.Draw(overlay)
     
-    for pred in predictions:
-        if pred['label'].lower() == 'fracture' and pred['score'] >= conf_threshold:
-            fractures_found = True
-            box = pred['box']
-            label = f"Fraktur ({pred['score']:.1%})"
-            color = "#2563eb" if pred['score'] > 0.7 else "#eab308"
-            
-            draw.rectangle(
-                [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
-                outline=color,
-                width=2
-            )
-            
-            text_bbox = draw.textbbox((box['xmin'], box['ymin']-15), label)
-            draw.rectangle(text_bbox, fill=color)
-            draw.text((box['xmin'], box['ymin']-15), label, fill="white")
+    x1, y1 = box['xmin'], box['ymin']
+    x2, y2 = box['xmax'], box['ymax']
     
-    return image if fractures_found else None
+    if score > 0.8:
+        fill_color = (255, 0, 0, 100)
+        border_color = (255, 0, 0, 255)
+    elif score > 0.6:
+        fill_color = (255, 165, 0, 100)
+        border_color = (255, 165, 0, 255)
+    else:
+        fill_color = (255, 255, 0, 100)
+        border_color = (255, 255, 0, 255)
+    
+    draw.rectangle([x1, y1, x2, y2], fill=fill_color)
+    draw.rectangle([x1, y1, x2, y2], outline=border_color, width=2)
+    
+    return overlay
 
-def analyze_images(images, conf_threshold=0.6):
-    models = load_models()
-    results = []
+def draw_boxes(image, predictions):
+    result_image = image.copy().convert('RGBA')
     
-    for img in images:
-        pil_img = Image.fromarray(img)
+    for pred in predictions:
+        box = pred['box']
+        score = pred['score']
+        
+        overlay = create_heatmap_overlay(image, box, score)
+        result_image = Image.alpha_composite(result_image, overlay)
         
-        # KnochenAuge Analysis
-        predictions = models["KnochenAuge"](pil_img)
-        fractures_found = any(p['label'].lower() == 'fracture' and p['score'] >= conf_threshold 
-                            for p in predictions)
+        draw = ImageDraw.Draw(result_image)
+        temp = 36.5 + (score * 2.5)
+        label = f"{translate_label(pred['label'])} ({score:.1%} • {temp:.1f}°C)"
         
-        if fractures_found:
-            # Draw boxes on image
-            result_image = draw_boxes(pil_img.copy(), predictions, conf_threshold)
-            
-            # Additional analyses
-            wachter_pred = models["KnochenWächter"](pil_img)[0]
-            meister_pred = models["RöntgenMeister"](pil_img)[0]
-            
-            if result_image:
-                results.append({
-                    "image": result_image,
-                    "knochen_wachter": f"KnochenWächter: {wachter_pred['score']:.1%}",
-                    "rontgen_meister": f"RöntgenMeister: {meister_pred['score']:.1%}"
-                })
-    
-    # Format results for display
-    if not results:
-        return None, "Keine Frakturen gefunden."
-    
-    output_images = [r["image"] for r in results]
-    analysis_text = "\n\n".join([
-        f"Bild {i+1}:\n{r['knochen_wachter']}\n{r['rontgen_meister']}"
-        for i, r in enumerate(results)
-    ])
-    
-    return output_images, analysis_text
+        text_bbox = draw.textbbox((box['xmin'], box['ymin']-20), label)
+        draw.rectangle(text_bbox, fill=(0, 0, 0, 180))
+        
+        draw.text(
+            (box['xmin'], box['ymin']-20),
+            label,
+            fill=(255, 255, 255, 255)
+        )
+    
+    return result_image
 
-# Interface configuration
+# Modèles chargés globalement
+models = load_models()
+
+def analyze_image(image, conf_threshold=0.60):
+    if image is None:
+        return None, "Bitte laden Sie ein Bild hoch."
+    
+    # Convertir en PIL Image si nécessaire
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Analyses
+    predictions_watcher = models["KnochenWächter"](image)
+    predictions_master = models["RöntgenMeister"](image)
+    predictions_locator = models["KnochenAuge"](image)
+    
+    has_fracture = False
+    max_fracture_score = 0
+    result_html = "<div style='background: #f8f9fa; padding: 20px; border-radius: 10px;'>"
+    
+    # KnochenWächter results
+    result_html += "<h3>🛡️ KnochenWächter</h3>"
+    for pred in predictions_watcher:
+        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
+        if pred['score'] >= conf_threshold and 'fracture' in pred['label'].lower():
+            has_fracture = True
+            max_fracture_score = max(max_fracture_score, pred['score'])
+        result_html += f"""
+            <div style='background: white; padding: 10px; margin: 5px 0; border-radius: 5px; border: 1px solid #e9ecef;'>
+                <span style='color: {confidence_color}; font-weight: 500;'>
+                    {pred['score']:.1%}
+                </span> - {translate_label(pred['label'])}
+            </div>
+        """
+    
+    # RöntgenMeister results
+    result_html += "<h3>🎓 RöntgenMeister</h3>"
+    for pred in predictions_master:
+        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
+        result_html += f"""
+            <div style='background: white; padding: 10px; margin: 5px 0; border-radius: 5px; border: 1px solid #e9ecef;'>
+                <span style='color: {confidence_color}; font-weight: 500;'>
+                    {pred['score']:.1%}
+                </span> - {translate_label(pred['label'])}
+            </div>
+        """
+    
+    # Probabilité
+    if max_fracture_score > 0:
+        no_fracture_prob = 1 - max_fracture_score
+        result_html += f"""
+            <h3>📊 Wahrscheinlichkeit</h3>
+            <div style='background: white; padding: 10px; margin: 5px 0; border-radius: 5px; border: 1px solid #e9ecef;'>
+                Knochenbruch: <strong style='color: #0066cc'>{max_fracture_score:.1%}</strong><br>
+                Kein Knochenbruch: <strong style='color: #ffa500'>{no_fracture_prob:.1%}</strong>
+            </div>
+        """
+    
+    result_html += "</div>"
+    
+    # Image processing
+    predictions = models["KnochenAuge"](image)
+    filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
+    if filtered_preds:
+        result_image = draw_boxes(image, filtered_preds)
+        return result_image, result_html
+    else:
+        return image, result_html
+
+# Interface Gradio
 css = """
 .gradio-container {
-    background-color: transparent !important;
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif;
+    background-color: #f0f2f5;
+}
+.gr-button {
+    background-color: #f8f9fa !important;
+    border: 1px solid #e9ecef !important;
+    color: #1a1a1a !important;
 }
-.dark {
-    background-color: #1f2937;
-    color: #f3f4f6;
+.gr-button:hover {
+    background-color: #e9ecef !important;
+    transform: translateY(-1px);
 }
-.light {
-    background-color: #ffffff;
-    color: #1f2937;
+.output-html {
+    background: white;
+    padding: 20px;
+    border-radius: 10px;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
 }
 """
 
 with gr.Blocks(css=css) as demo:
+    gr.Markdown("### 📤 Fraktur Detektion")
+    
     with gr.Row():
         with gr.Column(scale=1):
-            file_upload = gr.File(
-                label="Röntgenbilder hochladen",
-                file_types=["image"],
-                file_count="multiple"
-            )
-            conf_slider = gr.Slider(
+            input_image = gr.Image(type="pil", label="Röntgenbild hochladen")
+            conf_threshold = gr.Slider(
                 minimum=0.0,
                 maximum=1.0,
-                value=0.6,
+                value=0.60,
                 step=0.05,
                 label="Konfidenzschwelle"
             )
-            analyze_btn = gr.Button("Bilder analysieren", variant="primary")
+            analyze_button = gr.Button("Analysieren", variant="primary")
         
-        with gr.Column(scale=2):
-            gallery = gr.Gallery(label="Ergebnisse").style(grid=2)
-            analysis_output = gr.Textbox(label="KI-Analyse", lines=4)
-    
-    analyze_btn.click(
-        fn=analyze_images,
-        inputs=[file_upload, conf_slider],
-        outputs=[gallery, analysis_output]
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="pil", label="Analysiertes Bild")
+            output_html = gr.HTML(label="Ergebnisse")
+    
+    analyze_button.click(
+        fn=analyze_image,
+        inputs=[input_image, conf_threshold],
+        outputs=[output_image, output_html]
     )
 
-# Launch configuration
+# Lancement de l'interface
 demo.launch(
-    show_api=False,
-    share=False,
     server_name="0.0.0.0",
     server_port=7860,
-    show_error=True,
-    enable_queue=True
+    share=False,
+    favicon_path=None,
+    show_api=False,
+    show_error=False,
+    debug=False
 )