File size: 2,528 Bytes
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 74 |
from fastapi import FastAPI, File, UploadFile
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
app = FastAPI()
# Mount a static directory for serving images
app.mount("/static", StaticFiles(directory="static"), name="static")
# Templates for rendering HTML
templates = Jinja2Templates(directory="templates")
def predict_yolo(image_path):
# Assuming you have a YOLO API endpoint that accepts an image and returns predictions
with open(image_path, "rb") as file:
files = {"file": file}
# Load a model
model = YOLO('yolov8n.pt') # pretrained YOLOv8n model
# Run batched inference on a list of images
results = model(image_path) # return a list of Results objects
# Process results list
for result in results:
boxes = result.boxes # Boxes object for bbox outputs
# masks = result.masks # Masks object for segmentation masks outputs
# keypoints = result.keypoints # Keypoints object for pose outputs
# probs = result.probs # Probs object for classification outputs
predictions = boxes.json()
return predictions
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("/uploadfile/")
async def create_upload_file(file: UploadFile = File(...)):
contents = await file.read()
image = Image.open(io.BytesIO(contents))
# Save the image to a static directory
save_path = f"static/{file.filename}"
image.save(save_path)
# Perform YOLO prediction
predictions = predict_yolo(save_path)
# Draw bounding boxes on the image
image_np = np.array(image)
image_with_boxes = draw_boxes(image_np, predictions)
# Save the image with bounding boxes
image_with_boxes_path = f"static/{file.filename.split('.')[0]}_with_boxes.jpg"
cv2.imwrite(image_with_boxes_path, cv2.cvtColor(image_with_boxes, cv2.COLOR_RGB2BGR))
# Render the HTML with the image and bounding boxes
return templates.TemplateResponse("prediction.html", {"request": file, "image_path": image_with_boxes_path})
@app.get("/")
async def read_root():
return {"message": "Hello, this is a YOLO prediction API using FastAPI!"} |