fix

app.py

@@ -1,80 +1,103 @@
 import os
 import cv2
 import numpy as np
-from fastapi import FastAPI, WebSocket, status
-from onnxruntime import InferenceSession
-from ultralytics import YOLO
+from fastapi import FastAPI, WebSocket, status, UploadFile, File
+from fastapi.responses import JSONResponse
+from contextlib import asynccontextmanager
 import uvicorn
 import base64
 from typing import Tuple, List
+import onnxruntime as ort
+from ultralytics import YOLO
 
 # Configuration
 MODEL_PT_PATH = "model.pt"
 MODEL_ONNX_PATH = "model.onnx"
 INPUT_SIZE = 640
-CLASS_NAMES = ["class0", "class1"]  # Your class names
+CLASS_NAMES = ["class0", "class1"]  # Replace with your class names
 CONF_THRESHOLD = 0.5
 IOU_THRESHOLD = 0.45
 
-app = FastAPI(title="Object Detection API")
-
-# Load model once at startup
-@app.on_event("startup")
-async def load_model():
-    # Convert model if needed
+# --- Modern FastAPI Lifespan Setup (Replaces @app.on_event) ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Initialize and clean up model resources."""
+    # Convert PyTorch to ONNX if needed
     if not os.path.exists(MODEL_ONNX_PATH):
         print("Converting PyTorch model to ONNX...")
-        try:
-            model = YOLO(MODEL_PT_PATH)
-            model.export(
-                format="onnx",
-                imgsz=INPUT_SIZE,
-                opset=12,
-                simplify=True,
-                dynamic=False,
-                half=False
-            )
-            if os.path.exists("yolov8n.onnx"):
-                os.rename("yolov8n.onnx", MODEL_ONNX_PATH)
-        except Exception as e:
-            raise RuntimeError(f"ONNX conversion failed: {str(e)}")
-    
-    # Initialize ONNX runtime session
-    options = ort.SessionOptions()
-    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-    app.state.model = InferenceSession(
-        MODEL_ONNX_PATH,
-        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
-        sess_options=options
-    )
-    
+        model = YOLO(MODEL_PT_PATH)
+        model.export(
+            format="onnx",
+            imgsz=INPUT_SIZE,
+            opset=12,
+            simplify=True,
+            dynamic=False,
+            half=False
+        )
+        if os.path.exists("yolov8n.onnx"):
+            os.rename("yolov8n.onnx", MODEL_ONNX_PATH)
+
+    # Load ONNX model with GPU
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    app.state.model = ort.InferenceSession(MODEL_ONNX_PATH, providers=providers)
+
     # Warm-up
     dummy_input = np.random.randn(1, 3, INPUT_SIZE, INPUT_SIZE).astype(np.float32)
     app.state.model.run(None, {"images": dummy_input})
+    print("Model loaded and ready!")
+    
+    yield  # App runs here
+    
+    # Cleanup (optional)
+    print("Shutting down...")
 
+# Initialize FastAPI with lifespan
+app = FastAPI(title="YOLOv8 API", lifespan=lifespan)
+
+# --- Core Detection Functions (Same as Before) ---
 def preprocess_image(image: np.ndarray) -> Tuple[np.ndarray, float, Tuple[int, int]]:
-    """Preprocess image with letterboxing"""
+    """Resize and normalize image for YOLOv8 input."""
     h, w = image.shape[:2]
     scale = min(INPUT_SIZE / h, INPUT_SIZE / w)
     new_h, new_w = int(h * scale), int(w * scale)
-    
     resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    
+    # Letterboxing
     canvas = np.full((INPUT_SIZE, INPUT_SIZE, 3), 114, dtype=np.uint8)
     ph, pw = (INPUT_SIZE - new_h) // 2, (INPUT_SIZE - new_w) // 2
     canvas[ph:ph+new_h, pw:pw+new_w] = resized
