From cbca0c3584ad0ae9cfdb6893573edc400836b8a4 Mon Sep 17 00:00:00 2001 From: saas-backend-agent Date: Thu, 3 Sep 2026 20:41:11 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E5=88=86=E7=89=87=E6=8C=87?= =?UTF-8?q?=E7=BA=B9=E5=AD=98=E5=82=A8=E6=94=B9=E9=80=A0=20+=20=E5=AD=98?= =?UTF-8?q?=E9=87=8F=E6=8C=87=E7=BA=B9=E9=87=8D=E5=BB=BA=E8=84=9A=E6=9C=AC?= =?UTF-8?q?=20#1657?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 改动 ### 1. 新建 video_fingerprint_chunks 表(Migration 063) - 按时间分片存储 pHash + color_histogram - 索引:video_id, project_id, user_id ### 2. 新增 VideoFingerprintChunkModel - packages/adapters/sqlalchemy_impl/models.py ### 3. 改造 dedup.py 指纹计算 - compute_fingerprint() 改为按时间分片抽帧 - 短视频(≤60s):每 2s 一片 - 长视频(>60s):每 5s 一片 - VideoFingerprint 新增 chunks 字段(list of FingerprintChunk) - 向后兼容:keyframe_phashes/color_histograms 保留 - to_chunk_models() 方法转换为 SQLAlchemy Model - check_duplicate() 优先从分片表读取,回退到 JSON 字段 - check_duplicate_task() 写入分片表 ### 4. 改造 dedup_helpers.py - create_video_record_and_dedup() 同步写入分片表 ### 5. 存量指纹重建脚本 - apps/api/scripts/rebuild_fingerprint_chunks.py - 支持 --dry-run 和 --batch-size - 幂等:已有分片数据的视频跳过 ### 6. 单元测试(11 个) - 分片策略:60s→30片,120s→24片 - to_chunk_models() 输出正确 - _save_fingerprint_chunks 幂等性 - to_dict() 向后兼容 Closes #1657 --- .../063_add_video_fingerprint_chunks.py | 46 +++ .../api/scripts/rebuild_fingerprint_chunks.py | 173 ++++++++++ apps/worker/video_processing/dedup.py | 201 +++++++++-- apps/worker/video_processing/dedup_helpers.py | 8 + packages/adapters/sqlalchemy_impl/models.py | 17 + tests/unit/test_fingerprint_chunks.py | 313 ++++++++++++++++++ 6 files changed, 734 insertions(+), 24 deletions(-) create mode 100644 alembic/versions/063_add_video_fingerprint_chunks.py create mode 100644 apps/api/scripts/rebuild_fingerprint_chunks.py create mode 100644 tests/unit/test_fingerprint_chunks.py diff --git a/alembic/versions/063_add_video_fingerprint_chunks.py b/alembic/versions/063_add_video_fingerprint_chunks.py new file mode 100644 index 000000000..a629d5d57 --- /dev/null +++ b/alembic/versions/063_add_video_fingerprint_chunks.py @@ -0,0 +1,46 @@ +"""add video_fingerprint_chunks table for per-chunk fingerprint storage + +Revision ID: 063_fingerprint_chunks +Revises: 062_edit_plan_id +Create Date: 2026-09-03 +""" + +import sqlalchemy as sa + +from alembic import op + +revision = "063_fingerprint_chunks" +down_revision = "062_edit_plan_id" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "video_fingerprint_chunks", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("video_id", sa.String(36), nullable=False), + sa.Column("project_id", sa.String(36), nullable=False), + sa.Column("user_id", sa.String(36), nullable=False, server_default=""), + sa.Column("start_time_ms", sa.Integer, nullable=False), + sa.Column("end_time_ms", sa.Integer, nullable=False), + sa.Column("phash_binary", sa.String(16), nullable=False), + sa.Column("color_histogram", sa.JSON, nullable=False), + sa.Column("frame_count", sa.Integer, nullable=False, server_default="1"), + sa.Column( + "created_at", + sa.DateTime, + nullable=False, + server_default=sa.func.now(), + ), + ) + op.create_index("ix_vfc_video_id", "video_fingerprint_chunks", ["video_id"]) + op.create_index("ix_vfc_project_id", "video_fingerprint_chunks", ["project_id"]) + op.create_index("ix_vfc_user_id", "video_fingerprint_chunks", ["user_id"]) + + +def downgrade() -> None: + op.drop_index("ix_vfc_user_id", table_name="video_fingerprint_chunks") + op.drop_index("ix_vfc_project_id", table_name="video_fingerprint_chunks") + op.drop_index("ix_vfc_video_id", table_name="video_fingerprint_chunks") + op.drop_table("video_fingerprint_chunks") diff --git a/apps/api/scripts/rebuild_fingerprint_chunks.py b/apps/api/scripts/rebuild_fingerprint_chunks.py new file mode 100644 index 000000000..5568f82d1 --- /dev/null +++ b/apps/api/scripts/rebuild_fingerprint_chunks.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""存量指纹重建脚本 — 为已有视频生成 video_fingerprint_chunks 分片数据。 + +功能: +- 查询 generated_videos 中 video_fingerprint IS NOT NULL 但尚无分片数据的视频 +- 从 OSS 下载视频 → 用新的分片算法重新计算指纹 → 写入分片表 +- 支持 --dry-run(只打印不写入)和 --batch-size(默认 50) +- 幂等:已存在分片数据的视频跳过 + +用法: + # 预览(不写入) + python rebuild_fingerprint_chunks.py --dry-run + + # 执行重建 + python rebuild_fingerprint_chunks.py --batch-size 50 +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import tempfile + +# 确保可以 import worker_app 和 packages +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "worker")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", +) +logger = logging.getLogger("rebuild_fingerprint_chunks") + + +def find_videos_needing_rebuild(session, batch_size: int) -> list[dict]: + """查询需要重建分片指纹的视频。""" + from sqlalchemy import and_ + + from packages.adapters.sqlalchemy_impl.models import GeneratedVideoModel, VideoFingerprintChunkModel + + # 有 video_fingerprint 的视频 + has_fingerprint = GeneratedVideoModel.video_fingerprint.isnot(None) + has_fingerprint = and_(has_fingerprint, GeneratedVideoModel.video_fingerprint != "") + + # 排除已有分片数据的视频 + subq = session.query(VideoFingerprintChunkModel.video_id).distinct().subquery() + no_chunks = ~GeneratedVideoModel.id.in_(subq) + + videos = ( + session.query(GeneratedVideoModel) + .filter(and_(has_fingerprint, no_chunks)) + .order_by(GeneratedVideoModel.generated_at.desc()) + .limit(batch_size) + .all() + ) + + return [ + { + "id": v.id, + "project_id": v.project_id, + "user_id": v.user_id or "", + "duration": v.duration, + } + for v in videos + ] + + +def rebuild_one(video_info: dict, dry_run: bool = False) -> int: + """重建单个视频的分片数据。返回写入的 chunk 数量。""" + from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel + from packages.shared.storage import get_storage_service + from video_processing.dedup import VideoDeduplicator, _save_fingerprint_chunks + from worker_app.db import SessionLocal + + video_id = video_info["id"] + project_id = video_info["project_id"] + user_id = video_info["user_id"] + + if dry_run: + logger.info("[DRY-RUN] Would rebuild video %s (project=%s)", video_id, project_id) + return 0 + + session = SessionLocal() + temp_dir = tempfile.mkdtemp() + + try: + # 再次检查幂等性 + existing_count = ( + session.query(VideoFingerprintChunkModel).filter(VideoFingerprintChunkModel.video_id == video_id).count() + ) + if existing_count > 0: + logger.info("Video %s already has %d chunks, skipping", video_id, existing_count) + return 0 + + # 下载视频 + storage_service = get_storage_service() + local_path = os.path.join(temp_dir, f"{video_id}.mp4") + storage_key = f"projects/{project_id}/generated/{video_id}/{video_id}.mp4" + storage_service.download_file(storage_key, local_path) + + # 重新计算指纹 + deduplicator = VideoDeduplicator() + fingerprint = deduplicator.compute_fingerprint(local_path) + + # 写入分片表 + _save_fingerprint_chunks(fingerprint, video_id, project_id, user_id, session) + session.commit() + + chunk_count = len(fingerprint.chunks) + logger.info("Rebuilt %d chunks for video %s", chunk_count, video_id) + return chunk_count + + except Exception as e: + logger.error("Failed to rebuild video %s: %s", video_id, e) + session.rollback() + return -1 + finally: + session.close() + import shutil + + shutil.rmtree(temp_dir, ignore_errors=True) + + +def main(): + parser = argparse.ArgumentParser(description="存量指纹重建脚本") + parser.add_argument("--dry-run", action="store_true", help="只打印不写入") + parser.add_argument("--batch-size", type=int, default=50, help="每批处理数量(默认 50)") + parser.add_argument("--total-limit", type=int, default=0, help="总处理数量限制(0=不限制)") + args = parser.parse_args() + + from worker_app.db import SessionLocal + + session = SessionLocal() + + try: + videos = find_videos_needing_rebuild(session, args.batch_size) + logger.info("Found %d videos needing rebuild", len(videos)) + + if args.dry_run: + for v in videos: + logger.info("[DRY-RUN] Video %s | project=%s | duration=%.1fs", v["id"], v["project_id"], v["duration"]) + return + + total_chunks = 0 + processed = 0 + failed = 0 + + for v in videos: + if args.total_limit > 0 and processed >= args.total_limit: + break + + result = rebuild_one(v, dry_run=False) + if result < 0: + failed += 1 + else: + total_chunks += result + processed += 1 + + logger.info( + "Rebuild complete: processed=%d, chunks=%d, failed=%d", + processed, + total_chunks, + failed, + ) + + finally: + session.close() + + +if __name__ == "__main__": + main() diff --git a/apps/worker/video_processing/dedup.py b/apps/worker/video_processing/dedup.py index ecbeec613..1cbd4e123 100755 --- a/apps/worker/video_processing/dedup.py +++ b/apps/worker/video_processing/dedup.py @@ -4,8 +4,9 @@ import hashlib import logging import os import tempfile -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional +from uuid import uuid4 import cv2 import numpy as np @@ -15,10 +16,16 @@ from worker_app.celery_app import celery_app from worker_app.db import SessionLocal from packages.adapters.sqlalchemy_impl.generated_video_repository import SQLAlchemyGeneratedVideoRepository +from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel from packages.shared.storage import get_storage_service logger = logging.getLogger(__name__) +# 分片策略常量 +SHORT_VIDEO_CHUNK_SEC = 2 # ≤60秒视频,每 2 秒一个分片 +LONG_VIDEO_CHUNK_SEC = 5 # >60秒视频,每 5 秒一个分片 +SHORT_VIDEO_THRESHOLD_SEC = 60 + def compute_phash(image: np.ndarray, hash_size: int = 8) -> str: """计算图像的感知哈希(pHash),基于 DCT(离散余弦变换)。 @@ -80,6 +87,28 @@ def compute_color_histogram(image: np.ndarray, bins: int = 32) -> list[float]: return hist +def compute_chunk_interval(duration: float) -> float: + """根据视频时长返回分片间隔(秒)。 + + 短视频(≤60秒):每 2 秒一个分片 + 长视频(>60秒):每 5 秒一个分片 + """ + if duration <= SHORT_VIDEO_THRESHOLD_SEC: + return SHORT_VIDEO_CHUNK_SEC + return LONG_VIDEO_CHUNK_SEC + + +@dataclass +class FingerprintChunk: + """单个分片指纹数据。""" + + start_time_ms: int + end_time_ms: int + phash_binary: str + color_histogram: list[float] + frame_count: int = 1 + + @dataclass class VideoFingerprint: """Video fingerprint containing multiple similarity metrics.""" @@ -89,6 +118,7 @@ class VideoFingerprint: color_histograms: list[list[float]] duration: float resolution: tuple[int, int] + chunks: list[FingerprintChunk] = field(default_factory=list) def to_dict(self) -> dict: # 注意:color_histograms 里的值可能是 np.float32(来自 cv2.normalize), @@ -101,8 +131,37 @@ class VideoFingerprint: "color_histograms": native_histograms, "duration": float(self.duration), "resolution": [int(self.resolution[0]), int(self.resolution[1])], + "chunks": [ + { + "start_time_ms": c.start_time_ms, + "end_time_ms": c.end_time_ms, + "phash_binary": c.phash_binary, + "color_histogram": [float(v) for v in c.color_histogram], + "frame_count": c.frame_count, + } + for c in self.chunks + ], } + def to_chunk_models(self, video_id: str, project_id: str, user_id: str = "") -> list[VideoFingerprintChunkModel]: + """将分片数据转为 SQLAlchemy Model 列表,用于批量写入 video_fingerprint_chunks 表。""" + models = [] + for chunk in self.chunks: + models.append( + VideoFingerprintChunkModel( + id=uuid4().hex, + video_id=video_id, + project_id=project_id, + user_id=user_id, + start_time_ms=chunk.start_time_ms, + end_time_ms=chunk.end_time_ms, + phash_binary=chunk.phash_binary, + color_histogram=[float(v) for v in chunk.color_histogram], + frame_count=chunk.frame_count, + ) + ) + return models + class VideoDeduplicator: """Video deduplication using multiple fingerprint methods.""" @@ -111,7 +170,12 @@ class VideoDeduplicator: HISTOGRAM_THRESHOLD = 0.85 def compute_fingerprint(self, video_path: str) -> VideoFingerprint: - """Compute video fingerprint using MD5, pHash, and color histogram.""" + """Compute video fingerprint using MD5, pHash, and color histogram. + + 按时间分片抽帧:短视频(≤60s)每 2s 一片,长视频每 5s 一片。 + 每片取 1 帧计算 pHash + color_histogram。 + 同时保留 keyframe_phashes/color_histograms 聚合字段(向后兼容)。 + """ cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise RuntimeError(f"Cannot open video: {video_path}") @@ -123,42 +187,82 @@ class VideoDeduplicator: height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) md5_hash = hashlib.md5(usedforsecurity=False) - keyframe_phashes = [] - color_histograms = [] + chunks: list[FingerprintChunk] = [] - frame_interval = max(1, frame_count // 10) - for i in range(0, frame_count, frame_interval): - cap.set(cv2.CAP_PROP_POS_FRAMES, i) + # 分片间隔(秒) + chunk_interval_sec = compute_chunk_interval(duration) + chunk_interval_ms = int(chunk_interval_sec * 1000) + duration_ms = int(duration * 1000) + + # 遍历每个分片时间窗口,取 1 帧 + start_ms = 0 + while start_ms < duration_ms: + end_ms = min(start_ms + chunk_interval_ms, duration_ms) + # 定位到分片中点 + seek_ms = (start_ms + end_ms) / 2 + cap.set(cv2.CAP_PROP_POS_MSEC, seek_ms) ret, frame = cap.read() - if not ret: - continue + if ret: + # MD5 计算 + _, buffer = cv2.imencode(".jpg", frame) + md5_hash.update(buffer) - _, buffer = cv2.imencode(".jpg", frame) - md5_hash.update(buffer) + phash = compute_phash(frame) + hist = compute_color_histogram(frame) - keyframe_phashes.append(compute_phash(frame)) - color_histograms.append(compute_color_histogram(frame)) + chunks.append( + FingerprintChunk( + start_time_ms=start_ms, + end_time_ms=end_ms, + phash_binary=phash, + color_histogram=hist, + frame_count=1, + ) + ) + + start_ms = end_ms cap.release() + # 向后兼容:聚合 keyframe_phashes / color_histograms + keyframe_phashes = [c.phash_binary for c in chunks] + color_histograms = [c.color_histogram for c in chunks] + return VideoFingerprint( md5=md5_hash.hexdigest(), keyframe_phashes=keyframe_phashes, color_histograms=color_histograms, duration=duration, resolution=(width, height), + chunks=chunks, ) + def _get_existing_chunks(self, video_id: str, session: Session) -> list[dict]: + """从 video_fingerprint_chunks 表读取分片数据。返回空列表表示无分片数据。""" + rows = ( + session.query(VideoFingerprintChunkModel) + .filter(VideoFingerprintChunkModel.video_id == video_id) + .order_by(VideoFingerprintChunkModel.start_time_ms) + .all() + ) + return [ + { + "phash_binary": r.phash_binary, + "color_histogram": r.color_histogram, + "start_time_ms": r.start_time_ms, + "end_time_ms": r.end_time_ms, + } + for r in rows + ] + def check_duplicate(self, fingerprint: VideoFingerprint, project_id: str, session: Session) -> Optional[dict]: """检查视频是否与项目中已有视频重复。 - 判定逻辑(按优先级): - 1. MD5 精确匹配:完全一致则 similarity=1.0,立即返回 - 2. pHash 相似度:计算新视频每帧 phash 与已有视频每帧 phash 的最小汉明距离, - 取所有帧的平均值 avg_distance。若 avg_distance < PHASH_THRESHOLD(10), - 则判定为重复,similarity = 1.0 - (avg_distance / 64) + 查重逻辑: + 1. MD5 精确匹配 → similarity=1.0 + 2. pHash 相似度(优先从分片表读取,回退到 JSON 字段) - 注意:返回第一个通过阈值的匹配(非最优匹配)。 + 判定阈值:avg_distance < PHASH_THRESHOLD(10) Args: fingerprint: 待检测视频的指纹 @@ -182,8 +286,15 @@ class VideoDeduplicator: if fingerprint.md5 == ef.get("md5"): return {"duplicate": True, "duplicate_of": existing.id, "reason": "exact_md5_match", "similarity": 1.0} - # 感知哈希相似度 - existing_phashes = ef.get("keyframe_phashes", []) + # 优先从分片表读取已有视频的分片 phash + existing_phashes = [] + chunk_data = self._get_existing_chunks(existing.id, session) + if chunk_data: + existing_phashes = [c["phash_binary"] for c in chunk_data] + else: + # 回退:从 JSON 字段读取(存量旧视频) + existing_phashes = ef.get("keyframe_phashes", []) + if not existing_phashes: continue @@ -247,7 +358,14 @@ class VideoDeduplicator: "similarity": 1.0, } - existing_phashes = ef.get("keyframe_phashes", []) + # 优先从分片表读取 + existing_phashes = [] + chunk_data = self._get_existing_chunks(existing.id, session) + if chunk_data: + existing_phashes = [c["phash_binary"] for c in chunk_data] + else: + existing_phashes = ef.get("keyframe_phashes", []) + if not existing_phashes: continue @@ -372,8 +490,14 @@ class VideoDeduplicator: if fingerprint.md5 == ef.get("md5"): return 100.0 - # pHash 相似度 - existing_phashes = ef.get("keyframe_phashes", []) + # 优先从分片表读取 + existing_phashes = [] + chunk_data = self._get_existing_chunks(existing.id, session) + if chunk_data: + existing_phashes = [c["phash_binary"] for c in chunk_data] + else: + existing_phashes = ef.get("keyframe_phashes", []) + if not existing_phashes or not fingerprint.keyframe_phashes: continue @@ -388,6 +512,31 @@ class VideoDeduplicator: return round(max(max_similarity, 0.0), 2) +def _save_fingerprint_chunks( + fingerprint: VideoFingerprint, + video_id: str, + project_id: str, + user_id: str, + session: Session, +) -> None: + """将指纹分片数据批量写入 video_fingerprint_chunks 表。幂等:已有数据时跳过。""" + # 幂等检查:已有分片数据则跳过 + existing_count = ( + session.query(VideoFingerprintChunkModel).filter(VideoFingerprintChunkModel.video_id == video_id).count() + ) + if existing_count > 0: + logger.debug("Fingerprint chunks already exist for video %s (%d chunks), skipping", video_id, existing_count) + return + + if not fingerprint.chunks: + logger.warning("No chunks in fingerprint for video %s, skipping chunk save", video_id) + return + + chunk_models = fingerprint.to_chunk_models(video_id, project_id, user_id) + session.bulk_save_objects(chunk_models) + logger.info("Saved %d fingerprint chunks for video %s", len(chunk_models), video_id) + + @celery_app.task(bind=True, max_retries=3, name="worker.check_duplicate") def check_duplicate_task(self: Task, generated_video_id: str) -> dict: """Celery task to check if generated video is a duplicate.""" @@ -421,6 +570,10 @@ def check_duplicate_task(self: Task, generated_video_id: str) -> dict: video.duplicate_of = None video_repo.update(video) + + # 写入分片表 + _save_fingerprint_chunks(fingerprint, generated_video_id, video.project_id, video.user_id, session) + session.commit() logger.info(f"Duplicate check completed for video {generated_video_id}: is_duplicate={video.is_duplicate}") diff --git a/apps/worker/video_processing/dedup_helpers.py b/apps/worker/video_processing/dedup_helpers.py index efcc2072f..9b92a5da7 100755 --- a/apps/worker/video_processing/dedup_helpers.py +++ b/apps/worker/video_processing/dedup_helpers.py @@ -100,6 +100,14 @@ def create_video_record_and_dedup( generated_video.video_fingerprint = fingerprint.to_dict() + # 写入分片指纹表 + from video_processing.dedup import _save_fingerprint_chunks + + try: + _save_fingerprint_chunks(fingerprint, video_id, project_id, user_id, session) + except Exception as chunk_err: + logger.warning("Failed to save fingerprint chunks for %s: %s", video_id, chunk_err) + # (a) 历史成片查重 duplicate_result = deduplicator.check_duplicate(fingerprint, project_id, session) diff --git a/packages/adapters/sqlalchemy_impl/models.py b/packages/adapters/sqlalchemy_impl/models.py index af3712106..060b07480 100755 --- a/packages/adapters/sqlalchemy_impl/models.py +++ b/packages/adapters/sqlalchemy_impl/models.py @@ -620,3 +620,20 @@ class CoverTemplateModel(Base): config = Column(JSON, nullable=False, default=dict) created_at = Column(DateTime(timezone=True), nullable=False, default=lambda: datetime.now(timezone.utc)) updated_at = Column(DateTime(timezone=True), nullable=False, default=lambda: datetime.now(timezone.utc)) + + +class VideoFingerprintChunkModel(Base): + """分片视频指纹 — 每个视频按时间分片存储 pHash + color_histogram.""" + + __tablename__ = "video_fingerprint_chunks" + + id = Column(String(36), primary_key=True) + video_id = Column(String(36), nullable=False, index=True) + project_id = Column(String(36), nullable=False, index=True) + user_id = Column(String(36), nullable=False, index=True, default="") + start_time_ms = Column(Integer, nullable=False) + end_time_ms = Column(Integer, nullable=False) + phash_binary = Column(String(16), nullable=False) + color_histogram = Column(JSON, nullable=False) + frame_count = Column(Integer, nullable=False, default=1) + created_at = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc)) diff --git a/tests/unit/test_fingerprint_chunks.py b/tests/unit/test_fingerprint_chunks.py new file mode 100644 index 000000000..a10ad5260 --- /dev/null +++ b/tests/unit/test_fingerprint_chunks.py @@ -0,0 +1,313 @@ +"""分片指纹存储单元测试 — Issue #1657. + +覆盖: +- 分片策略:60秒视频 → 30片,120秒视频 → 24片 +- VideoFingerprint.to_chunk_models() 输出正确 +- _save_fingerprint_chunks 幂等性(已有数据跳过) +- to_dict() 向后兼容 +""" + +from __future__ import annotations + +import sys +from unittest.mock import MagicMock + + +def _mock_module(**attrs): + """Create a mock module with __spec__ to avoid AttributeError.""" + m = MagicMock() + m.__spec__ = None + for k, v in attrs.items(): + setattr(m, k, v) + return m + + +# ── Module-level setup: mock deps, import dedup, then restore sys.modules ── +_SAVED_MODULES_KEYS = set(sys.modules.keys()) +_SAVED_MODULES_VALUES = { + k: sys.modules.get(k) + for k in [ + "cv2", + "celery", + "sqlalchemy", + "sqlalchemy.orm", + "sqlalchemy.engine", + "sqlalchemy.ext", + "sqlalchemy.ext.declarative", + "worker_app.db", + "worker_app.celery_app", + "worker_app.core.config", + "packages.adapters.sqlalchemy_impl.session", + "packages.adapters.sqlalchemy_impl.generated_video_repository", + "packages.adapters.sqlalchemy_impl.models", + "packages.shared.config", + "packages.shared.storage", + ] +} + +# Set up mocks +sys.modules["cv2"] = _mock_module() + +_mock_celery = MagicMock() +_mock_celery.Task = MagicMock +_mock_celery.Celery = MagicMock +_mock_celery.__spec__ = None +sys.modules["celery"] = _mock_celery + +_mock_sqla = MagicMock() +_mock_sqla.__path__ = [] +_mock_sqla.__spec__ = None +sys.modules["sqlalchemy"] = _mock_sqla + +_mock_sqla_orm = MagicMock() +_mock_sqla_orm.__path__ = [] +_mock_sqla_orm.__spec__ = None +_mock_sqla_orm.Session = MagicMock +sys.modules["sqlalchemy.orm"] = _mock_sqla_orm +sys.modules["sqlalchemy.engine"] = _mock_module() +sys.modules["sqlalchemy.ext"] = _mock_module() +sys.modules["sqlalchemy.ext.declarative"] = _mock_module() + +sys.modules["worker_app.db"] = _mock_module(SessionLocal=MagicMock()) +sys.modules["worker_app.celery_app"] = _mock_module(celery_app=MagicMock()) +sys.modules["worker_app.core.config"] = _mock_module(get_settings=MagicMock(return_value=MagicMock())) + +sys.modules["packages.adapters.sqlalchemy_impl.session"] = _mock_module( + Base=MagicMock(), + build_engine=MagicMock(), + build_session_factory=MagicMock(), + ensure_database_exists=MagicMock(), + initialize_database=MagicMock(), +) +sys.modules["packages.adapters.sqlalchemy_impl.generated_video_repository"] = _mock_module() + + +# Mock VideoFingerprintChunkModel with class-level column attributes +class _FakeChunkModel: + video_id = MagicMock() + project_id = MagicMock() + user_id = MagicMock() + start_time_ms = MagicMock() + end_time_ms = MagicMock() + phash_binary = MagicMock() + color_histogram = MagicMock() + frame_count = MagicMock() + created_at = MagicMock() + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + +sys.modules["packages.adapters.sqlalchemy_impl.models"] = _mock_module( + VideoFingerprintChunkModel=_FakeChunkModel, +) +sys.modules["packages.shared.config"] = _mock_module(get_shared_settings=MagicMock(return_value=MagicMock())) +sys.modules["packages.shared.storage"] = _mock_module() + +# Import dedup while mocks are active +from video_processing.dedup import ( # noqa: E402 + FingerprintChunk, + VideoFingerprint, + _save_fingerprint_chunks, + compute_chunk_interval, +) + +# ── Restore sys.modules immediately after import ── +for _key in list(sys.modules.keys()): + if _key not in _SAVED_MODULES_KEYS: + del sys.modules[_key] +for _key, _value in _SAVED_MODULES_VALUES.items(): + if _value is not None: + sys.modules[_key] = _value + elif _key in sys.modules: + del sys.modules[_key] +del _SAVED_MODULES_KEYS, _SAVED_MODULES_VALUES, _key, _value + + +class TestChunkInterval: + """测试分片间隔策略。""" + + def test_short_video_interval(self): + """短视频(≤60秒)每 2 秒一个分片。""" + assert compute_chunk_interval(0) == 2 + assert compute_chunk_interval(30) == 2 + assert compute_chunk_interval(60) == 2 + + def test_long_video_interval(self): + """长视频(>60秒)每 5 秒一个分片。""" + assert compute_chunk_interval(61) == 5 + assert compute_chunk_interval(120) == 5 + assert compute_chunk_interval(300) == 5 + + def test_chunk_count_60s_video(self): + """60秒视频 → 30 片(60/2=30)。""" + duration = 60 + interval = compute_chunk_interval(duration) + expected_chunks = int(duration / interval) + assert expected_chunks == 30 + + def test_chunk_count_120s_video(self): + """120秒视频 → 24 片(120/5=24)。""" + duration = 120 + interval = compute_chunk_interval(duration) + expected_chunks = int(duration / interval) + assert expected_chunks == 24 + + +class TestVideoFingerprintToChunkModels: + """测试 VideoFingerprint.to_chunk_models() 输出。""" + + def test_to_chunk_models_output(self): + """to_chunk_models 返回正确的 Model 列表。""" + fp = VideoFingerprint( + md5="abc123", + keyframe_phashes=["a1b2", "c3d4"], + color_histograms=[[0.1] * 96, [0.2] * 96], + duration=10.0, + resolution=(1920, 1080), + chunks=[ + FingerprintChunk(start_time_ms=0, end_time_ms=2000, phash_binary="a1b2", color_histogram=[0.1] * 96), + FingerprintChunk(start_time_ms=2000, end_time_ms=4000, phash_binary="c3d4", color_histogram=[0.2] * 96), + ], + ) + + models = fp.to_chunk_models(video_id="v1", project_id="p1", user_id="u1") + + assert len(models) == 2 + assert models[0].video_id == "v1" + assert models[0].project_id == "p1" + assert models[0].user_id == "u1" + assert models[0].start_time_ms == 0 + assert models[0].end_time_ms == 2000 + assert models[0].phash_binary == "a1b2" + assert models[1].start_time_ms == 2000 + assert models[1].end_time_ms == 4000 + assert models[1].phash_binary == "c3d4" + + def test_to_chunk_models_empty_chunks(self): + """空 chunks 列表返回空 Model 列表。""" + fp = VideoFingerprint( + md5="abc", + keyframe_phashes=[], + color_histograms=[], + duration=0, + resolution=(0, 0), + chunks=[], + ) + + models = fp.to_chunk_models(video_id="v1", project_id="p1") + assert models == [] + + +class TestSaveFingerprintChunksIdempotent: + """测试 _save_fingerprint_chunks 幂等性。""" + + def test_save_skips_existing(self): + """已有分片数据时跳过写入。""" + fp = VideoFingerprint( + md5="abc", + keyframe_phashes=["a1b2"], + color_histograms=[[0.1] * 96], + duration=5.0, + resolution=(1920, 1080), + chunks=[ + FingerprintChunk(start_time_ms=0, end_time_ms=2000, phash_binary="a1b2", color_histogram=[0.1] * 96), + ], + ) + + session = MagicMock() + # Mock: 已有 1 条分片数据 + session.query.return_value.filter.return_value.count.return_value = 1 + + _save_fingerprint_chunks(fp, video_id="v1", project_id="p1", user_id="u1", session=session) + + # bulk_save_objects 不应被调用 + session.bulk_save_objects.assert_not_called() + + def test_save_writes_new(self): + """无分片数据时写入。""" + fp = VideoFingerprint( + md5="abc", + keyframe_phashes=["a1b2"], + color_histograms=[[0.1] * 96], + duration=5.0, + resolution=(1920, 1080), + chunks=[ + FingerprintChunk(start_time_ms=0, end_time_ms=2000, phash_binary="a1b2", color_histogram=[0.1] * 96), + ], + ) + + session = MagicMock() + # Mock: 无分片数据 + session.query.return_value.filter.return_value.count.return_value = 0 + + _save_fingerprint_chunks(fp, video_id="v1", project_id="p1", user_id="u1", session=session) + + # bulk_save_objects 应被调用一次 + session.bulk_save_objects.assert_called_once() + saved_models = session.bulk_save_objects.call_args[0][0] + assert len(saved_models) == 1 + assert saved_models[0].video_id == "v1" + assert saved_models[0].phash_binary == "a1b2" + + def test_save_skips_no_chunks(self): + """指纹无 chunks 时跳过。""" + fp = VideoFingerprint( + md5="abc", + keyframe_phashes=[], + color_histograms=[], + duration=0, + resolution=(0, 0), + chunks=[], + ) + + session = MagicMock() + session.query.return_value.filter.return_value.count.return_value = 0 + + _save_fingerprint_chunks(fp, video_id="v1", project_id="p1", user_id="u1", session=session) + + # bulk_save_objects 不应被调用 + session.bulk_save_objects.assert_not_called() + + +class TestFingerprintToDictBackwardCompat: + """测试 to_dict() 向后兼容性。""" + + def test_to_dict_includes_chunks(self): + """to_dict() 包含 chunks 字段。""" + fp = VideoFingerprint( + md5="abc123", + keyframe_phashes=["a1b2"], + color_histograms=[[0.1] * 96], + duration=5.0, + resolution=(1920, 1080), + chunks=[ + FingerprintChunk(start_time_ms=0, end_time_ms=2000, phash_binary="a1b2", color_histogram=[0.1] * 96), + ], + ) + + d = fp.to_dict() + + assert "chunks" in d + assert len(d["chunks"]) == 1 + assert d["chunks"][0]["start_time_ms"] == 0 + assert d["chunks"][0]["end_time_ms"] == 2000 + assert d["chunks"][0]["phash_binary"] == "a1b2" + + def test_to_dict_preserves_legacy_fields(self): + """to_dict() 保留 keyframe_phashes 和 color_histograms 字段。""" + fp = VideoFingerprint( + md5="abc", + keyframe_phashes=["a1b2", "c3d4"], + color_histograms=[[0.1] * 96, [0.2] * 96], + duration=10.0, + resolution=(1920, 1080), + ) + + d = fp.to_dict() + + assert "keyframe_phashes" in d + assert "color_histograms" in d + assert len(d["keyframe_phashes"]) == 2 + assert len(d["color_histograms"]) == 2 -- 2.54.0 From 244691d3358fca8f652d706b6e3e70042748656c Mon Sep 17 00:00:00 2001 From: CI Bot Date: Thu, 3 Sep 2026 12:50:38 +0000 Subject: [PATCH 2/2] style: auto-format with black + isort + prettier [skip ci-format-check] --- apps/api/scripts/rebuild_fingerprint_chunks.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/apps/api/scripts/rebuild_fingerprint_chunks.py b/apps/api/scripts/rebuild_fingerprint_chunks.py index 5568f82d1..922731e87 100644 --- a/apps/api/scripts/rebuild_fingerprint_chunks.py +++ b/apps/api/scripts/rebuild_fingerprint_chunks.py @@ -69,11 +69,12 @@ def find_videos_needing_rebuild(session, batch_size: int) -> list[dict]: def rebuild_one(video_info: dict, dry_run: bool = False) -> int: """重建单个视频的分片数据。返回写入的 chunk 数量。""" - from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel - from packages.shared.storage import get_storage_service from video_processing.dedup import VideoDeduplicator, _save_fingerprint_chunks from worker_app.db import SessionLocal + from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel + from packages.shared.storage import get_storage_service + video_id = video_info["id"] project_id = video_info["project_id"] user_id = video_info["user_id"] -- 2.54.0