Update app.py

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 from transformers import pipeline
 from PIL import Image, ImageDraw
 import numpy as np
-from PIL import ImageColor
 import colorsys
 
 st.set_page_config(
@@ -21,29 +20,22 @@ st.markdown("""
         padding-top: 0 !important;
         padding-bottom: 0 !important;
         max-width: 1400px !important;
-        margin: 0 auto !important;
     }
     
-    .main-container {
-        display: flex;
-        gap: 1rem;
-        padding: 1rem;
+    .upload-container {
         background: white;
+        padding: 1.5rem;
         border-radius: 10px;
         box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        margin: 1rem;
-    }
-    
-    .upload-section {
-        flex: 1;
-        padding: 1rem;
-        border-radius: 8px;
-        background: #f8f9fa;
+        margin-bottom: 1rem;
+        text-align: center;
     }
     
-    .result-section {
-        flex: 2;
-        padding: 1rem;
+    .results-container {
+        background: white;
+        padding: 1.5rem;
+        border-radius: 10px;
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
     }
     
     .result-box {
@@ -78,7 +70,7 @@ st.markdown("""
     }
     
     .stButton > button {
-        width: 100%;
+        width: 200px;
         background-color: #0066cc !important;
         color: white !important;
         border: none !important;
@@ -92,16 +84,12 @@ st.markdown("""
         transform: translateY(-1px);
     }
     
-    #MainMenu, footer, header {
+    #MainMenu, footer, header, [data-testid="stToolbar"] {
         display: none !important;
     }
     
     /* Hide deprecation warning */
-    [data-testid="stExpander"] {
-        display: none !important;
-    }
-    
-    .element-container:has(>.stAlert) {
+    [data-testid="stExpander"], .element-container:has(>.stAlert) {
         display: none !important;
     }
 </style>
@@ -119,11 +107,11 @@ def load_models():
 def translate_label(label):
     translations = {
         "fracture": "Knochenbruch",
-        "no fracture": "Kein Bruch",
+        "no fracture": "Kein Knochenbruch",
         "normal": "Normal",
         "abnormal": "Auffällig",
         "F1": "Knochenbruch",
-        "NF": "Kein Bruch"
+        "NF": "Kein Knochenbruch"
     }
     return translations.get(label.lower(), label)
 
@@ -131,27 +119,22 @@ def create_heatmap_overlay(image, box, score):
     overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
     draw = ImageDraw.Draw(overlay)
     
-    # Create gradient colors based on confidence
     def get_heatmap_color(value):
-        # Convert to HSV for better control
-        hue = (1 - value) * 0.3  # 0.3 = reddish, 0 = red
+        hue = (1 - value) * 0.3
         saturation = 0.8
         value = 0.9
-        # Convert back to RGB
         rgb = colorsys.hsv_to_rgb(hue, saturation, value)
         return tuple(int(x * 255) for x in rgb)
     
-    # Draw the heatmap with gradient
     x1, y1 = box['xmin'], box['ymin']
     x2, y2 = box['xmax'], box['ymax']
     
     steps = 20
     for i in range(steps):
-        alpha = int(255 * (1 - i/steps) * 0.6)  # Gradient transparency
+        alpha = int(255 * (1 - i/steps) * 0.6)
         color = get_heatmap_color(score)
         rect_color = color + (alpha,)
         
-        # Create shrinking rectangles for gradient effect
         shrink = i * ((x2-x1)/(steps*2))
         draw.rectangle([x1+shrink, y1+shrink, x2-shrink, y2-shrink], 
                       fill=rect_color)
@@ -159,19 +142,16 @@ def create_heatmap_overlay(image, box, score):
     return overlay
 
 def draw_boxes(image, predictions):
-    # Create a copy of the image to work with
     result_image = image.copy().convert('RGBA')
     
     for pred in predictions:
         box = pred['box']
         score = pred['score']
-        label = f"{translate_label(pred['label'])} ({score:.2%})"
+        label = f"{translate_label(pred['label'])} ({score:.1%})"
         
-        # Create and combine heatmap overlay
         heatmap = create_heatmap_overlay(image, box, score)
         result_image = Image.alpha_composite(result_image, heatmap)
         
-        # Draw border and label
         draw = ImageDraw.Draw(result_image)
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
@@ -179,7 +159,6 @@ def draw_boxes(image, predictions):
             width=2
         )
         
-        # Add label with background
         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")
@@ -189,80 +168,83 @@ def draw_boxes(image, predictions):
 def main():
     models = load_models()
     
-    # Initialize session state
-    if 'analyzed' not in st.session_state:
-        st.session_state.analyzed = False
-    
-    # Main container
-    st.markdown('<div class="main-container">', unsafe_allow_html=True)
-    
-    # Upload section
-    st.markdown('<div class="upload-section">', unsafe_allow_html=True)
-    st.markdown("### 📤 Röntgenbild Upload")
-    uploaded_file = st.file_uploader("", type=['png', 'jpg', 'jpeg'])
-    
-    conf_threshold = st.slider(
-        "Konfidenzschwelle",
-        min_value=0.0, max_value=1.0,
-        value=0.60, step=0.05
-    )
-    
-    analyze_button = st.button("Analysieren")
-    st.markdown('</div>', unsafe_allow_html=True)
-    
-    # Results section
-    st.markdown('<div class="result-section">', unsafe_allow_html=True)
-    
-    if uploaded_file and analyze_button:
-        st.session_state.analyzed = True
+    with st.container():
+        st.write("### 📤 Röntgenbild hochladen")
+        uploaded_file = st.file_uploader("", type=['png', 'jpg', 'jpeg'])
         
-        with st.spinner("Analysiere Bild..."):
+        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
+            )
+        with col2:
+            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.markdown("### 🎯 KI-Analyse")
+                st.write("#### 🤖 KI-Diagnose")
                 
-                # KnochenWächter results
-                st.markdown("#### 🛡️ KnochenWächter")
-                predictions = models["KnochenWächter"](image)
-                for pred in predictions:
+                # KnochenWächter
+                predictions_watcher = models["KnochenWächter"](image)
+                has_fracture = False
+                for pred in predictions_watcher:
                     if pred['score'] >= conf_threshold:
+                        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
+                        if 'fracture' in pred['label'].lower():
+                            has_fracture = True
                         st.markdown(f"""
                             <div class="result-box">
-                                <span style="color: {'#0066cc' if pred['score'] > 0.7 else '#ffa500'}; font-weight: 500;">
+                                <span style="color: {confidence_color}; font-weight: 500;">
                                     {pred['score']:.1%}
                                 </span> - {translate_label(pred['label'])}
                             </div>
                         """, unsafe_allow_html=True)
                 
-                # RöntgenMeister results
-                st.markdown("#### 🎓 RöntgenMeister")
-                predictions = models["RöntgenMeister"](image)
-                for pred in predictions:
+                # RöntgenMeister
+                predictions_master = models["RöntgenMeister"](image)
+                for pred in predictions_master:
                     if pred['score'] >= conf_threshold:
+                        confidence_color = '#0066cc' if pred['score'] > 0.7 else '#ffa500'
                         st.markdown(f"""
                             <div class="result-box">
-                                <span style="color: {'#0066cc' if pred['score'] > 0.7 else '#ffa500'}; font-weight: 500;">
+                                <span style="color: {confidence_color}; font-weight: 500;">
                                     {pred['score']:.1%}
                                 </span> - {translate_label(pred['label'])}
                             </div>
                         """, unsafe_allow_html=True)
+                
+                # Calculate and display fracture probability
+                fracture_prob = max((p['score'] for p in predictions_watcher 
+                                   if 'fracture' in p['label'].lower()), default=0)
+                no_fracture_prob = 1 - fracture_prob
+                
+                st.write("#### 📊 Wahrscheinlichkeit")
+                st.markdown(f"""
+                    <div class="result-box">
+                        Knochenbruch: <strong>{fracture_prob:.1%}</strong><br>
+                        Kein Knochenbruch: <strong>{no_fracture_prob:.1%}</strong>
+                    </div>
+                """, unsafe_allow_html=True)
             
             with col2:
-                st.markdown("### 🔍 Visualisierung")
+                st.write("#### 🎯 Visualisierung")
                 predictions = models["KnochenAuge"](image)
                 filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
                 
-                if filtered_preds:
+                if filtered_preds and has_fracture:
                     result_image = draw_boxes(image, filtered_preds)
                     st.image(result_image, use_container_width=True)
                 else:
                     st.image(image, use_container_width=True)
-    
-    st.markdown('</div>', unsafe_allow_html=True)
-    st.markdown('</div>', unsafe_allow_html=True)
 
 if __name__ == "__main__":
     main()