fix

app.py

@@ -1,11 +1,12 @@
 import os
 import cv2
 import numpy as np
-from fastapi import FastAPI, WebSocket, status, UploadFile, File
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, status
 from fastapi.responses import JSONResponse
 from contextlib import asynccontextmanager
 import uvicorn
 import base64
+from typing import List, Tuple
 import onnxruntime as ort
 from ultralytics import YOLO
 
@@ -13,79 +14,215 @@ from ultralytics import YOLO
 MODEL_PT_PATH = "model.pt"
 MODEL_ONNX_PATH = "model.onnx"
 INPUT_SIZE = 640
-CLASS_NAMES = ["class0", "class1"]  # Replace with your class names
+CLASS_NAMES = ["class0", "class1"]  # Replace with your actual class names
 CONF_THRESHOLD = 0.5
 IOU_THRESHOLD = 0.45
 
-# --- Modern FastAPI Lifespan Setup ---
+# --- Lifespan Management ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Initialize model on startup."""
+    """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...")
-        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)
+        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)
+            print("✅ ONNX conversion successful!")
+        except Exception as e:
+            raise RuntimeError(f"ONNX conversion failed: {str(e)}")
 
-    # Warm-up
+    # Initialize ONNX Runtime session with GPU
+    print("⚙️ Initializing ONNX Runtime session...")
+    providers = [
+        ('CUDAExecutionProvider', {
+            'device_id': 0,
+            'arena_extend_strategy': 'kNextPowerOfTwo',
+            'gpu_mem_limit': 2 * 1024 * 1024 * 1024,  # 2GB
+            'cudnn_conv_algo_search': 'HEURISTIC',
+            'do_copy_in_default_stream': True,
+        }),
+        'CPUExecutionProvider'
+    ]
+    
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    
+    app.state.model = ort.InferenceSession(
+        MODEL_ONNX_PATH,
+        providers=providers,
+        sess_options=sess_options
+    )
+    
+    # Warm-up run
+    print("🔥 Warming up model...")
     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 successfully!")
-    yield
+    print("🚀 Model ready for inference!")
+    
+    yield  # App runs here
+    
+    print("🛑 Cleaning up resources...")
 
 # --- Initialize FastAPI App ---
-app = FastAPI(lifespan=lifespan)  # Must be named 'app' for Hugging Face Spaces
+app = FastAPI(
+    title="YOLOv8 Object Detection API",
+    description="Real-time object detection with WebSocket and HTTP endpoints",
+    lifespan=lifespan
+)
+
+# --- Core Detection Functions ---
+def letterbox_image(image: np.ndarray) -> Tuple[np.ndarray, float, Tuple[int, int]]:
+    """Preprocess image with letterboxing for YOLOv8."""
+    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)
+    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)
+
+def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> List[int]:
+    """Non-Maximum Suppression to filter overlapping boxes."""
+    keep = []
+    if len(boxes) == 0:
+        return keep
+        
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = 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 = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        iou = inter / (areas[i] + areas[order[1:]] - inter)
+        inds = np.where(iou <= iou_threshold)[0]
+        order = order[inds + 1]
+    
+    return keep
+
+async def detect_objects(image: np.ndarray) -> dict:
+    """Run object detection pipeline."""
+    # Preprocess
+    input_tensor, scale, (pad_w, pad_h) = letterbox_image(image)
+    
+    # Inference
+    outputs = app.state.model.run(None, {"images": input_tensor})
+    predictions = np.squeeze(outputs[0]).T
+    scores = np.max(predictions[:, 4:], axis=1)
+    valid = scores > CONF_THRESHOLD
+    predictions = predictions[valid]
+    
+    if predictions.size == 0:
+        return {"detections": []}
 
-# --- Rest of your code (WebSocket and HTTP endpoints) ---
+    # Decode boxes
+    boxes = predictions[:, :4]
+    boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
+    boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
+    boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
+    boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
+    
+    # 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)
+    
+    # NMS
+    class_ids = np.argmax(predictions[:, 4:], axis=1)
+    keep = nms(boxes, scores[valid], IOU_THRESHOLD)
+    
+    # Format results
+    detections = []
+    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": [float(boxes[i][0]), float(boxes[i][1]), 
+                     float(boxes[i][2]), float(boxes[i][3])]
+        })
+    
+    return {"detections": detections}
+
+# --- API Endpoints ---
 @app.get("/")
 async def health_check():
-    return {"status": "OK", "message": "API is running"}
+    return {"status": "OK", "message": "Object Detection API is running"}
 
 @app.websocket("/ws/detect")
 async def websocket_detection(websocket: WebSocket):
     await websocket.accept()
     try:
         while True:
-            data = await websocket.receive_text()
-            _, encoded = data.split(",", 1)
-            image_bytes = base64.b64decode(encoded)
-            nparr = np.frombuffer(image_bytes, np.uint8)
-            image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            
-            # Process and return detections
-            results = await detect_objects(image)
-            await websocket.send_json(results)
+            try:
+                # Receive base64 image
+                data = await websocket.receive_text()
+                if not data.startswith("data:"):
+                    await websocket.send_json({"error": "Invalid image format"})
+                    continue
+                    
+                _, encoded = data.split(",", 1)
+                image_bytes = base64.b64decode(encoded)
+                nparr = np.frombuffer(image_bytes, np.uint8)
+                image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                
+                # Process and return detections
+                results = await detect_objects(image)
+                await websocket.send_json(results)
+                
+            except Exception as e:
+                print(f"⚠️ Processing error: {str(e)}")
+                await websocket.send_json({"error": str(e)})
+                continue
+                
+    except WebSocketDisconnect:
+        print("Client disconnected")
     except Exception as e:
-        print(f"WebSocket error: {e}")
-        await websocket.close(code=status.WS_1011_INTERNAL_ERROR)
+        print(f"WebSocket error: {str(e)}")
+    finally:
+        await websocket.close()
 
 @app.post("/detect")
 async def http_detect(image: UploadFile = File(...)):
-    """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)
-    return await detect_objects(img)
-
-# --- Helper Functions ---
-async def detect_objects(image: np.ndarray) -> dict:
-    """Your existing detection logic here"""
-    return {"detections": []}  # Replace with actual implementation
+    """Process single image via HTTP POST."""
+    try:
+        contents = await image.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        return await detect_objects(img)
+    except Exception as e:
+        return JSONResponse(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            content={"error": f"Image processing failed: {str(e)}"}
+        )
 
-# --- Hugging Face Spaces Requirement ---
-# The variable `app` must be defined at the top level
+# --- For Local Development ---
+if __name__ == "__main__":
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)