DataFlowing / processor /app /batch_fingerprint.py
quyu666's picture
sync processor from GitLab 2bbf9fe: UMI SLAM 动作模块、批次指纹与设备进度、导出/路由/前端更新
edae745
Raw History Blame Contribute Delete
7.01 kB
"""批次内容指纹 — 识别"换个包名重传同一份数据"。
采集端偶尔会把已经上传过的录制再推一遍,但**包名变了**(例如重装、
重开 App 后本地去重标记丢失,或包名规则从序号变成时间戳)。此时
``session.py`` 的 ``incoming_name in existing_ids`` 判不出重复,同一份
数据就会被当成新批次追加,审核页出现两条一模一样的数据,导出也会
重复计一份(实测:UMIGripper_AI 的 4 组批次视频字节级相同、
parquet 除 ``episode_index`` 外逐值相同)。
本模块给出一个与**包名无关**的内容身份:
- 每个视频文件的 md5(按 source_key 排序)
- canonical data parquet 的内容哈希(排除 ``episode_index`` ——
同一次录制被追加成不同 episode 时,只有这一列会变)
合成为一个短字符串 key,存到 ``state/batch_fingerprints.json``。
上传时先算新批次的 key,再按项目查表:命中且目标 episode 仍然存在,
就判定为重复上传。
**失败一律放行**:算不出指纹、索引读不到、目标 episode 已被删除 ——
任一情况都按"新批次"正常入库(即本模块不存在时的行为),绝不因为
指纹机制本身的问题丢数据。
"""
from __future__ import annotations
import hashlib
import json
import threading
from pathlib import Path
from typing import Any
from app.localstore import STATE_ROOT
INDEX_PATH = STATE_ROOT / "batch_fingerprints.json"
_LOCK = threading.Lock()
# 视频按块读;读取量固定,不随文件大小膨胀。
_CHUNK = 1 << 20
def _file_md5(path: Path) -> str | None:
digest = hashlib.md5()
try:
with open(path, "rb") as handle:
for block in iter(lambda: handle.read(_CHUNK), b""):
digest.update(block)
except OSError:
return None
return digest.hexdigest()
def _parquet_content_hash(path: Path) -> str | None:
"""canonical data parquet 的内容哈希,排除 ``episode_index``。
``episode_index`` 是同一次录制被追加成不同 episode 时唯一会变的列,
必须排除,否则重复上传算不出相同的 key。用 Arrow IPC 序列化取字节:
对嵌套 list 列(手套 250x250x3 力矩阵、手部关键点)比对 pandas 的
hash 更稳,且与行序无关的假设更少 —— 这里行序本身就是内容的一部分。
"""
try:
import pyarrow as pa
import pyarrow.parquet as pq
except ImportError:
return None
try:
table = pq.read_table(path)
if "episode_index" in table.column_names:
table = table.drop_columns(["episode_index"])
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, table.schema) as writer:
writer.write_table(table)
return hashlib.md5(sink.getvalue().to_pybytes()).hexdigest()
except Exception:
return None
def compute_fingerprint(root: Path) -> str | None:
"""Staging 目录(已 normalize)的内容指纹;算不出返回 None。"""
root = Path(root)
parts: list[str] = []
videos_root = root / "videos"
if videos_root.is_dir():
from app.lerobot_v21 import iter_video_streams
entries: list[tuple[str, str]] = []
for source, path in iter_video_streams(videos_root):
digest = _file_md5(Path(path))
if digest is None:
return None
entries.append((str(source), digest))
if not entries:
return None
parts.append("|".join(f"v:{s}={d}" for s, d in sorted(entries)))
data_files = sorted(
path for path in (root / "data").rglob("*.parquet")
if path.is_file() and "/meta/" not in path.as_posix().casefold()
)
if not data_files:
return None
digest = _parquet_content_hash(data_files[0])
if digest is None:
return None
parts.append(f"d:{digest}")
return ";".join(parts)
def _read_index() -> dict[str, dict[str, str]]:
try:
value = json.loads(INDEX_PATH.read_text(encoding="utf-8"))
except (OSError, ValueError):
return {}
return value if isinstance(value, dict) else {}
def _write_index(index: dict) -> None:
try:
INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)
temporary = INDEX_PATH.with_suffix(".json.tmp")
temporary.write_text(
json.dumps(index, ensure_ascii=False, indent=2), encoding="utf-8")
temporary.replace(INDEX_PATH)
except OSError:
pass
def lookup_duplicate(project_folder: str, fingerprint: str | None,
live_ids: set[str]) -> str | None:
"""返回与本次上传内容相同的、且**仍然存在**的 episode_id。
``live_ids`` 来自当前项目数据集的实际扫描结果:索引里指向已删除
episode 的陈旧条目在这里被自然过滤掉,因此"删除后重新上传同一条
录制"仍会正常入库,不会被误判成重复。
"""
if not fingerprint:
return None
with _LOCK:
index = _read_index()
entries = index.get(str(project_folder))
if not isinstance(entries, dict):
return None
episode_id = entries.get(str(fingerprint))
if not episode_id:
return None
return str(episode_id) if str(episode_id) in live_ids else None
def record_fingerprint(project_folder: str, fingerprint: str | None,
episode_id: str) -> None:
"""记录已提交批次的内容指纹(失败静默:只是失去一次去重机会)。"""
if not fingerprint:
return
with _LOCK:
index = _read_index()
if not isinstance(index, dict):
index = {}
entries = index.get(str(project_folder))
if not isinstance(entries, dict):
entries = {}
entries[str(fingerprint)] = str(episode_id)
# 同一 episode 换内容(同名重传)时清掉它名下的旧指纹,避免
# 索引里留下永远指不到实际内容的条目。
for key in [k for k, v in entries.items()
if str(v) == str(episode_id) and k != str(fingerprint)]:
entries.pop(key, None)
index[str(project_folder)] = entries
_write_index(index)
def forget_fingerprint(episode_id: str) -> None:
"""批次被删除后清除其指纹条目。"""
target = str(episode_id)
with _LOCK:
index = _read_index()
changed = False
for entries in index.values():
if not isinstance(entries, dict):
continue
for key in [k for k, v in entries.items() if str(v) == target]:
entries.pop(key, None)
changed = True
if changed:
_write_index(index)
def describe(project_folder: str) -> dict[str, Any]:
"""诊断用:该项目已记录的内容指纹。"""
with _LOCK:
index = _read_index()
entries = index.get(str(project_folder))
return entries if isinstance(entries, dict) else {}