pothole-api-v2 / api.py
Mohamed-salah-dev's picture
Update api.py
b2f7ecf verified
Raw
History Blame Contribute Delete
1.42 kB
import os
import gdown
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from ultralytics import YOLO
import cv2
import numpy as np
import uvicorn
app = FastAPI()
# الموديل بتاعك
FILE_ID = '1a5-mI8D6oT2K2r0Y-j0U-pE4k3t2r9mC'
drive_url = f'https://drive.google.com/uc?id={FILE_ID}'
model_path = 'best.pt'
if not os.path.exists(model_path):
gdown.download(drive_url, model_path, quiet=False)
model = YOLO(model_path)
@app.get("/")
def home():
return {"status": "Online", "model": "YOLOv8 Loaded"}
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
try:
contents = await file.read()
nparr = np.frombuffer(contents, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
results = model(img)
predictions = []
for result in results:
for box in result.boxes:
predictions.append({
"class": model.names[int(box.cls)],
"confidence": float(box.conf),
"bbox": [float(x) for x in box.xyxy[0]]
})
return JSONResponse(content={"predictions": predictions})
except Exception as e:
return JSONResponse(content={"error": str(e)}, status_code=500)
if __name__ == "__main__":
# لازم بورت 7860 عشان Hugging Face
uvicorn.run(app, host="0.0.0.0", port=7860)