pip install scipy

import cv2
import numpy as np
from scipy.spatial.transform import Rotation as R

class PreciseGeoCalculator:
    def __init__(self, camera_matrix, dist_coeffs):
        self.camera_matrix = camera_matrix
        self.dist_coeffs = dist_coeffs
    
    def pixel_to_geo(self, pixel_x, pixel_y, drone_pos, drone_attitude):
        """
        精确的像素坐标到地理坐标的转换
        drone_pos: [longitude, latitude, altitude]
        drone_attitude: [yaw, pitch, roll] (度)
        """
        # 1. 畸变校正
        undistorted_x, undistorted_y = self.undistort_pixel(pixel_x, pixel_y)
        
        # 2. 像素坐标转相机坐标（归一化平面）
        x_cam, y_cam = self.pixel_to_camera(undistorted_x, undistorted_y)
        
        # 3. 相机坐标转无人机坐标
        x_body, y_body, z_body = self.camera_to_body(x_cam, y_cam, drone_attitude)
        
        # 4. 无人机坐标转世界坐标（ENU坐标系）
        x_enu, y_enu, z_enu = self.body_to_enu(x_body, y_body, z_body, drone_attitude)
        
        # 5. ENU坐标转地理坐标
        target_lat, target_lon = self.enu_to_geo(
            x_enu, y_enu, 
            drone_pos[1], drone_pos[0]  # drone_lat, drone_lon
        )
        
        return target_lat, target_lon
    
    def undistort_pixel(self, pixel_x, pixel_y):
        """畸变校正"""
        point = np.array([[pixel_x, pixel_y]], dtype=np.float32)
        undistorted = cv2.undistortPoints(
            point, self.camera_matrix, self.dist_coeffs, P=self.camera_matrix
        )
        return undistorted[0][0]
    
    def pixel_to_camera(self, x, y):
        """像素坐标转相机坐标系"""
        fx, fy = self.camera_matrix[0, 0], self.camera_matrix[1, 1]
        cx, cy = self.camera_matrix[0, 2], self.camera_matrix[1, 2]
        
        x_cam = (x - cx) / fx
        y_cam = (y - cy) / fy
        
        return x_cam, y_cam
    
    def camera_to_body(self, x_cam, y_cam, attitude):
        """
        相机坐标系 → 无人机机体坐标系
        考虑相机安装角度（假设相机与机体坐标系对齐）
        """
        # 将归一化坐标转换为3D射线（假设深度为1）
        z_depth = 1.0  # 单位深度
        
        # 相机坐标系下的3D点
        x_cam_3d = x_cam * z_depth
        y_cam_3d = y_cam * z_depth
        z_cam_3d = z_depth
        
        # 如果相机有安装角度偏移，需要额外变换
        # 这里假设相机坐标系与机体坐标系对齐
        x_body = x_cam_3d
        y_body = y_cam_3d
        z_body = z_cam_3d
        
        return x_body, y_body, z_body
    
    def body_to_enu(self, x_body, y_body, z_body, attitude):
        """
        机体坐标系 → ENU坐标系 (East-North-Up)
        使用完整的旋转矩阵变换
        """
        yaw, pitch, roll = np.radians(attitude)
        
        # 构建旋转矩阵（ZYX顺序，即偏航-俯仰-横滚）
        # 从机体坐标系到ENU坐标系的变换
        cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
        cos_pitch, sin_pitch = np.cos(pitch), np.sin(pitch)
        cos_roll, sin_roll = np.cos(roll), np.sin(roll)
        
        # 旋转矩阵 R_enu_body
        R = np.array([
            [cos_yaw*cos_pitch, cos_yaw*sin_pitch*sin_roll - sin_yaw*cos_roll, cos_yaw*sin_pitch*cos_roll + sin_yaw*sin_roll],
            [sin_yaw*cos_pitch, sin_yaw*sin_pitch*sin_roll + cos_yaw*cos_roll, sin_yaw*sin_pitch*cos_roll - cos_yaw*sin_roll],
            [-sin_pitch,        cos_pitch*sin_roll,                           cos_pitch*cos_roll]
        ])
        
        # 坐标变换
        body_vec = np.array([x_body, y_body, z_body])
        enu_vec = R @ body_vec
        
        return enu_vec[0], enu_vec[1], enu_vec[2]
    
    def enu_to_geo(self, x_enu, y_enu, drone_lat, drone_lon):
        """
        ENU坐标转换为地理坐标
        """
        # 地球半径（米）
        R_earth = 6378137.0
        
        # 转换为经纬度偏移
        lat_offset = y_enu / (R_earth * np.pi / 180.0)
        lon_offset = x_enu / (R_earth * np.pi / 180.0 * np.cos(np.radians(drone_lat)))
        
        target_lat = drone_lat + lat_offset
        target_lon = drone_lon + lon_offset
        
        return target_lat, target_lon

# 使用示例
if __name__ == "__main__":
    calibration_data = np.load('camera_calibration.npz')
    camera_matrix = calibration_data['camera_matrix']
    dist_coeffs = calibration_data['dist_coeffs']
    # 使用示例
    calculator = PreciseGeoCalculator(camera_matrix, dist_coeffs)

    # 无人机6自由度信息
    drone_position = [116.4074, 39.9042, 100.0]  # [经度, 纬度, 高度]
    drone_attitude = [45.0, -10.0, 2.0]          # [偏航角, 俯仰角, 横滚角] (度)

    # 目标在图像中的像素坐标
    target_pixel = (960, 540)  # 图像中心

    # 计算精确的地理坐标
    target_lat, target_lon = calculator.pixel_to_geo(
        target_pixel[0], target_pixel[1],
        drone_position,
        drone_attitude
    )

    print(f"目标精确位置: 纬度 {target_lat:.8f}, 经度 {target_lon:.8f}")