51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
|
from fastapi import FastAPI, File, UploadFile, Request
|
||
|
|
from fastapi.responses import HTMLResponse
|
||
|
|
from fastapi.staticfiles import StaticFiles
|
||
|
|
from fastapi.templating import Jinja2Templates
|
||
|
|
import uvicorn
|
||
|
|
from ultralytics import YOLO
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from PIL import Image
|
||
|
|
import io
|
||
|
|
import os
|
||
|
|
import base64
|
||
|
|
|
||
|
|
app = FastAPI()
|
||
|
|
|
||
|
|
# Initialize templates
|
||
|
|
templates = Jinja2Templates(directory="templates")
|
||
|
|
|
||
|
|
# Load YOLO model
|
||
|
|
current_directory = os.getcwd()
|
||
|
|
model_path = current_directory + '\\datasets\\runs\\detect\\train8\\weights\\last.pt'
|
||
|
|
model = YOLO(model_path)
|
||
|
|
|
||
|
|
@app.get("/", response_class=HTMLResponse)
|
||
|
|
async def home(request: Request):
|
||
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
||
|
|
|
||
|
|
@app.post("/upload")
|
||
|
|
async def upload_image(file: UploadFile = File(...)):
|
||
|
|
# Read the uploaded file
|
||
|
|
contents = await file.read()
|
||
|
|
|
||
|
|
# Convert to numpy array
|
||
|
|
nparr = np.frombuffer(contents, np.uint8)
|
||
|
|
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||
|
|
|
||
|
|
# Run YOLO prediction
|
||
|
|
results = model(img)
|
||
|
|
|
||
|
|
# Draw bounding boxes
|
||
|
|
annotated_img = results[0].plot()
|
||
|
|
|
||
|
|
# Convert to base64
|
||
|
|
_, img_encoded = cv2.imencode('.jpg', annotated_img)
|
||
|
|
img_base64 = base64.b64encode(img_encoded).decode('utf-8')
|
||
|
|
|
||
|
|
return {"image": img_base64}
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
uvicorn.run(app, host="localhost", port=8000)
|