34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
import io
|
||
from PIL import Image
|
||
from ultralytics import YOLO
|
||
|
||
|
||
class ModelHandler:
|
||
def __init__(self):
|
||
"""加载 YOLOv11 模型"""
|
||
self.model = YOLO("/opt/nuclio/best.pt") # 确保路径正确
|
||
|
||
def infer(self, image_data, threshold=0.3):
|
||
"""
|
||
执行推理
|
||
:param image_data: 图片的二进制数据
|
||
:param threshold: 置信度阈值（默认0.3）
|
||
:return: 符合阈值的检测结果
|
||
"""
|
||
image = Image.open(io.BytesIO(image_data))
|
||
results = self.model(image)
|
||
|
||
detections = []
|
||
for result in results:
|
||
for box in result.boxes.data.tolist():
|
||
x1, y1, x2, y2, score, class_id = box
|
||
if score >= threshold: # 过滤低置信度目标
|
||
detections.append({
|
||
"confidence": score,
|
||
"label": self.model.names[int(class_id)],
|
||
"points": [x1, y1, x2, y2],
|
||
"type": "rectangle",
|
||
})
|
||
|
||
return detections
|