Spaces:
Sleeping
Sleeping
from fastapi import FastAPI, File, UploadFile | |
from fastapi.middleware.cors import CORSMiddleware | |
from pydantic import BaseModel | |
import numpy as np | |
import cv2 | |
from ultralytics import YOLO | |
from PIL import Image | |
import base64 | |
from io import BytesIO | |
app = FastAPI() | |
model = YOLO("yolov8n.pt") | |
origins = ["*"] | |
app.add_middleware( | |
CORSMiddleware, | |
allow_origins=origins, | |
allow_credentials=True, | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
async def detect_objects(file: UploadFile): | |
# Process the uploaded image for object detection | |
image_bytes = await file.read() | |
image = np.frombuffer(image_bytes, dtype=np.uint8) | |
image = cv2.imdecode(image, cv2.IMREAD_COLOR) | |
# Perform object detection with YOLOv8 | |
detections = model(image) | |
return detections[0].tojson() | |
class ImageData(BaseModel): | |
image: str # Data gambar dalam format base64 | |
async def upload_image(image_data: ImageData): | |
# Mengonversi base64 ke gambar | |
base64_data = image_data.image.split(',')[1] | |
image = Image.open(BytesIO(base64.b64decode(base64_data))) | |
detections = model(image) | |
return detections[0].tojson() | |