import json
import os
import numpy as np
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
from scipy.spatial.distance import directed_hausdorff
from PIL import Image, ImageDraw
# 轨迹点数据结构
class TrajectoryPoint:
def __init__(self, x, y, angle=None):
self.x = x
self.y = y
self.angle = angle
def to_array(self):
return np.array([self.x, self.y])
# 图像轨迹数据类
class ImageTrajectoryData:
def __init__(self, timestamp, image_path, roadpoints):
self.timestamp = datetime.strptime(timestamp, '%Y-%m-%d %H:%M:%S.%f')
self.image_path = image_path
self.roadpoints = [TrajectoryPoint(p[0], p[1], p[2] if len(p) > 2 else None)
for p in roadpoints]
self.label = "no" # 默认标签
def trajectory_array(self):
return np.array([p.to_array() for p in self.roadpoints])
# 轨迹分析器
class TrajectoryAnalyzer:
def __init__(self, jsonl_path, output_dir):
self.jsonl_path = jsonl_path
self.output_dir = output_dir
self.data = []
os.makedirs(output_dir, exist_ok=True)
def load_and_sort_data(self):
"""加载JSONL文件并按时间戳降序排序"""
with open(self.jsonl_path, 'r',encoding='utf-8') as f:
for line in f:
entry = json.loads(line)
img_data = ImageTrajectoryData(
entry['timestamp'],
entry['image_path'],
entry['mop_pnc_info']['roadpoints']
)
self.data.append(img_data)
# 按时间戳降序排序(最新在前)
self.data.sort(key=lambda x: x.timestamp, reverse=True)
print(f"Loaded {len(self.data)} records, sorted by timestamp (descending)")
def plot_trajectory(self, img_data, index):
"""绘制轨迹路线图并保存"""
img = Image.open(img_data.image_path)
draw = ImageDraw.Draw(img)
width, height = img.size
# 绘制轨迹点
for i, point in enumerate(img_data.roadpoints):
# 转换为图像坐标(假设轨迹点在0-1范围)
x = int(point.x * width)
y = int(point.y * height)
# 绘制点
radius = 5
draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill='red')
# 绘制方向(如果有角度)
if point.angle is not None:
length = 20
end_x = x + length * np.cos(point.angle)
end_y = y + length * np.sin(point.angle)
draw.line([x, y, end_x, end_y], fill='blue', width=2)
# 绘制标签
draw.text((10, 10), f"Label: {img_data.label}", fill='green')
draw.text((10, 30), f"Time: {img_data.timestamp}", fill='green')
# 保存图像
output_path = os.path.join(self.output_dir, f"traj_{index}.png")
img.save(output_path)
return output_path
def hausdorff_distance(self, traj1, traj2):
"""计算两条轨迹的Hausdorff距离"""
if len(traj1) == 0 or len(traj2) == 0:
return float('inf')
return max(directed_hausdorff(traj1, traj2)[0],
directed_hausdorff(traj2, traj1)[0])
def detect_trajectory_changes(self, threshold=0.1):
"""检测轨迹变化点并打标签"""
# 从最新向最旧遍历(索引0是最新)
for i in range(len(self.data) - 1):
current = self.data[i]
next_data = self.data[i + 1]
# 计算轨迹相似度
dist = self.hausdorff_distance(
current.trajectory_array(),
next_data.trajectory_array()
)
# 检测显著变化
if dist > threshold:
print(f"轨迹变化检测于 {current.timestamp}, 距离: {dist:.4f}")
# 标记变化点及前3秒
change_time = current.timestamp
for j in range(i, len(self.data)):
if (change_time - self.data[j].timestamp) <= timedelta(seconds=3):
self.data[j].label = "yes"
else:
break # 超过3秒停止
def process_and_save(self):
"""处理所有数据并保存结果"""
# 检测轨迹变化
self.detect_trajectory_changes()
# 绘制所有轨迹图并更新JSONL
output_jsonl = []
for i, img_data in enumerate(self.data):
# 绘制轨迹图
img_path = self.plot_trajectory(img_data, i)
# 更新JSONL条目
output_entry = {
"timestamp": img_data.timestamp.strftime('%Y-%m-%d %H:%M:%S.%f'),
"image_path": img_path,
"mop_pnc_info": {
"roadpoints": [[p.x, p.y, p.angle] for p in img_data.roadpoints],
"label": img_data.label # 添加标签
}
}
output_jsonl.append(json.dumps(output_entry))
# 保存更新后的JSONL文件
output_path = os.path.join(self.output_dir, "labeled_data.jsonl")
with open(output_path, 'w') as f:
f.write("\n".join(output_jsonl))
print(f"处理完成! 结果保存至: {output_path}")
# ===== 使用示例 =====
if __name__ == "__main__":
# 创建示例JSONL文件(实际应用中替换为真实文件路径)
jsonl_path = "240426F_gt_all_scenes_frames_lines_m5_hs_dz5_train_fisheye4_240729.jsonl"
output_dir = "./output"
# 创建分析器并运行
analyzer = TrajectoryAnalyzer(jsonl_path, output_dir)
analyzer.load_and_sort_data()
analyzer.process_and_save()
报错:
entry['mop_pnc_info']['roadpoints']
~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
TypeError: list indices must be integers or slices, not str
最新发布