Update app.py

app.py

@@ -1,28 +1,25 @@
 import streamlit as st
 from transformers import pipeline
-import torch
 from PIL import Image, ImageDraw
-import io
-import base64
-from fastapi import FastAPI, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-import numpy as np
-import json
-from starlette.responses import JSONResponse
-
-# FastAPI app
-app = FastAPI()
+import torch
 
-# Enable CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+st.set_page_config(
+    page_title="Knochenbrucherkennung",
+    layout="wide",
+    initial_sidebar_state="collapsed"
 )
 
-# Load models
+st.markdown("""
+<style>
+    .main > div {
+        padding: 2rem;
+        background: #f8f9fa;
+        border-radius: 1rem;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    }
+</style>
+""", unsafe_allow_html=True)
+
 @st.cache_resource
 def load_models():
     return {
@@ -32,78 +29,104 @@ def load_models():
             model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
     }
 
-models = load_models()
+def translate_label(label):
+    translations = {
+        "fracture": "Knochenbruch",
+        "no fracture": "Kein Bruch",
+        "normal": "Normal",
+        "abnormal": "Abnormal"
+    }
+    for eng, deu in translations.items():
+        if eng.lower() in label.lower():
+            return deu
+    return label
 
-def draw_boxes(image, predictions, threshold=0.6):
+def draw_boxes(image, predictions):
     draw = ImageDraw.Draw(image)
-    filtered_preds = [p for p in predictions if p['score'] >= threshold]
-    
-    for pred in filtered_preds:
+    for pred in predictions:
         box = pred['box']
-        label = f"{pred['label']} ({pred['score']:.2%})"
+        label = f"{translate_label(pred['label'])} ({pred['score']:.2%})"
         
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
-            outline="red",
+            outline="#FF6B6B",
             width=2
         )
         
-        draw.text((box['xmin'], box['ymin']), label, fill="red")
-    
-    return image, filtered_preds
+        text_bbox = draw.textbbox((box['xmin'], box['ymin']), label)
+        draw.rectangle(text_bbox, fill="#FF6B6B")
+        draw.text((box['xmin'], box['ymin']), label, fill="white")
+    return image
 
-# API Endpoint
-@app.post("/detect")
-async def detect_fracture(file: UploadFile = File(...), confidence: float = 0.6):
-    try:
-        # Read and process image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
-        
-        # Get predictions from all models
-        results = {}
-        
-        # Object detection models
-        detection_preds = models["D3STRON"](image)
-        result_image = image.copy()
-        result_image, filtered_detections = draw_boxes(result_image, detection_preds, confidence)
-        
-        # Save result image
-        img_byte_arr = io.BytesIO()
-        result_image.save(img_byte_arr, format='PNG')
-        img_byte_arr = img_byte_arr.getvalue()
-        img_b64 = base64.b64encode(img_byte_arr).decode()
-        
-        # Classification models
-        class_results = {
-            "Heem2": models["Heem2"](image),
-            "Nandodeomkar": models["Nandodeomkar"](image)
-        }
-        
-        return JSONResponse({
-            "success": True,
-            "detections": filtered_detections,
-            "classifications": class_results,
-            "image": img_b64
-        })
-        
-    except Exception as e:
-        return JSONResponse({
-            "success": False,
-            "error": str(e)
-        })
-
-# Streamlit UI
 def main():
-    st.title("🦴 Fraktur Detektion")
+    st.title("🦴 Knochenbrucherkennung System")
+    
+    models = load_models()
     
-    # UI elements...
-    uploaded_file = st.file_uploader("Röntgenbild hochladen", type=['png', 'jpg', 'jpeg'])
-    confidence = st.slider("Konfidenzschwelle", 0.0, 1.0, 0.6, 0.05)
+    with st.expander("⚙️ Einstellungen", expanded=True):
+        conf_threshold = st.slider(
+            "Konfidenzschwelle",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.60,
+            step=0.01
+        )
     
+    uploaded_file = st.file_uploader(
+        "Röntgenbild hochladen",
+        type=['png', 'jpg', 'jpeg'],
+        key="xray_upload"
+    )
+
     if uploaded_file:
-        # Process image and display results...
-        pass
+        col1, col2 = st.columns([1, 2])
+        
+        with col1:
+            image = Image.open(uploaded_file)
+            max_size = (250, 250)
+            image.thumbnail(max_size, Image.Resampling.LANCZOS)
+            st.image(image, caption="Original Röntgenbild", use_container_width=True)
+
+        with col2:
+            tab1, tab2 = st.tabs(["📊 Klassifizierung", "🔍 Erkennung"])
+            
+            with tab1:
+                for name in ["Heem2", "Nandodeomkar"]:
+                    with st.container():
+                        st.subheader(f"Modell: {name}")
+                        with st.spinner("Analyse läuft..."):
+                            predictions = models[name](image)
+                            for pred in predictions:
+                                if pred['score'] >= conf_threshold:
+                                    score_color = "green" if pred['score'] > 0.7 else "orange"
+                                    st.markdown(f"""
+                                        <div style='padding: 10px; border-radius: 5px; background-color: #f0f2f6;'>
+                                            <span style='color: {score_color}; font-weight: bold;'>
+                                                {pred['score']:.1%}
+                                            </span> - {translate_label(pred['label'])}
+                                        </div>
+                                    """, unsafe_allow_html=True)
+
+            with tab2:
+                st.subheader("Modell: D3STRON")
+                with st.spinner("Erkennung läuft..."):
+                    predictions = models["D3STRON"](image)
+                    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_container_width=True)
+                        
+                        for pred in filtered_preds:
+                            st.markdown(f"""
+                                <div style='padding: 8px; border-left: 4px solid #FF6B6B; 
+                                     margin: 5px 0; background-color: #f0f2f6;'>
+                                    {translate_label(pred['label'])}: {pred['score']:.1%}
+                                </div>
+                            """, unsafe_allow_html=True)
+                    else:
+                        st.info("Keine Erkennungen über dem Schwellenwert")
 
 if __name__ == "__main__":
     main()