Virtual camera

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()

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()
 

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,15 +73,20 @@ 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)
     kd_slider = p.addUserDebugParameter("D", 0, 6000, 50)
 
     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),

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

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