SlopeWORLDFrame

 avatar
unknown
python
10 months ago
9.3 kB
28
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