ocr_service.py

import asyncio
import uuid
import os
import time
from pathlib import Path
from typing import Optional
import cv2
import pytesseract
from PIL import Image
import tempfile


OCR_TEMP_DIR = Path("/tmp/video_ocr")
OCR_TEMP_DIR.mkdir(parents=True, exist_ok=True)


def extract_frames(
    video_path: str,
    output_dir: str,
    interval_ms: int = 1000,
    max_frames: int = 100,
) -> list[dict]:
    """
    Extract frames from video at specified intervals.
    
    Args:
        video_path: Path to video file
        output_dir: Directory to save frames
        interval_ms: Time interval between frames in milliseconds
        max_frames: Maximum number of frames to extract
    
    Returns:
        List of frame info dicts with path, timestamp_ms, frame_number
    """
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        raise RuntimeError(f"Failed to open video: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration_ms = (total_frames / fps * 1000) if fps > 0 else 0
    
    frames_info = []
    frame_interval = int(fps * interval_ms / 1000)
    frame_interval = max(frame_interval, 1)
    
    frame_count = 0
    saved_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        if frame_count % frame_interval == 0 and saved_count < max_frames:
            timestamp_ms = int(frame_count / fps * 1000)
            frame_path = os.path.join(output_dir, f"frame_{saved_count:04d}.jpg")
            cv2.imwrite(frame_path, frame)
            
            frames_info.append({
                "path": frame_path,
                "timestamp_ms": timestamp_ms,
                "frame_number": frame_count,
            })
            saved_count += 1
        
        frame_count += 1
        
        if saved_count >= max_frames:
            break
    
    cap.release()
    return frames_info


def preprocess_frame(frame_path: str, enhance: bool = True) -> Image.Image:
    """
    Preprocess frame for better OCR results.
    
    Args:
        frame_path: Path to frame image
        enhance: Whether to apply enhancement
    
    Returns:
        PIL Image ready for OCR
    """
    img = cv2.imread(frame_path)
    
    if enhance:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
        _, enhanced = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        img = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)
    
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return Image.fromarray(img_rgb)


async def perform_ocr_on_frame(
    frame_path: str,
    lang: str = "eng",
    config: str = "--oem 3 --psm 6",
) -> dict:
    """
    Perform OCR on a single frame (async wrapper).
    """
    return await asyncio.to_thread(perform_ocr_on_frame_sync, frame_path, lang, config)


def perform_ocr_on_frame_sync(
    frame_path: str,
    lang: str = "eng",
    config: str = "--oem 3 --psm 6",
) -> dict:
    """
    Perform OCR on a single frame (synchronous).
    
    Args:
        frame_path: Path to frame image
        lang: Tesseract language code
        config: Tesseract config string
    
    Returns:
        Dict with text, confidence, and word-level data
    """
    try:
        img = preprocess_frame(frame_path)
        
        data = pytesseract.image_to_data(img, lang=lang, config=config, output_type=pytesseract.Output.DICT)
        text = pytesseract.image_to_string(img, lang=lang, config=config)
        
        words = []
        for i in range(len(data["text"])):
            if int(data["conf"][i]) > 0:
                words.append({
                    "text": data["text"][i],
                    "confidence": int(data["conf"][i]),
                    "left": data["left"][i],
                    "top": data["top"][i],
                    "width": data["width"][i],
                    "height": data["height"][i],
                })
        
        avg_confidence = (
            sum(w["confidence"] for w in words) / len(words) if words else 0
        )
        
        return {
            "text": text.strip(),
            "confidence": round(avg_confidence, 2),
            "words": words,
            "full_data": data,
        }
    except Exception as e:
        return {
            "text": "",
            "confidence": 0,
            "words": [],
            "error": str(e),
        }


async def process_video_ocr(
    video_path: str,
    lang: str = "eng",
    interval_ms: int = 1000,
    max_frames: int = 100,
    config: str = "--oem 3 --psm 6",
) -> dict:
    """
    Process video for OCR - extract frames and perform OCR.
    
    Args:
        video_path: Path to video file
        lang: Tesseract language code
        interval_ms: Interval between frames in ms
        max_frames: Maximum frames to process
        config: Tesseract config
    
    Returns:
        Dict with all OCR results and metadata
    """
    start_time = time.time()
    
    session_id = str(uuid.uuid4())[:8]
    frame_dir = OCR_TEMP_DIR / f"frames_{session_id}"
    frame_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        frames_info = extract_frames(video_path, str(frame_dir), interval_ms, max_frames)
        
        results = []
        for i, frame_info in enumerate(frames_info):
            ocr_result = perform_ocr_on_frame_sync(
                frame_info["path"], lang, config
            )
            
            if ocr_result.get("text"):
                results.append({
                    "frame_index": i,
                    "timestamp_ms": frame_info["timestamp_ms"],
                    "frame_number": frame_info["frame_number"],
                    **ocr_result,
                })
        
        all_text = " ".join(r["text"] for r in results if r.get("text"))
        avg_confidence = (
            sum(r["confidence"] for r in results) / len(results) if results else 0
        )
        
        return {
            "success": True,
            "session_id": session_id,
            "video_path": video_path,
            "text": all_text,
            "average_confidence": round(avg_confidence, 2),
            "frames_processed": len(results),
            "total_frames_extracted": len(frames_info),
            "processing_time_ms": round((time.time() - start_time) * 1000, 2),
            "language": lang,
            "results": results,
        }
    
    finally:
        pass


async def get_video_duration(video_path: str) -> float:
    """Get video duration in milliseconds using ffprobe."""
    import subprocess
    
    cmd = [
        "ffprobe",
        "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        video_path,
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode == 0:
        return float(result.stdout.strip()) * 1000
    return 0.0


def cleanup_session(session_id: str):
    """Cleanup temp files for a session."""
    frame_dir = OCR_TEMP_DIR / f"frames_{session_id}"
    if frame_dir.exists():
        import shutil
        shutil.rmtree(frame_dir, ignore_errors=True)