Fix

app.py

@@ -1,169 +1,51 @@
 # app.py
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 import numpy as np
 import cv2
-import onnxruntime as ort
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware  # <-- Add this import
-from typing import List, Dict, Tuple
-import os
-from ultralytics import YOLO
+from inference import ObjectDetector
+from typing import List
 
 # Configuration
-MODEL_PT_PATH = "model.pt"
 MODEL_ONNX_PATH = "model.onnx"
-INPUT_SIZE = 640
 CLASS_NAMES = ["class0", "class1"]
-CONF_THRESHOLD = 0.5
-IOU_THRESHOLD = 0.45
+INPUT_SIZE = 640
+
+# Initialize detector
+detector = ObjectDetector(
+    model_path=MODEL_ONNX_PATH,
+    class_names=CLASS_NAMES,
+    input_size=INPUT_SIZE
+)
 
 # Initialize FastAPI
 app = FastAPI()
 
-# Add CORS middleware  # <-- Add this section
+# CORS configuration
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allows all origins (for testing only!)
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"],  # Allows all methods
-    allow_headers=["*"],  # Allows all headers
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-def load_onnx_model() -> ort.InferenceSession:
-    """Initialize ONNX runtime session"""
-    options = ort.SessionOptions()
-    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-    return ort.InferenceSession(
-        MODEL_ONNX_PATH,
-        providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
-        sess_options=options
-    )
-
-# Convert model if needed and load
-ort_session = load_onnx_model()
-
-# Warm-up run
-dummy_input = np.random.randn(1, 3, INPUT_SIZE, INPUT_SIZE).astype(np.float32)
-ort_session.run(None, {"images": dummy_input})
-
-# ================== Core Processing Functions ================== #
-def compute_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
-    """Compute Intersection over Union between a box and multiple boxes"""
-    xmin = np.maximum(box[0], boxes[:, 0])
-    ymin = np.maximum(box[1], boxes[:, 1])
-    xmax = np.minimum(box[2], boxes[:, 2])
-    ymax = np.minimum(box[3], boxes[:, 3])
-    
-    intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
-    box_area = (box[2] - box[0]) * (box[3] - box[1])
-    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-    
-    return intersection_area / (box_area + boxes_area - intersection_area + 1e-6)
-
-def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> List[int]:
-    """Non-Maximum Suppression implementation"""
-    sorted_indices = np.argsort(scores)[::-1]
-    keep_boxes = []
-    
-    while sorted_indices.size > 0:
-        box_id = sorted_indices[0]
-        keep_boxes.append(box_id)
-        
-        ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
-        keep_indices = np.where(ious < iou_threshold)[0]
-        sorted_indices = sorted_indices[keep_indices + 1]
-    
-    return keep_boxes
-
-def preprocess_image(image: np.ndarray) -> Tuple[np.ndarray, float, Tuple[int, int]]:
-    """Resize and normalize image with letterboxing"""
-    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
-    
-    # Normalize and transpose for ONNX
-    blob = canvas.astype(np.float32) / 255.0
-    return blob.transpose(2, 0, 1)[None, ...], scale, (pw, ph)
-
-def postprocess(
-    predictions: np.ndarray, 
-    original_shape: Tuple[int, int],
-    scale: float,
-    padding: Tuple[int, int]
-) -> List[Dict]:
-    """Process model outputs into final detections"""
-    predictions = np.squeeze(predictions).T
-    scores = np.max(predictions[:, 4:], axis=1)
-    valid = scores > CONF_THRESHOLD
-    predictions = predictions[valid]
-    
-    if predictions.size == 0:
-        return []
-
-    # Extract boxes
-    boxes = predictions[:, :4]
-    boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
-    boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
-    
-    # Adjust for letterbox
-    pad_w, pad_h = padding
-    boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
-    boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
-    
-    # Clip to image dimensions
-    h, w = original_shape
-    boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
-    boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
-    
-    # Get class IDs and apply NMS
-    class_ids = np.argmax(predictions[:, 4:], axis=1)
-    indices = nms(boxes, scores[valid], IOU_THRESHOLD)
-    
-    return [{
-        "class": CLASS_NAMES[int(class_ids[i])],
-        "confidence": float(scores[valid][i]),
-        "bbox": boxes[i].tolist(),  # [x1, y1, x2, y2]
-        "bbox_normalized": [  # [x_center, y_center, width, height] normalized
-            float((boxes[i][0] + boxes[i][2])/2 / w),
-            float((boxes[i][1] + boxes[i][3])/2 / h),
-            float((boxes[i][2] - boxes[i][0]) / w),
-            float((boxes[i][3] - boxes[i][1]) / h)
-        ]
-    } for i in indices]
-
-# ================== API Endpoint ================== #
-@app.post("/detect")
+@app.post("/detect", response_model=List[dict])
 async def detect_objects(file: UploadFile = File(...)):
     try:
