|
|
import cv2
|
|
|
from ultralytics import YOLO
|
|
|
import numpy as np
|
|
|
import base64
|
|
|
|
|
|
class DataProcessor:
|
|
|
def __init__(self, model_path='./../app/models/yolov9e_object_classification.pt'):
|
|
|
self.model = YOLO(model_path)
|
|
|
|
|
|
def __decode_image_from_base64(self, data, color=True):
|
|
|
decoded = base64.b64decode(data)
|
|
|
nparr = np.frombuffer(decoded, np.uint8)
|
|
|
return cv2.imdecode(nparr, cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE)
|
|
|
|
|
|
def inline_detection(self, data):
|
|
|
color_matrix = self.__decode_image_from_base64(data['rgb'], color=True)
|
|
|
depth_matrix = self.__decode_image_from_base64(data['depth'], color=False)
|
|
|
|
|
|
if depth_matrix is None or color_matrix is None:
|
|
|
print("Не удалось получить изображения. Проверьте подключение камеры.")
|
|
|
return [], None
|
|
|
|
|
|
results = self.model(color_matrix)
|
|
|
annotated_image = results[0].plot()
|
|
|
|
|
|
detections_info = []
|
|
|
for detection in results[0].boxes:
|
|
|
x1, y1, x2, y2 = map(int, detection.xyxy[0])
|
|
|
class_id = int(detection.cls[0])
|
|
|
class_name = self.model.names[class_id]
|
|
|
depth_values = depth_matrix[y1:y2, x1:x2]
|
|
|
mean_depth = np.mean(depth_values)
|
|
|
|
|
|
detections_info.append({
|
|
|
'class_id': class_id,
|
|
|
'class_name': class_name,
|
|
|
'bbox': [x1, y1, x2, y2],
|
|
|
'mean_depth': float(mean_depth)
|
|
|
})
|
|
|
|
|
|
cv2.putText(annotated_image, f'{class_name} {mean_depth:.2f}m', (x1, y1 - 10),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
|
|
return detections_info, annotated_image |
