You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

105 lines
4.1 KiB

from Neko.RealSense import RealSenseController
import cv2
from ultralytics import YOLO
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from datetime import datetime
class CameraProcessor:
def __init__(self, model_path='./models/yolov9e_object_classification.pt'):
self.controller = RealSenseController()
self.model = YOLO(model_path)
def inline_detection(self):
depth_matrix = self.controller.acquisition.get_depth_image()
color_matrix = self.controller.acquisition.get_color_image()
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
def plot_3d_bounding_boxes(self, detections_info):
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
for detection in detections_info:
x1, y1, x2, y2 = detection['bbox']
mean_depth = detection['mean_depth']
class_name = detection['class_name']
width = x2 - x1
height = y2 - y1
depth = 0.05
box_points = [
[x1, y1, mean_depth],
[x2, y1, mean_depth],
[x2, y2, mean_depth],
[x1, y2, mean_depth],
[x1, y1, mean_depth + depth],
[x2, y1, mean_depth + depth],
[x2, y2, mean_depth + depth],
[x1, y2, mean_depth + depth]
]
faces = [
[box_points[0], box_points[1], box_points[5], box_points[4]],
[box_points[3], box_points[2], box_points[6], box_points[7]],
[box_points[0], box_points[3], box_points[7], box_points[4]],
[box_points[1], box_points[2], box_points[6], box_points[5]],
[box_points[0], box_points[1], box_points[2], box_points[3]],
[box_points[4], box_points[5], box_points[6], box_points[7]]
]
box = Poly3DCollection(faces, facecolors='cyan', linewidths=1, edgecolors='r', alpha=0.25)
ax.add_collection3d(box)
ax.text((x1 + x2) / 2, (y1 + y2) / 2, mean_depth + depth, f'{class_name} {mean_depth:.2f}m',
color='blue', fontsize=8)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Depth (m)")
ax.set_title("3D Bounding Boxes with Depth Information")
all_x = [d['bbox'][0] for d in detections_info] + [d['bbox'][2] for d in detections_info]
all_y = [d['bbox'][1] for d in detections_info] + [d['bbox'][3] for d in detections_info]
all_z = [d['mean_depth'] for d in detections_info] + [d['mean_depth'] + 0.05 for d in detections_info]
ax.set_xlim(min(all_x) - 50, max(all_x) + 50)
ax.set_ylim(min(all_y) - 50, max(all_y) + 50)
ax.set_zlim(min(all_z) - 50, max(all_z) + 50)
plt.show()
def test_time_delta(self):
start_time = datetime.now()
detections_info, _ = self.inline_detection()
end_time = datetime.now()
delta_time = end_time - start_time
print(f"Время выполнения: {delta_time}")