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!"}