|
from fastapi import FastAPI, File, UploadFile, HTTPException, Form |
|
from fastapi.responses import HTMLResponse |
|
from fastapi.staticfiles import StaticFiles |
|
from fastapi.templating import Jinja2Templates |
|
from starlette.requests import Request |
|
from ultralytics import YOLO |
|
import requests |
|
from PIL import Image |
|
import numpy as np |
|
import cv2 |
|
import io |
|
import os |
|
import httpx |
|
|
|
app = FastAPI() |
|
|
|
|
|
app.mount("/tmp", StaticFiles(directory="/tmp"), name="static") |
|
|
|
|
|
templates = Jinja2Templates(directory="templates") |
|
|
|
|
|
def predict_yolo(image_path): |
|
|
|
model = YOLO('ultralyticsplus/yolov8s') |
|
|
|
|
|
results = model(image_path) |
|
|
|
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: |
|
|
|
async with httpx.AsyncClient() as client: |
|
response = await client.get(url) |
|
response.raise_for_status() |
|
content = response.content |
|
|
|
image = Image.open(io.BytesIO(content)) |
|
|
|
results = predict_yolo(image) |
|
|
|
render = render_result(model=model, image=image, result=results[0]) |
|
|
|
|
|
image_byte_array = io.BytesIO() |
|
image.save(image_byte_array, format="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!"} |
