From cbb7d413842e0f5df4d8793de9ef73844241df65 Mon Sep 17 00:00:00 2001 From: lcaonmst Date: Mon, 20 Dec 2021 22:44:30 +0100 Subject: [PATCH 1/6] Add virtual camera and taking photo via user interface --- main.py | 4 +++ virtualcam/virtualcam.py | 78 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) create mode 100644 virtualcam/virtualcam.py diff --git a/main.py b/main.py index be1bfd3..6753ff9 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 def get_mode(): @@ -20,6 +21,8 @@ def get_mode(): ball_controller, paddle, wind_controllers = init_env_and_load_assets(p) +virtualcam = VirtualCam(p, [1, 1, 1], 240, 240) + if keyboard_mode: # add rotation speed controller @@ -40,6 +43,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/virtualcam/virtualcam.py b/virtualcam/virtualcam.py new file mode 100644 index 0000000..7b24e2d --- /dev/null +++ b/virtualcam/virtualcam.py @@ -0,0 +1,78 @@ +from typing import List +from utils.button import Button +from math import sqrt +from src.vision.camera import AbstractCameraService + + +# class VirtualCam: +# def __init__(self, p, pos: List[float]): +# 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 + 1) +# self.take_photo_button = Button( +# p.addUserDebugParameter("Take a photo", 1, 0, 0)) +# self.width = 240 +# self.height = 240 + + +# def check_and_take_photo(self): +# if self.take_photo_button.was_clicked(): +# self.client.getCameraImage( +# width=self.width, +# height=self.height, +# viewMatrix=self.view_matrix, +# projectionMatrix=self.projection_matrix) + + +class VirtualCam(AbstractCameraService): + 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 + 1) + self.take_photo_button = Button( + p.addUserDebugParameter("Take a photo", 1, 0, 0)) + self.width = 240 + self.height = 240 + + + def shape(self): + return self.width, self.height + + + def take_photo(self): + width, height, rgbImg, depthImg, segImg = self.client.getCameraImage( + width=self.width, + height=self.height, + viewMatrix=self.view_matrix, + projectionMatrix=self.projection_matrix) + return rgbImg, depthImg, segImg + + + def check_and_take_photo(self): + if self.take_photo_button.was_clicked(): + self.take_photo() + + + def intrinsics(self): + return None + + + def pose(self): + return None From 61525988e34cb7ac60132abff4bda0ba82c4226d Mon Sep 17 00:00:00 2001 From: lcaonmst Date: Mon, 20 Dec 2021 22:52:14 +0100 Subject: [PATCH 2/6] Refactor --- virtualcam/virtualcam.py | 28 ---------------------------- 1 file changed, 28 deletions(-) diff --git a/virtualcam/virtualcam.py b/virtualcam/virtualcam.py index 7b24e2d..4e77a9d 100644 --- a/virtualcam/virtualcam.py +++ b/virtualcam/virtualcam.py @@ -4,34 +4,6 @@ from src.vision.camera import AbstractCameraService -# class VirtualCam: -# def __init__(self, p, pos: List[float]): -# 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 + 1) -# self.take_photo_button = Button( -# p.addUserDebugParameter("Take a photo", 1, 0, 0)) -# self.width = 240 -# self.height = 240 - - -# def check_and_take_photo(self): -# if self.take_photo_button.was_clicked(): -# self.client.getCameraImage( -# width=self.width, -# height=self.height, -# viewMatrix=self.view_matrix, -# projectionMatrix=self.projection_matrix) - - class VirtualCam(AbstractCameraService): def __init__(self, p, pos: List[float], width: int, height: int): super().__init__() From 15d6c69985b81a6e486438445dfd2e04ffcfb641 Mon Sep 17 00:00:00 2001 From: lcaonmst Date: Mon, 10 Jan 2022 23:34:07 +0100 Subject: [PATCH 3/6] Recognize ball and paddle in the camera photo and compute its depth --- main.py | 1 - virtualcam/virtualcam.py | 75 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 71 insertions(+), 5 deletions(-) diff --git a/main.py b/main.py index 6753ff9..3ac7f2f 100644 --- a/main.py +++ b/main.py @@ -16,7 +16,6 @@ def get_mode(): return args.mode == 'keyboard' - keyboard_mode = get_mode() ball_controller, paddle, wind_controllers = init_env_and_load_assets(p) diff --git a/virtualcam/virtualcam.py b/virtualcam/virtualcam.py index 4e77a9d..4e04c2f 100644 --- a/virtualcam/virtualcam.py +++ b/virtualcam/virtualcam.py @@ -2,7 +2,9 @@ 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 VirtualCam(AbstractCameraService): def __init__(self, p, pos: List[float], width: int, height: int): @@ -17,11 +19,16 @@ def __init__(self, p, pos: List[float], width: int, height: int): fov=45.0, aspect=1.0, nearVal=max(0, dist - 1), - farVal=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)) def shape(self): @@ -29,12 +36,16 @@ def shape(self): def take_photo(self): - width, height, rgbImg, depthImg, segImg = self.client.getCameraImage( + 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) - return rgbImg, depthImg, segImg + 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 def check_and_take_photo(self): @@ -48,3 +59,59 @@ def intrinsics(self): def pose(self): return None + + + def get_depth(self, w: int, h: int): + return 2*self.depth_img[h,w] - 1 + + + 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 + pix_pos = asarray([x, y, z, 1]) + position = matmul(self.transform_matrix, pix_pos) + return position / position[3] + + + def search_for_ball(self): + avg = 0 + cnt = 0 + center = [0, 0] + for h in range(0, self.last_height): + for w in range(0, self.last_width): + if 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: + self.rgb_img[h,w][:3] = [85, 255, 0] + center[0] += h + center[1] += w + real_depth = self.far * self.near / (self.far - (self.far - self.near) * self.depth_img[h,w]) + avg += real_depth + cnt += 1 + avg /= cnt + center[0] /= cnt + center[1] /= cnt + print("Average depth: ", avg) + print("Center: ", center) + return avg, center + + + def search_for_paddle(self): + avg = 0 + cnt = 0 + center = [0, 0] + for h in range(0, self.last_height): + for w in range(0, self.last_width): + if max(self.rgb_img[h,w][1:3]) == 0 and self.rgb_img[h,w][0] > 10: + self.rgb_img[h,w][:3] = [0, 68, 255] + center[0] += h + center[1] += w + real_depth = self.far * self.near / (self.far - (self.far - self.near) * self.depth_img[h,w]) + avg += real_depth + cnt += 1 + avg /= cnt + center[0] /= cnt + center[1] /= cnt + print("Average depth: ", avg) + print("Center: ", center) + return avg, center + From c8a7e065efa2d4c87aaed82d0fd34a90995dff71 Mon Sep 17 00:00:00 2001 From: lcaonmst Date: Sun, 13 Mar 2022 19:39:06 +0100 Subject: [PATCH 4/6] Add feature translating 2d picture point to 3d world coordinates and refactor --- virtualcam/virtualcam.py | 87 +++++++++++++++++++++++++++++----------- 1 file changed, 63 insertions(+), 24 deletions(-) diff --git a/virtualcam/virtualcam.py b/virtualcam/virtualcam.py index 4e04c2f..b2bdb6d 100644 --- a/virtualcam/virtualcam.py +++ b/virtualcam/virtualcam.py @@ -6,7 +6,12 @@ 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) @@ -31,87 +36,121 @@ def __init__(self, p, pos: List[float], width: int, height: int): 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()) + # 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 None + return projection_matrix + # Returns matrix changing the position and orientation of origin. def pose(self): - return None - - - def get_depth(self, w: int, h: int): - return 2*self.depth_img[h,w] - 1 + 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): - avg = 0 cnt = 0 center = [0, 0] for h in range(0, self.last_height): for w in range(0, self.last_width): - if 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: + if self.is_ball_pixel(h, w): self.rgb_img[h,w][:3] = [85, 255, 0] center[0] += h center[1] += w - real_depth = self.far * self.near / (self.far - (self.far - self.near) * self.depth_img[h,w]) - avg += real_depth cnt += 1 - avg /= cnt + if cnt == 0: + return None center[0] /= cnt center[1] /= cnt - print("Average depth: ", avg) - print("Center: ", center) - return avg, center + 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): - avg = 0 cnt = 0 center = [0, 0] for h in range(0, self.last_height): for w in range(0, self.last_width): - if max(self.rgb_img[h,w][1:3]) == 0 and self.rgb_img[h,w][0] > 10: + if self.is_paddle_pixel(h, w): self.rgb_img[h,w][:3] = [0, 68, 255] center[0] += h center[1] += w - real_depth = self.far * self.near / (self.far - (self.far - self.near) * self.depth_img[h,w]) - avg += real_depth cnt += 1 - avg /= cnt + if cnt == 0: + return None center[0] /= cnt center[1] /= cnt - print("Average depth: ", avg) - print("Center: ", center) - return avg, center + 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 From cd34bb51c10858727742b8e0873dba41728ce145 Mon Sep 17 00:00:00 2001 From: lcaonmst Date: Sun, 13 Mar 2022 21:07:56 +0100 Subject: [PATCH 5/6] Connect virtual camera to ball tracker --- main.py | 2 +- trackers/ball_tracker.py | 16 ++++++++++++++-- utils/environment.py | 5 +++-- utils/pid_performer.py | 5 +++-- 4 files changed, 21 insertions(+), 7 deletions(-) diff --git a/main.py b/main.py index 113da68..b8daa33 100644 --- a/main.py +++ b/main.py @@ -33,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: diff --git a/trackers/ball_tracker.py b/trackers/ball_tracker.py index 96a0fde..4f8057b 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..e926d94 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,7 @@ 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 +81,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..aa88c23 100644 --- a/utils/pid_performer.py +++ b/utils/pid_performer.py @@ -2,14 +2,15 @@ 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 From e3e9eb6e67b831742c66f7b8b97139dc778e9164 Mon Sep 17 00:00:00 2001 From: Mateusz Date: Wed, 16 Mar 2022 13:13:49 +0100 Subject: [PATCH 6/6] Reformat code --- trackers/ball_tracker.py | 2 +- utils/environment.py | 7 +- utils/pid_performer.py | 8 +- virtualcam/virtualcam.py | 289 ++++++++++++++++++++------------------- 4 files changed, 159 insertions(+), 147 deletions(-) diff --git a/trackers/ball_tracker.py b/trackers/ball_tracker.py index 4f8057b..e00367f 100644 --- a/trackers/ball_tracker.py +++ b/trackers/ball_tracker.py @@ -33,7 +33,7 @@ def get_error_vector(self) -> List[float]: for ball_pos, paddle_pos in zip(ball_pos, paddle_pos) ][:2] else: - raise OutOfRange + 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 e926d94..e00cdf3 100644 --- a/utils/environment.py +++ b/utils/environment.py @@ -73,7 +73,12 @@ def load_paddle(p): def init_standard_pid_tools( - p: pybullet, ball: ABCBall, paddle: ABCPaddle, max_angle: float, min_angle: float, virtualcam: VirtualCam + 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) diff --git a/utils/pid_performer.py b/utils/pid_performer.py index aa88c23..b7d68a9 100644 --- a/utils/pid_performer.py +++ b/utils/pid_performer.py @@ -7,7 +7,13 @@ class PidPerformer: - def __init__(self, pybullet_client: pybullet, ball: ABCBall, paddle: ABCPaddle, virtualcam: VirtualCam): + 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, virtualcam diff --git a/virtualcam/virtualcam.py b/virtualcam/virtualcam.py index b2bdb6d..3f12cb7 100644 --- a/virtualcam/virtualcam.py +++ b/virtualcam/virtualcam.py @@ -10,147 +10,148 @@ # 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 - + # 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