WorldFrame

import pyzed.sl as sl
import cv2
import numpy as np
from tqdm import tqdm

class SVOViewer:
    def __init__(self, svo_file_path):
        self.svo_file_path = svo_file_path
        self.zed = sl.Camera()
        self.clicked_points = [(1376, 1070), (1423, 791)]  # puntos por defecto
        self.point_cloud = sl.Mat()
        self.last_image = None
        self.image_raw = None
        self.runtime_params = self._get_runtime_params()
        self.init_params = self._get_init_params()

    def _get_init_params(self):
        init_params = sl.InitParameters()
        init_params.set_from_svo_file(self.svo_file_path)
        init_params.svo_real_time_mode = False
        init_params.camera_resolution = sl.RESOLUTION.HD2K
        init_params.depth_mode = sl.DEPTH_MODE.NEURAL_PLUS
        init_params.coordinate_units = sl.UNIT.METER
        init_params.coordinate_system = sl.COORDINATE_SYSTEM.RIGHT_HANDED_Z_UP
        init_params.depth_stabilization = 15
        init_params.depth_minimum_distance = 0.5
        init_params.depth_maximum_distance = 9.0
        init_params.enable_image_enhancement = True
        return init_params

    def _get_runtime_params(self):
        runtime_params = sl.RuntimeParameters()
        runtime_params.measure3D_reference_frame = sl.REFERENCE_FRAME.WORLD
        return runtime_params

    def open_camera(self):
        err = self.zed.open(self.init_params)
        if err != sl.ERROR_CODE.SUCCESS:
            raise RuntimeError(f"Error al abrir el archivo SVO: {self.svo_file_path}")

        tracking_params = sl.PositionalTrackingParameters()
        tracking_params.set_gravity_as_origin = True
        tracking_params.enable_imu_fusion = True
        tracking_params.mode = sl.POSITIONAL_TRACKING_MODE.GEN_3
        self.zed.enable_positional_tracking(tracking_params)

        self.zed.set_svo_position(25500)  # empezar en un punto medio si hace falta
        print("[INFO] Cámara abierta correctamente.")

    def _mouse_callback(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            if len(self.clicked_points) >= 2:
                self.clicked_points = []
            self.clicked_points.append((x, y))
            self._update_display()

    def _draw_text_with_background(self, image, text, org, font=cv2.FONT_HERSHEY_SIMPLEX,
                                   font_scale=0.75, text_color=(0, 255, 0), bg_color=(0, 0, 0),
                                   thickness=2, padding=5):
        (text_width, text_height), _ = cv2.getTextSize(text, font, font_scale, thickness)
        x, y = org
        cv2.rectangle(image,
                      (x - padding, y - text_height - padding),
                      (x + text_width + padding, y + padding),
                      bg_color, cv2.FILLED)
        cv2.putText(image, text, (x, y), font, font_scale, text_color, thickness, cv2.LINE_AA)

    def _update_display(self):
        if self.image_raw is None:
            return

        img_disp = self.image_raw.copy()
        points_3d = []
        points_rotated = []

        # Obtener rotación de la cámara respecto al mundo
        pose = sl.Pose()
        # tracking_state = self.zed.get_position(pose, sl.REFERENCE_FRAME.WORLD)
        # rotation_matrix = pose.get_rotation_matrix().r

        for idx, (x, y) in enumerate(self.clicked_points):
            point = self.point_cloud.get_value(x, y)
            xyz = point[1][:3]
            xyz_np = np.array(xyz)

            if np.isfinite(xyz_np).all():
                # xyz_rotated = np.dot(rotation_matrix, xyz_np)
                text_lines = [
                    f"P{idx+1}: px=({x},{y})",
                    f"   XYZ : X={xyz[0]:.2f} Y={xyz[1]:.2f} Z={xyz[2]:.2f} m",
                    # f"   Rot.: Xr={xyz_rotated[0]:.2f} Yr={xyz_rotated[1]:.2f} Zr={xyz_rotated[2]:.2f} m"
                ]
                points_3d.append(xyz_np)
                # points_rotated.append(xyz_rotated)
            else:
                text_lines = [
                    f"P{idx+1}: px=({x},{y})",
                    "   XYZ : X=NaN Y=NaN Z=NaN",
                    "   Rot.: Xr=NaN Yr=NaN Zr=NaN"
                ]

            cv2.circle(img_disp, (x, y), 6, (255, 0, 0), -1)
            cv2.putText(img_disp, f"P{idx+1}", (x + 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)

            line_spacing = 35
            block_spacing = 120
            base_y = 30 + idx * block_spacing

            for i, line in enumerate(text_lines):
                y_offset = base_y + i * line_spacing
                self._draw_text_with_background(img_disp, line, (30, y_offset),
                                                font_scale=0.75, text_color=(0, 255, 0))

        if len(self.clicked_points) == 2 and len(points_3d) == 2:
            cv2.line(img_disp, self.clicked_points[0], self.clicked_points[1], (0, 0, 255), 2)

            dz_orig = points_3d[1][2] - points_3d[0][2]  # diferencia en altura
            dy_orig = points_3d[1][1] - points_3d[0][1]  # diferencia en profundidad
            pendiente_orig = abs(dz_orig / dy_orig) if dy_orig != 0 else float('inf')

            # dz_rot = points_rotated[1][2] - points_rotated[0][2]
            # dx_rot = points_rotated[1][0] - points_rotated[0][0]
            # dy_rot = points_rotated[1][1] - points_rotated[0][1]
            # avance_rotado = np.sqrt(dx_rot ** 2 + dy_rot ** 2)
            # pendiente_rot = abs(dz_rot / avance_rotado) if avance_rotado != 0 else float('inf')

            pendiente_orig_pct = pendiente_orig * 100
            # pendiente_rot_pct = pendiente_rot * 100
            pendiente_orig_deg = np.degrees(np.arctan(pendiente_orig))
            # pendiente_rot_deg = np.degrees(np.arctan(pendiente_rot))

            # text_dist = f"Avance original: {dy_orig:.2f} m | rotado: {avance_rotado:.2f} m"
            text_pend = (
                f"Pendiente: {pendiente_orig:.2f} → {pendiente_orig_pct:.1f}% → {pendiente_orig_deg:.2f}°   |   "
                # f"Rotada: {pendiente_rot:.2f} → {pendiente_rot_pct:.1f}% → {pendiente_rot_deg:.2f}°"
            )

            # self._draw_text_with_background(img_disp, text_dist, (30, 280), font_scale=0.75, text_color=(255, 255, 0))
            self._draw_text_with_background(img_disp, text_pend, (30, 320), font_scale=0.75, text_color=(0, 255, 255))

        self.last_image = img_disp
        cv2.imshow("Ojo izquierdo", self.last_image)

    def run(self):
        self.open_camera()

        image = sl.Mat()
        cv2.namedWindow("Ojo izquierdo", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Ojo izquierdo", 900, 900)
        cv2.setMouseCallback("Ojo izquierdo", self._mouse_callback)

        frame_count = self.zed.get_svo_number_of_frames()
        print("Mostrando vídeo. Pulsa 'Enter' para avanzar, 'q' para salir.")

        for _ in tqdm(range(frame_count), desc="Procesando frames del video"):
            if self.zed.grab(self.runtime_params) != sl.ERROR_CODE.SUCCESS:
                break

            self.zed.retrieve_image(image, sl.VIEW.LEFT)
            self.zed.retrieve_measure(self.point_cloud, sl.MEASURE.XYZRGBA)

            img_cv2 = image.get_data()
            img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_RGBA2BGR)
            self.image_raw = img_cv2
            self._update_display()

            while True:
                key = cv2.waitKey(0) & 0xFF
                if key == 13:  # Enter
                    break
                elif key == ord('q'):
                    self.zed.close()
                    cv2.destroyAllWindows()
                    return

        self.zed.close()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    viewer = SVOViewer("/home/cristian-decode/Escritorio/Pendientes_Patinete/B/soncladera.svo2_1.svo2")
    viewer.run()