#!/usr/bin/env python """一次性迁移脚本:旧「每会话一个子目录」布局 → v1.1.0 任务级池化布局。 旧布局(egodata v1.0,每会话一个子目录): /_NNNNNN/ ├── metadata.json / timestamps.json / calibration/*.json ├── videos//chunk-0000/.mp4 ├── data//chunk-0000/chunk_000000.parquet(每传感器一文件) ├── depth//.mkv | NNNNNN.png | NNNNNN.bin └── meta/{info.json, stats.json, tasks.jsonl, episodes/{chunk_000000.parquet, chunk-000/file-000.parquet}} 新布局(v1.1.0 任务级池化;v1.1.2 起文件前缀 file- → episode-): / ├── videos/chunk-NNN//episode-NNN.ext # 每 episode 一文件 ├── data/chunk-NNN/episode-NNN.parquet # 单文件稀疏列 └── meta/{info.json, stats.json, tasks.jsonl, episodes/chunk-NNN/episode-NNN.parquet} 迁移内容(每会话): 1. 按旧 metadata.json episode_index/created_at 排序定新全局序号 N( 已存在池化 episode 之后续排) 2. 视频 move: videos//chunk-0000/.mp4 → 池化路径 3. 深度: mkv 直接 move;png16/bin 序列经 numpy+ffmpeg 无损合成 单流 FFV1 gray16le mkv(数据无损;显示层热力图由读取端按需生成) 4. data parquet 列合并: timestamps.json 为骨架(含 wall_time/ hardware_ns)+ 各传感器/IMU 按 frame_index 对齐(同帧多槽行取首行) + action/手部占位零 + 旧 status. 列(如有) 5. episodes 分片 append 一行(created_at/video_codec/drop_stats 取自 metadata.json,calibration 取 calibration/*.json 原文) 6. 任务级 info.json(并集合并)/ tasks.jsonl / stats.json(v1.1.1 自含 count 累加器,无 .stats_state 边车) 7. pipeline.db 改写(recording.file_path → 任务目录 + episode_index; upload_task.session_path 同);改前自动备份 8. keypoints_output//<旧会话名>/ → /episode_{N:06d}/ 9. 旧会话目录先**全量快照复制**到 data/recordings/_migrate_backup_/ 再移动媒体(快照在前是 restore 无损的前提),迁移完成删除残壳 (--no-backup 跳过快照=显式破坏性模式;manifest 记录全部映射) 用法: venv/bin/python scripts/migrate_pooled_storage.py --dry-run venv/bin/python scripts/migrate_pooled_storage.py [--backup] [--task 名称] venv/bin/python scripts/migrate_pooled_storage.py --validate venv/bin/python scripts/migrate_pooled_storage.py --restore [--task 名称] 限制(一次性工具,已知可接受): - stats 为累加器语义:--restore 不回退 stats 数值(若任务无剩余 episode 则删除 stats 文件;有剩余则保留并含已迁移数值) - 旧格式行无视频帧-传感器一一对应元信息:传感器/IMU 按 frame_index 挂到同帧首行(旧读取端同口径) - 深度合成单流 FFV1(非双流),显示端按需热力图 - 无 ffmpeg 时跳过 png16/bin 深度会话并警告(视频/数据仍迁移) - 快照复制需要 2× 磁盘临时余量;确认备份无误后再删除备份区 """ import argparse import json import os import re import shutil import sqlite3 import subprocess import sys import time sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import numpy as np from config import settings from core.helpers import (episode_chunk_file, POOLED_CHUNK_SIZE, pooled_video_path, pooled_data_parquet_path, pooled_episodes_path, _legacy_episodes_shard_path, pooled_info_path, pooled_stats_path, pooled_tasks_jsonl_path, load_stats_acc, merge_stat_block, acc_to_stats_json, recalc_stats, list_task_episodes) from core.egodata_writer import (_episode_rows_table, _read_episode_rows, _atomic_write_parquet, _LockedFile) try: import fcntl as _fcntl except ImportError: _fcntl = None try: import cv2 # png16 解码 except ImportError: cv2 = None MANIFEST_VERSION = 1 MANIFEST_NAME = ".migrate_manifest.json" RECORDINGS_DIR = os.path.join(settings.DATA_DIR, "recordings") DB_PATH = os.path.join(settings.DATA_DIR, "pipeline.db") DEFAULT_SENSOR_DIM = 256 # ── 日志 / 报告 ────────────────────────────────────────────── class Report: """迁移执行报告(dry-run 只打印,实跑落到 backup 根目录)。""" def __init__(self): self.sessions = [] self.warnings = [] def add(self, task, session, n, kind, detail, created, db_rows, warns): self.sessions.append({ "task": task, "session": session, "N": n, "kind": kind, "detail": detail, "created": created, "db_rows": db_rows, }) self.warnings.extend(warns) def print(self, label): print(f"\n===== 迁移报告 [{label}] =====") for s in self.sessions: print(f" [{s['kind']}] {s['task']} :: {s['session']} → " f"episode {s['N']} ({s['detail']})") for f in s["created"]: print(f" + {f}") for d in s["db_rows"]: print(f" db {d['table']} id={d['id']}") if self.warnings: print(" ⚠ 警告:") for w in self.warnings: print(f" - {w}") if not self.sessions: print(" (无可迁移/已回滚的会话)") # ── 会话发现 ──────────────────────────────────────────────── def task_dirs(root, task_filter): out = [] for name in sorted(os.listdir(root)): if name.startswith("_") or name.startswith("."): continue d = os.path.join(root, name) if not os.path.isdir(d): continue if task_filter and name not in task_filter: continue out.append((name, d)) return out def legacy_sessions(task_dir, task): """任务目录下的旧会话子目录(有 metadata.json 且非池化结构目录)。""" out = [] try: children = os.listdir(task_dir) except OSError: return out for name in children: d = os.path.join(task_dir, name) if not os.path.isdir(d): continue if name in ("videos", "data", "meta", "depth", "calibration") \ or name.startswith("chunk-") or name.startswith("_"): continue if os.path.isfile(os.path.join(d, "metadata.json")): out.append(d) out.sort(key=lambda d: _session_sort_key(d, task)) return out def _session_sort_key(session_dir, task): """排序键:旧 metadata.json 的 episode_index,回退 created_at/名字。""" try: with open(os.path.join(session_dir, "metadata.json"), "r", encoding="utf-8") as f: m = json.load(f) except (OSError, json.JSONDecodeError): m = {} return (int(m.get("episode_index", 0) or 0), float(m.get("created_at", 0) or 0), os.path.basename(session_dir)) def load_json(path, default=None): try: with open(path, "r", encoding="utf-8") as f: return json.load(f) except (OSError, json.JSONDecodeError): return default # ── 旧数据读取 ────────────────────────────────────────────── def read_ts_rows(session_dir): """timestamps.json → ([{frame_index, timestamp, wall_time, hardware_ns}], total_frames);缺失返回 (None, None)。""" ts = load_json(os.path.join(session_dir, "timestamps.json")) if not ts or not isinstance(ts.get("timestamps"), list): return None, None return ts["timestamps"], int(ts.get("total_frames") or 0) def read_sensor_parquet(session_dir, sensor): """data//chunk-0000/chunk_000000.parquet → {frame_index: row}""" for rel in (f"data/{sensor}/chunk-0000/chunk_000000.parquet", f"data/{sensor}/chunk_000000.parquet", f"data/{sensor}/chunk-000/file-000.parquet"): p = os.path.join(session_dir, rel) if os.path.isfile(p): import pyarrow.parquet as pq t = pq.read_table(p) idx = {} for r in t.to_pylist(): fi = int(r.get("frame_index", 0)) if fi not in idx: # 同帧多槽行取首行 idx[fi] = r return idx return None def find_legacy_media(session_dir, metadata): """定位旧视频/深度文件。 返回 ({cam: (src, kind)}, {depth_slot: (src, kind)}) kind ∈ {move, png16, bin} """ vids, depths = {}, {} for cam, info in (metadata.get("cameras") or {}).items(): if not isinstance(info, dict): continue ctype = info.get("type", "rgb") if ctype == "depth": src = _find_depth_source(session_dir, cam) if src: depths[cam] = src continue # rgb: videos//chunk-0000/.mp4(兼容无 chunk 层) for rel in (f"videos/{cam}/chunk-0000/{cam}.mp4", f"videos/{cam}/{cam}.mp4"): p = os.path.join(session_dir, rel) if os.path.isfile(p): vids[cam] = (p, "move") break if cam not in vids: # 兜底:videos/ 下任意 mp4(槽名可能与 metadata 名不同) d = os.path.join(session_dir, "videos", cam) if os.path.isdir(d): for root, _dirs, files in os.walk(d): for fn in sorted(files): if fn.endswith(".mp4"): vids[cam] = (os.path.join(root, fn), "move") break if cam in vids: break return vids, depths def _find_depth_source(session_dir, slot): mkv = os.path.join(session_dir, "depth", slot, f"{slot}.mkv") if os.path.isfile(mkv): return (mkv, "move") d = os.path.join(session_dir, "depth", slot) if not os.path.isdir(d): return None pngs = sorted(f for f in os.listdir(d) if f.endswith(".png") and f.split(".")[0].isdigit()) if pngs: return (os.path.join(d, pngs[0]), "png16") bins = sorted(f for f in os.listdir(d) if f.endswith(".bin") and f.split(".")[0].isdigit()) if bins: return (os.path.join(d, bins[0]), "bin") return None def depth_frame_files(session_dir, slot): """深度帧文件按数字文件名升序 → [(abs_path, frame_no)]""" d = os.path.join(session_dir, "depth", slot) if not os.path.isdir(d): return [] out = [] for fn in os.listdir(d): stem = fn.rsplit(".", 1)[0] if stem.isdigit() and fn.rsplit(".", 1)[1] in ("png", "bin"): out.append((os.path.join(d, fn), int(stem))) out.sort(key=lambda x: x[1]) return out # ── 深度合成(png16 / bin → 单流 FFV1 gray16le mkv) ───────── def find_ffmpeg(metadata): """ffmpeg 二进制:metadata 记录的 imageio_ffmpeg 路径 → PATH → None。""" vc = metadata.get("video_codec") or {} p = vc.get("ffmpeg") if p and os.path.isfile(p): return p try: import imageio_ffmpeg exe = imageio_ffmpeg.get_ffmpeg_exe() if exe and os.path.isfile(exe): return exe except Exception: pass return shutil.which("ffmpeg") def synthesize_depth_mkv(task_dir, session_dir, slot, metadata, out_path, warn): """png16/bin 序列 → 无损单流 FFV1 gray16le mkv(out_path)。""" frames = depth_frame_files(session_dir, slot) if not frames: warn(f"{slot}: 深度帧目录为空,跳过") return False ffmpeg = find_ffmpeg(metadata) if not ffmpeg: warn(f"{slot}: 无 ffmpeg,跳过深度合成({len(frames)} 帧未迁移)") return False cam = (metadata.get("cameras") or {}).get(slot) or {} fps = float(metadata.get("fps") or cam.get("fps") or 30.0) width = int(cam.get("width") or 0) height = int(cam.get("height") or 0) tmp_raw = os.path.join(task_dir, f".migrate_tmp_{slot}_N{int(time.time()*1000)%100000}.raw") try: # 逐帧解码/读入 memmap,避免整段驻留内存 first = None count = 0 ext = os.path.splitext(frames[0][0])[1].lower() if ext == ".png": if cv2 is None: warn(f"{slot}: 无 cv2(OpenCV),跳过 png16 深度合成") return False for path, _no in frames: img = cv2.imread(path, cv2.IMREAD_UNCHANGED) if img is None: continue if img.ndim == 3: img = img[..., 0] if img.dtype != np.uint16: warn(f"{slot}: png 非 uint16({img.dtype}),跳过深度合成") return False if first is None: first = img width, height = img.shape[1], img.shape[0] count += 1 else: # bin: raw uint16 if width <= 0 or height <= 0: warn(f"{slot}: bin 序列缺分辨率(metadata 无 width/height)," "跳过深度合成") return False first = np.zeros((height, width), dtype=np.uint16) count = len(frames) if first is None or count == 0: return False npx = width * height mm = np.memmap(tmp_raw, dtype=np.uint16, mode="w+", shape=(count * npx,)) i = 0 if ext == ".png": for path, _no in frames: img = cv2.imread(path, cv2.IMREAD_UNCHANGED) if img is None or img.ndim == 3: img = img[..., 0] if img is not None and img.ndim == 3 \ else np.zeros((height, width), np.uint16) mm[i * npx:(i + 1) * npx] = img.reshape(-1)[:npx] i += 1 else: for path, _no in frames: raw = np.fromfile(path, dtype=np.uint16) mm[i * npx:(i + 1) * npx] = raw[:npx] i += 1 mm.flush() del mm os.makedirs(os.path.dirname(out_path), exist_ok=True) cmd = [ffmpeg, "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "gray16le", "-s", f"{width}x{height}", "-framerate", f"{fps}", "-i", tmp_raw, "-c:v", "ffv1", "-pix_fmt", "gray16le", out_path] r = subprocess.run(cmd, capture_output=True, text=True) if r.returncode != 0: warn(f"{slot}: ffmpeg 合成失败: {r.stderr.strip()[:300]}") return False return True finally: try: os.remove(tmp_raw) except OSError: pass # ── data parquet 列合并 ───────────────────────────────────── def build_merged_table(meta, ts_rows, sensor_idx, imu_idx): """按新格式契约构建本 episode 的 data 表。 ts_rows 为骨架(一行一帧槽位行);sensor_idx/imu_idx 为 {frame_index: 旧行}(同帧多槽行取首行)。 """ import pyarrow as pa n_ep = int(meta.get("episode_index", 0)) # 由调用方覆盖为新 N sensor_dim = int((meta.get("sensor_dim") or DEFAULT_SENSOR_DIM) or DEFAULT_SENSOR_DIM) hand_dim = int(settings.HAND_POSE_DIM or 63) sensors = sorted(sensor_idx.keys()) cols = {} cols["episode_index"] = pa.array([n_ep] * len(ts_rows), pa.int64()) cols["frame_index"] = pa.array([int(r.get("frame_index", 0)) for r in ts_rows], pa.int64()) cols["timestamp"] = pa.array([float(r.get("timestamp", 0.0)) for r in ts_rows], pa.float32()) cols["task_index"] = pa.array([0] * len(ts_rows), pa.int64()) cols["wall_time"] = pa.array([float(r.get("wall_time", 0.0)) for r in ts_rows], pa.float64()) cols["hardware_ns"] = pa.array([int(r.get("hardware_ns", 0) or 0) for r in ts_rows], pa.int64()) cols["action"] = pa.array([[0.0]] * len(ts_rows), pa.list_(pa.float32(), 1)) for sn in sensors: idx = sensor_idx[sn] vals = [] for r in ts_rows: fi = int(r.get("frame_index", 0)) row = idx.get(fi) v = (row or {}).get(f"observation.{sn}") if v is None: v = [0.0] * sensor_dim v = [float(x) for x in v] if len(v) < sensor_dim: v = v + [0.0] * (sensor_dim - len(v)) vals.append(v[:sensor_dim]) cols[f"observation.{sn}"] = pa.array( vals, pa.list_(pa.float32(), sensor_dim)) if imu_idx is not None: # IMU 只挂在每帧首行(旧格式=stereo_left 行,左右目共享样本) ts_list, imu_list = [], [] seen = set() for r in ts_rows: fi = int(r.get("frame_index", 0)) row = imu_idx.get(fi) if fi in seen or row is None: ts_list.append([]) imu_list.append([]) continue seen.add(fi) ts_list.append([int(x) for x in (row.get("imu_ts_ns") or [])]) imu_list.append([[float(x) for x in s] for s in (row.get("observation.imu") or [])]) cols["imu_ts_ns"] = pa.array(ts_list, pa.list_(pa.int64())) cols["observation.imu"] = pa.array( imu_list, pa.list_(pa.list_(pa.float32(), 6))) for pose in (settings.HAND_POSE_LEFT, settings.HAND_POSE_RIGHT): cols[f"observation.{pose}"] = pa.array( [[0.0] * hand_dim] * len(ts_rows), pa.list_(pa.float32(), hand_dim)) # status.:旧 parquet 携带过设备状态列则保留(同帧取首行) dids = set() for idx in sensor_idx.values(): for row in (idx or {}).values(): for key in (row or {}).keys(): if isinstance(key, str) and key.startswith("status."): dids.add(key[len("status."):]) for did in sorted(dids): vals = [] for r in ts_rows: fi = int(r.get("frame_index", 0)) row = next((idx.get(fi) for idx in sensor_idx.values() if f"status.{did}" in (idx.get(fi) or {})), None) vals.append(str((row or {}).get(f"status.{did}", "connected"))) cols[f"status.{did}"] = pa.array(vals, pa.string()) return pa.table(cols), sensors # ── 统计(stats.json 自含累加器,v1.1.1) ────────────────── def episode_stat_blocks(table, sensors, has_imu): """本 episode 各列的统计块 {key: {count,sum,sum_sq,min,max}}。""" blocks = {} cols = {c: table.column(c).to_pylist() for c in table.column_names} for sn in sensors: name = f"observation.{sn}" vals = [np.asarray(v, dtype=np.float64) for v in cols[name] if v and any(x != 0 for x in v)] if not vals: continue arr = np.stack(vals) blocks[name] = { "count": int(arr.shape[0]), "sum": arr.sum(axis=0).tolist(), "sum_sq": (arr * arr).sum(axis=0).tolist(), "min": arr.min(axis=0).tolist(), "max": arr.max(axis=0).tolist(), } if has_imu: samples = [np.asarray(s, dtype=np.float64) for v in cols["observation.imu"] for s in v] if samples: arr = np.stack(samples) blocks["observation.imu"] = { "count": int(arr.shape[0]), "sum": arr.sum(axis=0).tolist(), "sum_sq": (arr * arr).sum(axis=0).tolist(), "min": arr.min(axis=0).tolist(), "max": arr.max(axis=0).tolist(), } return blocks def merge_state_and_write(task_dir, blocks): """把各 episode 统计块合并进 stats.json(v1.1.1 自含累加器,同写入器)。""" acc, need_recalc = load_stats_acc(task_dir) if need_recalc: acc = recalc_stats(task_dir) for key, blk in blocks.items(): dim = len(blk["sum"]) merge_stat_block(acc, key, blk, dim) _atomic_write_json(pooled_stats_path(task_dir), acc_to_stats_json(acc)) def _atomic_write_json(path, obj): os.makedirs(os.path.dirname(path), exist_ok=True) tmp = path + ".tmp" with open(tmp, "w", encoding="utf-8") as f: json.dump(obj, f, indent=2, ensure_ascii=False) os.replace(tmp, path) # ── info.json 合成(任务级并集合并,同写入器语义) ────────── def build_info(task_dir, task, sessions_meta): """sessions_meta: [(metadata, sensors, has_imu, calib_dict)] 按 N 升序。""" path = pooled_info_path(task_dir) base = load_json(path, {}) if not isinstance(base, dict): base = {} merged_features = dict(base.get("features") or {}) merged_cameras = dict(base.get("cameras") or {}) merged_exts = dict(base.get("video_extensions") or {}) merged_calib = dict(base.get("calibration") or {}) sensors_union, has_imu_any = set(), False last = None for meta, sensors, has_imu, calib in sessions_meta: last = meta sensors_union |= set(sensors) has_imu_any = has_imu_any or has_imu for sn in sorted(sensors): merged_features[f"observation.{sn}"] = { "dtype": "float32", "shape": [16, 16]} if has_imu: merged_features["observation.imu"] = { "dtype": "float32", "shape": [6]} merged_features["action"] = {"dtype": "float32", "shape": [1]} for cam, info in (meta.get("cameras") or {}).items(): if not isinstance(info, dict): continue if info.get("type") == "depth": merged_exts[cam] = "mkv" continue entry = {"height": int(info.get("height", 0)), "width": int(info.get("width", 0))} cam_fps = float(info.get("fps") or meta.get("fps") or 0.0) if cam_fps and cam_fps != float(meta.get("fps") or 0.0): entry["fps"] = cam_fps merged_cameras[cam] = entry merged_exts.update({cam: "mp4" for cam in merged_cameras}) merged_calib.update(calib or {}) devices, device_names = [], {} if last: devices = [{"key": d.get("key", ""), "kind": d.get("kind", ""), "name": d.get("name", ""), "slots": list(d.get("slots") or [])} for d in (last.get("devices") or [])] for d in (last.get("devices") or []): if d.get("name"): for slot in (d.get("slots") or []): device_names[slot] = d["name"] info = dict(base) info.update({ "format": "pooled_episodes_v1", "chunks_size": POOLED_CHUNK_SIZE, "data_path": "data/chunk-{c:03d}/episode-{f:03d}.parquet", "video_path": "videos/chunk-{c:03d}/{image_key}/episode-{f:03d}.{ext}", "episodes_path": "meta/episodes/chunk-{c:03d}/episode-{f:03d}.parquet", "codebase_version": (last or {}).get("codebase_version", "v3.0"), "app_version": f"{settings.APP_VERSION}-migrated", "fps": float((last or {}).get("fps") or 30.0), "video": len(merged_cameras) > 0, "task_name": task, "features": merged_features, "cameras": merged_cameras, "devices": devices, "device_names": device_names, "sensors": sorted(sensors_union), "sensor_dim": int((last or {}).get("sensor_dim") or DEFAULT_SENSOR_DIM), "created_at": max([float((m or {}).get("created_at") or 0.0) for m, *_ in sessions_meta] or [time.time()]), "calibration": merged_calib, "total_episodes": len(list_task_episodes(task_dir)), "video_extensions": merged_exts, }) return info # ── episodes 分片 append ─────────────────────────────────── def _task_lock(task_dir): lock_path = os.path.join(task_dir, "meta", "episodes", ".lock") os.makedirs(os.path.dirname(lock_path), exist_ok=True) fh = open(lock_path, "a+") if _fcntl is not None: _fcntl.flock(fh.fileno(), _fcntl.LOCK_EX) return _LockedFile(fh) def append_episode_row(task_dir, n, row): """每段一个文件、单行原子写(与 data/videos 同编号);旧分片里 的同号行一并删除,避免双份。""" path = pooled_episodes_path(task_dir, n) os.makedirs(os.path.dirname(path), exist_ok=True) with _task_lock(task_dir): _atomic_write_parquet(_episode_rows_table([row]), path) cidx, _ = episode_chunk_file(n) legacy = _legacy_episodes_shard_path(task_dir, cidx) if legacy != path and os.path.isfile(legacy): rows = _read_episode_rows(legacy) kept = [r for r in rows if int(r.get("episode_index", 0)) != n] if len(kept) != len(rows): if kept: _atomic_write_parquet(_episode_rows_table(kept), legacy) else: os.remove(legacy) def remove_episode_row(task_dir, n): """移除 episode N 的 episodes 元数据(每段文件 → os.remove; 旧分片回退 → 读-改-写删行)。N=1 时分片与每段文件同路径, 先按内容判每段文件才删,避免误删整个旧分片。""" path = pooled_episodes_path(task_dir, n) cidx, _ = episode_chunk_file(n) legacy = _legacy_episodes_shard_path(task_dir, cidx) with _task_lock(task_dir): removed_path = False if os.path.isfile(path): rows = _read_episode_rows(path) if len(rows) <= 1 and (not rows or int(rows[0].get("episode_index", 0)) == n): os.remove(path) # 每段文件(或旧分片只剩该行) removed_path = True if os.path.isfile(legacy) and not (removed_path and legacy == path): rows = _read_episode_rows(legacy) kept = [r for r in rows if int(r.get("episode_index", 0)) != n] if len(kept) != len(rows): if kept: _atomic_write_parquet(_episode_rows_table(kept), legacy) else: os.remove(legacy) # 分片空 → 删除文件本身 # ── DB 改写 ──────────────────────────────────────────────── def db_rows_for_path(db_path, col, old_path): """返回 {table: [(id, 现值)]}——只匹配旧会话目录绝对路径。""" if not os.path.isfile(db_path): return [], [] con = sqlite3.connect(db_path) try: rec = con.execute( "SELECT id, file_path FROM recording " "WHERE file_path=? AND episode_index=0", (old_path,)).fetchall() up = con.execute( "SELECT id, session_path FROM upload_task " "WHERE session_path=?", (old_path,)).fetchall() except sqlite3.OperationalError: rec, up = [], [] finally: con.close() return [{"id": r[0], "table": "recording", "old": r[1]} for r in rec], [{"id": r[0], "table": "upload_task", "old": r[1]} for r in up] def db_update_episode(db_path, rows, task_dir, n, task): con = sqlite3.connect(db_path) try: for r in rows: if r["table"] == "recording": con.execute("UPDATE recording SET file_path=?, " "episode_index=? WHERE id=?", (task_dir, n, r["id"])) else: con.execute("UPDATE upload_task SET session_path=?, " "session_name=?, episode_index=? WHERE id=?", (task_dir, task, n, r["id"])) con.commit() finally: con.close() def db_restore_episode(db_path, rows): con = sqlite3.connect(db_path) try: for r in rows: if r["table"] == "recording": con.execute("UPDATE recording SET file_path=?, " "episode_index=0 WHERE id=?", (r["old"], r["id"])) else: con.execute("UPDATE upload_task SET session_path=?, " "episode_index=0 WHERE id=?", (r["old"], r["id"])) con.commit() finally: con.close() # ── 单个会话迁移 ─────────────────────────────────────────── def migrate_session(task_dir, task, session_dir, n, backup_root, report, dry_run, backup, db_path): warns = [] meta = load_json(os.path.join(session_dir, "metadata.json"), {}) if not meta: return "skip", warns ts_rows, total_frames = read_ts_rows(session_dir) if ts_rows is None or not ts_rows: # 无 timestamps.json:从传感器 parquet 并集重建骨架(wall_time 近似) import pyarrow.parquet as pq skeleton = {} for sub in os.listdir(os.path.join(session_dir, "data")): idx = read_sensor_parquet(session_dir, sub) for fi, r in (idx or {}).items(): skeleton.setdefault(fi, { "frame_index": fi, "timestamp": float(r.get("timestamp", 0.0)), "wall_time": float(meta.get("created_at", time.time())) + float(r.get("timestamp", 0.0)), "hardware_ns": int(r.get("hardware_ns", 0) or 0), }) ts_rows = [skeleton[k] for k in sorted(skeleton)] warns.append("无 timestamps.json,骨架由 data parquet 重建" "(wall_time 近似)") if not ts_rows: return "skip", ["无 timestamps.json 且无 data parquet,跳过"] vids, depths = find_legacy_media(session_dir, meta) # 防撞护栏:任何计划目标已存在 → 整会话跳过,绝不覆盖池化已有文件 # (部分迁移/并发重跑时 N 可能撞号,文件级检查是最强不变量) existing = [] for cam in sorted(vids): if os.path.exists(pooled_video_path(task_dir, cam, n, ext="mp4")): existing.append(os.path.relpath( pooled_video_path(task_dir, cam, n, ext="mp4"), task_dir)) for slot in sorted(depths): if os.path.exists(pooled_video_path(task_dir, slot, n, ext="mkv")): existing.append(os.path.relpath( pooled_video_path(task_dir, slot, n, ext="mkv"), task_dir)) if os.path.exists(pooled_data_parquet_path(task_dir, n)): existing.append(os.path.relpath( pooled_data_parquet_path(task_dir, n), task_dir)) # episodes 行撞号(restore 半途失败残留行 → 追加会造重复行; # list_task_episodes 只扫 data/videos 文件,看不到 episodes 行) # 每段文件与旧分片两条路径都要查(N=1 时两者同路径,set 去重) import pyarrow.parquet as pq cidx, _ = episode_chunk_file(n) for cand in {pooled_episodes_path(task_dir, n), _legacy_episodes_shard_path(task_dir, cidx)}: if os.path.isfile(cand): try: if any(r["episode_index"] == n for r in pq.read_table( cand, columns=["episode_index"]).to_pylist()): existing.append(f"meta/episodes 行 {n}") except Exception: pass if existing: return "skip", [f"防撞护栏:池化目标已存在,跳过本会话绝不覆盖:" f"{existing[:4]}"] # ① 完整快照先行(媒体 move 出目录之前),否则 restore 会丢数据 backup_path = os.path.join(backup_root, task, os.path.basename(session_dir)) if not dry_run and backup: os.makedirs(os.path.dirname(backup_path), exist_ok=True) try: shutil.copytree(session_dir, backup_path) except OSError as e: return "skip", [f"快照失败({e}),本会话不动"] if not os.path.isfile(os.path.join(backup_path, "metadata.json")): shutil.rmtree(backup_path, ignore_errors=True) return "skip", ["快照缺 metadata.json(校验失败),本会话不动"] sensors_found = {} imu_idx = None data_dir = os.path.join(session_dir, "data") if os.path.isdir(data_dir): for sub in sorted(os.listdir(data_dir)): if sub == "imu": imu_idx = read_sensor_parquet(session_dir, "imu") continue idx = read_sensor_parquet(session_dir, sub) if idx: # 全零传感器视为从未出现(旧格式零填充=缺席) present = False for r in idx.values(): v = r.get(f"observation.{sub}") if v and any(float(x) != 0 for x in v): present = True break if present: sensors_found[sub] = idx else: warns.append(f"传感器 {sub} 全零,视为缺席不写列") created = [] def plan_create(rel): return os.path.join(task_dir, rel) # 视频 move for cam, (src, kind) in sorted(vids.items()): dst = pooled_video_path(task_dir, cam, n, ext="mp4") created.append(dst) if not dry_run: os.makedirs(os.path.dirname(dst), exist_ok=True) _move_file(src, dst) for cam, info in (meta.get("cameras") or {}).items(): if (isinstance(info, dict) and info.get("type", "rgb") != "depth" and cam not in vids): warns.append(f"摄像头 {cam} 无视频文件(旧目录缺 mp4)") # 深度 move / 合成 depth_done = {} for slot, (src, kind) in sorted(depths.items()): dst = pooled_video_path(task_dir, slot, n, ext="mkv") created.append(dst) if kind == "move": if not dry_run: os.makedirs(os.path.dirname(dst), exist_ok=True) _move_file(src, dst) depth_done[slot] = "mkv" else: # png16 / bin if dry_run: depth_done[slot] = f"{kind}→mkv" continue ok = synthesize_depth_mkv(task_dir, session_dir, slot, meta, dst, warns.append) depth_done[slot] = "mkv" if ok else "skipped" # data parquet import pyarrow as pa table, sensors = build_merged_table(meta, ts_rows, sensors_found, imu_idx) table = table.set_column( 0, "episode_index", pa.array([n] * table.num_rows, pa.int64())) data_path = pooled_data_parquet_path(task_dir, n) created.append(data_path) if not dry_run: os.makedirs(os.path.dirname(data_path), exist_ok=True) _atomic_write_parquet(table, data_path) # episodes 行 calib = {} cal_dir = os.path.join(session_dir, "calibration") if os.path.isdir(cal_dir): for fn in sorted(os.listdir(cal_dir)): if fn.endswith(".json"): c = load_json(os.path.join(cal_dir, fn)) if isinstance(c, dict): calib[fn[:-5]] = c video_codec = dict(meta.get("video_codec") or {}) video_codec.pop("ffmpeg", None) # 机器相关绝对路径不入库 fps = float(meta.get("fps") or 30.0) ts_vals = [float(r.get("timestamp", 0.0)) for r in ts_rows] ep_row = { "episode_index": n, "task_index": 0, "start_frame_index": int(min(r["frame_index"] for r in ts_rows)), "end_frame_index": int(max(r["frame_index"] for r in ts_rows)), "length": int(total_frames or len({int(r["frame_index"]) for r in ts_rows})), "created_at": float(meta.get("created_at") or time.time()), # 真实墙钟时长(多槽行共享时间轴,不随行数翻倍) "duration_sec": float(max(0.0, max(ts_vals) - min(ts_vals)) + 1.0 / fps), "drop_stats": json.dumps(meta.get("drop_stats") or {}, ensure_ascii=False), "video_codec": json.dumps(video_codec, ensure_ascii=False), "calibration": json.dumps(calib, ensure_ascii=False), } if not dry_run: append_episode_row(task_dir, n, ep_row) # DB 改写(先于目录移动,崩溃后重跑可从 manifest 续) db_rec, db_up = db_rows_for_path(db_path, "file_path", session_dir) if not dry_run and (db_rec or db_up): db_update_episode(db_path, db_rec + db_up, task_dir, n, task) # keypoints_output 重键 kp_moves = [] kp_base = os.path.join(settings.DATA_DIR, "keypoints_output", task) kp_old = os.path.join(kp_base, os.path.basename(session_dir)) kp_new = os.path.join(kp_base, f"episode_{n:06d}") if os.path.isdir(kp_old): kp_moves.append((kp_old, kp_new)) if not dry_run: os.makedirs(os.path.dirname(kp_new), exist_ok=True) _move_file(kp_old, kp_new) # ② 媒体已 move 入池化(快照在前,--no-backup 时无快照);删除残壳 # 防重跑重复发现(--no-backup 属显式破坏性模式) if not dry_run: shutil.rmtree(session_dir, ignore_errors=True) report.add(task, os.path.basename(session_dir), n, "migrated" if not dry_run else "planned", f"视频 {len(vids)} 路 / 深度 {list(depth_done.values())} / " f"数据 {table.num_rows} 行 x {len(sensors)} 传感器", created, db_rec + db_up, warns) return { "session": os.path.basename(session_dir), "N": n, "backup_path": backup_path, "created_files": created, "db_rows": db_rec + db_up, "keypoints": [{"from": a, "to": b} for a, b in kp_moves], "metadata": meta, "sensors": sensors, "has_imu": imu_idx is not None and any((r.get("observation.imu") or []) for r in imu_idx.values()), "calib": calib, "warnings": warns, }, warns def _move_file(src, dst): try: os.rename(src, dst) except OSError: shutil.move(src, dst) # ── 任务级收尾 ───────────────────────────────────────────── def finalize_task(task_dir, task, session_entries, dry_run): """info.json / tasks.jsonl / stats 合成(实跑时)。""" metas = [(e["metadata"], e["sensors"], e["has_imu"], e["calib"]) for e in session_entries if e.get("metadata")] info = build_info(task_dir, task, metas) blocks = {} if not dry_run: for e in session_entries: import pyarrow.parquet as pq p = pooled_data_parquet_path(task_dir, e["N"]) if os.path.isfile(p): t = pq.read_table(p) blocks.update(episode_stat_blocks( t, e["sensors"], e["has_imu"])) _atomic_write_json(pooled_info_path(task_dir), info) if not os.path.isfile(pooled_tasks_jsonl_path(task_dir)): with open(pooled_tasks_jsonl_path(task_dir), "w", encoding="utf-8") as f: json.dump({"task_index": 0, "task": task}, f, ensure_ascii=False) f.write("\n") merge_state_and_write(task_dir, blocks) return info # ── manifest ─────────────────────────────────────────────── def manifest_path(task_dir): return os.path.join(task_dir, "meta", MANIFEST_NAME) def load_manifest(task_dir): m = load_json(manifest_path(task_dir), {}) if not isinstance(m, dict): m = {} return m def save_manifest(task_dir, manifest): _atomic_write_json(manifest_path(task_dir), manifest) # ── 回滚 ─────────────────────────────────────────────────── def restore_task(task_dir, task, manifest, dry_run, report): """按 manifest 回滚一个任务的全部迁移。""" entries = manifest.get("sessions") or [] for e in reversed(entries): n = int(e["N"]) for f in e.get("created_files") or []: if os.path.isfile(f): if not dry_run: os.remove(f) remove_episode_row(task_dir, n) if not dry_run else None if not dry_run and (e.get("db_rows")): db_restore_episode(DB_PATH, e["db_rows"]) for kp in e.get("keypoints") or []: if os.path.isdir(kp["to"]): if not dry_run: os.makedirs(os.path.dirname(kp["from"]), exist_ok=True) _move_file(kp["to"], kp["from"]) bp = e.get("backup_path") if bp and os.path.isdir(bp): if not dry_run: os.makedirs(task_dir, exist_ok=True) _move_file(bp, os.path.join( task_dir, os.path.basename(bp))) report.add(task, e.get("session", "?"), n, "restored" if not dry_run else "planned-restore", f"backup={bp}", [], [], e.get("warnings") or []) # 备份根目录若已清空则一并回收 if not dry_run: for entry in entries: bp = entry.get("backup_path") if bp: root = os.path.dirname(os.path.dirname(bp)) if os.path.isdir(root) and not os.listdir(root): _rmdir_empty(root) # 任务清空则回收 info/stats/tasks;有剩余 episode 则保留(stats 不回退) if not dry_run and not list_task_episodes(task_dir): for p in (pooled_info_path(task_dir), pooled_stats_path(task_dir), pooled_tasks_jsonl_path(task_dir)): try: os.remove(p) except OSError: pass # 池化结构空目录自底向上回收(不动已恢复的旧会话目录) if not dry_run: for sub in ("videos", "data"): for root, dirs, _files in os.walk( os.path.join(task_dir, sub), topdown=False): _rmdir_empty(root) ep_dir = os.path.join(task_dir, "meta", "episodes") for root, _dirs, _files in os.walk(ep_dir, topdown=False): # 注意:不能用 walk 的 dirs 判断——那是下探时快照, # 子目录在自底向上过程中已被删,快照仍非空会挡住清理 leftovers = [e for e in os.listdir(root) if e != ".lock"] if not leftovers: try: os.remove(os.path.join(root, ".lock")) except OSError: pass _rmdir_empty(root) _rmdir_empty(os.path.join(task_dir, "meta")) if not dry_run: try: os.remove(manifest_path(task_dir)) except OSError: pass _rmdir_empty(os.path.join(task_dir, "meta")) def _rmdir_empty(path): try: os.rmdir(path) except OSError: pass # ── 校验 ─────────────────────────────────────────────────── def _episode_video_files(task_dir, n): """episode N 的全部视频文件路径列表(任意 key/扩展名)。""" _, fidx = episode_chunk_file(n) out = [] vroot = os.path.join(task_dir, "videos") if os.path.isdir(vroot): for cname in os.listdir(vroot): cdir = os.path.join(vroot, cname) if not os.path.isdir(cdir): continue for key in os.listdir(cdir): kd = os.path.join(cdir, key) if not os.path.isdir(kd): continue for fn in os.listdir(kd): # v1.1.2 起 episode- 前缀;file- 前缀为更早数据(容忍) if (fn.startswith(f"episode-{fidx:03d}.") or fn.startswith(f"file-{fidx:03d}.")): out.append(os.path.join(kd, fn)) return out def validate_task(task_dir, task): """读取端自检:episodes 行 ↔ data parquet ↔ 视频文件组 ↔ info.json。""" print(f"── 校验 {task} ──") import pyarrow.parquet as pq ok = True info = load_json(pooled_info_path(task_dir), {}) if info.get("format") != "pooled_episodes_v1": print(" ✗ info.json 缺失或非池化格式") return False episodes = [] dupes = set() ep_dir = os.path.join(task_dir, "meta", "episodes") if not os.path.isdir(ep_dir): print(" ✗ meta/episodes 缺失") return False for c in sorted(os.listdir(ep_dir)): if not c.startswith("chunk-"): continue d = os.path.join(ep_dir, c) for f in sorted(os.listdir(d)): if not f.endswith(".parquet"): continue p = os.path.join(d, f) rows = pq.read_table(p).to_pylist() # 每段文件不变量:单行且 episode-{f}(旧数据 file-{f})与 # episode_index 同编号 if len(rows) == 1: n = int(rows[0]["episode_index"]) fm = re.match(r"^(?:episode|file)-(\d+)\.parquet$", f) fidx = int(fm.group(1)) if fm else -1 if fidx != (n - 1) % POOLED_CHUNK_SIZE: print(f" ✗ {os.path.relpath(p, task_dir)}: " f"episode-{fidx:03d} 与 episode {n} 编号不符") ok = False for r in rows: n = int(r["episode_index"]) if any(int(x["episode_index"]) == n for x in episodes): dupes.add(n) episodes.append(r) if dupes: print(f" ✗ episode 行重复(新旧布局并存): {sorted(dupes)}") ok = False print(f" episodes 行: {len(episodes)}") for r in episodes: n = int(r["episode_index"]) data = pooled_data_parquet_path(task_dir, n) if not os.path.isfile(data): print(f" ✗ episode {n}: 缺 data parquet") ok = False continue t = pq.read_table(data) cols = set(t.column_names) need = {"episode_index", "frame_index", "timestamp", "task_index", "wall_time", "hardware_ns", "action"} miss = need - cols if miss: print(f" ✗ episode {n}: 缺列 {sorted(miss)}") ok = False else: print(f" ✓ episode {n}: {t.num_rows} 行, 列 " f"{sorted(c for c in cols if c.startswith('observation') or c.startswith('status'))}") # 视频:episode 自己的文件组(key 可跨会话改名,不能要求 # 每个 key 每个 episode 都有文件) if not _episode_video_files(task_dir, n): print(f" ✗ episode {n}: 无任何视频文件") ok = False # 每个声明的 key 至少在某个 episode 有文件(任务级并集口径) for key, ext in (info.get("video_extensions") or {}).items(): if not any(os.path.isfile(pooled_video_path( task_dir, key, int(r["episode_index"]), ext=ext)) for r in episodes): print(f" ✗ 视频 key {key} ({ext}) 全任务无文件") ok = False if os.path.isfile(pooled_stats_path(task_dir)): print(f" ✓ stats.json 存在: {sorted(load_json(pooled_stats_path(task_dir), {}))}") return ok def split_episodes(task_dir, task, dry_run): """旧分片(每 chunk 一个多行文件)→ 每段一个文件;幂等。 - 多行文件按 episode_index 拆成每段单行文件(每段文件已存在 则跳过,不覆盖——每段文件为准)。 - v1.1.2 起每段文件前缀 episode-,与旧分片名 file-{chunk} 不再 重名(v1.1.1 及更早两者同路径,写回源路径特判已无必要但保留)。 - 拆完删除旧分片;单行文件不动。 - 拆分过则刷新 info.json episodes_path 模板(旧任务写的是旧模板)。 """ ep_dir = os.path.join(task_dir, "meta", "episodes") if not os.path.isdir(ep_dir): print(f" {task}: 无 meta/episodes,跳过") return changed = False for c in sorted(os.listdir(ep_dir)): if not c.startswith("chunk-"): continue d = os.path.join(ep_dir, c) for fn in sorted(os.listdir(d)): if not fn.endswith(".parquet"): continue p = os.path.join(d, fn) rows = _read_episode_rows(p) if len(rows) <= 1: continue # 每段新布局或只剩一行,不动 dests = {} for r in rows: n = int(r.get("episode_index", 0)) if n: dests[pooled_episodes_path(task_dir, n)] = r wrote = 0 for dst, r in dests.items(): if dst == p: continue # 源路径自身最后写 if os.path.isfile(dst): continue # 每段文件已存在(幂等,不覆盖) if dry_run: wrote += 1 continue os.makedirs(os.path.dirname(dst), exist_ok=True) _atomic_write_parquet(_episode_rows_table([r]), dst) wrote += 1 if dry_run: print(f" [plan] {task}: {os.path.relpath(p, task_dir)} " f"{len(rows)} 行 → {len(dests)} 个每段文件") continue if p in dests: # chunk-0:N=1 的行写回源路径(分片变成每段文件) _atomic_write_parquet(_episode_rows_table([dests[p]]), p) else: os.remove(p) changed = True print(f" {task}: {os.path.relpath(p, task_dir)} {len(rows)} 行 " f"→ {len(dests)} 个每段文件") if changed and not dry_run: info_path = pooled_info_path(task_dir) info = load_json(info_path, {}) if isinstance(info, dict): info["episodes_path"] = ("meta/episodes/chunk-{c:03d}/" "episode-{f:03d}.parquet") _atomic_write_json(info_path, info) if dry_run: print(f" [plan] {task}: 以上为计划,未落盘") # ── 主流程 ───────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser(description="旧会话 → v1.1.0 任务池化布局") ap.add_argument("--task", action="append", default=[], help="只处理指定任务目录名(可多次)") ap.add_argument("--dry-run", action="store_true", help="只打印迁移计划,不落盘") ap.add_argument("--backup", action="store_true", default=True, help="旧会话目录移入备份区(默认开)") ap.add_argument("--no-backup", dest="backup", action="store_false", help="旧会话目录不备份(不推荐)") ap.add_argument("--force", action="store_true", help="任务已有 manifest 时仍处理剩余旧会话") ap.add_argument("--restore", action="store_true", help="按 manifest 回滚迁移") ap.add_argument("--validate", action="store_true", help="读取端自检(迁移后校验报告)") ap.add_argument("--split-episodes", action="store_true", help="旧 episodes 分片(每 chunk 多行)拆成每段一文件(幂等)") ap.add_argument("--recordings-root", default=None, help="测试用:覆盖录制根目录(默认 data/recordings)") ap.add_argument("--db-path", default=None, help="测试用:覆盖 pipeline.db 路径") args = ap.parse_args() global RECORDINGS_DIR, DB_PATH if args.recordings_root: RECORDINGS_DIR = args.recordings_root if args.db_path: DB_PATH = args.db_path report = Report() run_ts = time.strftime("%Y%m%d_%H%M%S") backup_root = os.path.join(RECORDINGS_DIR, f"_migrate_backup_{run_ts}") db_backup_path = os.path.join( settings.DATA_DIR, f"pipeline.db.migrate_backup_{run_ts}") db_needs_backup = False if args.restore: for task, task_dir in task_dirs(RECORDINGS_DIR, args.task): manifest = load_manifest(task_dir) if manifest.get("sessions"): restore_task(task_dir, task, manifest, args.dry_run, report) report.print("restore") return 0 if args.validate: all_ok = True for task, task_dir in task_dirs(RECORDINGS_DIR, args.task): if os.path.isdir(os.path.join(task_dir, "meta", "episodes")): all_ok = validate_task(task_dir, task) and all_ok print() print("PASS: 校验全部通过" if all_ok else "FAIL: 存在校验失败项") return 0 if all_ok else 1 if args.split_episodes: for task, task_dir in task_dirs(RECORDINGS_DIR, args.task): split_episodes(task_dir, task, args.dry_run) print() print("PASS: 旧分片拆分完成(幂等,可重复跑)") return 0 for task, task_dir in task_dirs(RECORDINGS_DIR, args.task): sessions = legacy_sessions(task_dir, task) if not sessions: continue manifest = load_manifest(task_dir) if manifest.get("sessions") and not args.force: report.warnings.append( f"{task}: 已有 manifest(迁移过),--force 才会重扫") continue done = {os.path.basename(e["backup_path"]) for e in (manifest.get("sessions") or []) if e.get("backup_path")} todo = [s for s in sessions if os.path.basename(s) not in done] if not todo: continue n_next = max(list_task_episodes(task_dir) or [0]) + 1 entries = manifest.get("sessions") or [] for session_dir in todo: db_rec, db_up = db_rows_for_path( DB_PATH, "file_path", session_dir) if db_rec or db_up: db_needs_backup = True if not args.dry_run and not os.path.isfile(db_backup_path): shutil.copy2(DB_PATH, db_backup_path) print(f"pipeline.db 已备份: {db_backup_path}") entry, warns = migrate_session( task_dir, task, session_dir, n_next, backup_root, report, args.dry_run, args.backup, DB_PATH) if entry == "skip": report.warnings.extend(warns) continue entries.append(entry) n_next += 1 finalize_task(task_dir, task, entries, args.dry_run) manifest = { "version": MANIFEST_VERSION, "run_ts": run_ts, "db_backup": db_backup_path if db_needs_backup else None, "sessions": entries, } if not args.dry_run: save_manifest(task_dir, manifest) report.print("dry-run" if args.dry_run else "migrate") return 0 if __name__ == "__main__": sys.exit(main())