Spaces:
Runtime error
Runtime error
File size: 1,801 Bytes
7b6b122 f863e99 7b6b122 f863e99 7b6b122 a9284f8 f863e99 7b6b122 f863e99 7b6b122 |
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 |
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import torch
import pandas as pd
import uvicorn
import json
from PIL import Image, UnidentifiedImageError
import io
import os
from datetime import datetime
import os
os.environ['TORCH_HOME'] = '/tmp/torch' # Set writable cache for PyTorch Hub
# Load YOLOv5s pretrained on COCO
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# Folder to store annotated images
SAVE_DIR = "detections"
os.makedirs(SAVE_DIR, exist_ok=True)
app = FastAPI(title="COCO Object Detection API")
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
try:
# Read uploaded file into bytes
image_bytes = await file.read()
# Open with Pillow safely
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
except UnidentifiedImageError:
return JSONResponse(
content={"error": "Unrecognized or invalid image file format."},
status_code=400
)
# Run inference
results = model(image)
# Save annotated image locally
plotted_image = results.render()[0] # numpy array with boxes
pil_image = Image.fromarray(plotted_image)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
save_path = os.path.join(SAVE_DIR, f"detection_{timestamp}.jpg")
pil_image.save(save_path, format="JPEG")
print (f"Saved annotated image to {save_path}")
# Convert results to JSON
df: pd.DataFrame = results.pandas().xyxy[0]
json_data = json.loads(df.to_json(orient="records"))
return JSONResponse(content=json_data)
@app.get("/")
def root():
return {"message": "Send POST /predict with an image file to get detections."}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8080)
|