File size: 2,210 Bytes
7200a76 c52b1e1 b57bea3 7200a76 c52b1e1 0802393 c52b1e1 8b7fae6 d89619d 8b7fae6 d89619d c52b1e1 d89619d 8bc9319 d89619d c52b1e1 d89619d bde6aa6 4114862 c52b1e1 |
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 63 64 65 66 67 68 69 70 71 72 73 |
from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from ultralytics import YOLO
import requests
from PIL import Image
import numpy as np
import cv2
import io
import os
import httpx
app = FastAPI()
# Mount a static directory for serving images
app.mount("/tmp", StaticFiles(directory="/tmp"), name="static")
# Templates for rendering HTML
templates = Jinja2Templates(directory="templates")
def predict_yolo(image_path):
# Load a model
model = YOLO('ultralyticsplus/yolov8s')
# Run batched inference on a list of images
results = model(image_path) # return a list of Results objects
return results
def draw_boxes(image, boxes):
for box in boxes:
x, y, w, h = box["bbox"]
cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 0), 2)
return image
@app.post("/", response_class=HTMLResponse)
async def detect_yolo(request: Request, url: str = Form(...)):
try:
# Download the image from the specified URL
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status() # Raise an exception if there is an error in the request
content = response.content
image = Image.open(io.BytesIO(content))
results = predict_yolo(image)
render = render_result(model=model, image=image, result=results[0])
# Save the modified image to a byte sequence
image_byte_array = io.BytesIO()
image.save(image_byte_array, format="PNG")
# Return the image as a response with content type "image/png"
return templates.TemplateResponse("result.html", {"request": request, "image": base64.b64encode(image_byte_array.getvalue()).decode()})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
@app.get("/test")
async def read_root():
return {"message": "TEST"}
@app.get("/")
async def read_root():
return {"message": "Hello, this is a YOLO prediction API using FastAPI!"} |