xiaoxia-saas/apps/worker/video_processing/dedup.py

"""Video deduplication module - compute fingerprints and detect duplicates."""
import hashlib
import json
import logging
import os
import subprocess
import tempfile
from dataclasses import dataclass
from typing import Optional

import cv2
import numpy as np
from celery import Task
from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.generated_video_repository import SQLAlchemyGeneratedVideoRepository
from packages.shared.storage import get_storage_service

from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

logger = logging.getLogger(__name__)


def compute_phash(image: np.ndarray, hash_size: int = 8) -> str:
    """Compute perceptual hash of an image using DCT."""
    # Resize to 32x32 for DCT
    resized = cv2.resize(image, (hash_size * 4, hash_size * 4))
    gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY).astype(np.float32)
    # Apply 2D DCT
    dct = cv2.dct(gray)
    # Take top-left 8x8 low-frequency components
    dct_low = dct[:hash_size, :hash_size]
    # Compute median (excluding DC component at [0,0])
    dct_low[0, 0] = 0
    median = np.median(dct_low)
    # Generate hash based on comparison with median
    diff = (dct_low > median).astype(int)
    hash_str = "".join(str(b) for row in diff for b in row)
    return hex(int(hash_str, 2))[2:]


def hamming_distance(hash1: str, hash2: str) -> int:
    """Calculate Hamming distance between two hex hashes."""
    h1, h2 = int(hash1, 16), int(hash2, 16)
    return bin(h1 ^ h2).count("1")


def compute_color_histogram(image: np.ndarray, bins: int = 32) -> list[float]:
    """Compute color histogram for an image."""
    hist = []
    for i in range(3):
        h = cv2.calcHist([image], [i], None, [bins], [0, 256])
        h = cv2.normalize(h, h).flatten()
        hist.extend(h)
    return hist


@dataclass
class VideoFingerprint:
    """Video fingerprint containing multiple similarity metrics."""
    md5: str
    keyframe_phashes: list[str]
    color_histograms: list[list[float]]
    duration: float
    resolution: tuple[int, int]

    def to_dict(self) -> dict:
        return {"md5": self.md5, "keyframe_phashes": self.keyframe_phashes, "color_histograms": self.color_histograms, "duration": self.duration, "resolution": list(self.resolution)}


class VideoDeduplicator:
    """Video deduplication using multiple fingerprint methods."""

    PHASH_THRESHOLD = 10
    HISTOGRAM_THRESHOLD = 0.85

    def compute_fingerprint(self, video_path: str) -> VideoFingerprint:
        """Compute video fingerprint using MD5, pHash, and color histogram."""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open video: {video_path}")
        
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = frame_count / fps if fps > 0 else 0
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        md5_hash = hashlib.md5()
        keyframe_phashes = []
        color_histograms = []
        
        frame_interval = max(1, frame_count // 10)
        for i in range(0, frame_count, frame_interval):
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if not ret:
                continue
            
            _, buffer = cv2.imencode(".jpg", frame)
            md5_hash.update(buffer)
            
            keyframe_phashes.append(compute_phash(frame))
            color_histograms.append(compute_color_histogram(frame))
        
        cap.release()
        
        return VideoFingerprint(md5=md5_hash.hexdigest(), keyframe_phashes=keyframe_phashes, color_histograms=color_histograms, duration=duration, resolution=(width, height))

    def check_duplicate(self, fingerprint: VideoFingerprint, project_id: str, session: Session) -> Optional[dict]:
        """Check if video is duplicate of existing one. Returns duplicate info if found."""
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        existing_videos = video_repo.list_by_project(project_id)
        
        for existing in existing_videos:
            if not existing.video_fingerprint:
                continue
            
            ef = existing.video_fingerprint
            
            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", [])
            if existing_phashes:
                total_distance = 0
                min_distances = []
                for phash in fingerprint.keyframe_phashes:
                    distances = [hamming_distance(phash, ep) for ep in existing_phashes]
                    min_distances.append(min(distances))
                avg_distance = sum(min_distances) / len(min_distances) if min_distances else 100
                
                if avg_distance < self.PHASH_THRESHOLD:
                    return {"duplicate": True, "duplicate_of": existing.id, "reason": "phash_similar", "similarity": 1.0 - (avg_distance / 64)}
        
        return None


@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."""
    session = SessionLocal()
    temp_dir = tempfile.mkdtemp()
    
    try:
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        storage_service = get_storage_service()
        deduplicator = VideoDeduplicator()
        
        video = video_repo.get(generated_video_id)
        if video is None:
            raise ValueError(f"Generated video {generated_video_id} not found")
        
        local_path = os.path.join(temp_dir, f"{generated_video_id}.mp4")
        storage_key = video.file_url.split("/")[-1]
        storage_service.download_file(f"workspaces/{video.workspace_id}/projects/{video.project_id}/generated/{generated_video_id}/{generated_video_id}.mp4", local_path)
        
        fingerprint = deduplicator.compute_fingerprint(local_path)
        
        duplicate_result = deduplicator.check_duplicate(fingerprint, video.project_id, session)
        
        video.video_fingerprint = fingerprint.to_dict()
        if duplicate_result:
            video.is_duplicate = True
            video.duplicate_of = duplicate_result["duplicate_of"]
        else:
            video.is_duplicate = False
            video.duplicate_of = None
        
        video_repo.update(video)
        session.commit()
        
        logger.info(f"Duplicate check completed for video {generated_video_id}: is_duplicate={video.is_duplicate}")
        
        return {"ok": True, "video_id": generated_video_id, "is_duplicate": video.is_duplicate, "duplicate_of": video.duplicate_of, "fingerprint": fingerprint.to_dict()}
    except Exception as e:
        logger.error(f"Duplicate check failed for {generated_video_id}: {str(e)}")
        session.rollback()
        raise self.retry(exc=e, countdown=60)
    finally:
        session.close()
        import shutil
        shutil.rmtree(temp_dir, ignore_errors=True)