Spaces:
Runtime error
Runtime error
File size: 1,773 Bytes
baa5379 9e17458 baa5379 9e17458 baa5379 9e17458 baa5379 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
# main.py
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
import uvicorn
import shutil
import os
import uuid
import cv2
import numpy as np
import base64
from yolo_numbering import predict_yolo as predict_yolo_numbering
from yolo_anomaly import predict_yolo_anomaly
app = FastAPI()
UPLOAD_DIR = "/tmp/uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)
@app.get("/")
async def root():
return {"message": "API is up and running!"}
@app.post("/predict/")
async def predict_endpoint(
file: UploadFile = File(...),
model: str = Form(...), # either "detectron" or "yolo"
):
# Save uploaded image
file_ext = file.filename.split('.')[-1]
filename = f"{uuid.uuid4()}.{file_ext}"
file_path = os.path.join(UPLOAD_DIR, filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
try:
if model == "Anomaly":
result_img, predictions = predict_yolo_anomaly(file_path)
elif model == "Numbering":
result_img, predictions = predict_yolo_numbering(file_path)
else:
return JSONResponse({"error": "Invalid model choice"}, status_code=400)
# Encode image to bytes (optional)
_, img_encoded = cv2.imencode(".jpg", result_img)
img_bytes = img_encoded.tobytes()
img_b64 = base64.b64encode(img_bytes).decode('utf-8')
return {
"predictions": predictions,
"image_base64": img_b64
}
except Exception as e:
return JSONResponse({"error": str(e)}, status_code=500)
finally:
os.remove(file_path) # Clean up uploaded file
# Only for testing locally
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000) |