Update app.py

app.py

@@ -1,7 +1,9 @@
 import streamlit as st
 from transformers import pipeline
 from PIL import Image, ImageDraw
-import torch
+import numpy as np
+from PIL import ImageColor
+import colorsys
 
 st.set_page_config(
     page_title="Fraktur Detektion",
@@ -11,34 +13,37 @@ st.set_page_config(
 
 st.markdown("""
 <style>
-    /* Base styles */
     .stApp {
         background: #f0f2f5 !important;
     }
     
     .block-container {
-        padding: 1rem !important;
+        padding-top: 0 !important;
+        padding-bottom: 0 !important;
         max-width: 1400px !important;
         margin: 0 auto !important;
     }
     
-    /* Custom containers */
-    .center-upload {
+    .main-container {
+        display: flex;
+        gap: 1rem;
+        padding: 1rem;
         background: white;
-        padding: 2rem;
         border-radius: 10px;
         box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        margin-bottom: 2rem;
-        text-align: center;
+        margin: 1rem;
     }
     
-    .analysis-container {
-        background: white;
-        padding: 1.5rem;
-        border-radius: 10px;
-        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        margin-bottom: 1rem;
-        animation: slideIn 0.5s ease-out;
+    .upload-section {
+        flex: 1;
+        padding: 1rem;
+        border-radius: 8px;
+        background: #f8f9fa;
+    }
+    
+    .result-section {
+        flex: 2;
+        padding: 1rem;
     }
     
     .result-box {
@@ -49,13 +54,11 @@ st.markdown("""
         border: 1px solid #e9ecef;
     }
     
-    /* Text styles */
     h1, h2, h3, h4, p {
         color: #1a1a1a !important;
         margin: 0.5rem 0 !important;
     }
     
-    /* Image styles */
     .stImage {
         background: white;
         padding: 0.5rem;
@@ -64,38 +67,18 @@ st.markdown("""
     }
     
     .stImage > img {
-        max-height: 250px !important;
+        max-height: 300px !important;
         width: auto !important;
         margin: 0 auto !important;
         display: block !important;
     }
     
-    /* Animations */
-    @keyframes slideIn {
-        from {
-            opacity: 0;
-            transform: translateY(-10px);
-        }
-        to {
-            opacity: 1;
-            transform: translateY(0);
-        }
-    }
-    
-    /* Hide unnecessary elements */
-    #MainMenu, footer {
-        display: none !important;
-    }
-    
-    /* Custom columns spacing */
-    [data-testid="column"] {
-        padding: 0.5rem !important;
-        background: transparent !important;
+    [data-testid="stFileUploader"] {
+        width: 100% !important;
     }
     
-    /* Button styling */
     .stButton > button {
-        width: 200px;
+        width: 100%;
         background-color: #0066cc !important;
         color: white !important;
         border: none !important;
@@ -108,6 +91,19 @@ st.markdown("""
         background-color: #0052a3 !important;
         transform: translateY(-1px);
     }
+    
+    #MainMenu, footer, header {
+        display: none !important;
+    }
+    
+    /* Hide deprecation warning */
+    [data-testid="stExpander"] {
+        display: none !important;
+    }
+    
+    .element-container:has(>.stAlert) {
+        display: none !important;
+    }
 </style>
 """, unsafe_allow_html=True)
 
@@ -125,109 +121,148 @@ def translate_label(label):
         "fracture": "Knochenbruch",
         "no fracture": "Kein Bruch",
         "normal": "Normal",
-        "abnormal": "Auffällig"
+        "abnormal": "Auffällig",
+        "F1": "Knochenbruch",
+        "NF": "Kein Bruch"
     }
     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)
+    
+    # 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
+        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
+        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)
+    
+    return overlay
+
 def draw_boxes(image, predictions):
-    draw = ImageDraw.Draw(image)
+    # Create a copy of the image to work with
+    result_image = image.copy().convert('RGBA')
+    
     for pred in predictions:
         box = pred['box']
-        label = f"{translate_label(pred['label'])} ({pred['score']:.2%})"
-        color = "#0066cc" if pred['score'] > 0.7 else "#ffa500"
+        score = pred['score']
+        label = f"{translate_label(pred['label'])} ({score:.2%})"
         
+        # 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'])],
-            outline=color,
+            outline="#FFFFFF",
             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")
-    return image
+        # 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")
+    
+    return result_image
 
 def main():
     models = load_models()
     
-    # Initial upload section
-    st.markdown('<div class="center-upload">', unsafe_allow_html=True)
+    # 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_files = st.file_uploader("", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)
+    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,
-        key='confidence'
+        value=0.60, step=0.05
     )
     
-    analyze_button = st.button("Analysieren", key='analyze')
+    analyze_button = st.button("Analysieren")
     st.markdown('</div>', unsafe_allow_html=True)
-
-    # Analysis section
-    if analyze_button and uploaded_files:
-        for idx, uploaded_file in enumerate(uploaded_files):
-            st.markdown(f'<div class="analysis-container">', 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.spinner("Analysiere Bild..."):
+            image = Image.open(uploaded_file)
             
-            with st.spinner("Analysiere Bild..."):
-                image = Image.open(uploaded_file)
-                
-                # Create three columns
-                col1, col2, col3 = st.columns(3)
-                
-                # Column 1: Original Image
-                with col1:
-                    st.markdown("### 🖼️ Original")
-                    st.image(image, use_column_width=True)
+            col1, col2 = st.columns(2)
+            
+            with col1:
+                st.markdown("### 🎯 KI-Analyse")
                 
-                # Column 2: AI Analysis
-                with col2:
-                    st.markdown("### 🤖 KI-Analyse")
-                    
-                    # KnochenWächter results
-                    predictions = models["KnochenWächter"](image)
-                    st.markdown("#### 🛡️ KnochenWächter")
-                    for pred in predictions:
-                        if pred['score'] >= conf_threshold:
-                            st.markdown(f"""
-                                <div class="result-box">
-                                    <span style='color: {"#0066cc" if pred["score"] > 0.7 else "#ffa500"}; font-weight: 500;'>
-                                        {pred['score']:.1%}
-                                    </span> - {translate_label(pred['label'])}
-                                </div>
-                            """, unsafe_allow_html=True)
-                    
-                    # RöntgenMeister results
-                    predictions = models["RöntgenMeister"](image)
-                    st.markdown("#### 🎓 RöntgenMeister")
-                    for pred in predictions:
-                        if pred['score'] >= conf_threshold:
-                            st.markdown(f"""
-                                <div class="result-box">
-                                    <span style='color: {"#0066cc" if pred["score"] > 0.7 else "#ffa500"}; font-weight: 500;'>
-                                        {pred['score']:.1%}
-                                    </span> - {translate_label(pred['label'])}
-                                </div>
-                            """, unsafe_allow_html=True)
+                # KnochenWächter results
+                st.markdown("#### 🛡️ KnochenWächter")
+                predictions = models["KnochenWächter"](image)
+                for pred in predictions:
+                    if pred['score'] >= conf_threshold:
+                        st.markdown(f"""
+                            <div class="result-box">
+                                <span style="color: {'#0066cc' if pred['score'] > 0.7 else '#ffa500'}; font-weight: 500;">
+                                    {pred['score']:.1%}
+                                </span> - {translate_label(pred['label'])}
+                            </div>
+                        """, unsafe_allow_html=True)
                 
-                # Column 3: Localization (only if fracture detected)
-                with col3:
-                    predictions = models["KnochenAuge"](image)
-                    has_fracture = any(
-                        p['score'] >= conf_threshold and 'fracture' in p['label'].lower() 
-                        for p in predictions
-                    )
-                    
-                    if has_fracture:
-                        st.markdown("### 🔍 Fraktur Lokalisation")
-                        filtered_preds = [p for p in predictions if p['score'] >= conf_threshold]
-                        if filtered_preds:
-                            result_image = image.copy()
-                            result_image = draw_boxes(result_image, filtered_preds)
-                            st.image(result_image, use_column_width=True)
+                # RöntgenMeister results
+                st.markdown("#### 🎓 RöntgenMeister")
+                predictions = models["RöntgenMeister"](image)
+                for pred in predictions:
+                    if pred['score'] >= conf_threshold:
+                        st.markdown(f"""
+                            <div class="result-box">
+                                <span style="color: {'#0066cc' if pred['score'] > 0.7 else '#ffa500'}; font-weight: 500;">
+                                    {pred['score']:.1%}
+                                </span> - {translate_label(pred['label'])}
+                            </div>
+                        """, unsafe_allow_html=True)
             
-            st.markdown('</div>', unsafe_allow_html=True)
+            with col2:
+                st.markdown("### 🔍 Visualisierung")
+                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)
+                    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()