-        # Read and validate image
         if not file.content_type.startswith("image/"):
-            raise HTTPException(status_code=400, detail="Invalid file type")
+            raise HTTPException(400, "Invalid file type")
 
         image_data = await file.read()
         image = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
+        
         if image is None:
-            raise HTTPException(status_code=400, detail="Invalid image data")
+            raise HTTPException(400, "Invalid image data")
+
+        detections = detector.predict(image)
+        return detections
 
-        # Preprocess
-        input_tensor, scale, padding = preprocess_image(image)
-        
-        # Inference
-        outputs = ort_session.run(None, {"images": input_tensor})
-        
-        # Post-process
-        detections = postprocess(outputs[0], image.shape[:2], scale, padding)
-        
-        return {"detections": detections}
-    
     except HTTPException as he:
         raise he
     except Exception as e:
-        return {"error": str(e), "detail": "Internal processing error"}
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+        return {"error": str(e)}

inference.py

@@ -0,0 +1,112 @@
+# inference.py
+import numpy as np
+import cv2
+import onnxruntime as ort
+from typing import List, Dict, Tuple
+
+class ObjectDetector:
+    def __init__(self, model_path: str, class_names: List[str], input_size: int = 640):
+        self.class_names = class_names
+        self.input_size = input_size
+        self.session = self._load_model(model_path)
+        self._warmup()
+
+    def _load_model(self, model_path: str) -> ort.InferenceSession:
+        options = ort.SessionOptions()
+        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        return ort.InferenceSession(
+            model_path,
+            providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
+            sess_options=options
+        )
+
+    def _warmup(self):
+        dummy_input = np.random.randn(1, 3, self.input_size, self.input_size).astype(np.float32)
+        self.session.run(None, {"images": dummy_input})
+
+    @staticmethod
+    def compute_iou(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+        xmin = np.maximum(box[0], boxes[:, 0])
+        ymin = np.maximum(box[1], boxes[:, 1])
+        xmax = np.minimum(box[2], boxes[:, 2])
+        ymax = np.minimum(box[3], boxes[:, 3])
+        
+        intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
+        box_area = (box[2] - box[0]) * (box[3] - box[1])
+        boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+        return intersection_area / (box_area + boxes_area - intersection_area + 1e-6)
+
+    @staticmethod
+    def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float) -> List[int]:
+        sorted_indices = np.argsort(scores)[::-1]
+        keep_boxes = []
+        
+        while sorted_indices.size > 0:
+            box_id = sorted_indices[0]
+            keep_boxes.append(box_id)
+            ious = ObjectDetector.compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
+            keep_indices = np.where(ious < iou_threshold)[0]
+            sorted_indices = sorted_indices[keep_indices + 1]
+        return keep_boxes
+
+    def preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, float, Tuple[int, int]]:
+        h, w = image.shape[:2]
+        scale = min(self.input_size / h, self.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((self.input_size, self.input_size, 3), 114, dtype=np.uint8)
+        ph, pw = (self.input_size - new_h) // 2, (self.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 postprocess(
+        self,
+        predictions: np.ndarray,
+        original_shape: Tuple[int, int],
+        scale: float,
+        padding: Tuple[int, int],
+        conf_threshold: float = 0.5,
+        iou_threshold: float = 0.45
+    ) -> List[Dict]:
+        predictions = np.squeeze(predictions).T
+        scores = np.max(predictions[:, 4:], axis=1)
+        valid = scores > conf_threshold
+        predictions = predictions[valid]
+        
+        if predictions.size == 0:
+            return []
+
+        boxes = predictions[:, :4]
+        boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
+        boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
+        
+        pad_w, pad_h = padding
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
+        
+        h, w = original_shape
+        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
+        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
+        
+        class_ids = np.argmax(predictions[:, 4:], axis=1)
+        indices = self.nms(boxes, scores[valid], iou_threshold)
+        
+        return [{
+            "class": self.class_names[int(class_ids[i])],
+            "confidence": float(scores[valid][i]),
+            "bbox": boxes[i].tolist(),
+            "bbox_normalized": [
+                float((boxes[i][0] + boxes[i][2])/2 / w),
+                float((boxes[i][1] + boxes[i][3])/2 / h),
+                float((boxes[i][2] - boxes[i][0]) / w),
+                float((boxes[i][3] - boxes[i][1]) / h)
+            ]
+        } for i in indices]
+
+    def predict(self, image: np.ndarray) -> List[Dict]:
+        input_tensor, scale, padding = self.preprocess(image)
+        outputs = self.session.run(None, {"images": input_tensor})
+        return self.postprocess(outputs[0], image.shape[:2], scale, padding)