SlopeWORLDFrame

 avatar
unknown
python
3 months ago
9.3 kB
9
Indexable
import pyzed.sl as sl
import cv2
import numpy as np
import struct
import laspy
import open3d as o3d
from tqdm import tqdm

def color_decode(x):
    x = x[0] if isinstance(x, np.ndarray) else x
    packed = struct.pack('f', x)
    r, g, b, a = struct.unpack('BBBB', packed)
    return r / 255.0, g / 255.0, b / 255.0

def pc_to_las(xyz, rgba, out_filename):
    numpy_matrix = np.asarray(xyz, dtype=np.float64)
    rgb = np.apply_along_axis(lambda x: color_decode(x), 2, rgba[..., np.newaxis])
    points = np.nan_to_num(numpy_matrix, nan=0, posinf=0, neginf=0)

    pc = o3d.geometry.PointCloud()
    pc.points = o3d.utility.Vector3dVector(points.reshape(-1, 3))
    pc.colors = o3d.utility.Vector3dVector(rgb.reshape(-1, 3))
    points = np.asarray(pc.points)
    colors = np.asarray(pc.colors)
    colors_scaled = (colors * 65535).astype(np.uint16)

    header = laspy.LasHeader(version="1.4", point_format=2)
    las = laspy.LasData(header)
    las.x = points[:, 0]
    las.y = points[:, 1]
    las.z = points[:, 2]
    las.red = colors_scaled[:, 0]
    las.green = colors_scaled[:, 1]
    las.blue = colors_scaled[:, 2]

    try:
        las.write(out_filename)
    except Exception as e:
        print(f"Error exportando {out_filename}: {e}")

class SVOViewer:
    def __init__(self, svo_file_path):
        self.svo_file_path = svo_file_path
        self.zed = sl.Camera()
        self.clicked_points = [(1096, 1116), (1081, 657)]  # puntos por defecto
        self.point_cloud = sl.Mat()
        self.last_image = None
        self.image_raw = None
        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 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(0)
        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 = []

        pose = sl.Pose()
        tracking_state = self.zed.get_position(pose, sl.REFERENCE_FRAME.WORLD)
        print(f"Tengo estado {tracking_state}")
        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]
            dy_orig = points_3d[1][1] - points_3d[0][1]
            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 _save_pointclouds_to_las(self):
        print("[INFO] Exportando nubes de puntos...")

        rgba_mat = self.point_cloud.get_data()
        xyz = rgba_mat[:, :, :3]
        rgba = rgba_mat[:, :, 3]

        pose = sl.Pose()
        self.zed.get_position(pose, sl.REFERENCE_FRAME.WORLD)
        rotation_matrix = pose.get_rotation_matrix().r
        rotated_xyz = np.dot(xyz.reshape(-1, 3), rotation_matrix.T).reshape(xyz.shape)

        pc_to_las(xyz, rgba, "pointcloud_original.las")
        pc_to_las(rotated_xyz, rgba, "pointcloud_rotated.las")

        print("[INFO] Nubes exportadas correctamente.")

    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, 'S' para guardar, 'q' para salir.")

        for _ in tqdm(range(frame_count), desc="Procesando frames del video"):
            if self.zed.grab() != 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
                elif key == ord('s'):
                    self._save_pointclouds_to_las()

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

if __name__ == "__main__":
    viewer = SVOViewer("/home/cristian-decode/Escritorio/Pendientes_Patinete/D/slope.svo2")
    viewer.run()
Editor is loading...
Leave a Comment