# 版本: Python 3.10+
# 依赖: torch, opencv-python, grounding-dino, segment-anything

import torch
import cv2
import numpy as np
from typing import Optional

try:
    from groundingdino.util.inference import load_model as load_dino, predict as dino_predict
    from segment_anything import sam_model_registry, SamPredictor
except ImportError:
    print("请安装: pip install grounding-dino segment-anything")

class GroundingSAMPipeline:
    """Grounding-DINO + SAM 实例分割管道"""

    def __init__(
        self,
        dino_model: str = "GroundingDINO/groundingdinov5-x.yaml",
        sam_model: str = "vit_b",
        sam_ckpt: str = "sam_vit_b.pth",
        device: str = "cuda"
    ):
        self.device = device
        self.dino = load_dino(dino_model, dino_model.replace("yaml", "pth"))
        self.dino.to(device)
        self.sam = sam_model_registry[sam_model](checkpoint=sam_ckpt)
        self.sam.to(device)
        self.sam_predictor = SamPredictor(self.sam)

    def segment(
        self,
        image: np.ndarray,
        text_prompt: str,
        box_threshold: float = 0.35,
        text_threshold: float = 0.25
    ) -> tuple:
        """语义引导实例分割

        参数:
            image: BGR图像 [H, W, 3]
            text_prompt: 文本描述，如 "cup. bottle. bowl."
            box_threshold: 检测框置信度阈值
            text_threshold: 文本匹配阈值

        返回:
            masks: 二值掩码列表 [mask1, mask2, ...] 每个shape [H, W]
            boxes: 检测框列表 [[x1,y1,x2,y2], ...]
        """
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        boxes, logits, phrases = dino_predict(
            self.dino,
            image,
            text_prompt,
            box_threshold,
            text_threshold
        )

        self.sam_predictor.set_image(image_rgb)
        masks = []

        for box in boxes:
            x1, y1, x2, y2 = box
            input_box = np.array([x1, y1, x2, y2])

            mask, score, _ = self.sam_predictor.predict(
                point_coords=None,
                point_labels=None,
                box=input_box[None, :],
                multimask_output=False
            )
            masks.append(mask[0].astype(np.uint8))

        return masks, boxes.cpu().numpy()

    def select_best_mask(
        self,
        masks: list,
        priority: str = "largest_visible"
    ) -> int:
        """多物体选择策略

        priority: "nearest" | "largest_visible" | "least_occluded"
        """
        if len(masks) == 0:
            return -1
        if len(masks) == 1:
            return 0

        if priority == "largest_visible":
            areas = [np.sum(m > 0) for m in masks]
            return int(np.argmax(areas))
        elif priority == "nearest":
            centroids_y = [np.mean(np.where(m > 0)[0]) for m in masks]
            return int(np.argmin(centroids_y))
        else:
            return 0

# 使用示例
pipeline = GroundingSAMPipeline()

image = cv2.imread("scene.jpg")
masks, boxes = pipeline.segment(image, text_prompt="cup. bottle.")

best_idx = pipeline.select_best_mask(masks, priority="largest_visible")
best_mask = masks[best_idx]

print(f"检测到 {len(masks)} 个物体，选择第 {best_idx} 个（面积最大）")
print(f"掩码非零像素数: {np.sum(best_mask > 0)}")

# 版本: Python 3.10+
# 依赖: numpy, opencv-python, torch

import numpy as np
import cv2
import torch

try:
    from transformers import DPTForDepthEstimation, DPTImageProcessor
except ImportError:
    print("请安装: pip install transformers")

