Rune Harlyk
58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
import pybullet as pb
|
|
import sys
|
|
|
|
sys.path.append("../../../")
|
|
|
|
from raspberry_pi.src.Camera.CameraBase import CameraBase
|
|
|
|
viewMatrix = pb.computeViewMatrix(
|
|
cameraEyePosition=[0, 0, 1],
|
|
cameraTargetPosition=[0, 0, 0],
|
|
cameraUpVector=[0, 1, 0])
|
|
projectionMatrix = pb.computeProjectionMatrixFOV(
|
|
fov=160, aspect=1.0, nearVal=0.1, farVal=150.0)
|
|
|
|
class PyBulletCamera(CameraBase):
|
|
def __init__(self, view_matrix=viewMatrix, projection_matrix=projectionMatrix, width=640, height=480):
|
|
"""
|
|
Initializes the camera with PyBullet view and projection matrices.
|
|
:param view_matrix: The view matrix for the camera in PyBullet.
|
|
:param projection_matrix: The projection matrix for the camera.
|
|
:param width: Width of the camera image.
|
|
:param height: Height of the camera image.
|
|
"""
|
|
self.view_matrix = view_matrix
|
|
self.projection_matrix = projection_matrix
|
|
self.width = width
|
|
self.height = height
|
|
|
|
def get_image(self):
|
|
"""
|
|
Captures an image from the PyBullet simulation.
|
|
:return: An image frame captured from PyBullet.
|
|
"""
|
|
# Capture an image from the simulation
|
|
_, _, px, _, _ = pb.getCameraImage(
|
|
width=self.width, height=self.height,
|
|
viewMatrix=self.view_matrix,
|
|
projectionMatrix=self.projection_matrix)
|
|
|
|
# Convert PyBullet's RGBA image to an OpenCV BGR image
|
|
rgb_image = np.array(px, dtype=np.uint8)
|
|
rgb_image = np.reshape(rgb_image, (self.height, self.width, -1))
|
|
bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGBA2BGR)
|
|
|
|
return bgr_image
|
|
|
|
def update_view_matrix(self, model_id):
|
|
position, orientation = pb.getBasePositionAndOrientation(model_id)
|
|
position = [position[0], position[1]-0.1, position[2]]
|
|
orientation = [orientation[0], orientation[1], orientation[2]]
|
|
|
|
# Calculate the camera's view matrix based on the new position and orientation
|
|
camera_eye_pos = position # This should be adjusted based on your specific relative position requirements
|
|
camera_target_pos = [position[0] + orientation[0], position[1] + orientation[1], position[2] + orientation[2]]
|
|
self.view_matrix = pb.computeViewMatrix(camera_eye_pos, camera_target_pos, [0, 1, 0])
|