File size: 3,313 Bytes
c52b1e1 b57bea3 c52b1e1 0802393 c52b1e1 8b7fae6 51a67f0 8b7fae6 c52b1e1 9cd86d8 c52b1e1 f746146 c52b1e1 f746146 c52b1e1 4114862 c52b1e1 51a67f0 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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
import os
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('/usr/src/ultralytics/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"/tmp/{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"/tmp/{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("/test")
async def read_root():
return {"message": "TEST"}
@app.get("/")
async def read_root():
# Download image content from URL
image_url = "https://s1.eestatic.com/2018/04/10/actualidad/actualidad_298735505_130172549_1000x500.jpg"
response = requests.get(image_url)
response.raise_for_status()
# Open the image using PIL
image = Image.open(io.BytesIO(response.content))
# Perform YOLO prediction
predictions = predict_yolo(image_url)
# 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"/tmp/{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})
return {"message": "Hello, this is a YOLO prediction API using FastAPI!"} |