将检测任务迁移python

python-inference-service/app/main.py

@@ -0,0 +1,164 @@
+import os
+import base64
+import cv2
+import json
+import numpy as np
+from typing import Dict, List
+from fastapi import FastAPI, HTTPException, File, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
+
+from app.models import Detection, DetectionRequest, DetectionResponse, ModelInfo, ModelsResponse
+from app.detector import ModelManager
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Python Model Inference Service",
+    description="API for object detection using Python models",
+    version="1.0.0"
+)
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize model manager
+model_manager = None
+
+# Load models from configuration
+@app.on_event("startup")
+async def startup_event():
+    global model_manager
+    model_manager = ModelManager()
+    
+    # Look for models.json configuration file
+    models_json_path = os.getenv("MODELS_JSON", os.path.join(os.path.dirname(__file__), "..", "models", "models.json"))
+    
+    if os.path.exists(models_json_path):
+        try:
+            with open(models_json_path, "r") as f:
+                models_config = json.load(f)
+            model_manager.load(models_config)
+            print(f"Loaded model configuration from {models_json_path}")
+        except Exception as e:
+            print(f"Failed to load models from {models_json_path}: {str(e)}")
+    else:
+        print(f"Models configuration not found: {models_json_path}")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    global model_manager
+    if model_manager:
+        model_manager.close()
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "ok"}
+
+@app.get("/api/models", response_model=ModelsResponse)
+async def get_models():
+    """Get available models"""
+    global model_manager
+    
+    if not model_manager:
+        raise HTTPException(status_code=500, detail="Model manager not initialized")
+    
+    detectors = model_manager.all()
+    models = []
+    
+    for detector in detectors:
+        model_info = ModelInfo(
+            name=detector.model_name,
+            path=getattr(detector, 'model_path', ''),
+            size=[detector.input_width, detector.input_height],
+            backend="Python",
+            loaded=True
+        )
+        models.append(model_info)
+    
+    return ModelsResponse(models=models)
+
+@app.post("/api/detect", response_model=DetectionResponse)
+async def detect(request: DetectionRequest):
+    """Detect objects in an image"""
+    global model_manager
+    
+    if not model_manager:
+        raise HTTPException(status_code=500, detail="Model manager not initialized")
+    
+    # Get detector for requested model
+    detector = model_manager.get(request.model_name)
+    if not detector:
+        raise HTTPException(status_code=404, detail=f"Model not found: {request.model_name}")
+    
+    # Decode base64 image
+    try:
+        # Remove data URL prefix if present
+        if "base64," in request.image_data:
+            image_data = request.image_data.split("base64,")[1]
+        else:
+            image_data = request.image_data
+            
+        # Decode base64 image
+        image_bytes = base64.b64decode(image_data)
+        nparr = np.frombuffer(image_bytes, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
+        if image is None:
+            raise HTTPException(status_code=400, detail="Invalid image data")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to decode image: {str(e)}")
+    
+    # Run detection
+    detections, inference_time = detector.detect(image)
+    
+    return DetectionResponse(
+        model_name=request.model_name,
+        detections=detections,
+        inference_time=inference_time
+    )
+
+@app.post("/api/detect/file", response_model=DetectionResponse)
+async def detect_file(
+    model_name: str,
+    file: UploadFile = File(...)
+):
+    """Detect objects in an uploaded image file"""
+    global model_manager
+    
+    if not model_manager:
+        raise HTTPException(status_code=500, detail="Model manager not initialized")
+    
+    # Get detector for requested model
+    detector = model_manager.get(model_name)
+    if not detector:
+        raise HTTPException(status_code=404, detail=f"Model not found: {model_name}")
+    
+    # Read uploaded file
+    try:
+        contents = await file.read()
+        nparr = np.frombuffer(contents, np.uint8)
+        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
+        if image is None:
+            raise HTTPException(status_code=400, detail="Invalid image data")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to process image: {str(e)}")
+    
+    # Run detection
+    detections, inference_time = detector.detect(image)
+    
+    return DetectionResponse(
+        model_name=model_name,
+        detections=detections,
+        inference_time=inference_time
+    )
+
+if __name__ == "__main__":
+    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)