diff --git a/main.py b/main.py index ede799d..b8daa33 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ import pybullet as p from utils.environment import init_env_and_load_assets, update_wind_controllers +from virtualcam.virtualcam import VirtualCam from utils.pid_performer import PidPerformer @@ -22,6 +23,7 @@ def get_mode(): keyboard_mode = mode == "keyboard" ball_controller, ball, paddle, wind_controllers = init_env_and_load_assets(p) +virtualcam = VirtualCam(p, [1, 1, 1], 240, 240) if keyboard_mode: # add rotation speed controller @@ -31,7 +33,7 @@ def get_mode(): paddle.create_joint_controllers() if pid_flag: - pid_performer = PidPerformer(p, ball, paddle) + pid_performer = PidPerformer(p, ball, paddle, virtualcam) while True: @@ -48,6 +50,7 @@ def get_mode(): if ball_controller.should_throw_ball(): ball_controller.throw_ball(paddle.get_center_position()) + virtualcam.check_and_take_photo() update_wind_controllers(p, *wind_controllers) p.stepSimulation() diff --git a/trackers/ball_tracker.py b/trackers/ball_tracker.py index 96a0fde..e00367f 100644 --- a/trackers/ball_tracker.py +++ b/trackers/ball_tracker.py @@ -2,17 +2,19 @@ from ball.abc_ball import ABCBall from paddle.abc_paddle import ABCPaddle +from virtualcam.virtualcam import VirtualCam SMOOTHER_SIZE = 7 N = 3 class BallTracker: - def __init__(self, ball: ABCBall, paddle: ABCPaddle): + def __init__(self, ball: ABCBall, paddle: ABCPaddle, virtualcam: VirtualCam): self.ball = ball self.paddle = paddle + self.virtualcam = virtualcam - def get_error_vector(self) -> List[float]: + def get_error_vector2(self) -> List[float]: ball_pos = self.ball.get_position() if self.paddle.check_if_in_range(ball_pos): paddle_pos = self.paddle.get_center_position() @@ -23,6 +25,16 @@ def get_error_vector(self) -> List[float]: else: raise OutOfRange + def get_error_vector(self) -> List[float]: + ball_pos, paddle_pos = self.virtualcam.get_objects_location() + if self.paddle.check_if_in_range(ball_pos): + return [ + ball_pos - paddle_pos + for ball_pos, paddle_pos in zip(ball_pos, paddle_pos) + ][:2] + else: + raise OutOfRange + def get_ball_position(self) -> List[float]: return self.ball.get_position() diff --git a/utils/environment.py b/utils/environment.py index 491807d..e00cdf3 100644 --- a/utils/environment.py +++ b/utils/environment.py @@ -12,6 +12,7 @@ from paddle.paddle import Paddle from trackers.ball_tracker import BallTracker from utils.button import Button +from virtualcam.virtualcam import VirtualCam G = 9.81 BASE_PLANE_POSITION = [0, 0, -0.1] @@ -72,7 +73,12 @@ def load_paddle(p): def init_standard_pid_tools( - p: pybullet, ball: ABCBall, paddle: ABCPaddle, max_angle: float, min_angle: float + p: pybullet, + ball: ABCBall, + paddle: ABCPaddle, + max_angle: float, + min_angle: float, + virtualcam: VirtualCam, ) -> Tuple[Dict[str, float], Button, PIDBalancer]: kp_slider = p.addUserDebugParameter("P", 0, 500, 60) ki_slider = p.addUserDebugParameter("I", 0, 50, 1) @@ -80,7 +86,7 @@ def init_standard_pid_tools( set_pid_button = Button(p.addUserDebugParameter("Change PID", 1, 0, 0)) - engine_tracker = BallTracker(ball, paddle) + engine_tracker = BallTracker(ball, paddle, virtualcam) pid_controller = PIDController( p.readUserDebugParameter(kp_slider), diff --git a/utils/pid_performer.py b/utils/pid_performer.py index 2b5ff22..b7d68a9 100644 --- a/utils/pid_performer.py +++ b/utils/pid_performer.py @@ -2,14 +2,21 @@ from pid.pid_balancer import OUT_OF_RANGE from paddle.abc_paddle import ABCPaddle from utils.environment import init_standard_pid_tools +from virtualcam.virtualcam import VirtualCam import pybullet class PidPerformer: - def __init__(self, pybullet_client: pybullet, ball: ABCBall, paddle: ABCPaddle): + def __init__( + self, + pybullet_client: pybullet, + ball: ABCBall, + paddle: ABCPaddle, + virtualcam: VirtualCam, + ): self.pybullet_client = pybullet_client self.pid_sliders, self.pid_button, self.pid_balancer = init_standard_pid_tools( - pybullet_client, ball, paddle, 55, -55 + pybullet_client, ball, paddle, 55, -55, virtualcam ) self.paddle = paddle self.pid_balancer.controller.debug = True diff --git a/virtualcam/virtualcam.py b/virtualcam/virtualcam.py new file mode 100644 index 0000000..3f12cb7 --- /dev/null +++ b/virtualcam/virtualcam.py @@ -0,0 +1,157 @@ +from typing import List +from utils.button import Button +from math import sqrt +from src.vision.camera import AbstractCameraService +from numpy import asarray, matmul +from numpy.linalg import inv +from PIL import Image + + +# Class represents a virtual camera. +class VirtualCam(AbstractCameraService): + + # Initialize the virtual camera by pybullet client, + # camera position and the shape of photos. + def __init__(self, p, pos: List[float], width: int, height: int): + super().__init__() + dist = sqrt(pos[0] ** 2 + pos[1] ** 2) + self.client = p + self.view_matrix = p.computeViewMatrix( + cameraEyePosition=pos, + cameraTargetPosition=[0, 0, 0.5], + cameraUpVector=[pos[0], pos[1], pos[2] + 1], + ) + self.projection_matrix = p.computeProjectionMatrixFOV( + fov=45.0, aspect=1.0, nearVal=max(0, dist - 1), farVal=dist + 3 + ) + self.near = max(0, dist - 1) + self.far = dist + 3 + self.take_photo_button = Button( + p.addUserDebugParameter("Take a photo", 1, 0, 0) + ) + self.width = 240 + self.height = 240 + projection_matrix = asarray(self.projection_matrix).reshape([4, 4], order="F") + view_matrix = asarray(self.view_matrix).reshape([4, 4], order="F") + self.transform_matrix = inv(matmul(projection_matrix, view_matrix)) + + # Returns the shape of the photo. + def shape(self): + return self.width, self.height + + # Takes and returns photos taken by virtual camera. + def take_photo(self): + ( + self.last_width, + self.last_height, + self.rgb_img, + self.depth_img, + self.seg_img, + ) = self.client.getCameraImage( + width=self.width, + height=self.height, + viewMatrix=self.view_matrix, + projectionMatrix=self.projection_matrix, + ) + # print(self.search_for_ball()) + # print(self.search_for_paddle()) + im = Image.fromarray(self.rgb_img) + im.save("image.png") + return self.rgb_img, self.depth_img, self.seg_img + + # Takes the photo if button is clicked. + def check_and_take_photo(self): + if self.take_photo_button.was_clicked(): + self.take_photo() + self.get_objects_location(take_photo=False) + + # Returns matrix transposing 3d world coordinates with changed origin to 2d picture. + def intrinsics(self): + return projection_matrix + + # Returns matrix changing the position and orientation of origin. + def pose(self): + return view_matrix + + # Translates pixel from the previous picture into 3d world coordinates. + # https://stackoverflow.com/questions/59128880/getting-world-coordinates-from-opengl-depth-buffer + def translate_to_origin_frame(self, w: int, h: int): + x = (2 * w - self.last_width) / self.last_width + y = -(2 * h - self.last_height) / self.last_height + z = 2 * float(self.depth_img[h, w]) - 1 + + print("pixel position: ", w, h) + print("vector position: ", x, y, z) + pix_pos = asarray([x, y, z, 1]) + position = matmul(self.transform_matrix, pix_pos) + return position / position[3] + + # Checks whether pixel from the previous picture belongs to the ball. + def is_ball_pixel(self, h: int, w: int): + return ( + max(self.rgb_img[h, w][:3]) - min(self.rgb_img[h, w][:3]) < 2 + and float(self.depth_img[h, w]) < 0.9 + and min(self.rgb_img[h, w][:3]) < 225 + and min(self.rgb_img[h, w][:3]) > 50 + ) + + # Checks whether pixel from the previous picture belongs to the paddle. + def is_paddle_pixel(self, h: int, w: int): + return max(self.rgb_img[h, w][1:3]) == 0 and self.rgb_img[h, w][0] > 10 + + # Finds the center of the ball in the previous picture (in terms of the pixels). + # Returns none if have not found it. + def search_for_ball(self): + cnt = 0 + center = [0, 0] + for h in range(0, self.last_height): + for w in range(0, self.last_width): + if self.is_ball_pixel(h, w): + self.rgb_img[h, w][:3] = [85, 255, 0] + center[0] += h + center[1] += w + cnt += 1 + if cnt == 0: + return None + center[0] /= cnt + center[1] /= cnt + return center + + # Finds the center of the paddle in the previous picture (in terms of the pixels). + # Return none if have not found it. + def search_for_paddle(self): + cnt = 0 + center = [0, 0] + for h in range(0, self.last_height): + for w in range(0, self.last_width): + if self.is_paddle_pixel(h, w): + self.rgb_img[h, w][:3] = [0, 68, 255] + center[0] += h + center[1] += w + cnt += 1 + if cnt == 0: + return None + center[0] /= cnt + center[1] /= cnt + return center + + # Finds the 3d world coordinates of the ball and paddle in the previous picture. + # If take_photo is set, takes a photo before + def get_objects_location(self, take_photo=True): + if take_photo: + self.take_photo() + # print("Searching ball...") + center_ball = self.search_for_ball() + # print("Searching paddle...") + center_paddle = self.search_for_paddle() + # print("ball: ") + pos_ball = self.translate_to_origin_frame( + int(center_ball[1]), int(center_ball[0]) + ) + # print("paddle: ") + pos_paddle = self.translate_to_origin_frame( + int(center_paddle[1]), int(center_paddle[0]) + ) + print("Ball position: ", pos_ball) + print("Paddle position: ", pos_paddle) + return center_ball, center_paddle