Update app.py

app.py

@@ -1,17 +1,26 @@
 import streamlit as st
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse
 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
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # FastAPI app
-app = FastAPI()
+app = FastAPI(
+    title="Fracture Detection API",
+    description="API for detecting fractures in X-ray images using multiple ML models",
+    version="1.0.0"
+)
 
 # Enable CORS
 app.add_middleware(
@@ -20,21 +29,46 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
+    expose_headers=["*"]
 )
 
-# Load models
+# Load models with caching
 @st.cache_resource
 def load_models():
-    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")
-    }
+    logger.info("Loading ML 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:
+        logger.error(f"Error loading models: {str(e)}")
+        raise
 
-models = load_models()
+# Initialize models
+try:
+    models = load_models()
+    logger.info("Models loaded successfully")
+except Exception as e:
+    logger.error(f"Failed to load models: {str(e)}")
+    models = None
 
 def draw_boxes(image, predictions, threshold=0.6):
+    """
+    Draw bounding boxes and labels on the image for detected fractures.
+    
+    Args:
+        image (PIL.Image): Input image
+        predictions (list): List of predictions from the model
+        threshold (float): Confidence threshold for filtering predictions
+    
+    Returns:
+        tuple: (annotated image, filtered predictions)
+    """
     draw = ImageDraw.Draw(image)
     filtered_preds = [p for p in predictions if p['score'] >= threshold]
     
@@ -42,68 +76,202 @@ def draw_boxes(image, predictions, threshold=0.6):
         box = pred['box']
         label = f"{pred['label']} ({pred['score']:.2%})"
         
+        # Draw bounding box
         draw.rectangle(
             [(box['xmin'], box['ymin']), (box['xmax'], box['ymax'])],
             outline="red",
             width=2
         )
         
-        draw.text((box['xmin'], box['ymin']), label, fill="red")
+        # Draw label
+        draw.text(
+            (box['xmin'], box['ymin'] - 10),
+            label,
+            fill="red"
+        )
     
     return image, filtered_preds
 
-# API Endpoint
-@app.post("/detect")
-async def detect_fracture(file: UploadFile = File(...), confidence: float = 0.6):
+def process_image(image, confidence_threshold):
+    """
+    Process an image through all models and return combined results.
+    
+    Args:
+        image (PIL.Image): Input image
+        confidence_threshold (float): Confidence threshold for filtering predictions
+    
+    Returns:
+        dict: Combined results from all models
+    """
     try:
-        # Read and process image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
-        
-        # Get predictions from all models
-        results = {}
-        
-        # Object detection models
+        # Object detection
         detection_preds = models["D3STRON"](image)
         result_image = image.copy()
-        result_image, filtered_detections = draw_boxes(result_image, detection_preds, confidence)
+        result_image, filtered_detections = draw_boxes(
+            result_image,
+            detection_preds,
+            confidence_threshold
+        )
         
-        # Save result image
+        # Save annotated 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)
-        }
+        # Classification results
+        class_results = {}
+        
+        # Heem2 model
+        try:
+            heem2_result = models["Heem2"](image)
+            class_results["Heem2"] = heem2_result
+        except Exception as e:
+            logger.error(f"Error in Heem2 model: {str(e)}")
+            class_results["Heem2"] = {"error": str(e)}
         
-        return JSONResponse({
+        # Nandodeomkar model
+        try:
+            nando_result = models["Nandodeomkar"](image)
+            class_results["Nandodeomkar"] = nando_result
+        except Exception as e:
+            logger.error(f"Error in Nandodeomkar model: {str(e)}")
+            class_results["Nandodeomkar"] = {"error": str(e)}
+        
+        return {
             "success": True,
             "detections": filtered_detections,
             "classifications": class_results,
             "image": img_b64
-        })
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise
+
+# API Endpoints
+@app.post("/detect")
+@app.post("/api/predict")
+async def detect_fracture(
+    file: UploadFile = File(...),
+    confidence: float = Form(default=0.6)
+):
+    """
+    Endpoint for fracture detection in X-ray images.
+    
+    Args:
+        file (UploadFile): Uploaded image file
+        confidence (float): Confidence threshold for predictions
+    
+    Returns:
+        JSONResponse: Detection results including annotated image
+    """
+    logger.info(f"Received request with confidence threshold: {confidence}")
+    
+    try:
+        # Validate confidence threshold
+        if not 0 <= confidence <= 1:
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "success": False,
+                    "error": "Confidence threshold must be between 0 and 1"
+                }
+            )
         
+        # Read and validate image
+        contents = await file.read()
+        try:
+            image = Image.open(io.BytesIO(contents))
+        except Exception as e:
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "success": False,
+                    "error": f"Invalid image file: {str(e)}"
+                }
+            )
+        
+        # Process image
+        try:
+            results = process_image(image, confidence)
+            logger.info("Image processed successfully")
+            return JSONResponse(content=results)
+            
+        except Exception as e:
+            logger.error(f"Error processing image: {str(e)}")
+            return JSONResponse(
+                status_code=500,
+                content={
+                    "success": False,
+                    "error": f"Error processing image: {str(e)}"
+                }
+            )
+            
     except Exception as e:
-        return JSONResponse({
-            "success": False,
-            "error": str(e)
-        })
+        logger.error(f"Unexpected error: {str(e)}")
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "error": f"Unexpected error: {str(e)}"
+            }
+        )
 
 # Streamlit UI
 def main():
-    st.title("🦴 Fraktur Detektion")
+    st.title("🦴 Fracture Detection System")
+    st.write("Upload an X-ray image to detect potential fractures")
     
-    # 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)
+    # File uploader
+    uploaded_file = st.file_uploader(
+        "Upload X-ray image",
+        type=['png', 'jpg', 'jpeg']
+    )
     
-    if uploaded_file:
-        # Process image and display results...
-        pass
+    # Confidence threshold slider
+    confidence = st.slider(
+        "Confidence Threshold",
+        min_value=0.0,
+        max_value=1.0,
+        value=0.6,
+        step=0.05
+    )
+    
+    if uploaded_file is not None:
+        # Display original image
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Original Image", use_column_width=True)
+        
+        if st.button("Analyze Image"):
+            try:
+                # Process image
+                results = process_image(image, confidence)
+                
+                if results["success"]:
+                    # Display results
+                    st.success("Analysis completed successfully!")
+                    
+                    # Show annotated image
+                    annotated_image = Image.open(io.BytesIO(base64.b64decode(results["image"])))
+                    st.image(annotated_image, caption="Detected Fractures", use_column_width=True)
+                    
+                    # Show detections
+                    if results["detections"]:
+                        st.subheader("Detected Fractures")
+                        for det in results["detections"]:
+                            st.write(f"- {det['label']}: {det['score']:.2%} confidence")
+                    
+                    # Show classifications
+                    st.subheader("Classification Results")
+                    for model, preds in results["classifications"].items():
+                        st.write(f"**{model} Model:**")
+                        st.json(preds)
+                else:
+                    st.error("Analysis failed. Please try again.")
+                    
+            except Exception as e:
+                st.error(f"Error during analysis: {str(e)}")
 
 if __name__ == "__main__":
     main()