Update app.py

app.py

@@ -3,31 +3,33 @@ from transformers import pipeline
 from PIL import Image, ImageDraw
 import numpy as np
 import colorsys
-from streamlit.web.server.server import Server
-import streamlit.components.v1 as components
+import os
 
-# Add WebSocket headers configuration
-if Server.get_current():
-    Server.get_current()._websocket_headers = {
-        "Access-Control-Allow-Origin": "*",
-        "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
-        "Access-Control-Allow-Headers": "Content-Type"
-    }
+# Configuration des en-têtes pour WebSocket
+os.environ['STREAMLIT_SERVER_WEBSOCKET_HEADERS'] = '{"Access-Control-Allow-Origin": "*"}'
+os.environ['STREAMLIT_SERVER_ENABLE_CORS'] = 'true'
 
+# Configuration de la page
 st.set_page_config(
     page_title="Fraktur Detektion",
     layout="wide",
     initial_sidebar_state="collapsed"
 )
 
-# Add Edge WebSocket compatibility
-components.html("""
-    <script>
-        if (window.WebSocket && navigator.userAgent.indexOf("Edge") > -1) {
-            window.WebSocket = window.WebSocket || window.MozWebSocket;
-        }
-    </script>
-""", height=0)
+# Script pour gérer les WebSockets sur Edge
+st.markdown("""
+<script>
+    if (window.WebSocket && navigator.userAgent.indexOf("Edge") > -1) {
+        const originalWebSocket = window.WebSocket;
+        window.WebSocket = function(url, protocols) {
+            if (url.includes('_stcore/stream')) {
+                url = url.replace('wss://', 'ws://');
+            }
+            return new originalWebSocket(url, protocols);
+        };
+    }
+</script>
+""", unsafe_allow_html=True)
 
 st.markdown("""
 <style>
@@ -115,10 +117,28 @@ st.markdown("""
     [data-testid="stExpander"], .element-container:has(>.stAlert) {
         display: none !important;
     }
+
+    /* Fix for WebSocket connection issues */
+    iframe {
+        visibility: hidden;
+    }
 </style>
+
+<script>
+    // Fix for WebSocket connection
+    window.addEventListener('load', function() {
+        setTimeout(function() {
+            const frames = document.getElementsByTagName('iframe');
+            for (let frame of frames) {
+                frame.style.visibility = 'visible';
+            }
+        }, 1000);
+    });
+</script>
 """, unsafe_allow_html=True)
 
-@st.cache_resource
+# Configuration du cache pour les modèles
+@st.cache_resource(show_spinner=False)
 def load_models():
     return {
         "KnochenAuge": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
@@ -145,21 +165,17 @@ def create_heatmap_overlay(image, box, score):
     x1, y1 = box['xmin'], box['ymin']
     x2, y2 = box['xmax'], box['ymax']
     
-    # Couleur basée sur le score
     if score > 0.8:
-        fill_color = (255, 0, 0, 100)  # Rouge
+        fill_color = (255, 0, 0, 100)
         border_color = (255, 0, 0, 255)
     elif score > 0.6:
-        fill_color = (255, 165, 0, 100)  # Orange
+        fill_color = (255, 165, 0, 100)
         border_color = (255, 165, 0, 255)
     else:
-        fill_color = (255, 255, 0, 100)  # Jaune
+        fill_color = (255, 255, 0, 100)
         border_color = (255, 255, 0, 255)
     
-    # Rectangle semi-transparent
     draw.rectangle([x1, y1, x2, y2], fill=fill_color)
-    
-    # Bordure
     draw.rectangle([x1, y1, x2, y2], outline=border_color, width=2)
     
     return overlay
@@ -171,20 +187,16 @@ def draw_boxes(image, predictions):
         box = pred['box']
         score = pred['score']
         
-        # Création de l'overlay
         overlay = create_heatmap_overlay(image, box, score)
         result_image = Image.alpha_composite(result_image, overlay)
         
-        # Ajout du texte
         draw = ImageDraw.Draw(result_image)
         temp = 36.5 + (score * 2.5)
         label = f"{translate_label(pred['label'])} ({score:.1%} • {temp:.1f}°C)"
         
-        # Fond noir pour le texte
         text_bbox = draw.textbbox((box['xmin'], box['ymin']-20), label)
         draw.rectangle(text_bbox, fill=(0, 0, 0, 180))
         
-        # Texte en blanc
         draw.text(
             (box['xmin'], box['ymin']-20),
             label,
@@ -194,100 +206,99 @@ def draw_boxes(image, predictions):
     return result_image
 
 def main():
-    models = load_models()
-    
-    with st.container():
-        st.write("### 📤 Röntgenbild hochladen")
-        uploaded_file = st.file_uploader("Bild auswählen", type=['png', 'jpg', 'jpeg'], label_visibility="collapsed")
+    try:
+        models = load_models()
         
-        col1, col2 = st.columns([2, 1])
-        with col1:
-            conf_threshold = st.slider(
-                "Konfidenzschwelle",
-                min_value=0.0, max_value=1.0,
-                value=0.60, step=0.05,
-                label_visibility="visible"
-            )
-        with col2:
-            analyze_button = st.button("Analysieren")
-
-    if uploaded_file and analyze_button:
-        with st.spinner("Bild wird analysiert..."):
-            image = Image.open(uploaded_file)
-            results_container = st.container()
+        with st.container():
+            st.write("### 📤 Röntgenbild hochladen")
+            uploaded_file = st.file_uploader("Bild auswählen", type=['png', 'jpg', 'jpeg'], label_visibility="collapsed")
             
-            predictions_watcher = models["KnochenWächter"](image)
-            predictions_master = models["RöntgenMeister"](image)
-            predictions_locator = models["KnochenAuge"](image)
-            
-            has_fracture = False
-            max_fracture_score = 0
-            filtered_locations = [p for p in predictions_locator 
-                                if p['score'] >= conf_threshold]
-            
-            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'])
-            
-            with results_container:
-                st.write("### 🔍 Analyse Ergebnisse")
-                col1, col2 = st.columns(2)
+            col1, col2 = st.columns([2, 1])
+            with col1:
+                conf_threshold = st.slider(
+                    "Konfidenzschwelle",
+                    min_value=0.0, max_value=1.0,
+                    value=0.60, step=0.05,
+                    label_visibility="visible"
+                )
+            with col2:
+                analyze_button = st.button("Analysieren")
+
+        if uploaded_file and analyze_button:
+            with st.spinner("Bild wird analysiert..."):
+                image = Image.open(uploaded_file)
+                results_container = st.container()
                 
-                with col1:
-                    st.write("#### 🤖 KI-Diagnose")
-                    
-                    st.markdown("#### 🛡️ KnochenWächter")
-                    # Afficher tous les résultats de KnochenWächter
-                    for pred in predictions_watcher:
-                        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
-                        label_lower = pred['label'].lower()
-                        # Mettre à jour max_fracture_score seulement pour les fractures
-                        if pred['score'] >= conf_threshold and 'fracture' in label_lower:
-                            has_fracture = True
-                            max_fracture_score = max(max_fracture_score, pred['score'])
-                        # Afficher tous les résultats
-                        st.markdown(f"""
-                            <div class="result-box" style="color: #1a1a1a;">
-                                <span style="color: {confidence_color}; font-weight: 500;">
-                                    {pred['score']:.1%}
-                                </span> - {translate_label(pred['label'])}
-                            </div>
-                        """, unsafe_allow_html=True)
-                    
-                    st.markdown("#### 🎓 RöntgenMeister")
-                    # Afficher tous les résultats de RöntgenMeister
-                    for pred in predictions_master:
-                        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
-                        st.markdown(f"""
-                            <div class="result-box" style="color: #1a1a1a;">
-                                <span style="color: {confidence_color}; font-weight: 500;">
-                                    {pred['score']:.1%}
-                                </span> - {translate_label(pred['label'])}
-                            </div>
-                        """, unsafe_allow_html=True)
-                    
-                    if max_fracture_score > 0:
-                        st.write("#### 📊 Wahrscheinlichkeit")
-                        no_fracture_prob = 1 - max_fracture_score
-                        st.markdown(f"""
-                            <div class="result-box" style="color: #1a1a1a;">
-                                Knochenbruch: <strong style="color: #0066cc">{max_fracture_score:.1%}</strong><br>
-                                Kein Knochenbruch: <strong style="color: #ffa500">{no_fracture_prob:.1%}</strong>
-                            </div>
-                        """, unsafe_allow_html=True)
+                predictions_watcher = models["KnochenWächter"](image)
+                predictions_master = models["RöntgenMeister"](image)
+                predictions_locator = models["KnochenAuge"](image)
+                
+                has_fracture = False
+                max_fracture_score = 0
+                filtered_locations = [p for p in predictions_locator 
+                                    if p['score'] >= conf_threshold]
                 
-                with col2:
-                    predictions = models["KnochenAuge"](image)
-                    filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
+                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'])
+                
+                with results_container:
+                    st.write("### 🔍 Analyse Ergebnisse")
+                    col1, col2 = st.columns(2)
+                    
+                    with col1:
+                        st.write("#### 🤖 KI-Diagnose")
+                        
+                        st.markdown("#### 🛡️ KnochenWächter")
+                        for pred in predictions_watcher:
+                            confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
+                            label_lower = pred['label'].lower()
+                            if pred['score'] >= conf_threshold and 'fracture' in label_lower:
+                                has_fracture = True
+                                max_fracture_score = max(max_fracture_score, pred['score'])
+                            st.markdown(f"""
+                                <div class="result-box" style="color: #1a1a1a;">
+                                    <span style="color: {confidence_color}; font-weight: 500;">
+                                        {pred['score']:.1%}
+                                    </span> - {translate_label(pred['label'])}
+                                </div>
+                            """, unsafe_allow_html=True)
+                        
+                        st.markdown("#### 🎓 RöntgenMeister")
+                        for pred in predictions_master:
+                            confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
+                            st.markdown(f"""
+                                <div class="result-box" style="color: #1a1a1a;">
+                                    <span style="color: {confidence_color}; font-weight: 500;">
+                                        {pred['score']:.1%}
+                                    </span> - {translate_label(pred['label'])}
+                                </div>
+                            """, unsafe_allow_html=True)
+                        
+                        if max_fracture_score > 0:
+                            st.write("#### 📊 Wahrscheinlichkeit")
+                            no_fracture_prob = 1 - max_fracture_score
+                            st.markdown(f"""
+                                <div class="result-box" style="color: #1a1a1a;">
+                                    Knochenbruch: <strong style="color: #0066cc">{max_fracture_score:.1%}</strong><br>
+                                    Kein Knochenbruch: <strong style="color: #ffa500">{no_fracture_prob:.1%}</strong>
+                                </div>
+                            """, unsafe_allow_html=True)
                     
-                    if filtered_preds:
-                        st.write("#### 🎯 Fraktur Lokalisation")
-                        result_image = draw_boxes(image, filtered_preds)
-                        st.image(result_image, use_container_width=True)
-                    else:
-                        st.write("#### 🖼️ Röntgenbild")
-                        st.image(image, use_container_width=True)
+                    with col2:
+                        predictions = models["KnochenAuge"](image)
+                        filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
+                        
+                        if filtered_preds:
+                            st.write("#### 🎯 Fraktur Lokalisation")
+                            result_image = draw_boxes(image, filtered_preds)
+                            st.image(result_image, use_container_width=True)
+                        else:
+                            st.write("#### 🖼️ Röntgenbild")
+                            st.image(image, use_container_width=True)
+    except Exception as e:
+        st.error(f"Ein Fehler ist aufgetreten: {str(e)}")
 
 if __name__ == "__main__":
     main()