File size: 4,019 Bytes
4674250 | 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 97 98 99 100 | import io
import base64
import numpy as np
import cv2
from PIL import Image as PILImage
from ultralytics import YOLO
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from matplotlib import cm
from fastapi.middleware.cors import CORSMiddleware
app = FastAPI()
# Add CORS middleware to allow requests from anywhere
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Allow all origins
allow_credentials=True,
allow_methods=["*"], # Allow all HTTP methods
allow_headers=["*"], # Allow all headers
)
class ImageRequest(BaseModel):
image: str # Base64-encoded image string
class ObjectDetectionSystem:
def __init__(self, model_path='Model/yolov8l.pt'):
self.yolo_model = YOLO(model_path)
self.color_map = cm.get_cmap("tab20", len(self.yolo_model.names)) # Generate distinct colors
def detect_objects(self, image_base64: str, confidence_threshold: float = 0.5):
try:
# Decode Base64 and convert to NumPy array
image_data = base64.b64decode(image_base64)
pil_image = PILImage.open(io.BytesIO(image_data)).convert("RGB")
image_np = np.array(pil_image)
# Resize image for consistent processing
input_size = 640 # Example size for YOLO models
height, width, _ = image_np.shape
scale = input_size / max(height, width)
resized_image = cv2.resize(image_np, (int(width * scale), int(height * scale)))
# Perform object detection
results = self.yolo_model(resized_image)
# Draw results on image
output_image = resized_image.copy()
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
confidence = float(box.conf[0])
class_id = int(box.cls[0])
class_name = self.yolo_model.names[class_id]
if confidence > confidence_threshold:
# Generate color for this class
color = tuple(int(c * 255) for c in self.color_map(class_id)[:3])
# Draw bounding box
cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
# Add label with confidence
label = f'{class_name} ({confidence:.2f})'
label_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
label_y = max(y1, label_size[1] + 10)
cv2.rectangle(output_image, (x1, label_y - label_size[1] - 10),
(x1 + label_size[0], label_y + baseline - 10), color, -1)
cv2.putText(output_image, label, (x1, label_y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
# Convert back to Base64
output_pil = PILImage.fromarray(output_image)
buffered = io.BytesIO()
output_pil.save(buffered, format="PNG")
encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
return encoded_image
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# Initialize the detector with a more advanced model
detector = ObjectDetectionSystem('Model/yolov8l.pt')
@app.post("/detect")
async def detect_objects(request: ImageRequest):
try:
result_image = detector.detect_objects(request.image)
return {"processed_image": result_image}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def root():
return {"message": "Object Detection API is running"}
|