372 lines
15 KiB
Python
372 lines
15 KiB
Python
import os
|
||
import glob
|
||
import json
|
||
import numpy as np
|
||
import open3d as o3d
|
||
from typing import Dict, List, Optional, Union
|
||
from dataclasses import dataclass, field
|
||
from tqdm import tqdm
|
||
|
||
@dataclass
|
||
class FrameData:
|
||
"""单帧原始数据，完全未经过数学转换"""
|
||
frame_index: int # 帧编号（从0开始）
|
||
lidar_points: np.ndarray # (N, 3) float32，原始点云坐标
|
||
image_paths: Dict[str, str] # 视角名 -> 图像路径
|
||
camera_raw: Dict[str, dict] # 视角名称 -> {
|
||
# "projection_matrix": List[float], # 4x4 投影矩阵
|
||
# "image_size": List[int] # [width, height]
|
||
# }
|
||
labels: List[dict] # 原始3D框标注，每个框包含：
|
||
# - "num": int # 框ID
|
||
# - "color": str
|
||
# - "class": str # 原始类别名
|
||
# - "points": List[List[float]] # 8个点的xyz坐标
|
||
# - "center": List[float] # [x,y,z]
|
||
# - "rotateZ": float # 绕Z轴旋转角，弧度
|
||
# # 正方向：从X正半轴转向Y正半轴（逆时针）
|
||
# - "dx": float
|
||
# - "dy": float
|
||
# - "dz": float
|
||
contain_labels: bool = False # 是否包含标注框
|
||
|
||
@dataclass
|
||
class SampleData:
|
||
"""一个 sample 文件夹包含的全部帧"""
|
||
sample_id: str # 数字ID
|
||
frames: Dict[int, FrameData] # 帧索引 -> FrameData
|
||
|
||
class MyDataset:
|
||
f"""
|
||
私有数据集加载。点云坐标系（世界坐标系）Y轴正方向为前，X轴正方向为右，Z轴正方向为上。
|
||
数据组织规范：
|
||
root_folder/
|
||
sample{id}/
|
||
├── pcd_sequence{id}/ # 点云帧
|
||
├── 0/, 1/, .../ # 各帧图像
|
||
├── test{id}.json # 相机参数
|
||
└── test{id}-mark.json # 3D标注
|
||
"""
|
||
def __init__(self, root_folder: str):
|
||
self.root = root_folder
|
||
self._samples: Dict[str, SampleData] = {}
|
||
|
||
self.color2class = {
|
||
"#5414ED": "car",
|
||
"#F6EE64": "pick-up-truck",
|
||
"#F6A087": "small-truck",
|
||
"#BC4EF1": "truck",
|
||
"#4E9AF1": "bus",
|
||
"#F1A94E": "special-vehicle",
|
||
"#E1DFDD": "ignore",
|
||
"#F91906": "tricyclist-withrider",
|
||
"#FA5F51": "tricyclist-withoutrider",
|
||
"#B8CB30": "bicycle-withrider",
|
||
"#E6FD4E": "bicycle-withoutrider",
|
||
"#876363": "people",
|
||
"#2CBDF5": "crowd-people",
|
||
"#C9F52C": "crowd-bicycle",
|
||
"#DC6788": "crowd-car",
|
||
"#6EC913": "traffic-cone",
|
||
"#0DDE69": "plastic-barrier",
|
||
"#8260D2": "crash-barrels",
|
||
"#F1D1D1": "warning-triangle",
|
||
"#FE6DF4": "crowd-traffic-cone",
|
||
"#D1AA35": "crowd-plastic-barrier",
|
||
"#3BE8D0": "crowd-crash-barrels",
|
||
"#2B7567": "crowd-warning-triangle"
|
||
}
|
||
self.class2color = {cls: color for color, cls in self.color2class.items()}
|
||
|
||
self.cam2filename = {
|
||
"front_120": "scanofilm_surround_front_120_8M.jpg",
|
||
"front_left_100": "scanofilm_surround_front_left_100_2M.jpg",
|
||
"front_right_100": "scanofilm_surround_front_right_100_2M.jpg",
|
||
"rear_100": "scanofilm_surround_rear_100_2M.jpg",
|
||
"rear_left_100": "scanofilm_surround_rear_left_100_2M.jpg",
|
||
"rear_right_100": "scanofilm_surround_rear_right_100_2M.jpg",
|
||
}
|
||
|
||
self._load_all_samples()
|
||
# ------------------------------------------------------------------
|
||
# 路径与ID辅助方法
|
||
# ------------------------------------------------------------------
|
||
@staticmethod
|
||
def _extract_id(folder_name: str) -> str:
|
||
"""从文件夹名提取连续数字作为sample_id"""
|
||
return ''.join(c for c in os.path.basename(folder_name) if c.isdigit())
|
||
|
||
def _find_sample_folders(self) -> List[str]:
|
||
"""返回所有 sample{id} 文件夹路径"""
|
||
pattern = os.path.join(self.root, "sample*")
|
||
return [f for f in glob.glob(pattern) if os.path.isdir(f)]
|
||
|
||
def _find_paths_in_sample(self, sample_folder: str, sample_id: str) -> dict:
|
||
"""返回 sample 内部各资源路径字典"""
|
||
pcd_folder = os.path.join(sample_folder, f"pcd_sequence{sample_id}")
|
||
img_folders = []
|
||
if os.path.isdir(pcd_folder):
|
||
pcd_files = glob.glob(os.path.join(pcd_folder, "*.pcd"))
|
||
num_frames = len(pcd_files)
|
||
img_folders = [os.path.join(sample_folder, str(i)) for i in range(num_frames)]
|
||
page_json = os.path.join(sample_folder, f"test{sample_id}.json")
|
||
mark_json = os.path.join(sample_folder, f"test{sample_id}-mark.json")
|
||
return {
|
||
"pcd_folder": pcd_folder,
|
||
"img_folders": img_folders,
|
||
"page_json": page_json,
|
||
"mark_json": mark_json
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
# 原始数据读取（无转换）
|
||
# ------------------------------------------------------------------
|
||
def _read_lidar_points(self, pcd_path: str) -> np.ndarray:
|
||
try:
|
||
# Convert path to bytes for better Chinese path support on Windows
|
||
pcd_path_bytes = pcd_path.encode('utf-8') if isinstance(pcd_path, str) else pcd_path
|
||
pcd = o3d.io.read_point_cloud(
|
||
pcd_path_bytes,
|
||
remove_nan_points=True,
|
||
remove_infinite_points=True,
|
||
format="auto"
|
||
)
|
||
pcd_np = np.asarray(pcd.points, dtype=np.float32) # [N, 3]
|
||
return pcd_np
|
||
|
||
except Exception as e:
|
||
print(f"Error loading point cloud {pcd_path}: {e}")
|
||
return None
|
||
|
||
def _load_all_lidar_frames(self, pcd_folder: str) -> Dict[int, np.ndarray]:
|
||
idx2pcd_np = {}
|
||
pcd_files = glob.glob(os.path.join(pcd_folder, "*.pcd"))
|
||
for pcd_file in pcd_files:
|
||
idx = int(os.path.basename(pcd_file).split(".")[0].split("n")[-1])
|
||
pcd_np = self._read_lidar_points(pcd_file)
|
||
if pcd_np is not None:
|
||
idx2pcd_np[idx] = pcd_np
|
||
return idx2pcd_np
|
||
|
||
def _load_image_paths(self, img_folders: List[str]) -> Dict[int, Dict[str, str]]:
|
||
"""返回 帧索引 -> {视角名: 图像路径}"""
|
||
idx2imgs = {}
|
||
for folder in img_folders:
|
||
if not os.path.isdir(folder):
|
||
continue
|
||
idx = int(os.path.basename(folder))
|
||
paths = {}
|
||
for cam, fname in self.cam2filename.items():
|
||
full_path = os.path.join(folder, fname)
|
||
assert os.path.isfile(full_path), f"Image file not found: {full_path}"
|
||
paths[cam] = full_path
|
||
if paths:
|
||
idx2imgs[idx] = paths
|
||
return idx2imgs
|
||
|
||
def _load_camera_raw(self, page_json_path: str) -> Dict[int, Dict[str, dict]]:
|
||
"""返回 帧索引 -> {视角名: {"projection_matrix": list, "image_size": list}}"""
|
||
if not os.path.isfile(page_json_path):
|
||
return {}
|
||
with open(page_json_path, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
try:
|
||
extend_source = data['data']['files'][0]['extendSources']
|
||
except (KeyError, IndexError):
|
||
return {}
|
||
|
||
idx2cam = {}
|
||
for d in extend_source:
|
||
page_element = d.get('pageElement', {})
|
||
if "sensor" not in page_element:
|
||
continue
|
||
# 提取帧索引
|
||
try:
|
||
idx = int(d['fileName'].split('.')[0])
|
||
except:
|
||
continue
|
||
sensor = page_element['sensor']
|
||
|
||
cam_name = sensor.replace("ofilm_surround_", "")
|
||
for suffix in ["_2M", "_8M"]:
|
||
if cam_name.endswith(suffix):
|
||
cam_name = cam_name[:-len(suffix)]
|
||
break
|
||
|
||
if cam_name not in self.cam2filename:
|
||
continue
|
||
|
||
matrix = page_element.get('mtx', [])
|
||
image_size = page_element.get('imageSize', [])
|
||
if not matrix or not image_size:
|
||
continue
|
||
|
||
matrix = np.array(matrix, dtype=np.float64)
|
||
assert matrix.shape == (3, 4), f"Invalid projection matrix shape: {matrix.shape}"
|
||
cam_info = {
|
||
"projection_matrix": matrix,
|
||
"image_size": image_size
|
||
}
|
||
if idx not in idx2cam:
|
||
idx2cam[idx] = {}
|
||
idx2cam[idx][cam_name] = cam_info
|
||
return idx2cam
|
||
|
||
def extract_text_label_from_json(self, json_path: str) -> List[Dict]:
|
||
with open(json_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)['data']['list'][0]
|
||
|
||
type_dict = data['rate']
|
||
idx2rate = {}
|
||
for key, value in type_dict.items():
|
||
float_value = float(value['accuracy'].replace('%', '')) / 100.0
|
||
idx2rate[key] = float_value
|
||
|
||
# 找出最大准确率对应的索引
|
||
max_idx = max(idx2rate, key=idx2rate.get)
|
||
max_rate = idx2rate[max_idx]
|
||
|
||
if max_rate == 1.0:
|
||
result = data['result']
|
||
label_list = result[max_idx]
|
||
|
||
return label_list
|
||
else:
|
||
return []
|
||
|
||
def process_text_label(self, json_path: str):
|
||
label_list = self.extract_text_label_from_json(json_path)
|
||
|
||
if len(label_list) == 0:
|
||
return {}
|
||
|
||
idx2label: Dict[int, List[Dict]] = {}
|
||
for label in label_list:
|
||
idx2label[label['index']] = self.process_label_for_one_frame(label['value'])
|
||
return idx2label
|
||
|
||
def process_label_for_one_frame(self, label_list: List[Dict]) -> List[Dict]:
|
||
if len(label_list) == 0:
|
||
return []
|
||
|
||
# there could be multiple 3D bounding boxes in one frame
|
||
labels = []
|
||
for label in label_list:
|
||
new_label = {}
|
||
|
||
# num is the id of the 3D bounding box
|
||
new_label['num'] = label['num']
|
||
|
||
# label
|
||
new_label['label'] = label['label']
|
||
|
||
# class
|
||
new_label['class'] = label['label']['class-name']
|
||
|
||
# is_moving
|
||
new_label['is_moving'] = not bool(label['label']['static'])
|
||
|
||
# isolation
|
||
# True: 被建筑物，护栏，绿化带等隔离开
|
||
# False: 可⾏驶区域
|
||
new_label['isolation'] = bool(label['label']['isolation'])
|
||
|
||
# color
|
||
new_label['color'] = self.class2color[new_label['class']]
|
||
|
||
# corners
|
||
assert len(label['newPoints']) == 8, f"Points number for one 3D bounding box is 8, but got {len(label['newPoints'])}."
|
||
eight_points = []
|
||
for i in range(8):
|
||
eight_points.append(
|
||
[
|
||
float(label['newPoints'][i]['x']),
|
||
float(label['newPoints'][i]['y']),
|
||
float(label['newPoints'][i]['z'])
|
||
]
|
||
)
|
||
new_label['points'] = eight_points
|
||
|
||
# center
|
||
new_label['center'] = [
|
||
float(label['x']),
|
||
float(label['y']),
|
||
float(label['z'])
|
||
]
|
||
|
||
# rotate - 只保留绕Z轴的旋转
|
||
# 单位：弧度
|
||
# 正方向：从X轴正半轴转向Y轴正半轴（逆时针）
|
||
new_label['rotateZ'] = float(label['rotateZ'])
|
||
|
||
new_label['dx'] = float(label['width'])
|
||
new_label['dy'] = float(label['height'])
|
||
new_label['dz'] = float(label['depth'])
|
||
|
||
labels.append(new_label)
|
||
return labels
|
||
|
||
# ------------------------------------------------------------------
|
||
# Sample 级组装
|
||
# ------------------------------------------------------------------
|
||
def _load_one_sample(self, sample_folder: str) -> Optional[SampleData]:
|
||
"""加载单个 sample 文件夹，返回 SampleData 对象"""
|
||
sample_id = self._extract_id(sample_folder)
|
||
paths = self._find_paths_in_sample(sample_folder, sample_id)
|
||
|
||
# 读取各模态原始数据
|
||
lidar_dict = self._load_all_lidar_frames(paths["pcd_folder"])
|
||
img_dict = self._load_image_paths(paths["img_folders"])
|
||
cam_dict = self._load_camera_raw(paths["page_json"])
|
||
label_dict = self.process_text_label(paths["mark_json"])
|
||
if not (len(label_dict) == len(lidar_dict) == len(img_dict) == len(cam_dict)):
|
||
print(f"Sample {sample_id}: Mismatch in frame counts between lidar, images, cameras, and labels.")
|
||
return None
|
||
|
||
frames = {}
|
||
for idx in range(len(lidar_dict)):
|
||
frames[idx] = FrameData(
|
||
frame_index=idx,
|
||
lidar_points=lidar_dict[idx],
|
||
image_paths=img_dict[idx],
|
||
camera_raw=cam_dict[idx],
|
||
labels=label_dict[idx],
|
||
contain_labels=len(label_dict[idx]) > 0
|
||
)
|
||
if not frames:
|
||
return None
|
||
return SampleData(sample_id=sample_id, frames=frames)
|
||
|
||
def _load_all_samples(self):
|
||
"""遍历所有 sample 文件夹并加载"""
|
||
sample_folders = self._find_sample_folders()
|
||
process_bar = tqdm(sample_folders, desc="Loading samples")
|
||
for sf in process_bar:
|
||
sample = self._load_one_sample(sf)
|
||
if sample is not None:
|
||
self._samples[sample.sample_id] = sample
|
||
|
||
print(f"Loaded {len(self._samples)} samples.")
|
||
|
||
# ------------------------------------------------------------------
|
||
# 公开接口
|
||
# ------------------------------------------------------------------
|
||
@property
|
||
def sample_ids(self) -> List[str]:
|
||
return list(self._samples.keys())
|
||
|
||
def get_sample(self, sample_id: str) -> Optional[SampleData]:
|
||
return self._samples.get(sample_id)
|
||
|
||
def __len__(self) -> int:
|
||
return len(self._samples)
|
||
|
||
def __getitem__(self, idx: int) -> SampleData:
|
||
"""通过整数索引访问 sample（按sample_id排序）"""
|
||
sids = sorted(self.sample_ids)
|
||
if idx < 0 or idx >= len(sids):
|
||
raise IndexError(f"Index {idx} out of range [0, {len(sids)-1}]")
|
||
return self._samples[sids[idx]]
|
||
|
||
if __name__ == "__main__":
|
||
dataset = MyDataset(root_folder="sample") |