Spaces:
Sleeping
Sleeping
File size: 4,054 Bytes
b165696 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
OCR-based subtitle extractor — fallback when no subtitle file is available.
Captures a frame at the midpoint of each Whisper audio segment,
crops the bottom 15% (where subtitles typically appear),
and runs Tesseract OCR with the appropriate Indic language pack.
"""
from __future__ import annotations
import logging
from pathlib import Path
from typing import Callable
import cv2
import numpy as np
logger = logging.getLogger(__name__)
SUBTITLE_REGION_RATIO = 0.15 # bottom 15% of frame
class OCRExtractor:
def __init__(
self,
language: str = "hi",
progress_hook: Callable | None = None,
):
"""
Parameters
----------
language : str
yt-dlp / BCP-47 language code (e.g. 'hi', 'kn', 'en').
Mapped to Tesseract language pack internally.
"""
from config.settings import SUPPORTED_LANGUAGES
lang_info = SUPPORTED_LANGUAGES.get(language, SUPPORTED_LANGUAGES["en"])
self._tess_lang = lang_info["tesseract"]
self._progress_hook = progress_hook
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def extract_from_video(
self,
video_path: str,
audio_segments: list[dict],
) -> list[dict]:
"""
For each audio segment, OCR the frame at its midpoint.
Returns
-------
list of { "start": float, "end": float, "text": str }
"""
import pytesseract
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise RuntimeError(f"Cannot open video: {video_path}")
fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
results = []
for i, seg in enumerate(audio_segments):
midpoint = (seg["start"] + seg["end"]) / 2.0
frame_no = int(midpoint * fps)
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no)
ok, frame = cap.read()
if not ok:
logger.warning("Could not read frame %d (t=%.2fs)", frame_no, midpoint)
results.append({"start": seg["start"], "end": seg["end"], "text": ""})
continue
cropped = self._crop_subtitle_region(frame)
preprocessed = self._preprocess(cropped)
text = self._run_ocr(pytesseract, preprocessed)
results.append(
{"start": seg["start"], "end": seg["end"], "text": text.strip()}
)
if self._progress_hook and i % 10 == 0:
self._progress_hook(i / len(audio_segments))
cap.release()
logger.info("OCR extracted text for %d segments", len(results))
return results
# ------------------------------------------------------------------
# Internals
# ------------------------------------------------------------------
def _crop_subtitle_region(self, frame: np.ndarray) -> np.ndarray:
h, w = frame.shape[:2]
crop_top = int(h * (1 - SUBTITLE_REGION_RATIO))
return frame[crop_top:h, 0:w]
def _preprocess(self, img: np.ndarray) -> np.ndarray:
"""Improve OCR accuracy: grayscale → denoise → threshold."""
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Upscale for better OCR
scaled = cv2.resize(gray, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
# Adaptive threshold works better than simple binary for varied backgrounds
thresh = cv2.adaptiveThreshold(
scaled, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
)
return thresh
def _run_ocr(self, pytesseract, img: np.ndarray) -> str:
config = (
f"--oem 3 --psm 6 -l {self._tess_lang}"
" -c tessedit_char_blacklist=|"
)
try:
return pytesseract.image_to_string(img, config=config)
except Exception as exc:
logger.warning("Tesseract error: %s", exc)
return ""
|