| from fastapi import FastAPI, File, UploadFile |
| from fastapi.middleware.cors import CORSMiddleware |
| import uvicorn |
| import cv2 |
| import numpy as np |
| from ultralytics import YOLO |
|
|
| app = FastAPI() |
|
|
| |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_credentials=True, |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| model = YOLO("helmet.pt") |
|
|
| @app.post("/predict") |
| async def predict(file: UploadFile = File(...)): |
| image_bytes = await file.read() |
| np_img = np.frombuffer(image_bytes, np.uint8) |
| img = cv2.imdecode(np_img, cv2.IMREAD_COLOR) |
|
|
| results = model(img)[0] |
|
|
| detections = [] |
| for box in results.boxes: |
| cls = int(box.cls[0]) |
| conf = float(box.conf[0]) |
| label = model.names[cls] |
| detections.append({ |
| "label": label, |
| "confidence": round(conf, 3) |
| }) |
|
|
| return { |
| "count": len(detections), |
| "detections": detections |
| } |
|
|
| if __name__ == "__main__": |
| uvicorn.run(app, host="0.0.0.0", port=7860) |
|
|
