-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtraj_to_chunk.py
More file actions
60 lines (51 loc) · 1.56 KB
/
traj_to_chunk.py
File metadata and controls
60 lines (51 loc) · 1.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import json
import numpy as np
from tqdm import tqdm
# ======================
# 配置
# ======================
INPUT_JSONL = "/home/star/project/qwen3_sft_jyh/LLaMA-Factory/data/llamafactory_traj_vl_all_filtered_3s5hz16_nodup.jsonl"
OUTPUT_JSONL = "/home/star/project/qwen3_sft_jyh/cluster/traj_chunks_0p6s_dxdyphi_navigation.jsonl"
# 固定 5 段
SEGMENTS = [
(0, 3),
(3, 6),
(6, 9),
(9, 12),
(12, 15),
]
def compute_dxdyphi(traj, i0, i1):
p0 = traj[i0]
p1 = traj[i1]
dx = p1[0] - p0[0]
dy = p1[1] - p0[1]
dphi = p1[2] - p0[2]
return dx, dy, dphi
count = 0
with open(INPUT_JSONL, "r") as fin, open(OUTPUT_JSONL, "w") as fout:
for line in tqdm(fin, desc="Processing trajectories"):
j = json.loads(line)
# 解析轨迹
traj_str = j["conversations"][1]["value"]
traj = json.loads(traj_str)["traj_future_3s"]
traj = np.array(traj, dtype=float) # [16,3]
chunks = []
for seg_id, (i0, i1) in enumerate(SEGMENTS):
dx, dy, dphi = compute_dxdyphi(traj, i0, i1)
chunks.append({
"seg_id": seg_id,
"dx": float(dx),
"dy": float(dy),
"dphi": float(dphi),
})
out = {
"id": j["id"],
"scene": j["scene"],
"time_str": j["time_str"],
"flag": j.get("flag", None),
"chunks": chunks
}
fout.write(json.dumps(out) + "\n")
count += 1
print(f"Done. Processed {count} trajectories.")
print(f"Saved to: {OUTPUT_JSONL}")