-    
     blob = canvas.astype(np.float32) / 255.0
     return blob.transpose(2, 0, 1)[None, ...], scale, (pw, ph)
 
-async def process_image(image: np.ndarray) -> dict:
-    """Process image and return detection results"""
+def nms(boxes, scores, iou_threshold):
+    """Non-Maximum Suppression to filter overlapping boxes."""
+    keep = []
+    if len(boxes) == 0:
+        return keep
+    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+    
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        w, h = np.maximum(0.0, xx2 - xx1), np.maximum(0.0, yy2 - yy1)
+        iou = (w * h) / (areas[i] + areas[order[1:]] - w * h)
+        inds = np.where(iou <= iou_threshold)[0]
+        order = order[inds + 1]
+    return keep
+
+async def detect_objects(image: np.ndarray) -> dict:
+    """Run YOLOv8 inference and return detections."""
     # Preprocess
-    input_tensor, scale, padding = preprocess_image(image)
+    input_tensor, scale, (pad_w, pad_h) = preprocess_image(image)
     
     # Inference
     outputs = app.state.model.run(None, {"images": input_tensor})
-    
-    # Post-process
     predictions = np.squeeze(outputs[0]).T
     scores = np.max(predictions[:, 4:], axis=1)
     valid = scores > CONF_THRESHOLD
@@ -82,76 +105,60 @@ async def process_image(image: np.ndarray) -> dict:
     
     if predictions.size == 0:
         return {"detections": []}
-    
-    # Convert boxes
+
+    # Decode boxes
     boxes = predictions[:, :4]
     boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
     boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
-    
-    # Adjust coordinates
-    pad_w, pad_h = padding
     boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
     boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
     
-    # Clip boxes
+    # Clip to image bounds
     h, w = image.shape[:2]
     boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
     boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
     
-    # Apply NMS
+    # NMS
     class_ids = np.argmax(predictions[:, 4:], axis=1)
-    indices = nms(boxes, scores[valid], IOU_THRESHOLD)
+    keep = nms(boxes, scores[valid], IOU_THRESHOLD)
     
     # Format results
     detections = []
-    for i in indices:
+    for i in keep:
         detections.append({
             "class_id": int(class_ids[i]),
             "class_name": CLASS_NAMES[class_ids[i]],
             "confidence": float(scores[valid][i]),
-            "bbox": {
-                "x1": float(boxes[i][0]),
-                "y1": float(boxes[i][1]),
-                "x2": float(boxes[i][2]),
-                "y2": float(boxes[i][3])
-            }
+            "bbox": [float(x) for x in boxes[i]]  # [x1, y1, x2, y2]
         })
     
     return {"detections": detections}
 
+# --- API Endpoints ---
 @app.websocket("/ws/detect")
-async def websocket_endpoint(websocket: WebSocket):
+async def websocket_detection(websocket: WebSocket):
+    """Real-time detection via WebSocket."""
     await websocket.accept()
     try:
         while True:
-            # Receive base64 image
             data = await websocket.receive_text()
-            header, encoded = data.split(",", 1)
+            _, encoded = data.split(",", 1)
             image_bytes = base64.b64decode(encoded)
-            
-            # Convert to numpy array
             nparr = np.frombuffer(image_bytes, np.uint8)
             image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
             
-            # Process and return results
-            results = await process_image(image)
+            results = await detect_objects(image)
             await websocket.send_json(results)
-            
     except Exception as e:
         print(f"WebSocket error: {e}")
         await websocket.close(code=status.WS_1011_INTERNAL_ERROR)
 
 @app.post("/detect")
 async def http_detect(image: UploadFile = File(...)):
-    # Read and decode image
+    """HTTP endpoint for single-image detection."""
     contents = await image.read()
     nparr = np.frombuffer(contents, np.uint8)
     img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    
-    # Process and return results
-    return await process_image(img)
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    return await detect_objects(img)