Spaces:
Running
Running
Download processor/app/batch_fingerprint.py from Stouching-Robotics/DataFlowing: direct link, hf CLI and curl.
- Browser
- Download file 7.01 kB
-
https://huggingface.co/spaces/Stouching-Robotics/DataFlowing/resolve/main/processor/app/batch_fingerprint.py
- Command line
-
hf download hf://spaces/Stouching-Robotics/DataFlowing/processor/app/batch_fingerprint.py
-
curl -L -o batch_fingerprint.py https://huggingface.co/spaces/Stouching-Robotics/DataFlowing/resolve/main/processor/app/batch_fingerprint.py
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 {} | |