Update app.py

app.py

@@ -119,49 +119,81 @@ def create_heatmap_overlay(image, box, score):
     overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
     draw = ImageDraw.Draw(overlay)
     
-    def get_heatmap_color(value):
-        hue = (1 - value) * 0.3
-        saturation = 0.8
-        value = 0.9
-        rgb = colorsys.hsv_to_rgb(hue, saturation, value)
-        return tuple(int(x * 255) for x in rgb)
+    def get_temp_color(value):
+        # Rouge pour valeurs élevées, bleu pour faibles
+        if value > 0.8:
+            return (255, 0, 0)  # Rouge vif
+        elif value > 0.6:
+            return (255, 69, 0)  # Rouge-orange
+        elif value > 0.4:
+            return (255, 165, 0)  # Orange
+        else:
+            return (255, 255, 0)  # Jaune
     
     x1, y1 = box['xmin'], box['ymin']
     x2, y2 = box['xmax'], box['ymax']
+    width = x2 - x1
+    height = y2 - y1
     
-    steps = 20
+    # Créer un effet de gradient radial
+    steps = 30
     for i in range(steps):
-        alpha = int(255 * (1 - i/steps) * 0.6)
-        color = get_heatmap_color(score)
-        rect_color = color + (alpha,)
+        alpha = int(255 * (1 - (i / steps)) * 0.7)
+        base_color = get_temp_color(score)
+        color = base_color + (alpha,)
         
-        shrink = i * ((x2-x1)/(steps*2))
-        draw.rectangle([x1+shrink, y1+shrink, x2-shrink, y2-shrink], 
-                      fill=rect_color)
+        # Effet radial
+        shrink_x = (i * width) / (steps * 2)
+        shrink_y = (i * height) / (steps * 2)
+        
+        draw.rectangle(
+            [x1 + shrink_x, y1 + shrink_y, x2 - shrink_x, y2 - shrink_y],
+            fill=color,
+            outline=None
+        )
+    
+    # Ajouter une bordure plus visible
+    border_color = get_temp_color(score) + (200,)
+    draw.rectangle([x1, y1, x2, y2], outline=border_color, width=2)
     
     return overlay
 
 def draw_boxes(image, predictions):
     result_image = image.copy().convert('RGBA')
     
-    for pred in predictions:
+    # Trier les prédictions par score pour afficher les plus fortes en dernier
+    sorted_predictions = sorted(predictions, key=lambda x: x['score'])
+    
+    for pred in sorted_predictions:
         box = pred['box']
         score = pred['score']
-        label = f"{translate_label(pred['label'])} ({score:.1%})"
         
+        # Créer et appliquer la carte thermique
         heatmap = create_heatmap_overlay(image, box, score)
         result_image = Image.alpha_composite(result_image, heatmap)
         
+        # Ajouter le label avec la température
         draw = ImageDraw.Draw(result_image)
-        draw.rectangle(
-            [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
-            outline="#FFFFFF",
-            width=2
+        temp = 36.5 + (score * 2.5)  # Simulation de température: 36.5°C - 39°C
+        label = f"{translate_label(pred['label'])} ({score:.1%}) • {temp:.1f}°C"
+        
+        # Background du texte avec dégradé
+        text_bbox = draw.textbbox((box['xmin'], box['ymin']-25), label)
+        padding = 3
+        text_bbox = (
+            text_bbox[0]-padding, text_bbox[1]-padding,
+            text_bbox[2]+padding, text_bbox[3]+padding
         )
+        draw.rectangle(text_bbox, fill="#000000CC")
         
-        text_bbox = draw.textbbox((box['xmin'], box['ymin']-20), label)
-        draw.rectangle(text_bbox, fill="#000000AA")
-        draw.text((box['xmin'], box['ymin']-20), label, fill="white")
+        # Texte
+        draw.text(
+            (box['xmin'], box['ymin']-25),
+            label,
+            fill="#FFFFFF",
+            stroke_width=1,
+            stroke_fill="#000000"
+        )
     
     return result_image
 
@@ -184,22 +216,37 @@ def main():
             analyze_button = st.button("Analysieren")
 
     if uploaded_file and analyze_button:
-        with st.spinner("Bild wird analysiert..."):
-            image = Image.open(uploaded_file)
-            
-            st.write("### 🔍 Analyse Ergebnisse")
-            
-            col1, col2 = st.columns(2)
-            
-            with col1:
-                st.write("#### 🤖 KI-Diagnose")
+                    with st.spinner("Bild wird analysiert..."):
+                image = Image.open(uploaded_file)
+                results_container = st.container()
                 
-                # KnochenWächter
+                # Récupération de toutes les prédictions d'abord
                 predictions_watcher = models["KnochenWächter"](image)
+                predictions_master = models["RöntgenMeister"](image)
+                predictions_locator = models["KnochenAuge"](image)
+                
+                # Filtrage et traitement des résultats
                 has_fracture = False
                 max_fracture_score = 0
-
-                st.write("#### 🛡️ KnochenWächter")
+                filtered_locations = [p for p in predictions_locator 
+                                    if p['score'] >= conf_threshold 
+                                    and 'fracture' in p['label'].lower()]
+                
+                for pred in predictions_watcher:
+                    if pred['score'] >= conf_threshold and 'fracture' in pred['label'].lower():
+                        has_fracture = True
+                        max_fracture_score = max(max_fracture_score, pred['score'])
+                
+                # Affichage des résultats
+                with results_container:
+                    st.write("### 🔍 Analyse Ergebnisse")
+                    col1, col2 = st.columns(2)
+                    
+                    with col1:
+                        st.write("#### 🤖 KI-Diagnose")
+                        
+                        # KnochenWächter
+                        st.write("##### 🛡️ KnochenWächter")
                 for pred in predictions_watcher:
                     if pred['score'] >= conf_threshold:
                         confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'