Update app.py

app.py

@@ -1,26 +1,17 @@
 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
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from starlette.responses import JSONResponse
 
 # FastAPI app
-app = FastAPI(
-    title="Fracture Detection API",
-    description="API for detecting fractures in X-ray images using multiple ML models",
-    version="1.0.0"
-)
+app = FastAPI()
 
 # Enable CORS
 app.add_middleware(
@@ -29,46 +20,21 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-    expose_headers=["*"]
 )
 
-# Load models with caching
+# Load models
 @st.cache_resource
 def load_models():
-    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
+    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
-try:
-    models = load_models()
-    logger.info("Models loaded successfully")
-except Exception as e:
-    logger.error(f"Failed to load models: {str(e)}")
-    models = None
+models = load_models()
 
 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]
     
@@ -76,202 +42,68 @@ 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 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 combined results.
-    
-    Args:
-        image (PIL.Image): Input image
-        confidence_threshold (float): Confidence threshold for filtering predictions
-    
-    Returns:
-        dict: Combined results from all models
-    """
+# 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 annotated image
+        # 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 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)}
-        
-        # 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)}
+        # Classification models
+        class_results = {
+            "Heem2": models["Heem2"](image),
+            "Nandodeomkar": models["Nandodeomkar"](image)
+        }
         
-        return {
+        return JSONResponse({
             "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:
-        logger.error(f"Unexpected error: {str(e)}")
-        return JSONResponse(
-            status_code=500,
-            content={
-                "success": False,
-                "error": f"Unexpected error: {str(e)}"
-            }
-        )
+        return JSONResponse({
+            "success": False,
+            "error": str(e)
+        })
 
 # Streamlit UI
 def main():
-    st.title("🦴 Fracture Detection System")
-    st.write("Upload an X-ray image to detect potential fractures")
+    st.title("🦴 Fraktur Detektion")
     
-    # File uploader
-    uploaded_file = st.file_uploader(
-        "Upload X-ray image",
-        type=['png', 'jpg', 'jpeg']
-    )
+    # 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)
     
-    # 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 uploaded_file:
+        # Process image and display results...
+        pass
 
 if __name__ == "__main__":
     main()