Update app.py

app.py

@@ -1,23 +1,13 @@
 import streamlit as st
-
-# Set page config must be the first Streamlit command
-st.set_page_config(
-    page_title="Fracture Detection System",
-    page_icon="🦴",
-    layout="wide"
-)
-
-import base64
-from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
 from transformers import pipeline
 import torch
 from PIL import Image, ImageDraw
 import io
-from threading import Thread
-import uvicorn
-import json
+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
@@ -32,27 +22,19 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load models with caching
+# Load models
 @st.cache_resource
 def load_models():
-    try:
-        return {
-            "D3STRON": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
-            "Heem2": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
-            "Nandodeomkar": pipeline("image-classification", 
-                model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
-        }
-    except Exception as e:
-        st.error(f"Error loading models: {str(e)}")
-        return None
+    return {
+        "D3STRON": pipeline("object-detection", model="D3STRON/bone-fracture-detr"),
+        "Heem2": pipeline("image-classification", model="Heem2/bone-fracture-detection-using-xray"),
+        "Nandodeomkar": pipeline("image-classification", 
+            model="nandodeomkar/autotrain-fracture-detection-using-google-vit-base-patch-16-54382127388")
+    }
 
-# Initialize models
 models = load_models()
 
 def draw_boxes(image, predictions, threshold=0.6):
-    """
-    Draw bounding boxes on the image for fracture detections
-    """
     draw = ImageDraw.Draw(image)
     filtered_preds = [p for p in predictions if p['score'] >= threshold]
     
@@ -60,148 +42,68 @@ def draw_boxes(image, predictions, threshold=0.6):
         box = pred['box']
         label = f"{pred['label']} ({pred['score']:.2%})"
         
-        # Draw rectangle
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
             outline="red",
             width=2
         )
         
-        # Draw label
-        draw.text(
-            (box['xmin'], box['ymin'] - 10),
-            label,
-            fill="red"
-        )
+        draw.text((box['xmin'], box['ymin']), label, fill="red")
     
     return image, filtered_preds
 
-def process_image(image, confidence_threshold):
-    """
-    Process an image through all models and return results
-    """
+# API Endpoint
+@app.post("/detect")
+async def detect_fracture(file: UploadFile = File(...), confidence: float = 0.6):
     try:
-        # Object detection
+        # 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_threshold)
+        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()
-        result_base64 = base64.b64encode(img_byte_arr).decode()
+        img_b64 = base64.b64encode(img_byte_arr).decode()
         
-        # Classifications
+        # Classification models
         class_results = {
             "Heem2": models["Heem2"](image),
             "Nandodeomkar": models["Nandodeomkar"](image)
         }
         
-        return {
+        return JSONResponse({
             "success": True,
             "detections": filtered_detections,
             "classifications": class_results,
-            "image": result_base64
-        }
-        
-    except Exception as e:
-        return {
-            "success": False,
-            "error": str(e)
-        }
-
-# FastAPI endpoint
-@app.post("/api/predict")
-async def predict(request: Request):
-    try:
-        # Read JSON request body
-        body = await request.json()
-        
-        # Extract base64 image and confidence threshold
-        image_base64 = body['data'][0]
-        confidence_threshold = float(body['data'][1])
-        
-        # Decode base64 image
-        image_bytes = base64.b64decode(image_base64)
-        image = Image.open(io.BytesIO(image_bytes))
-        
-        # Process image
-        result = process_image(image, confidence_threshold)
-        
-        return JSONResponse(result)
+            "image": img_b64
+        })
         
     except Exception as e:
         return JSONResponse({
             "success": False,
             "error": str(e)
-        }, status_code=500)
+        })
 
-# Streamlit interface
-def streamlit_interface():
-    st.title("🦴 Système de Détection de Fractures")
-    
-    # File uploader
-    uploaded_file = st.file_uploader(
-        "Upload X-ray Image",
-        type=['png', 'jpg', 'jpeg'],
-        help="Upload an X-ray image for fracture detection"
-    )
+# Streamlit UI
+def main():
+    st.title("🦴 Fraktur Detektion")
     
-    # Confidence threshold slider
-    confidence = st.slider(
-        "Confidence Threshold",
-        min_value=0.0,
-        max_value=1.0,
-        value=0.6,
-        step=0.05,
-        help="Adjust the confidence threshold for detection"
-    )
+    # 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)
     
     if uploaded_file:
-        # Display original image
-        col1, col2 = st.columns(2)
-        
-        with col1:
-            st.subheader("Original X-ray")
-            image = Image.open(uploaded_file)
-            st.image(image, use_column_width=True)
-        
-        if st.button("Analyze"):
-            with st.spinner('Analyzing image...'):
-                try:
-                    # Process image
-                    results = process_image(image, confidence)
-                    
-                    if results["success"]:
-                        with col2:
-                            st.subheader("Detection Results")
-                            # Display processed image
-                            result_image = Image.open(io.BytesIO(base64.b64decode(results["image"])))
-                            st.image(result_image, use_column_width=True)
-                        
-                        # Display detections
-                        st.subheader("Detected Fractures:")
-                        for detection in results["detections"]:
-                            st.write(f"- {detection['label']}: {detection['score']:.2%}")
-                        
-                        # Display classifications
-                        st.subheader("Classification Results:")
-                        st.json(results["classifications"])
-                    else:
-                        st.error("Error processing image: " + results.get("error", "Unknown error"))
-                
-                except Exception as e:
-                    st.error(f"Error during analysis: {str(e)}")
-
-def run_fastapi():
-    """Run the FastAPI server"""
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+        # Process image and display results...
+        pass
 
 if __name__ == "__main__":
-    # Start FastAPI in a separate thread
-    api_thread = Thread(target=run_fastapi, daemon=True)
-    api_thread.start()
-    
-    # Run Streamlit interface
-    streamlit_interface()
+    main()