| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312 |
- #!/usr/bin/env python3
- # -*- coding: utf-8 -*-
- """
- 3D 场景可视化脚本
- 功能:
- 1. RANSAC 平面拟合地面和天花板
- 2. 基于地面拟合结果过滤门(门底部距地面>0.1 米过滤)
- 3. 将门和拍摄点位平移到天花板高度进行可视化
- 4. 入户门用特殊颜色标记
- 5. 输出 PLY 文件,支持在 3D 查看器中查看
- """
- import os
- import sys
- import json
- import argparse
- from pathlib import Path
- from dataclasses import dataclass
- from typing import Dict, List, Optional, Tuple
- import cv2
- import numpy as np
- import open3d as o3d
- from tqdm import tqdm
- from ultralytics import YOLOE
- from camera_spherical import Intrinsic_Spherical_NP
- from scipy.spatial import Delaunay
- # ============================================================================
- # 数据类
- # ============================================================================
- @dataclass
- class PoseData:
- """位姿数据"""
- uuid: str
- rotation: Dict[str, float] # w, x, y, z
- translation: Dict[str, float] # x, y, z
- @dataclass
- class Door3D:
- """3D 门实例"""
- id: int
- center: np.ndarray # 中心坐标 [x, y, z]
- bbox_8points: np.ndarray # 8 个角点 [[x,y,z], ...] (8x3)
- source_detections: List[Dict] # 来源检测信息
- source_uuid: Optional[str] = None # 来源点位的 UUID(用于获取 puck_z)
- # ============================================================================
- # 场景可视化处理器
- # ============================================================================
- class SceneVisualizer:
- """
- 场景 3D 可视化处理器:
- - RANSAC 平面拟合地面/天花板
- - 门过滤(基于地面距离)
- - 门和点位平移到天花板高度
- - 输出可视化 PLY 文件
- """
- # 门检测类别
- DOOR_CLASSES = [
- "door", "indoor door", "exterior door",
- "wooden door", "metal door", "glass door", "double door",
- "single door", "open door", "closed door"
- ]
- def __init__(
- self,
- scene_folder: str,
- model_path: str = "yoloe-26x-seg.pt",
- conf: float = 0.35,
- iou: float = 0.45,
- voxel_size: float = 0.03,
- depth_scale: float = 256.0,
- depth_min: float = 0.02,
- # 地面/天花板拟合参数
- ground_ransac_dist: float = 0.05, # RANSAC 内点距离阈值
- ground_ransac_prob: float = 0.99, # RANSAC 置信度
- ceiling_percentile: float = 95.0, # 天花板点分位数
- # 门过滤参数
- door_ground_dist: float = 0.1, # 门底部距地面最大距离
- door_height_min: float = 1.0,
- door_height_max: float = 3.0,
- door_width_min: float = 0.3,
- door_width_max: float = 3.0,
- door_thickness_max: float = 0.5,
- # 3D 门合并参数
- merge_iou_thresh: float = 0.1, # 降低 IoU 阈值,增加合并敏感度
- merge_dist_thresh: float = 0.3, # 中心距离阈值 (米) - 同一物理门的检测应该很近
- merge_z_overlap_thresh: float = 0.5, # Z 方向重叠度阈值 (0-1,0=不重叠,1=完全重叠)
- # 入户门参数
- entrance_method: str = "score", # "score" 或 "json"
- entrance_json_path: Optional[str] = None,
- ):
- """
- 初始化场景可视化处理器
- Args:
- scene_folder: 场景文件夹路径
- model_path: YOLOE 模型路径
- conf: 检测置信度阈值
- iou: NMS IoU 阈值
- voxel_size: 点云体素下采样尺寸
- depth_scale: 深度图缩放因子
- depth_min: 最小有效深度
- ground_ransac_dist: RANSAC 地面拟合距离阈值
- ground_ransac_prob: RANSAC 置信度
- ceiling_percentile: 天花板点分位数(取最高的 x% 点拟合天花板)
- door_ground_dist: 门底部距地面最大距离(超过则过滤)
- door_height_min/max: 门高度范围
- door_width_min/max: 门宽度范围
- door_thickness_max: 门最大厚度
- merge_iou_thresh: 3D 门合并 IoU 阈值
- merge_dist_thresh: 3D 门合并中心距离阈值
- entrance_method: 入户门判定方法 ("score" 或 "json")
- entrance_json_path: 入户门 JSON 路径(当 method=json 时使用)
- """
- self.scene_folder = Path(scene_folder)
- self.conf = conf
- self.iou = iou
- self.voxel_size = voxel_size
- self.depth_scale = depth_scale
- self.depth_min = depth_min
- # 地面/天花板参数
- self.ground_ransac_dist = ground_ransac_dist
- self.ground_ransac_prob = ground_ransac_prob
- self.ceiling_percentile = ceiling_percentile
- # 门过滤参数
- self.door_ground_dist = door_ground_dist
- self.door_height_min = door_height_min
- self.door_height_max = door_height_max
- self.door_width_min = door_width_min
- self.door_width_max = door_width_max
- self.door_thickness_max = door_thickness_max
- # 3D 门合并参数
- self.merge_iou_thresh = merge_iou_thresh
- self.merge_dist_thresh = merge_dist_thresh
- self.merge_z_overlap_thresh = merge_z_overlap_thresh
- # 地面 Z 值向下兼容:合并时取最低的地面 Z 值
- self.ground_z_fallback_low = True
- # 入户门参数
- self.entrance_method = entrance_method
- self.entrance_json_path = entrance_json_path
- self.entrance_door_id = None
- self.estimated_entrance_position = None # 当没有检测到门时的估计位置
- # 地面参数(在拟合时设置)
- self.ground_d = None
- self.ground_z_from_puck = None # 从 puck 参数估计的地面 Z
- # 子目录
- self.rgb_folder = self.scene_folder / "pano_img"
- self.depth_folder = self.scene_folder / "depth_img"
- self.pose_file = self.scene_folder / "vision.txt"
- # 输出目录
- self.output_folder = self.scene_folder / "output"
- # 加载位姿
- self.poses = self._load_poses()
- # 加载 YOLOE 模型
- print(f"加载 YOLOE 模型:{model_path}")
- self.model = YOLOE(model_path)
- self.model.set_classes(self.DOOR_CLASSES)
- print(f"检测类别:{self.DOOR_CLASSES}")
- def _load_poses(self) -> Dict[str, PoseData]:
- """从 vision.txt 加载位姿信息"""
- if not self.pose_file.exists():
- raise FileNotFoundError(f"位姿文件不存在:{self.pose_file}")
- with open(self.pose_file, 'r') as f:
- data = json.load(f)
- poses = {}
- self.puck_z_dict = {} # 保存每个点位的 puck_z
- for loc in data.get('sweepLocations', []):
- uuid = str(loc['uuid'])
- poses[uuid] = PoseData(
- uuid=uuid,
- rotation=loc['pose']['rotation'],
- translation=loc['pose']['translation']
- )
- # 保存每个点位的 puck_z
- if 'puck' in loc and 'z' in loc['puck']:
- self.puck_z_dict[uuid] = loc['puck']['z']
- print(f"加载 {len(poses)} 个拍摄点位")
- # 计算整体地面 Z(用于没有 puck 数据时的回退)
- if self.puck_z_dict:
- puck_z_values = list(self.puck_z_dict.values())
- self.ground_z_from_puck = np.median(puck_z_values)
- print(f"从 puck 参数估计地面 Z (中位数): {self.ground_z_from_puck:.4f}m")
- print(f"puck_z 范围:[{min(puck_z_values):.4f}, {max(puck_z_values):.4f}]")
- else:
- self.ground_z_from_puck = None
- print("⚠️ 未找到 puck 参数,将使用 RANSAC 拟合地面")
- return poses
- def _build_pose_matrix(self, pose: PoseData) -> np.ndarray:
- """构建 4x4 位姿变换矩阵"""
- R = o3d.geometry.get_rotation_matrix_from_quaternion(
- np.array([pose.rotation['w'], pose.rotation['x'],
- pose.rotation['y'], pose.rotation['z']])
- )
- t = np.array([
- pose.translation['x'],
- pose.translation['y'],
- pose.translation['z']
- ])
- T = np.eye(4)
- T[:3, :3] = R
- T[:3, 3] = t
- return T
- def _mask_to_3d_points(
- self,
- mask: np.ndarray,
- depth: np.ndarray,
- pose_matrix: np.ndarray
- ) -> Optional[np.ndarray]:
- """
- 将 2D mask 映射到世界坐标系 3D 点
- Args:
- mask: 二值 mask (H, W)
- depth: 深度图 (H, W)
- pose_matrix: 4x4 位姿矩阵
- Returns:
- 世界坐标系下的 3D 点 (N, 3)
- """
- H, W = depth.shape
- sph = Intrinsic_Spherical_NP(W, H)
- # 获取 mask 内的像素
- ys, xs = np.where(mask > 0)
- if len(xs) == 0:
- return None
- # 有效深度掩码
- valid = depth[ys, xs] > self.depth_min
- if not np.any(valid):
- return None
- xs, ys = xs[valid], ys[valid]
- depths = depth[ys, xs]
- # 计算方向向量
- bx, by, bz = sph.bearing([xs.astype(np.float64), ys.astype(np.float64)])
- bx, by, bz = np.array(bx), np.array(by), np.array(bz)
- # 相机坐标系
- pts_cam = np.stack([bx * depths, by * depths, bz * depths], axis=1)
- # Z 轴 180 度翻转
- R_z180 = np.diag([-1.0, -1.0, 1.0])
- pts_cam = pts_cam @ R_z180.T
- # 世界坐标系
- pts_w = (pose_matrix[:3, :3] @ pts_cam.T).T + pose_matrix[:3, 3]
- return pts_w
- def _extract_mask_contours(self, masks) -> Tuple[List[List[List[float]]], List[np.ndarray]]:
- """
- 从 YOLOE mask 结果提取轮廓
- Args:
- masks: YOLOE masks (H, W, N)
- Returns:
- (轮廓列表,对应 mask 数组)
- """
- contours = []
- mask_arrays = []
- if masks is None:
- return contours, mask_arrays
- masks_np = masks.cpu().numpy()
- for i in range(masks_np.shape[0]):
- mask = masks_np[i]
- # 二值化
- mask_bin = (mask > 0.5).astype(np.uint8) * 255
- # 提取轮廓
- cnts, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
- if not cnts:
- continue
- # 只保留面积最大的轮廓
- largest = max(cnts, key=cv2.contourArea)
- if len(largest) >= 3:
- # 简化轮廓
- epsilon = 0.02 * cv2.arcLength(largest, True)
- approx = cv2.approxPolyDP(largest, epsilon, True)
- contour = approx.reshape(-1, 2).astype(float).tolist()
- contours.append(contour)
- mask_arrays.append((mask_bin > 0).astype(np.uint8))
- else:
- mask_arrays.append((mask_bin > 0).astype(np.uint8))
- return contours, mask_arrays
- def detect_single_image(
- self,
- img_path: str,
- depth: np.ndarray,
- pose_matrix: np.ndarray,
- save_path: Optional[str] = None
- ) -> Tuple[List[np.ndarray], List[float]]:
- """
- 检测单张图像并提取 mask 3D 点
- Returns:
- (mask_3d_points 列表,scores 列表)
- """
- results = self.model.predict(
- img_path,
- imgsz=(1024, 2048),
- conf=self.conf,
- iou=self.iou,
- max_det=50,
- augment=True,
- retina_masks=True,
- half=False,
- verbose=False,
- )
- result = results[0]
- mask_3d_points = []
- scores = []
- if result.masks is not None:
- masks = result.masks.data
- contours, mask_arrays = self._extract_mask_contours(masks)
- scores = result.boxes.conf.cpu().numpy().tolist()
- H, W = depth.shape
- for mask_bin in mask_arrays:
- mask_resized = cv2.resize(mask_bin, (W, H), interpolation=cv2.INTER_NEAREST)
- pts_3d = self._mask_to_3d_points(mask_resized, depth, pose_matrix)
- if pts_3d is not None and len(pts_3d) > 10:
- mask_3d_points.append(pts_3d)
- if save_path:
- os.makedirs(os.path.dirname(save_path), exist_ok=True)
- result.save(save_path)
- return mask_3d_points, scores
- def _axis_aligned_bbox(self, points: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- """计算轴对齐包围盒 (min, max)"""
- lo = np.min(points, axis=0)
- hi = np.max(points, axis=0)
- return lo, hi
- def _filter_outliers(self, points: np.ndarray, std_thresh: float = 3.0) -> np.ndarray:
- """
- 过滤 3D 点云中的离群点
- Args:
- points: 3D 点 (N, 3)
- std_thresh: 标准差阈值(默认 3 倍标准差)
- Returns:
- 过滤后的点
- """
- if len(points) < 10:
- return points
- # 计算每个维度的均值和标准差
- mean = np.mean(points, axis=0)
- std = np.std(points, axis=0)
- # 保留在 3 倍标准差内的点
- mask = np.all(np.abs(points - mean) < std_thresh * std, axis=1)
- filtered = points[mask]
- if len(filtered) < len(points) * 0.5:
- # 如果过滤掉的点超过一半,返回原始点
- return points
- return filtered
- def _filter_door_points_by_depth(self, points: np.ndarray) -> np.ndarray:
- """
- 根据深度一致性过滤门的 3D 点
- 门的特点是:在门平面法向量方向上很薄
- 使用 PCA 找出最薄的方向,然后过滤掉超出阈值的点
- Args:
- points: 3D 点 (N, 3)
- Returns:
- 过滤后的点
- """
- if len(points) < 50:
- return points
- # PCA 分析
- centered = points - np.mean(points, axis=0)
- cov = np.cov(centered.T)
- eigenvalues, eigenvectors = np.linalg.eigh(cov)
- # 特征值从大到小排序
- idx = np.argsort(eigenvalues)[::-1]
- eigenvalues = eigenvalues[idx]
- eigenvectors = eigenvectors[:, idx]
- # 最小特征值对应的特征向量是门的法向量(厚度方向)
- normal = eigenvectors[:, 2] # 最小特征值对应的方向
- # 将点投影到法向量方向
- projected = np.dot(points, normal)
- # 过滤掉超出 2 倍标准差的点
- mean_proj = np.mean(projected)
- std_proj = np.std(projected)
- mask = np.abs(projected - mean_proj) < 2.0 * std_proj
- filtered = points[mask]
- if len(filtered) < len(points) * 0.5:
- return points
- return filtered
- def _compute_door_dimensions_pca(self, points: np.ndarray) -> Tuple[float, float, float]:
- """
- 使用 PCA 计算门的实际尺寸
- 门的特点是:在一个方向上很薄(厚度),在另外两个方向上较大(高度和宽度)
- Args:
- points: 3D 点 (N, 3)
- Returns:
- (width, thickness, height) 门的宽度、厚度、高度
- """
- if len(points) < 10:
- size = points.max(axis=0) - points.min(axis=0)
- return size[0], size[1], size[2]
- # PCA 分析
- centered = points - np.mean(points, axis=0)
- cov = np.cov(centered.T)
- eigenvalues, eigenvectors = np.linalg.eigh(cov)
- # 特征值从大到小排序
- idx = np.argsort(eigenvalues)[::-1]
- eigenvalues = eigenvalues[idx]
- eigenvectors = eigenvectors[:, idx]
- # 将点投影到主成分方向
- projected = centered @ eigenvectors
- # 计算每个方向的尺寸
- dims = projected.max(axis=0) - projected.min(axis=0)
- # 门的特点:厚度方向应该是最小的维度
- # 假设门的厚度是最小的维度
- thickness = dims.min()
- height = dims.max() # 高度是最大的维度
- width = dims.sum() - thickness - height # 宽度是中间的维度
- return width, thickness, height
- def _bbox_8corners(self, bbox_min: np.ndarray, bbox_max: np.ndarray) -> np.ndarray:
- """从 bbox min/max 获取 8 个角点"""
- cx, cy, cz = bbox_min
- ex, ey, ez = bbox_max
- return np.array([
- [cx, cy, cz], [ex, cy, cz], [ex, ey, cz], [cx, ey, cz],
- [cx, cy, ez], [ex, cy, ez], [ex, ey, ez], [cx, ey, ez],
- ])
- def _bbox_iou_3d(self, b1, b2) -> float:
- """3D IoU 计算"""
- lo = np.maximum(b1[0], b2[0])
- hi = np.minimum(b1[1], b2[1])
- inter = np.prod(np.maximum(hi - lo, 0))
- vol1 = np.prod(b1[1] - b1[0])
- vol2 = np.prod(b2[1] - b2[0])
- union = vol1 + vol2 - inter
- return inter / union if union > 0 else 0.0
- def _merge_3d_doors(self, door_candidates: List[Dict]) -> List[Door3D]:
- """
- 使用并查集合并 3D 门
- 合并策略:
- 1. 来自不同点位 Detection 才考虑合并(同一图像的可能是多个门)
- 2. 中心距离接近(同一物理门在不同视角下应该位置相近)
- 3. Z 方向高度重叠(同一门的高度范围应该相近)
- 4. IoU 或距离满足其一即可(放宽条件)
- 地面 Z 值向下兼容:
- - 对于不同点位的重叠或相邻很近的门,合并后地面的 Z 值以最低的为准
- - 例如:a 点位地面 z=-1.45, b 点位地面 z=-1.5,合并后以 -1.5 为准
- """
- if not door_candidates:
- return []
- n = len(door_candidates)
- if n == 1:
- d = door_candidates[0]
- return [Door3D(
- id=0,
- center=(d['bbox_min'] + d['bbox_max']) / 2,
- bbox_8points=self._bbox_8corners(d['bbox_min'], d['bbox_max']),
- source_detections=[d['source']],
- source_uuid=d.get('source_uuid')
- )]
- # 并查集
- parent = list(range(n))
- def find(x):
- if parent[x] != x:
- parent[x] = find(parent[x])
- return parent[x]
- def union(x, y):
- px, py = find(x), find(y)
- if px != py:
- parent[px] = py
- # 构建连通关系
- for i in range(n):
- for j in range(i + 1, n):
- ci = (door_candidates[i]['bbox_min'] + door_candidates[i]['bbox_max']) / 2
- cj = (door_candidates[j]['bbox_min'] + door_candidates[j]['bbox_max']) / 2
- dist = np.linalg.norm(ci - cj)
- # 检查是否来自同一张图像
- same_image = door_candidates[i]['source']['image'] == door_candidates[j]['source']['image']
- # 计算 Z 方向重叠度
- z_min_i, z_max_i = door_candidates[i]['bbox_min'][2], door_candidates[i]['bbox_max'][2]
- z_min_j, z_max_j = door_candidates[j]['bbox_min'][2], door_candidates[j]['bbox_max'][2]
- z_overlap_min = max(z_min_i, z_min_j)
- z_overlap_max = min(z_max_i, z_max_j)
- z_intersection = max(0, z_overlap_max - z_overlap_min)
- z_union = max(z_max_i, z_max_j) - min(z_min_i, z_min_j)
- z_overlap_ratio = z_intersection / z_union if z_union > 0 else 0
- # 计算 3D IoU
- iou = self._bbox_iou_3d(
- (door_candidates[i]['bbox_min'], door_candidates[i]['bbox_max']),
- (door_candidates[j]['bbox_min'], door_candidates[j]['bbox_max'])
- )
- # 合并条件:
- # 1. 来自不同图像:距离近 OR IoU 高,且 Z 重叠
- # 2. 来自同一图像:IoU 高且 Z 重叠(可能是同一个门在不同位置被检测到)
- if same_image:
- # 同一图像:IoU 高且 Z 重叠
- should_merge = (iou > 0.05 and z_overlap_ratio > 0.5)
- else:
- # 不同图像:(距离近 OR IoU 高) 且 Z 重叠
- should_merge = (
- (dist < self.merge_dist_thresh or iou > self.merge_iou_thresh) and
- z_overlap_ratio > self.merge_z_overlap_thresh
- )
- if should_merge:
- union(i, j)
- # 按连通分量分组
- from collections import defaultdict
- groups = defaultdict(list)
- for i in range(n):
- groups[find(i)].append(door_candidates[i])
- # 合并每个组
- doors = []
- for door_id, members in enumerate(groups.values()):
- if not members:
- continue
- # 原始合并逻辑:取 min 和 max
- # bbox_min = np.min([m['bbox_min'] for m in members], axis=0)
- # bbox_max = np.max([m['bbox_max'] for m in members], axis=0)
- # 地面 Z 值向下兼容:
- # 1. X, Y 方向保持不变(取最小/最大)
- # 2. Z 方向(地面)取最低值(最负的),门顶保持最高值
- all_bbox_mins = [m['bbox_min'] for m in members]
- all_bbox_maxs = [m['bbox_max'] for m in members]
- # X, Y 方向:取最小/最大
- merged_min_xy = np.min(all_bbox_mins, axis=0)[:2] # [x_min, y_min]
- merged_max_xy = np.max(all_bbox_maxs, axis=0)[:2] # [x_max, y_max]
- # Z 方向:地面 Z 向下兼容(取最小值/最负值),门顶取最大值
- z_mins = [b[2] for b in all_bbox_mins] # 所有门的底部 Z
- z_maxs = [b[2] for b in all_bbox_maxs] # 所有门的顶部 Z
- # 地面 Z 值向下兼容:取最低的 Z 值(最负的值)
- merged_z_min = np.min(z_mins)
- merged_z_max = np.max(z_maxs)
- bbox_min = np.array([merged_min_xy[0], merged_min_xy[1], merged_z_min])
- bbox_max = np.array([merged_max_xy[0], merged_max_xy[1], merged_z_max])
- sources = [m['source'] for m in members]
- # 使用第一个成员的 source_uuid(同一门的多个检测来自相近位置)
- source_uuid = members[0].get('source_uuid')
- doors.append(Door3D(
- id=door_id,
- center=(bbox_min + bbox_max) / 2,
- bbox_8points=self._bbox_8corners(bbox_min, bbox_max),
- source_detections=sources,
- source_uuid=source_uuid
- ))
- return doors
- def _fit_ground_plane_ransac(self, pc: o3d.geometry.PointCloud) -> Tuple[np.ndarray, float]:
- """
- 使用 RANSAC 拟合地面平面
- 策略:
- 1. 先用 Z 坐标最低的 5% 的点作为初始地面点
- 2. 对这些点使用 RANSAC 平面拟合
- 3. 确保法向量朝上(Z 轴正方向)
- Args:
- pc: 点云
- Returns:
- (平面法向量,平面到原点距离 d,满足 ax+by+cz+d=0)
- """
- print("\n=== RANSAC 地面拟合 ===")
- points = np.asarray(pc.points)
- # 方法 1:先用 Z 坐标最低的点作为初始种子
- z_coords = points[:, 2]
- z_threshold = np.percentile(z_coords, 5) # 最低的 5%
- ground_candidates = points[z_coords <= z_threshold]
- print(f"地面候选点数量:{len(ground_candidates)} (Z <= {z_threshold:.3f}m)")
- if len(ground_candidates) < 100:
- print("⚠️ 地面候选点过少,使用全局 RANSAC")
- # 回退到全局 RANSAC
- plane_model, inliers = pc.segment_plane(
- distance_threshold=self.ground_ransac_dist,
- ransac_n=3,
- num_iterations=1000
- )
- else:
- # 对地面候选点使用 PCA 拟合平面
- centered = ground_candidates - np.mean(ground_candidates, axis=0)
- cov = np.cov(centered.T)
- eigenvalues, eigenvectors = np.linalg.eigh(cov)
- # 最小特征值对应的特征向量是法向量
- normal = eigenvectors[:, 0] # 最小特征值
- mean_point = np.mean(ground_candidates, axis=0)
- # 平面方程:a(x-x0) + b(y-y0) + c(z-z0) = 0
- # 即:ax + by + cz + d = 0, d = -(ax0 + by0 + cz0)
- a, b, c = normal
- d = -np.dot(normal, mean_point)
- plane_model = [a, b, c, d]
- inliers = list(range(len(ground_candidates)))
- a, b, c, d = plane_model
- normal = np.array([a, b, c])
- # 确保法向量朝上(Z 轴正方向)
- if normal[2] < 0:
- normal = -normal
- d = -d
- inlier_ratio = len(inliers) / len(pc.points)
- print(f"地面平面方程:{normal[0]:.4f}x + {normal[1]:.4f}y + {normal[2]:.4f}z + {d:.4f} = 0")
- print(f"地面法向量:[{normal[0]:.4f}, {normal[1]:.4f}, {normal[2]:.4f}]")
- print(f"内点数量:{len(inliers)} / {len(pc.points)} ({inlier_ratio:.2%})")
- return normal, d
- def _fit_ceiling_plane(self, pc: o3d.geometry.PointCloud, ground_normal: np.ndarray, ground_d: float) -> Tuple[np.ndarray, float, float]:
- """
- 拟合天花板平面
- Args:
- pc: 点云
- ground_normal: 地面法向量
- ground_d: 地面平面方程的 d 参数
- Returns:
- (天花板平面法向量,平面到原点距离 d, 地面到天花板距离)
- """
- print("\n=== 天花板平面拟合 ===")
- points = np.asarray(pc.points)
- # 计算每个点沿法向量方向到地面的距离
- # distance = p·n + d_ground
- distances_to_ground = np.dot(points, ground_normal) + ground_d
- # 过滤掉异常高点(高于 5 米的点可能是噪声)
- # 正常房间高度在 2.5-4 米之间
- valid_mask = distances_to_ground < 5.0
- valid_points = points[valid_mask]
- valid_distances = distances_to_ground[valid_mask]
- if len(valid_points) < 1000:
- print(f"⚠️ 有效点过少 ({len(valid_points)} 个),使用所有点")
- valid_points = points
- valid_distances = distances_to_ground
- # 取最高的点拟合天花板(85%-95% 分位数之间,避免噪声)
- height_percentiles = np.percentile(valid_distances, [85, 90, 95])
- print(f"高度分位数:85%={height_percentiles[0]:.3f}m, 90%={height_percentiles[1]:.3f}m, 95%={height_percentiles[2]:.3f}m")
- # 使用 90%-95% 之间的点拟合天花板
- ceiling_threshold = height_percentiles[1] # 90% 分位数
- ceiling_mask = valid_distances >= ceiling_threshold
- ceiling_points = valid_points[ceiling_mask]
- if len(ceiling_points) < 100:
- print(f"⚠️ 天花板点过少 ({len(ceiling_points)} 个),降低阈值到 85%")
- ceiling_threshold = height_percentiles[0]
- ceiling_mask = valid_distances >= ceiling_threshold
- ceiling_points = valid_points[ceiling_mask]
- # 天花板法向量与地面相同(平行平面)
- ceiling_normal = ground_normal.copy()
- # 计算天花板到原点的距离
- # d_ceiling = -mean(p·n) for ceiling points
- d_ceiling = -np.mean(np.dot(ceiling_points, ceiling_normal))
- # 计算地面到天花板的距离(用天花板点的平均高度)
- floor_ceiling_dist = np.mean(valid_distances[ceiling_mask])
- ceiling_cloud = o3d.geometry.PointCloud()
- ceiling_cloud.points = o3d.utility.Vector3dVector(ceiling_points)
- print(f"天花板点数量:{len(ceiling_points)}")
- print(f"天花板平面方程:{ceiling_normal[0]:.4f}x + {ceiling_normal[1]:.4f}y + {ceiling_normal[2]:.4f}z + {d_ceiling:.4f} = 0")
- print(f"地面到天花板距离:{floor_ceiling_dist:.3f} 米")
- return ceiling_normal, d_ceiling, floor_ceiling_dist
- def _point_to_plane_distance(self, point: np.ndarray, normal: np.ndarray, d: float) -> float:
- """计算点到平面的距离"""
- return abs(np.dot(point, normal) + d)
- def _denoise_xy_projection(
- self,
- points: np.ndarray,
- grid_size: float = 0.15,
- min_points_per_cell: int = 5
- ) -> np.ndarray:
- """
- 投影到 XY 平面后进行网格滤波去噪
- 原理:
- 1. 将点云投影到 XY 平面
- 2. 划分网格,统计每个网格内的点数
- 3. 只保留点数足够的网格中的点
- 4. 这样可以去除孤立的噪声点,同时保留建筑轮廓
- Args:
- points: 3D 点 (N, 3)
- grid_size: 网格尺寸(米),默认 0.15m
- min_points_per_cell: 每个网格最少点数,低于此值的网格视为噪声
- Returns:
- 去噪后的点 (M, 3)
- """
- if len(points) < 10:
- return points
- xy = points[:, :2]
- z = points[:, 2]
- # 网格化
- grid_coords = np.floor(xy / grid_size).astype(int)
- # 统计每个网格的点数
- from collections import Counter, defaultdict
- cell_counts = Counter(map(tuple, grid_coords))
- # 保留点数足够的网格
- valid_cells = {
- cell for cell, count in cell_counts.items()
- if count >= min_points_per_cell
- }
- # 创建掩码
- mask = np.array([tuple(gc) in valid_cells for gc in grid_coords])
- if np.sum(mask) < len(points) * 0.3:
- # 如果过滤掉的点超过 70%,降低阈值
- print(f"⚠️ 去噪过滤过多 ({np.sum(mask)}/{len(points)}),降低阈值重试")
- valid_cells = {
- cell for cell, count in cell_counts.items()
- if count >= max(1, min_points_per_cell - 2)
- }
- mask = np.array([tuple(gc) in valid_cells for gc in grid_coords])
- filtered = points[mask]
- print(f"XY 投影去噪:{len(points)} → {len(filtered)} 点 ({len(filtered)/len(points)*100:.1f}%)")
- return filtered
- def _compute_xy_contour(self, points: np.ndarray) -> Tuple[np.ndarray, Optional[Delaunay]]:
- """
- 计算点云 XY 投影的轮廓(凹包/Alpha Shape)
- Args:
- points: 3D 点 (N, 3),应该已经去过噪
- Returns:
- (轮廓点集,Delaunay 三角网)
- 轮廓点集形状 (K, 2),Delaunay 三角网用于点在多边形内判断
- """
- from scipy.spatial import Delaunay
- if len(points) < 4:
- return points[:, :2], None
- xy = points[:, :2]
- # 使用 Alpha Shape 提取轮廓
- alpha = 2.0 # alpha 半径,越大越接近凸包
- # 计算 Delaunay 三角网
- tri = Delaunay(xy)
- # 提取 Alpha Shape 边界
- # 策略:移除边长超过阈值的三角形
- centers = []
- for i in range(tri.npoints):
- # 获取第 i 个三角形的三个顶点
- idx = tri.simplices[i]
- p0, p1, p2 = xy[idx[0]], xy[idx[1]], xy[idx[2]]
- # 计算三角形外接圆半径
- a = np.linalg.norm(p1 - p2)
- b = np.linalg.norm(p0 - p2)
- c = np.linalg.norm(p0 - p1)
- s = (a + b + c) / 2
- area = np.sqrt(max(0, s * (s - a) * (s - b) * (s - c)))
- if area < 1e-10:
- continue
- R = (a * b * c) / (4 * area)
- if R < alpha:
- center = (p0 + p1 + p2) / 3
- centers.append(center)
- if len(centers) < 3:
- # Alpha Shape 失败,回退到凸包
- print("⚠️ Alpha Shape 失败,使用凸包")
- from scipy.spatial import ConvexHull
- hull = ConvexHull(xy)
- return xy[hull.vertices], tri
- # 从中心点集提取边界
- centers = np.array(centers)
- from scipy.spatial import ConvexHull
- hull = ConvexHull(centers)
- contour = centers[hull.vertices]
- return contour, tri
- def _point_in_contour(self, point: np.ndarray, contour: np.ndarray, tri: Optional[Delaunay]) -> bool:
- """
- 判断点是否在轮廓内
- Args:
- point: 2D 点 [x, y]
- contour: 轮廓点集 (K, 2)
- tri: Delaunay 三角网(用于判断)
- Returns:
- 是否在轮廓内
- """
- if tri is None:
- # 没有三角网,使用射线法判断
- from matplotlib.path import Path
- path = Path(contour)
- return path.contains_point(point)
- # 使用 Delaunay 三角网判断
- # 如果点在凸包内,simplex 返回 >= 0
- simplex = tri.find_simplex(point)
- return simplex >= 0
- def _compute_door_boundary_score(
- self,
- door: Door3D,
- contour: np.ndarray,
- tri
- ) -> float:
- """
- 计算门到轮廓边界的距离评分
- 入户门应该在建筑轮廓的边缘,距离边界越近评分越高
- Args:
- door: 3D 门对象
- contour: 轮廓点集 (K, 2)
- tri: Delaunay 三角网
- Returns:
- 边界评分 (0-30 分)
- """
- door_xy = door.center[:2]
- # 计算门中心到轮廓边界的最短距离
- # 遍历轮廓的每条边
- min_dist = float('inf')
- n = len(contour)
- for i in range(n):
- p1 = contour[i]
- p2 = contour[(i + 1) % n]
- dist = self._point_to_line_segment_distance(door_xy, p1, p2)
- min_dist = min(min_dist, dist)
- if min_dist == float('inf'):
- return 0.0
- # 距离转评分(越近分数越高)
- # 0 米 = 30 分,2 米 = 0 分
- score = max(0, 30 - min_dist * 15)
- return score
- def _point_to_line_segment_distance(
- self,
- point: np.ndarray,
- line_start: np.ndarray,
- line_end: np.ndarray
- ) -> float:
- """
- 计算点到线段的最短距离
- Args:
- point: 2D 点 [x, y]
- line_start: 线段起点 [x1, y1]
- line_end: 线段终点 [x2, y2]
- Returns:
- 最短距离
- """
- p = np.array(point)
- a = np.array(line_start)
- b = np.array(line_end)
- ab = b - a
- ap = p - a
- # 投影长度
- t = np.dot(ap, ab) / np.dot(ab, ab)
- t = np.clip(t, 0, 1)
- # 投影点
- proj = a + t * ab
- return np.linalg.norm(p - proj)
- def _filter_door_by_ground(self, door: Door3D, ground_normal: np.ndarray, ground_d: float) -> Tuple[bool, str]:
- """
- 基于地面距离过滤门
- Args:
- door: 3D 门对象
- ground_normal: 地面法向量
- ground_d: 地面平面方程的 d 参数
- Returns:
- (是否通过,拒绝原因)
- """
- # 计算门底部到地面的距离
- # 门的底部是 bbox 中沿法向量方向最低的点
- points = door.bbox_8points
- # 计算每个点到地面的距离
- distances = np.array([self._point_to_plane_distance(p, ground_normal, ground_d) for p in points])
- # 门底部距离 = 最小距离(最接近地面的点)
- door_bottom_dist = distances.min()
- if door_bottom_dist > self.door_ground_dist:
- return False, f"门底部距地面 {door_bottom_dist:.3f}m > {self.door_ground_dist}m"
- return True, ""
- def _filter_door_by_ground_puck(self, door: Door3D) -> Tuple[bool, str]:
- """
- 基于 puck 参数检查门的地面距离
- 使用门对应来源点位的 puck_z
- Args:
- door: 3D 门对象
- Returns:
- (是否通过,拒绝原因)
- """
- # 门底部 Z 坐标(bbox 中 Z 最小的点)
- door_bottom_z = door.bbox_8points[:, 2].min()
- # 获取门对应来源点位的 puck_z
- puck_z = None
- if door.source_uuid and door.source_uuid in self.puck_z_dict:
- puck_z = self.puck_z_dict[door.source_uuid]
- elif self.ground_z_from_puck is not None:
- # 回退到中位数
- puck_z = self.ground_z_from_puck
- if puck_z is None:
- return True, "" # 没有 puck 数据,不检查
- # 门底部到 puck 地面的距离
- dist = door_bottom_z - puck_z
- # 允许 0.2m 的误差(门可能离地有一点距离,或者有门槛)
- if abs(dist) > 0.2:
- return False, f"门底部距地面 {dist:.3f}m (puck_z={puck_z:.3f})"
- return True, ""
- def _filter_door_by_properties(self, door: Door3D) -> Tuple[bool, List[str]]:
- """根据物理特性过滤门
- 注意:Z 轴是高度方向(向上)
- 当前只检查尺寸是否为空,不对长宽高具体数值做限制
- """
- reasons = []
- size = door.bbox_8points.max(axis=0) - door.bbox_8points.min(axis=0)
- height = size[2] # Z 方向是高度
- width = max(size[0], size[1]) # 宽度是 X/Y 中较大的(门的平面方向)
- thickness = min(size[0], size[1]) # 厚度是 X/Y 中较小的(门的深度方向)
- # 只检查尺寸是否有效,不对具体数值做限制
- if height <= 0:
- reasons.append(f"高度无效 ({height:.2f}m)")
- if width <= 0:
- reasons.append(f"宽度无效 ({width:.2f}m)")
- return len(reasons) == 0, reasons
- def _estimate_entrance_from_poses(
- self,
- pc: o3d.geometry.PointCloud,
- contour: Optional[np.ndarray] = None,
- tri = None
- ) -> Optional[np.ndarray]:
- """
- 当没有检测到门时,基于点位信息估计入户门位置
- 核心假设:
- 1. 入户门在建筑边缘 → 点位到轮廓边界的距离越近越可能是入户门
- 2. 室内点位被墙壁遮挡 → 可见性不一致的点位可能在室内
- 3. 入户门通常在建筑一端 → 距离中心远的点位更可能是入户方向
- 评分策略:
- - 边界距离评分 (70%): 点位 XY 投影到轮廓边界的距离
- - 可见性一致性评分 (过滤): 过滤掉明显在室内的点位
- - 距离中心评分 (30%): 点位到几何中心的距离
- Args:
- pc: 场景点云
- contour: 轮廓点集 (可选,None 则内部计算)
- tri: Delaunay 三角网 (可选)
- Returns:
- 估计的入户门位置 (x, y, z),无法估计时返回 None
- """
- if not self.poses:
- print("⚠️ 没有点位数据,无法估计入户门")
- return None
- print("\n=== 基于点位信息估计入户门 ===")
- # ========== 步骤 1:准备数据 ==========
- # 重新加载 vision.txt 获取完整的点位信息(包括 visibles)
- vision_file = self.scene_folder / "vision.txt"
- if not vision_file.exists():
- print("⚠️ vision.txt 不存在,无法获取可见性信息")
- return None
- with open(vision_file, 'r') as f:
- vision_data = json.load(f)
- # 构建点位查找表
- pose_lookup = {}
- for loc in vision_data.get('sweepLocations', []):
- uuid = str(loc['uuid'])
- pose_lookup[uuid] = {
- 'id': loc['id'],
- 'pose': loc['pose'],
- 'puck': loc.get('puck', {}),
- 'visibles': loc.get('visibles', []),
- 'position': np.array([
- loc['pose']['translation']['x'],
- loc['pose']['translation']['y'],
- loc['pose']['translation']['z']
- ])
- }
- # 收集所有点位信息
- pose_info = []
- for uuid, data in pose_lookup.items():
- pose_info.append({
- 'uuid': uuid,
- 'id': data['id'],
- 'position': data['position'],
- 'position_xy': data['position'][:2],
- 'visibles': set(data['visibles']),
- 'puck_z': data['puck'].get('z', 0)
- })
- # ========== 步骤 2:计算轮廓(用于边界距离评分) ==========
- if contour is None:
- # 使用 XY 投影去噪后的点云计算轮廓
- points_xy_denoised = self._denoise_xy_projection(
- np.asarray(pc.points),
- grid_size=0.15,
- min_points_per_cell=5
- )
- contour, tri = self._compute_xy_contour(points_xy_denoised)
- print(f"轮廓点数量:{len(contour)}")
- # ========== 步骤 3:可见性一致性过滤 ==========
- # 过滤掉明显在室内的点位
- # 规则:如果 A 对 B 可见,A 对 C 不可见,但 AC 距离 < AB 距离 → A 在室内(被墙遮挡)
- filtered_poses = []
- filtered_out = []
- for i, pose_a in enumerate(pose_info):
- is_indoor = False
- for j, pose_b in enumerate(pose_info):
- if i == j:
- continue
- # 检查 A 对 B 是否可见
- b_id = pose_b['id']
- a_visible_to_b = b_id in pose_a['visibles']
- # 计算 AB 距离
- dist_ab = np.linalg.norm(pose_a['position_xy'] - pose_b['position_xy'])
- # 检查其他点位 C
- for k, pose_c in enumerate(pose_info):
- if k == i or k == j:
- continue
- c_id = pose_c['id']
- a_visible_to_c = c_id in pose_a['visibles']
- dist_ac = np.linalg.norm(pose_a['position_xy'] - pose_c['position_xy'])
- # 规则:A 对 B 可见,A 对 C 不可见,但 AC < AB → A 在室内
- if a_visible_to_b and not a_visible_to_c and dist_ac < dist_ab:
- is_indoor = True
- filtered_out.append((pose_a, f"遮挡不一致:可见 B({b_id}) 不可见 C({c_id}), 但 AC({dist_ac:.2f}) < AB({dist_ab:.2f})"))
- break
- if is_indoor:
- break
- if not is_indoor:
- filtered_poses.append(pose_a)
- print(f"\n可见性过滤:{len(pose_info)} → {len(filtered_poses)} 个点位")
- if filtered_out:
- print("被过滤点位 (可能在室内):")
- for p, reason in filtered_out[:5]: # 只显示前 5 个
- print(f" 点位 {p['uuid']}: {reason}")
- # 如果没有点位通过过滤,回退到所有点位
- if not filtered_poses:
- print("⚠️ 所有点位都被过滤,使用全部点位")
- filtered_poses = pose_info
- # ========== 步骤 4:计算几何中心 ==========
- all_positions_xy = np.array([p['position_xy'] for p in filtered_poses])
- centroid = np.mean(all_positions_xy, axis=0)
- # ========== 步骤 5:计算每个点位的评分 ==========
- # 边界距离评分 (70%) + 距离中心评分 (30%)
- # 计算每个点到轮廓边界的距离
- boundary_distances = []
- for p in filtered_poses:
- dist = self._compute_distance_to_contour_boundary(p['position_xy'], contour)
- boundary_distances.append(dist)
- # 计算每个点到中心的距离
- center_distances = np.linalg.norm(all_positions_xy - centroid, axis=1)
- # 归一化
- max_boundary_dist = max(boundary_distances) if boundary_distances else 1.0
- max_center_dist = center_distances.max() if center_distances.max() > 0 else 1.0
- # 计算评分
- print("\n点位评分详情:")
- best_score = -float('inf')
- best_pose = None
- for i, p in enumerate(filtered_poses):
- # 边界距离评分 (70%): 越靠近边界分越高
- if max_boundary_dist > 0:
- boundary_score = (1 - boundary_distances[i] / max_boundary_dist) * 70
- else:
- boundary_score = 35 # 所有点都在边界上
- # 距离中心评分 (30%): 越远离中心分越高
- if max_center_dist > 0:
- center_score = (center_distances[i] / max_center_dist) * 30
- else:
- center_score = 15
- total_score = boundary_score + center_score
- print(f" 点位 {p['uuid']} (ID={p['id']}): 边界={boundary_score:.1f} + 中心距={center_score:.1f} = {total_score:.1f} "
- f"(边界距离={boundary_distances[i]:.2f}m, 中心距={center_distances[i]:.2f}m)")
- if total_score > best_score:
- best_score = total_score
- best_pose = p
- print(f"\n选择点位 {best_pose['uuid']} (综合评分={best_score:.1f})")
- return best_pose['position']
- def _compute_distance_to_contour_boundary(
- self,
- point_xy: np.ndarray,
- contour: np.ndarray
- ) -> float:
- """
- 计算点到轮廓边界的最短距离
- Args:
- point_xy: 2D 点 [x, y]
- contour: 轮廓点集 (K, 2)
- Returns:
- 最短距离
- """
- min_dist = float('inf')
- n = len(contour)
- for i in range(n):
- p1 = contour[i]
- p2 = contour[(i + 1) % n]
- dist = self._point_to_line_segment_distance(point_xy, p1, p2)
- min_dist = min(min_dist, dist)
- return min_dist
- def _compute_distance_to_pointcloud_boundary(
- self,
- point: np.ndarray,
- pc: o3d.geometry.PointCloud
- ) -> float:
- """
- 计算点到点云边界的距离
- Args:
- point: 3D 点 [x, y, z]
- pc: 点云
- Returns:
- 到边界的距离
- """
- # 投影到 XY 平面
- points_xy = np.asarray(pc.points)[:, :2]
- point_xy = point[:2]
- # 计算点云边界(凸包)
- from scipy.spatial import ConvexHull
- hull = ConvexHull(points_xy)
- hull_indices = hull.vertices
- hull_points = points_xy[hull_indices]
- # 计算点到凸包边界的最短距离
- min_dist = float('inf')
- n = len(hull_points)
- for i in range(n):
- p1 = hull_points[i]
- p2 = hull_points[(i + 1) % n]
- dist = self._point_to_line_segment_distance(point_xy, p1, p2)
- min_dist = min(min_dist, dist)
- return min_dist
- def _identify_entrance_door(
- self,
- doors: List[Door3D],
- pc: o3d.geometry.PointCloud,
- contour: Optional[np.ndarray] = None,
- tri = None
- ) -> int:
- """
- 识别入户门 ID
- Args:
- doors: 门列表
- pc: 场景点云
- contour: 点云轮廓(可选,用于边界评分)
- tri: Delaunay 三角网(可选,用于点在多边形内判断)
- Returns:
- 入户门在 doors 列表中的索引,-1 表示无法确定
- """
- if self.entrance_method == "json":
- # 从 JSON 文件读取入户门 ID
- if self.entrance_json_path is None:
- print("⚠️ entrance_json_path 未指定,无法使用 json 方法")
- return -1
- with open(self.entrance_json_path, 'r') as f:
- data = json.load(f)
- entrance_door_id = data.get('entrance_door_id')
- if entrance_door_id is not None:
- self.entrance_door_id = entrance_door_id
- print(f"\n从 JSON 读取入户门 ID: {entrance_door_id}")
- # 找到对应的门(通过中心坐标匹配)
- if entrance_door_id is not None:
- # 尝试匹配 door_id
- for i, door in enumerate(doors):
- # 检查是否有 source 包含 door_id 信息
- for src in door.source_detections:
- if src.get('door_id') == entrance_door_id:
- print(f"匹配到门 {i}")
- return i
- # 如果找不到,返回中心最接近 world_position 的门
- world_position = data.get('world_position')
- if world_position is not None:
- wp = np.array(world_position)
- min_dist = float('inf')
- best_idx = -1
- for i, door in enumerate(doors):
- dist = np.linalg.norm(door.center - wp)
- if dist < min_dist:
- min_dist = dist
- best_idx = i
- if best_idx >= 0:
- print(f"通过 world_position 匹配到门 {best_idx} (距离={min_dist:.3f}m)")
- return best_idx
- return -1
- else:
- # 使用评分方法
- if len(doors) == 0:
- # ========== 没有检测到门时的回退策略 ==========
- print("\n⚠️ 未检测到任何门,使用点位信息估计入户门位置")
- estimated_entrance = self._estimate_entrance_from_poses(
- pc=pc,
- contour=contour,
- tri=tri
- )
- if estimated_entrance is not None:
- print(f"估计入户门位置:{estimated_entrance}")
- # 创建一个虚拟的入户门标记(不加入 doors 列表,仅用于可视化)
- self.estimated_entrance_position = estimated_entrance
- return -1
- if len(doors) == 1:
- return 0
- # 计算轮廓(用于边界评分)
- use_contour = False
- if contour is None:
- # 使用 XY 投影去噪后的点云计算轮廓
- points_xy_denoised = self._denoise_xy_projection(
- np.asarray(pc.points),
- grid_size=0.15,
- min_points_per_cell=5
- )
- contour, tri = self._compute_xy_contour(points_xy_denoised)
- print(f"轮廓点数量:{len(contour)}")
- use_contour = True
- # 计算每个门的评分
- all_centers = np.array([d.center for d in doors])
- best_score = -1
- best_idx = 0
- print("\n入户门评分详情:")
- for i, door in enumerate(doors):
- size = door.bbox_8points.max(axis=0) - door.bbox_8points.min(axis=0)
- height = size[2] # Z 方向是高度
- width = max(size[0], size[1]) # X/Y 中较大的是宽度
- # 尺寸评分(理想高 2.1m 宽 1.0m)
- height_score = max(0, 15 - abs(height - 2.1) * 10)
- width_score = max(0, 15 - abs(width - 1.0) * 12)
- size_score = height_score + width_score
- # 边缘位置评分(基于到其他门的距离)
- dists_to_others = np.linalg.norm(all_centers - door.center, axis=1)
- avg_dist = dists_to_others.mean()
- edge_score = min(25, avg_dist * 10)
- # 多视角支持
- source_count = len(door.source_detections)
- view_score = min(10, source_count * 3)
- # 边界评分(新增:入户门应该在建筑轮廓边缘)
- boundary_score = 0.0
- if use_contour and contour is not None:
- boundary_score = self._compute_door_boundary_score(door, contour, tri)
- total = size_score + edge_score + view_score + boundary_score
- # 打印每个门的评分详情
- print(f" 门{i}: 尺寸={size_score:.1f} + 边缘={edge_score:.1f} + 视角={view_score:.1f} + 边界={boundary_score:.1f} = {total:.1f}")
- if total > best_score:
- best_score = total
- best_idx = i
- print(f"\n入户门选择:门 {best_idx} (得分={best_score:.1f})")
- return best_idx
- def _rgb_depth_to_pointcloud(
- self,
- rgb: np.ndarray,
- depth: np.ndarray,
- pose_matrix: np.ndarray
- ) -> o3d.geometry.PointCloud:
- """将 RGB-D 转换为世界坐标系点云"""
- H, W = depth.shape
- sph = Intrinsic_Spherical_NP(W, H)
- px, py = np.meshgrid(np.arange(W), np.arange(H))
- px_flat = px.flatten().astype(np.float64)
- py_flat = py.flatten().astype(np.float64)
- bx, by, bz = sph.bearing([px_flat, py_flat])
- bx, by, bz = np.array(bx), np.array(by), np.array(bz)
- mask = depth.flatten() > self.depth_min
- d = depth.flatten()[mask]
- if len(d) == 0:
- return o3d.geometry.PointCloud()
- pts_cam = np.stack([bx[mask] * d, by[mask] * d, bz[mask] * d], axis=1)
- R_z180 = np.diag([-1.0, -1.0, 1.0])
- pts_cam = pts_cam @ R_z180.T
- pts_w = (pose_matrix[:3, :3] @ pts_cam.T).T + pose_matrix[:3, 3]
- if rgb.shape[:2] != depth.shape:
- rgb_d = cv2.resize(rgb, (W, H), interpolation=cv2.INTER_LINEAR)
- else:
- rgb_d = rgb
- colors = rgb_d.reshape(-1, 3)[mask].astype(np.float64) / 255.0
- pc = o3d.geometry.PointCloud()
- pc.points = o3d.utility.Vector3dVector(pts_w)
- pc.colors = o3d.utility.Vector3dVector(colors)
- return pc
- def _translate_point_to_ceiling(
- self,
- point: np.ndarray,
- floor_ceiling_dist: float,
- ground_normal: np.ndarray,
- ground_d: float
- ) -> np.ndarray:
- """
- 将点沿法向量方向平移,平移距离 = 地面到天花板的距离
- Args:
- point: 原始点
- floor_ceiling_dist: 地面到天花板的距离
- ground_normal: 地面法向量(单位向量)
- ground_d: 地面平面方程的 d
- Returns:
- 平移后的点
- """
- # 平移向量 = floor_ceiling_dist * normal
- # 这样门整体会向上移动层高的距离
- # 门的底部会从地面移到天花板高度
- # 门的顶部会在天花板以上(保持门的原始高度)
- translation = floor_ceiling_dist * ground_normal
- return point + translation
- def _create_door_visualization(
- self,
- door: Door3D,
- ground_normal: np.ndarray,
- ground_d: float,
- floor_ceiling_dist: float,
- is_entrance: bool = False,
- translate_to_ceiling: bool = True
- ) -> o3d.geometry.PointCloud:
- """
- 创建门的可视化点云(带边框)
- Args:
- door: 3D 门对象
- ground_normal: 地面法向量
- ground_d: 地面平面方程的 d
- floor_ceiling_dist: 地面到天花板的距离
- is_entrance: 是否入户门
- translate_to_ceiling: 是否平移到天花板
- Returns:
- 门点云(带颜色)
- """
- # 门是平面矩形,bbox_8points 中:
- # - 索引 0-3:底面 4 个点(Z 最小,靠近地面)
- # - 索引 4-7:顶面 4 个点(Z 最大,靠近门顶)
- # 投影到天花板只需要 4 个点,取底面 4 个点或顶面 4 个点即可
- # 直接取底面 4 个点(索引 0-3)
- rect_points = door.bbox_8points[:4] # 底面 4 个点构成矩形
- # 生成矩形 4 条边上的点
- edge_points = []
- edges = [(0, 1), (1, 2), (2, 3), (3, 0)] # 4 条边
- for i, j in edges:
- p_start = rect_points[i]
- p_end = rect_points[j]
- # 在边上生成 15 个点(更密集的轮廓)
- for t in np.linspace(0, 1, 15):
- edge_points.append(p_start + t * (p_end - p_start))
- points = np.array(edge_points)
- if translate_to_ceiling:
- # 将每个点平移到天花板高度
- points = np.array([
- self._translate_point_to_ceiling(p, floor_ceiling_dist, ground_normal, ground_d)
- for p in points
- ])
- # 创建点云
- pc = o3d.geometry.PointCloud()
- pc.points = o3d.utility.Vector3dVector(points)
- # 设置颜色:入户门红色,其他门绿色
- if is_entrance:
- colors = np.array([[1.0, 0.0, 0.0]] * len(points)) # 红色
- else:
- colors = np.array([[0.0, 1.0, 0.0]] * len(points)) # 绿色
- pc.colors = o3d.utility.Vector3dVector(colors)
- return pc
- def _create_door_visualization_with_color(
- self,
- door: Door3D,
- ground_normal: np.ndarray,
- ground_d: float,
- floor_ceiling_dist: float,
- color: np.ndarray,
- translate_to_ceiling: bool = True
- ) -> o3d.geometry.PointCloud:
- """
- 创建门的可视化点云(带边框)- 指定颜色版本
- Args:
- door: 3D 门对象
- ground_normal: 地面法向量
- ground_d: 地面平面方程的 d
- floor_ceiling_dist: 地面到天花板的距离
- color: RGB 颜色 [R, G, B]
- translate_to_ceiling: 是否平移到天花板
- Returns:
- 门点云(带颜色)
- """
- # 门是平面矩形,bbox_8points 中:
- # - 索引 0-3:底面 4 个点(Z 最小,靠近地面)
- # - 索引 4-7:顶面 4 个点(Z 最大,靠近门顶)
- # 投影到天花板只需要 4 个点,取底面 4 个点或顶面 4 个点即可
- # 直接取底面 4 个点(索引 0-3)
- rect_points = door.bbox_8points[:4] # 底面 4 个点构成矩形
- # 生成矩形 4 条边上的点
- edge_points = []
- edges = [(0, 1), (1, 2), (2, 3), (3, 0)] # 4 条边
- for i, j in edges:
- p_start = rect_points[i]
- p_end = rect_points[j]
- # 在边上生成 15 个点(更密集的轮廓)
- for t in np.linspace(0, 1, 15):
- edge_points.append(p_start + t * (p_end - p_start))
- points = np.array(edge_points)
- if translate_to_ceiling:
- # 将每个点平移到天花板高度
- points = np.array([
- self._translate_point_to_ceiling(p, floor_ceiling_dist, ground_normal, ground_d)
- for p in points
- ])
- # 创建点云
- pc = o3d.geometry.PointCloud()
- pc.points = o3d.utility.Vector3dVector(points)
- # 设置颜色
- colors = np.array([color] * len(points))
- pc.colors = o3d.utility.Vector3dVector(colors)
- return pc
- def _create_pose_visualization(
- self,
- pose: PoseData,
- ground_normal: np.ndarray,
- ground_d: float,
- floor_ceiling_dist: float,
- translate_to_ceiling: bool = True
- ) -> o3d.geometry.PointCloud:
- """
- 创建拍摄点位的可视化
- Args:
- pose: 位姿数据
- ground_normal: 地面法向量
- ground_d: 地面平面方程的 d
- floor_ceiling_dist: 地面到天花板的距离
- translate_to_ceiling: 是否平移到天花板
- Returns:
- 点位点云(蓝色)
- """
- point = np.array([
- pose.translation['x'],
- pose.translation['y'],
- pose.translation['z']
- ])
- if translate_to_ceiling:
- point = self._translate_point_to_ceiling(point, floor_ceiling_dist, ground_normal, ground_d)
- # 创建点云(单个点)
- pc = o3d.geometry.PointCloud()
- pc.points = o3d.utility.Vector3dVector([point])
- pc.colors = o3d.utility.Vector3dVector([[0.0, 0.0, 1.0]]) # 蓝色
- return pc
- def process_and_visualize(
- self,
- output_ply_path: str,
- translate_to_ceiling: bool = True
- ):
- """
- 处理场景并生成可视化 PLY 文件
- Args:
- output_ply_path: 输出 PLY 文件路径
- translate_to_ceiling: 是否将门和点位平移到天花板
- """
- # 创建输出目录
- self.output_folder.mkdir(parents=True, exist_ok=True)
- rgb_files = sorted(
- self.rgb_folder.glob("*.jpg"),
- key=lambda x: int(x.stem)
- )
- if not rgb_files:
- raise FileNotFoundError(f"在 {self.rgb_folder} 中未找到 RGB 图像")
- print(f"找到 {len(rgb_files)} 张全景图,开始处理...")
- # ========== 第一步:收集所有点云和门检测 ==========
- combined_pc = o3d.geometry.PointCloud()
- door_candidates = []
- for rgb_file in tqdm(rgb_files, desc="检测门"):
- idx = rgb_file.stem
- pose = self.poses.get(idx)
- if pose is None:
- continue
- depth_path = self.depth_folder / f"{idx}.png"
- if not depth_path.exists():
- continue
- rgb = cv2.cvtColor(cv2.imread(str(rgb_file)), cv2.COLOR_BGR2RGB)
- depth = cv2.imread(str(depth_path), cv2.IMREAD_UNCHANGED).astype(np.float32) / self.depth_scale
- pose_matrix = self._build_pose_matrix(pose)
- # YOLOE 检测
- mask_3d_points, scores = self.detect_single_image(
- str(rgb_file), depth, pose_matrix
- )
- # 收集 3D 门候选
- for i, pts_3d in enumerate(mask_3d_points):
- if len(pts_3d) > 10:
- # 先过滤离群点
- pts_3d_filtered = self._filter_outliers(pts_3d, std_thresh=2.0)
- # 再根据深度一致性过滤
- pts_3d_filtered = self._filter_door_points_by_depth(pts_3d_filtered)
- bbox_min, bbox_max = self._axis_aligned_bbox(pts_3d_filtered)
- door_candidates.append({
- 'bbox_min': bbox_min,
- 'bbox_max': bbox_max,
- 'points_3d': pts_3d_filtered,
- 'source': {
- 'image': rgb_file.name,
- 'score': scores[i] if i < len(scores) else 0.0
- },
- 'source_uuid': idx # 保存来源点位 UUID
- })
- # RGB-D 转点云
- pc = self._rgb_depth_to_pointcloud(rgb, depth, pose_matrix)
- combined_pc += pc
- # 下采样
- print(f"\n融合前点数:{len(combined_pc.points)}")
- combined_pc = combined_pc.voxel_down_sample(self.voxel_size)
- print(f"融合后点数:{len(combined_pc.points)}")
- # ========== 第二步:地面和天花板拟合 ==========
- ground_normal, ground_d = self._fit_ground_plane_ransac(combined_pc)
- self.ground_d = ground_d # 保存用于后续计算
- ceiling_normal, ceiling_d, floor_ceiling_dist = self._fit_ceiling_plane(
- combined_pc, ground_normal, ground_d
- )
- # ========== 第三步:3D 门合并和过滤 ==========
- print(f"\n3D 门候选:{len(door_candidates)}")
- # 打印每个候选的尺寸
- for i, c in enumerate(door_candidates):
- size = c['bbox_max'] - c['bbox_min']
- print(f" 候选 {i}: 尺寸=[{size[0]:.3f}, {size[1]:.3f}, {size[2]:.3f}], 中心={((c['bbox_min']+c['bbox_max'])/2).round(3)}")
- doors_3d = self._merge_3d_doors(door_candidates)
- print(f"合并后门数量:{len(doors_3d)}")
- # 过滤 - 只检查物理特性,不检查地面距离
- # 地面距离只用于入户门候选筛选
- print("\n过滤 3D 门 (只检查物理特性)...")
- valid_doors = [] # 通过物理特性过滤的门
- filtered_doors = [] # 未通过物理特性过滤的门
- for door in doors_3d:
- # 物理特性过滤(高度、宽度、厚度)
- passed_prop, prop_reasons = self._filter_door_by_properties(door)
- if not passed_prop:
- filtered_doors.append((door, prop_reasons))
- continue
- valid_doors.append(door)
- print(f"通过物理特性过滤:{len(valid_doors)} 个门")
- if len(filtered_doors) > 0:
- print(f"被过滤(物理特性):{len(filtered_doors)} 个门")
- for door, reasons in filtered_doors:
- print(f" 门{door.id} (中心={door.center.round(2)}):")
- for reason in reasons:
- print(f" - {reason}")
- # 检查地面距离,用于入户门候选筛选
- print("\n检查地面距离 (用于入户门候选筛选)...")
- ground_valid_doors = [] # 通过地面距离过滤的门(入户门候选)
- # 优先使用 puck 参数进行地面距离检查
- use_puck = (self.ground_z_from_puck is not None)
- for door in valid_doors:
- if use_puck:
- passed_ground, ground_reason = self._filter_door_by_ground_puck(door)
- else:
- passed_ground, ground_reason = self._filter_door_by_ground(
- door, ground_normal, ground_d
- )
- if passed_ground:
- ground_valid_doors.append(door)
- else:
- print(f" 门{door.id} (中心={door.center.round(2)}): {ground_reason}")
- if use_puck:
- print(f"通过地面距离过滤(puck_z={self.ground_z_from_puck:.3f}):{len(ground_valid_doors)} 个门")
- # 调试输出:显示每个门使用的 puck_z
- for door in valid_doors:
- puck_z_used = None
- if door.source_uuid and door.source_uuid in self.puck_z_dict:
- puck_z_used = self.puck_z_dict[door.source_uuid]
- print(f" 门{door.id}: source_uuid={door.source_uuid}, puck_z={puck_z_used:.3f}")
- else:
- print(f" 门{door.id}: 使用中位数 puck_z={self.ground_z_from_puck:.3f}")
- else:
- print(f"通过地面距离过滤(RANSAC):{len(ground_valid_doors)} 个门")
- # ========== 第四步:识别入户门 ==========
- print("\n" + "=" * 40)
- print("入户门识别")
- print("=" * 40)
- # 入户门候选只从通过地面距离过滤的门中选择
- entrance_idx = self._identify_entrance_door(
- ground_valid_doors,
- combined_pc,
- contour=None, # None 表示在方法内部计算
- tri=None
- )
- # ========== 第五步:创建可视化 ==========
- print("\n创建可视化点云...")
- # 1. 原始场景点云(降采样用于可视化)
- vis_pc = combined_pc.voxel_down_sample(0.05)
- # 2. 门的可视化点云 - 显示所有门(包括被过滤的)
- door_vis_points = []
- door_vis_colors = []
- # 颜色定义
- COLOR_ENTRANCE = np.array([1.0, 0.0, 0.0]) # 红色 - 入户门
- COLOR_VALID = np.array([0.0, 1.0, 0.0]) # 绿色 - 有效门(通过物理特性过滤)
- COLOR_FILTERED = np.array([1.0, 1.0, 0.0]) # 黄色 - 被过滤的门
- print(f"\n可视化门:")
- # 找到入户门在 valid_doors 中的索引
- entrance_door_in_valid = None
- if entrance_idx >= 0 and entrance_idx < len(ground_valid_doors):
- entrance_door = ground_valid_doors[entrance_idx]
- # 找到这个门在 valid_doors 中的索引
- for i, door in enumerate(valid_doors):
- if door.id == entrance_door.id:
- entrance_door_in_valid = i
- break
- # 显示有效门
- for i, door in enumerate(valid_doors):
- is_entrance = (i == entrance_door_in_valid)
- if is_entrance:
- color = COLOR_ENTRANCE
- label = f"入户门 (红色)"
- else:
- color = COLOR_VALID
- label = f"有效门 (绿色)"
- door_pc = self._create_door_visualization_with_color(
- door, ground_normal, ground_d, floor_ceiling_dist, color,
- translate_to_ceiling=translate_to_ceiling
- )
- pts = np.asarray(door_pc.points)
- door_vis_points.append(pts)
- door_vis_colors.append(np.asarray(door_pc.colors))
- print(f" {label}: 中心={door.center.round(3)}, 点数={len(pts)}")
- # 显示被过滤的门(黄色)
- for door, reasons in filtered_doors:
- door_pc = self._create_door_visualization_with_color(
- door, ground_normal, ground_d, floor_ceiling_dist, COLOR_FILTERED,
- translate_to_ceiling=translate_to_ceiling
- )
- pts = np.asarray(door_pc.points)
- door_vis_points.append(pts)
- door_vis_colors.append(np.asarray(door_pc.colors))
- print(f" 被过滤的门 (黄色): 中心={door.center.round(3)}, 点数={len(pts)}")
- for reason in reasons:
- print(f" - {reason}")
- if door_vis_points:
- door_vis_pc = o3d.geometry.PointCloud()
- door_vis_pc.points = o3d.utility.Vector3dVector(np.vstack(door_vis_points))
- door_vis_pc.colors = o3d.utility.Vector3dVector(np.vstack(door_vis_colors))
- # 3. 拍摄点位可视化
- pose_vis_points = []
- pose_vis_colors = []
- for pose in self.poses.values():
- pose_pc = self._create_pose_visualization(
- pose, ground_normal, ground_d, floor_ceiling_dist,
- translate_to_ceiling=translate_to_ceiling
- )
- pose_vis_points.append(np.asarray(pose_pc.points))
- pose_vis_colors.append(np.asarray(pose_pc.colors))
- if pose_vis_points:
- pose_vis_pc = o3d.geometry.PointCloud()
- pose_vis_pc.points = o3d.utility.Vector3dVector(np.vstack(pose_vis_points))
- pose_vis_pc.colors = o3d.utility.Vector3dVector(np.vstack(pose_vis_colors))
- # 4. 估计的入户门位置可视化(当没有检测到门时)
- entrance_vis_pc = None
- if self.estimated_entrance_position is not None:
- print("\n可视化估计的入户门位置...")
- entrance_point = self.estimated_entrance_position.copy()
- # 平移到天花板高度
- if translate_to_ceiling:
- entrance_point = self._translate_point_to_ceiling(
- entrance_point, floor_ceiling_dist, ground_normal, ground_d
- )
- # 创建一个较大的红色球体标记
- from scipy.spatial import ConvexHull
- # 生成球面点
- phi = np.linspace(0, np.pi, 10)
- theta = np.linspace(0, 2 * np.pi, 20)
- phi, theta = np.meshgrid(phi, theta)
- r = 0.15 # 球体半径
- x = entrance_point[0] + r * np.sin(phi) * np.cos(theta)
- y = entrance_point[1] + r * np.sin(phi) * np.sin(theta)
- z = entrance_point[2] + r * np.cos(phi)
- sphere_points = np.column_stack([x.flatten(), y.flatten(), z.flatten()])
- entrance_vis_pc = o3d.geometry.PointCloud()
- entrance_vis_pc.points = o3d.utility.Vector3dVector(sphere_points)
- entrance_colors = np.array([[1.0, 0.0, 0.0]] * len(sphere_points)) # 红色
- entrance_vis_pc.colors = o3d.utility.Vector3dVector(entrance_colors)
- print(f" 估计入户门位置 (红色球体): {entrance_point.round(3)}")
- # ========== 第六步:合并并保存 ==========
- # 创建一个完整的场景点云
- full_scene = o3d.geometry.PointCloud()
- # 添加原始点云(灰色)
- full_scene += vis_pc
- # 添加门点云
- if door_vis_points:
- full_scene += door_vis_pc
- # 添加拍摄点位点云
- if pose_vis_points:
- full_scene += pose_vis_pc
- # 添加估计的入户门位置可视化
- if entrance_vis_pc is not None:
- full_scene += entrance_vis_pc
- # 保存 PLY 文件
- # 创建输出目录
- os.makedirs(os.path.dirname(output_ply_path), exist_ok=True)
- o3d.io.write_point_cloud(output_ply_path, full_scene)
- print(f"\n可视化结果已保存:{output_ply_path}")
- # ========== 打印汇总 ==========
- print("\n" + "=" * 50)
- print("可视化汇总")
- print("=" * 50)
- print(f"原始点云点数:{len(vis_pc.points)}")
- print(f"有效门数量:{len(valid_doors)}")
- if entrance_idx >= 0:
- print(f"入户门索引:{entrance_idx}")
- elif self.estimated_entrance_position is not None:
- print(f"估计入户门位置:{self.estimated_entrance_position.round(3)} (红色球体标记)")
- else:
- print("入户门:未确定")
- print(f"拍摄点位数量:{len(self.poses)}")
- print(f"门点云点数:{sum(len(p) for p in door_vis_points) if door_vis_points else 0}")
- print(f"点位点云点数:{sum(len(p) for p in pose_vis_points) if pose_vis_points else 0}")
- if entrance_vis_pc is not None:
- print(f"估计入户门标记点数:{len(entrance_vis_pc.points)}")
- print(f"总点数:{len(full_scene.points)}")
- if translate_to_ceiling:
- print(f"\n门和点位已平移到天花板高度 (平移距离={floor_ceiling_dist:.3f}m)")
- else:
- print("\n门和点位保持原始位置")
- print("\n颜色说明:")
- print(" - 原始点云:RGB 颜色")
- print(" - 入户门:红色(门边框 / 估计位置球体)")
- print(" - 有效门:绿色(通过物理特性过滤)")
- print(" - 被过滤的门:黄色(物理特性不符合)")
- print(" - 拍摄点位:蓝色")
- print("=" * 50)
- return full_scene, valid_doors, entrance_idx
- # ============================================================================
- # 主函数
- # ============================================================================
- def main():
- parser = argparse.ArgumentParser(
- description="3D 场景可视化 - RANSAC 地面拟合 + 门平移到天花板",
- formatter_class=argparse.RawDescriptionHelpFormatter,
- epilog="""
- 示例:
- # 处理单个场景
- python visualize_scene_3d.py -s scene0001 --vis-ply
- # 指定入户门 JSON 文件
- python visualize_scene_3d.py -s scene0001 --vis-ply \\
- --entrance-method json \\
- --entrance-json /path/to/entrance.json
- # 调整地面拟合参数
- python visualize_scene_3d.py -s scene0001 --vis-ply \\
- --ground-ransac-dist 0.03 \\
- --ceiling-percentile 98
- # 不平移到天花板(保持原始位置)
- python visualize_scene_3d.py -s scene0001 --vis-ply --no-translate
- # 调整门过滤阈值
- python visualize_scene_3d.py -s scene0001 --vis-ply \\
- --door-ground-dist 0.15 \\
- --door-height-min 1.2
- # ========== 批处理模式 ==========
- # 处理 ten_pano_demo 目录下所有场景
- python visualize_scene_3d.py --batch-base /path/to/ten_pano_demo --vis-ply
- # 指定批处理输出目录
- python visualize_scene_3d.py --batch-base /path/to/ten_pano_demo \\
- --batch-output /path/to/output --vis-ply
- # 批处理 + 不平移
- python visualize_scene_3d.py --batch-base /path/to/ten_pano_demo \\
- --vis-ply --no-translate
- """
- )
- parser.add_argument("--scene", "-s", type=str, default="scene0001",
- help="场景文件夹")
- parser.add_argument("--model", "-m", type=str, default="yoloe-26x-seg.pt",
- help="YOLOE 模型路径")
- parser.add_argument("--conf", type=float, default=0.35,
- help="置信度阈值")
- parser.add_argument("--iou", type=float, default=0.45,
- help="NMS IoU 阈值")
- parser.add_argument("--voxel-size", type=float, default=0.03,
- help="点云体素大小")
- parser.add_argument("--depth-scale", type=float, default=256.0,
- help="深度图缩放因子")
- # 地面/天花板参数
- parser.add_argument("--ground-ransac-dist", type=float, default=0.05,
- help="RANSAC 地面拟合距离阈值")
- parser.add_argument("--ceiling-percentile", type=float, default=95.0,
- help="天花板点分位数")
- # 门过滤参数
- parser.add_argument("--door-ground-dist", type=float, default=0.1,
- help="门底部距地面最大距离")
- parser.add_argument("--door-height-min", type=float, default=1.0,
- help="门最小高度")
- parser.add_argument("--door-height-max", type=float, default=3.0,
- help="门最大高度")
- parser.add_argument("--door-width-min", type=float, default=0.3,
- help="门最小宽度")
- parser.add_argument("--door-width-max", type=float, default=3.0,
- help="门最大宽度")
- parser.add_argument("--door-thickness-max", type=float, default=0.5,
- help="门最大厚度")
- # 3D 门合并参数
- parser.add_argument("--merge-iou-thresh", type=float, default=0.1,
- help="门合并 IoU 阈值(默认 0.1,降低敏感度)")
- parser.add_argument("--merge-dist-thresh", type=float, default=0.3,
- help="门合并中心距离阈值(默认 0.3 米)")
- parser.add_argument("--merge-z-overlap-thresh", type=float, default=0.5,
- help="门合并 Z 方向重叠度阈值(默认 0.5,0-1 范围)")
- # 入户门参数
- parser.add_argument("--entrance-method", type=str, default="score",
- choices=["score", "json"],
- help="入户门判定方法")
- parser.add_argument("--entrance-json", type=str, default=None,
- help="入户门 JSON 文件路径")
- # 可视化参数
- parser.add_argument("--vis-ply", action="store_true",
- help="保存可视化 PLY 文件")
- parser.add_argument("--vis-output", type=str, default=None,
- help="输出 PLY 文件路径")
- parser.add_argument("--no-translate", action="store_true",
- help="不平移到天花板,保持原始位置")
- # 批处理参数
- parser.add_argument("--batch-base", type=str, default=None,
- help="批处理基础目录(如 ten_pano_demo),自动处理所有子场景")
- parser.add_argument("--batch-output", type=str, default=None,
- help="批处理输出基础目录,默认在基础目录下创建 output 文件夹")
- args = parser.parse_args()
- # ========== 批处理模式 ==========
- if args.batch_base:
- process_all_scenes(args)
- return
- # ========== 单场景模式 ==========
- if not Path(args.scene).exists():
- print(f"❌ 场景文件夹不存在:{args.scene}")
- return
- # 确定输出路径
- if args.vis_output is None:
- output_ply = Path(args.scene) / "output" / "visualization.ply"
- else:
- output_ply = Path(args.vis_output)
- processor = SceneVisualizer(
- scene_folder=args.scene,
- model_path=args.model,
- conf=args.conf,
- iou=args.iou,
- voxel_size=args.voxel_size,
- depth_scale=args.depth_scale,
- ground_ransac_dist=args.ground_ransac_dist,
- ceiling_percentile=args.ceiling_percentile,
- door_ground_dist=args.door_ground_dist,
- door_height_min=args.door_height_min,
- door_height_max=args.door_height_max,
- door_width_min=args.door_width_min,
- door_width_max=args.door_width_max,
- door_thickness_max=args.door_thickness_max,
- merge_iou_thresh=args.merge_iou_thresh,
- merge_dist_thresh=args.merge_dist_thresh,
- merge_z_overlap_thresh=args.merge_z_overlap_thresh,
- entrance_method=args.entrance_method,
- entrance_json_path=args.entrance_json,
- )
- if args.vis_ply:
- processor.process_and_visualize(
- output_ply_path=str(output_ply),
- translate_to_ceiling=not args.no_translate
- )
- else:
- print("⚠️ 未指定 --vis-ply,不进行可视化")
- print("请使用 --vis-ply 参数来生成可视化 PLY 文件")
- def process_all_scenes(args):
- """
- 批处理模式:处理基础目录下的所有场景
- 目录结构假设:
- base_dir/
- ├── depth/
- │ ├── scene1/
- │ │ └── depthmap/
- │ └── scene2/
- ├── rgb/
- │ ├── scene1/
- │ └── scene2/
- └── vision/
- ├── scene1/
- └── scene2/
- """
- from pathlib import Path
- from tqdm import tqdm
- base_dir = Path(args.batch_base)
- if not base_dir.exists():
- print(f"❌ 批处理基础目录不存在:{base_dir}")
- return
- # 确定输出目录
- if args.batch_output:
- output_base = Path(args.batch_output)
- else:
- output_base = base_dir / "output"
- # 获取所有场景(从 depth 目录获取)
- depth_dir = base_dir / "depth"
- if not depth_dir.exists():
- print(f"❌ depth 目录不存在:{depth_dir}")
- return
- scene_names = sorted([d.name for d in depth_dir.iterdir() if d.is_dir()])
- if not scene_names:
- print(f"⚠️ 未找到任何场景")
- return
- print(f"找到 {len(scene_names)} 个场景:{scene_names}")
- print(f"输出目录:{output_base}")
- print("=" * 60)
- # 创建输出基础目录
- output_base.mkdir(parents=True, exist_ok=True)
- for i, scene_name in enumerate(tqdm(scene_names, desc="总体进度")):
- print(f"\n[{i+1}/{len(scene_names)}] 处理场景:{scene_name}")
- print("-" * 40)
- # 构建各路径
- scene_rgb_dir = base_dir / "rgb" / scene_name
- scene_depth_dir = base_dir / "depth" / scene_name / "depthmap"
- scene_vision_dir = base_dir / "vision" / scene_name
- scene_vision_file = scene_vision_dir / "vision.txt"
- # 检查路径是否存在
- if not scene_rgb_dir.exists():
- print(f"⚠️ 跳过 {scene_name}: rgb 目录不存在")
- continue
- if not scene_depth_dir.exists():
- print(f"⚠️ 跳过 {scene_name}: depth 目录不存在")
- continue
- if not scene_vision_file.exists():
- print(f"⚠️ 跳过 {scene_name}: vision.txt 不存在")
- continue
- # 输出路径
- scene_output_dir = output_base / scene_name
- output_ply = scene_output_dir / "visualization.ply"
- # SceneVisualizer 期望的目录结构:
- # scene_folder/
- # pano_img/*.jpg
- # depth_img/depthmap/*.png
- # vision.txt
- # 所以需要创建一个临时结构或直接传递 scene_depth_dir 的父目录的父目录
- # 但实际上我们应该使用 scene_depth_dir 作为 depth 目录
- # 修改策略:传递 depth/scene_name 目录,因为它包含 depthmap
- # 而 rgb 和 vision 需要单独指定 - 但 SceneVisualizer 不支持
- # 解决方案:使用符号链接或者修改 SceneVisualizer 支持自定义子目录名
- # 简单方案:创建一个临时目录结构
- # 创建临时目录结构
- import tempfile
- import shutil
- temp_dir = tempfile.mkdtemp(prefix=f"scene_{scene_name}_")
- temp_scene = Path(temp_dir) / scene_name
- temp_scene.mkdir(parents=True, exist_ok=True)
- # 创建符号链接
- # SceneVisualizer 期望的结构:
- # scene_folder/pano_img/*.jpg
- # scene_folder/depth_img/{idx}.png
- # scene_folder/vision.txt
- (temp_scene / "pano_img").symlink_to(scene_rgb_dir)
- (temp_scene / "depth_img").symlink_to(scene_depth_dir) # 直接指向 depthmap 目录
- (temp_scene / "vision.txt").symlink_to(scene_vision_file)
- try:
- # 创建处理器
- processor = SceneVisualizer(
- scene_folder=str(temp_scene),
- model_path=args.model,
- conf=args.conf,
- iou=args.iou,
- voxel_size=args.voxel_size,
- depth_scale=args.depth_scale,
- ground_ransac_dist=args.ground_ransac_dist,
- ceiling_percentile=args.ceiling_percentile,
- door_ground_dist=args.door_ground_dist,
- door_height_min=args.door_height_min,
- door_height_max=args.door_height_max,
- door_width_min=args.door_width_min,
- door_width_max=args.door_width_max,
- door_thickness_max=args.door_thickness_max,
- merge_iou_thresh=args.merge_iou_thresh,
- merge_dist_thresh=args.merge_dist_thresh,
- merge_z_overlap_thresh=args.merge_z_overlap_thresh,
- entrance_method=args.entrance_method,
- entrance_json_path=args.entrance_json,
- )
- # 处理并可视化
- if args.vis_ply:
- try:
- processor.process_and_visualize(
- output_ply_path=str(output_ply),
- translate_to_ceiling=not args.no_translate
- )
- print(f"✓ 保存:{output_ply}")
- except Exception as e:
- print(f"❌ 处理失败:{e}")
- else:
- print("⚠️ 未指定 --vis-ply,跳过可视化")
- finally:
- # 清理临时目录
- import shutil
- shutil.rmtree(temp_dir, ignore_errors=True)
- if __name__ == "__main__":
- main()
|