mrisdi/object-detection-api

1from fastapi import FastAPI, File, UploadFile
2from fastapi.middleware.cors import CORSMiddleware
3from pydantic import BaseModel
4import numpy as np
5import cv2
6from ultralytics import YOLO
7from PIL import Image
8import base64
9from io import BytesIO
10
11app = FastAPI()
12model = YOLO("yolov8n.pt")
13
14origins = ["*"]
15
16app.add_middleware(
17    CORSMiddleware,
18    allow_origins=origins,
19    allow_credentials=True,
20    allow_methods=["*"],
21    allow_headers=["*"],
22)
23
24@app.post("/detect/")
25async def detect_objects(file: UploadFile):
26 # Process the uploaded image for object detection
27 image_bytes = await file.read()
28 image = np.frombuffer(image_bytes, dtype=np.uint8)
29 image = cv2.imdecode(image, cv2.IMREAD_COLOR)
30 
31 # Perform object detection with YOLOv8
32 detections = model(image)
33 
34 return detections[0].tojson()
35
36class ImageData(BaseModel):
37    image: str  # Data gambar dalam format base64
38    
39@app.post("/uploadimage")
40async def upload_image(image_data: ImageData):
41    # Mengonversi base64 ke gambar
42    base64_data = image_data.image.split(',')[1]
43    image = Image.open(BytesIO(base64.b64decode(base64_data)))
44    detections = model(image)
45    
46    return detections[0].tojson()
47