Spaces:
Runtime error
Runtime error
File size: 1,671 Bytes
a86a345 |
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 |
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io
import cv2
import torch
from ultralytics import YOLO
import numpy as np
from typing import List
# Initialize the FastAPI app
app = FastAPI()
# Load YOLOv8 model
model = YOLO("best_yolov8m_model.pt")
# Class names for algae
class_names = ["blue-green-algae", "brown-algae", "red-algae"]
def predict(image: Image.Image):
# Convert image to OpenCV format
image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
# Perform inference
results = model(image_cv)
# Process the results
detected_objects = []
for result in results:
for box in result.boxes:
class_id = int(box.cls)
conf = box.conf
bbox = box.xyxy
detected_objects.append({
"class": class_names[class_id],
"confidence": float(conf),
"bbox": bbox.tolist()
})
return detected_objects, results
@app.post("/predict/")
async def predict_api(file: UploadFile = File(...)):
# Read image from the uploaded file
image = Image.open(io.BytesIO(await file.read()))
# Run inference
detections, results = predict(image)
# Annotate the image with bounding boxes
result_image = results[0].plot() # Plot the bounding boxes
# Convert to bytes to return as response
_, encoded_image = cv2.imencode('.jpg', result_image)
result_image_bytes = encoded_image.tobytes()
return {
"image": result_image_bytes, # Encoded result image with bounding boxes
"detections": detections # List of detected algae types and bounding boxes
}
|