Update app.py

app.py

@@ -1,25 +1,28 @@
 import streamlit as st
 from transformers import pipeline
-from PIL import Image, ImageDraw
 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
 
-st.set_page_config(
-    page_title="Knochenbrucherkennung",
-    layout="wide",
-    initial_sidebar_state="collapsed"
-)
+# FastAPI app
+app = FastAPI()
 
-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)
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
+# Load models
 @st.cache_resource
 def load_models():
     return {
@@ -29,104 +32,78 @@ def load_models():
             model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
     }
 
-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
+models = load_models()
 
-def draw_boxes(image, predictions):
+def draw_boxes(image, predictions, threshold=0.6):
     draw = ImageDraw.Draw(image)
-    for pred in predictions:
+    filtered_preds = [p for p in predictions if p['score'] >= threshold]
+    
+    for pred in filtered_preds:
         box = pred['box']
-        label = f"{translate_label(pred['label'])} ({pred['score']:.2%})"
+        label = f"{pred['label']} ({pred['score']:.2%})"
         
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
-            outline="#FF6B6B",
+            outline="red",
             width=2
         )
         
-        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
+        draw.text((box['xmin'], box['ymin']), label, fill="red")
+    
+    return image, filtered_preds
 
+# 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("🦴 Knochenbrucherkennung System")
-    
-    models = load_models()
+    st.title("🦴 Fraktur Detektion")
     
-    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
-        )
+    # 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)
     
-    uploaded_file = st.file_uploader(
-        "Röntgenbild hochladen",
-        type=['png', 'jpg', 'jpeg'],
-        key="xray_upload"
-    )
-
     if uploaded_file:
-        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")
+        # Process image and display results...
+        pass
 
 if __name__ == "__main__":
     main()