class DepthEstimator:
    """单目深度估计 + 点云融合，使用MiDaS/DPT深度估计网络"""

    def __init__(self, model_name: str = "Intel/dpt-hybrid-midas"):
        self.processor = DPTImageProcessor.from_pretrained(model_name)
        self.model = DPTForDepthEstimation.from_pretrained(model_name)

    def estimate_depth(self, image: np.ndarray) -> np.ndarray:
        inputs = self.processor(images=image, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
            depth = outputs.predicted_depth

        depth = torch.nn.functional.interpolate(
            depth.unsqueeze(1),
            size=image.shape[:2],
            mode="bicubic",
            align_corners=False
        ).squeeze().numpy()

        depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
        return depth

    def depth_to_pointcloud(
        self,
        rgb: np.ndarray,
        depth: np.ndarray,
        K: np.ndarray,
        T_cam_world: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """深度图 → 点云

        参数:
            rgb: RGB图像 [H, W, 3]
            depth: 深度图 [H, W] 归一化到0~1
            K: 相机内参 [3, 3]
            T_cam_world: 相机外参，默认单位阵

        返回:
            points: 点云 [N, 3] (x,y,z)
        """
        H, W = depth.shape
        if T_cam_world is None:
            T_cam_world = np.eye(4)

        depth_real = depth * 5.0  # 假设深度范围0~5m

        u_coords, v_coords = np.meshgrid(np.arange(W), np.arange(H))
        z = depth_real[v_coords.flatten(), u_coords.flatten()]

        x = (u_coords.flatten() - K[0, 2]) * z / K[0, 0]
        y = (v_coords.flatten() - K[1, 2]) * z / K[1, 1]

        points = np.stack([x, y, z], axis=1)

        mask = z > 0.01
        points = points[mask]

        return points

    def mask_pointcloud(
        self,
        points: np.ndarray,
        mask: np.ndarray,
        rgb: np.ndarray,
        K: np.ndarray
    ) -> tuple:
        """根据二值掩码提取物体点云"""
        H, W = mask.shape
        u_coords, v_coords = np.meshgrid(np.arange(W), np.arange(H))
        valid_mask = mask.flatten() > 0

        depth = np.zeros(H * W)
        depth[valid_mask] = 1.0

        x = (u_coords.flatten() - K[0, 2]) * depth / K[0, 0]
        y = (v_coords.flatten() - K[1, 2]) * depth / K[1, 1]
        z = depth

        obj_points_raw = np.stack([x, y, z], axis=1)
        obj_points = obj_points_raw[valid_mask]

        obj_rgb = rgb.transpose(2, 0, 1).reshape(3, -1)
        obj_colors = obj_rgb[:, valid_mask].T

        return obj_points, obj_colors

# 使用示例
estimator = DepthEstimator()
rgb = cv2.imread("scene.jpg")
depth = estimator.estimate_depth(rgb)

K = np.array([[535.4, 0, 320.1],
              [0, 539.2, 240.1],
              [0, 0, 1]])
points = estimator.depth_to_pointcloud(rgb, depth, K)
print(f"场景点云: {points.shape}")

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

class PoseModelRefiner:
    """3D重建模型修正

    1. 质心校准: 点云质心→物体几何中心
    2. 去对称性处理: 对称物体(如圆柱体)需要先验对齐
    """

    def calibrate_centroid(self, point_cloud: np.ndarray) -> tuple:
        """质心校准

        问题: 深度估计的零点通常在相机位置，
             导致重建的点云质心≠物体几何中心

        解决: 假设物体在某个高度范围内(桌面上),
             校准质心z坐标
        """
        centroid = np.mean(point_cloud, axis=0)

        z_min = point_cloud[:, 2].min()
        z_max = point_cloud[:, 2].max()
        z_geo_center = (z_min + z_max) / 2

        calibrated = point_cloud.copy()
        calibrated[:, 2] += (z_geo_center - centroid[2])

        T_calib = np.eye(4)
        T_calib[:3, 3] = centroid

        return calibrated, T_calib

    def remove_symmetry(
        self,
        point_cloud: np.ndarray,
        symmetry_type: str = "bilateral_z"
    ) -> tuple:
        """去对称性处理

        对称物体(如杯子、瓶子)绕某轴旋转时视觉外观不变,
        需要额外的先验(如物体朝向)来消歧

        参数:
            symmetry_type: "bilateral_z"(绕Z轴对称) |
                          "planar"(平面对称) |
                          "none"(无对称)
        """
        if symmetry_type == "none":
            return point_cloud, np.eye(4)

        if symmetry_type == "bilateral_z":
            T_align = np.eye(4)
            T_align[2, 3] = 0.0
            return point_cloud, T_align

        return point_cloud, np.eye(4)

    def refine_model(
        self,
        point_cloud: np.ndarray,
        symmetry_type: str = "bilateral_z"
    ) -> tuple:
        refined_cloud, T1 = self.calibrate_centroid(point_cloud)
        refined_cloud, T2 = self.remove_symmetry(refined_cloud, symmetry_type)
        return refined_cloud, T1 @ T2

# 版本: Python 3.10+
# 依赖: numpy, opencv-python, scipy

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

class PoseEstimator:
    """EPnP + RANSAC 6DoF位姿估计

    输入: 2D图像点(物体掩码边缘/关键点) + 3D模型点(重建的物体点云)
    输出: 旋转矩阵R[3x3] + 平移向量t[3,] (物体相对相机)
    """

    def __init__(self, K: np.ndarray):
        self.K = K  # 相机内参

    def extract_2d_points_from_mask(
        self,
        mask: np.ndarray,
        num_points: int = 100
    ) -> np.ndarray:
        """从二值掩码中提取2D关键点，使用掩码边缘采样"""
        edges = cv2.Canny((mask * 255).astype(np.uint8), 100, 200)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        all_points = []
        for contour in contours:
            all_points.extend(contour.squeeze().tolist())

        all_points = np.array(all_points)
        if len(all_points) == 0:
            y_coords, x_coords = np.where(mask > 0)
            all_points = np.stack([x_coords, y_coords], axis=1)

        if len(all_points) > num_points:
            indices = np.random.choice(len(all_points), num_points, replace=False)
            all_points = all_points[indices]

        return all_points.astype(np.float32)

    def estimate_pose_ransac(
        self,
        points_3d: np.ndarray,
        points_2d: np.ndarray,
        ransac_threshold: float = 3.0
    ) -> tuple:
        """RANSAC + EPnP 鲁棒位姿估计

        参数:
            points_3d: 3D点云 [N, 3] (物体坐标系)
            points_2d: 对应2D像素坐标 [N, 2]
            ransac_threshold: RANSAC重投影误差阈值(像素)

        返回:
            R: 旋转矩阵 [3, 3]
            t: 平移向量 [3,]
            inliers: 内点数量
        """
        if len(points_3d) < 4:
            raise ValueError("至少需要4个对应点")

        success, rvec, t, inliers = cv2.solvePnPRansac(
            objectPoints=points_3d.astype(np.float32),
            imagePoints=points_2d.astype(np.float32),
            cameraMatrix=self.K,
            distCoeffs=None,
            reprojectionError=ransac_threshold,
            confidence=0.99,
            iterationsCount=1000
        )

        if not success:
            raise RuntimeError("PnP求解失败")

        R_mat = cv2.Rodrigues(rvec)[0]

        return R_mat, t.squeeze(), len(inliers)

    def compute_reprojection_error(
        self,
        points_3d: np.ndarray,
        R: np.ndarray,
        t: np.ndarray,
        points_2d: np.ndarray
    ) -> float:
        """计算重投影误差"""
        points_3d_h = np.concatenate([points_3d, np.ones((len(points_3d), 1))], axis=1).T
        proj = self.K @ (R @ points_3d[:3].T + t.reshape(3, 1))
        proj_2d = (proj[:2] / proj[2:]).T

        errors = np.linalg.norm(proj_2d - points_2d, axis=1)
        return float(np.mean(errors))

    def pose_to_transform(self, R: np.ndarray, t: np.ndarray) -> np.ndarray:
        """R,t → 4x4齐次变换矩阵"""
        T = np.eye(4)
        T[:3, :3] = R
        T[:3, 3] = t
        return T

# 使用示例
K = np.array([[535.4, 0, 320.1],
              [0, 539.2, 240.1],
              [0, 0, 1]])

estimator = PoseEstimator(K)

points_3d = np.random.randn(500, 3) * 0.1
points_2d = estimator.extract_2d_points_from_mask(best_mask, num_points=100)

R, t, inliers = estimator.estimate_pose_ransac(points_3d, points_2d, ransac_threshold=3.0)
T_cam_obj = estimator.pose_to_transform(R, t)

# 物体相对相机坐标系下的位置
print(f"物体位置(相机坐标系): t={t}")
print(f"物体朝向(欧拉角): {R.as_euler('xyz', degrees=True)}")
print(f"内点比例: {inliers/len(points_2d)*100:.1f}%")