Spaces:
Running on Zero
Running on Zero
Download app/core/ocr_worker_cli.py from hoangtaiii/DRIPPY4: direct link, hf CLI and curl.
- Browser
- Download file 29.6 kB
-
https://huggingface.co/spaces/hoangtaiii/DRIPPY4/resolve/main/app/core/ocr_worker_cli.py
- Command line
-
hf download hf://spaces/hoangtaiii/DRIPPY4/app/core/ocr_worker_cli.py
-
curl -L -o ocr_worker_cli.py https://huggingface.co/spaces/hoangtaiii/DRIPPY4/resolve/main/app/core/ocr_worker_cli.py
29.6 kB
| import sys | |
| import os | |
| import cv2 | |
| import time | |
| import difflib | |
| import argparse | |
| import traceback | |
| import json | |
| from pathlib import Path | |
| def register_ocr_gpu_dll_paths(): | |
| import os | |
| from pathlib import Path | |
| root = Path(__file__).resolve().parents[2] | |
| candidates = [ | |
| root / "env" / "Lib" / "site-packages" / "nvidia" / "cudnn" / "bin", | |
| root / "env" / "Lib" / "site-packages" / "nvidia" / "cublas" / "bin", | |
| root / "env" / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin", | |
| root / "env" / "Lib" / "site-packages" / "nvidia" / "cuda_nvrtc" / "bin", | |
| root / "env" / "Lib" / "site-packages" / "paddle" / "libs", | |
| ] | |
| added = [] | |
| for p in candidates: | |
| if p.exists(): | |
| try: | |
| os.add_dll_directory(str(p)) | |
| except Exception as e: | |
| print(f"[OCR DLL WARN] add_dll_directory failed for {p}: {e}") | |
| added.append(str(p)) | |
| if added: | |
| os.environ["PATH"] = os.pathsep.join(added) + os.pathsep + os.environ.get("PATH", "") | |
| print("[OCR DLL] registered paths:") | |
| for p in added: | |
| print(f" - {p}") | |
| register_ocr_gpu_dll_paths() | |
| # Enforce UTF-8 for Windows console | |
| if sys.platform == 'win32': | |
| try: | |
| if hasattr(sys.stdout, 'reconfigure'): | |
| sys.stdout.reconfigure(encoding='utf-8') | |
| if hasattr(sys.stderr, 'reconfigure'): | |
| sys.stderr.reconfigure(encoding='utf-8') | |
| except Exception: | |
| pass | |
| def get_similarity(s1, s2): | |
| if not s1 or not s2: | |
| return 0.0 | |
| return difflib.SequenceMatcher(None, s1, s2).ratio() | |
| def _is_cjk_char(ch): | |
| return "\u3400" <= ch <= "\u9fff" | |
| def _ocr_garbage_reason(text, max_frame_chars=120): | |
| text = " ".join(str(text or "").split()).strip() | |
| if not text: | |
| return "empty" | |
| if len(text) > int(max_frame_chars): | |
| return "frame_text_too_long" | |
| chars = [c for c in text if not c.isspace()] | |
| if not chars: | |
| return "empty" | |
| noise_chars = set("□√×÷′²③④") | |
| noise_ratio = sum(1 for c in chars if c in noise_chars) / len(chars) | |
| if noise_ratio > 0.25: | |
| return "ocr_symbol_noise" | |
| cjk_chars = [c for c in chars if _is_cjk_char(c)] | |
| if len(cjk_chars) >= 40: | |
| counts = {} | |
| for c in cjk_chars: | |
| counts[c] = counts.get(c, 0) + 1 | |
| dominant_ratio = max(counts.values()) / len(cjk_chars) | |
| unique_ratio = len(counts) / len(cjk_chars) | |
| if dominant_ratio > 0.35: | |
| return "dominant_cjk_repeat" | |
| if len(cjk_chars) >= 80 and unique_ratio < 0.15: | |
| return "low_cjk_diversity" | |
| return None | |
| def _atomic_write_json(path, payload): | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp_path = path.with_name(path.name + ".tmp") | |
| with open(tmp_path, "w", encoding="utf-8") as f: | |
| json.dump(payload, f, ensure_ascii=False, indent=2) | |
| os.replace(tmp_path, path) | |
| def _atomic_write_text(path, content): | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp_path = path.with_name(path.name + ".tmp") | |
| with open(tmp_path, "w", encoding="utf-8") as f: | |
| f.write(content) | |
| os.replace(tmp_path, path) | |
| def _extract_ocr_lines(results, min_confidence=0.0): | |
| lines = [] | |
| for page in (results or []): | |
| if not page: | |
| continue | |
| for line in page: | |
| if not line or len(line) <= 1 or not line[1]: | |
| continue | |
| try: | |
| text = str(line[1][0] or "").strip() | |
| score = float(line[1][1]) | |
| except Exception: | |
| continue | |
| if text and score >= float(min_confidence): | |
| lines.append((text, score)) | |
| return lines | |
| def _find_cjk_font(): | |
| candidates = [ | |
| Path("C:/Windows/Fonts/msyh.ttc"), | |
| Path("C:/Windows/Fonts/msyhbd.ttc"), | |
| Path("C:/Windows/Fonts/simhei.ttf"), | |
| Path("C:/Windows/Fonts/simsun.ttc"), | |
| ] | |
| for path in candidates: | |
| if path.exists(): | |
| return path | |
| return None | |
| def _build_ocr_probe_image(): | |
| from PIL import Image, ImageDraw, ImageFont | |
| import numpy as np | |
| text = "现在主要是拉尔夫劳伦贸易" | |
| image = Image.new("RGB", (760, 128), (22, 22, 22)) | |
| draw = ImageDraw.Draw(image) | |
| font_path = _find_cjk_font() | |
| if not font_path: | |
| raise RuntimeError("No Windows CJK font found for OCR quality probe.") | |
| font = ImageFont.truetype(str(font_path), 46) | |
| draw.text((28, 32), text, font=font, fill=(245, 245, 245)) | |
| return np.array(image), text | |
| def _run_ocr_quality_probe(reader, max_frame_chars=160): | |
| probe_img, expected_text = _build_ocr_probe_image() | |
| results = reader.ocr(probe_img, cls=True) | |
| lines = _extract_ocr_lines(results, min_confidence=0.1) | |
| joined = " ".join(text for text, _score in lines).strip() | |
| joined = " ".join(joined.split()) | |
| expected_chars = set("现在主要拉尔夫劳伦贸易") | |
| hit_count = sum(1 for ch in expected_chars if ch in joined) | |
| reason = _ocr_garbage_reason(joined, max_frame_chars=max_frame_chars) | |
| ok = bool(joined) and hit_count >= 4 and not reason and len(lines) <= 8 | |
| details = { | |
| "ok": ok, | |
| "expected": expected_text, | |
| "recognized": joined[:240], | |
| "line_count": len(lines), | |
| "hit_count": hit_count, | |
| "reject_reason": reason, | |
| "sample_lines": [ | |
| {"text": text, "confidence": round(float(score), 4)} | |
| for text, score in lines[:5] | |
| ], | |
| } | |
| return ok, details | |
| def format_srt_time(seconds): | |
| h = int(seconds // 3600) | |
| m = int((seconds % 3600) // 60) | |
| s = int(seconds % 60) | |
| ms = int((seconds - int(seconds)) * 1000) | |
| return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}" | |
| def is_cuda_fully_functional(): | |
| try: | |
| import torch | |
| if not torch.cuda.is_available(): | |
| return False | |
| device_idx = torch.cuda.current_device() | |
| gpu_major, gpu_minor = torch.cuda.get_device_capability(device_idx) | |
| arch_list = torch.cuda.get_arch_list() | |
| if not arch_list: | |
| return False | |
| supported_caps = [] | |
| for arch in arch_list: | |
| digits = ''.join(c for c in arch if c.isdigit()) | |
| if digits: | |
| val = int(digits) | |
| major = val // 10 | |
| minor = val % 10 | |
| supported_caps.append((major, minor)) | |
| if not supported_caps: | |
| return False | |
| max_supported_major = max(major for major, minor in supported_caps) | |
| if gpu_major > max_supported_major: | |
| return False | |
| import torch.nn as nn | |
| conv = nn.Conv2d(1, 1, 3).cuda() | |
| x = torch.randn(1, 1, 8, 8).cuda() | |
| y = conv(x) | |
| float(y.to("cpu")[0, 0, 0, 0]) | |
| return True | |
| except Exception: | |
| return False | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Standalone PaddleOCR Subtitle Extractor CLI") | |
| parser.add_argument("--video", required=True, help="Path to input video file") | |
| parser.add_argument("--output", required=True, help="Path to output SRT file") | |
| parser.add_argument("--region", required=False, help="Selected subtitle region as 'x,y,w,h,orig_w,orig_h'") | |
| parser.add_argument("--device", default="auto", choices=["cuda", "cpu", "auto"], help="Device to run inference on") | |
| parser.add_argument("--prefer-gpu", default="true", help="Prefer GPU execution (true/false)") | |
| parser.add_argument("--allow-cpu-fallback", default="true", help="Allow fallback to CPU (true/false)") | |
| parser.add_argument("--gpu-mem-mb", type=int, default=3072, help="GPU memory allocation limit in MB") | |
| parser.add_argument("--sample-interval", type=float, default=0.25, help="Sample interval in seconds") | |
| parser.add_argument("--min-line-confidence", type=float, default=0.45, help="Drop OCR text lines below this confidence") | |
| parser.add_argument("--max-frame-text-chars", type=int, default=120, help="Reject OCR frame text longer than this") | |
| parser.add_argument("--max-block-duration", type=float, default=6.5, help="Maximum subtitle block duration before forced split") | |
| args = parser.parse_args() | |
| video_path = Path(args.video) | |
| output_srt_path = Path(args.output) | |
| if not video_path.exists(): | |
| print(f"Error: Video file not found at {video_path}", file=sys.stderr) | |
| sys.exit(2) | |
| prefer_gpu = args.prefer_gpu.lower() in ("true", "1", "yes", "t") | |
| allow_cpu_fallback = args.allow_cpu_fallback.lower() in ("true", "1", "yes", "t") | |
| # ── Step 1: Set environment variables BEFORE any paddle/paddleocr import ── | |
| import os | |
| os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python") | |
| os.environ["FLAGS_use_mkldnn"] = "0" | |
| print("[OCR ENV] PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python") | |
| try: | |
| import google.protobuf | |
| pb_ver = google.protobuf.__version__ | |
| except Exception: | |
| pb_ver = "Unknown" | |
| print(f"[OCR ENV] protobuf version: {pb_ver}") | |
| # ── Step 2: Import paddle (lightweight) and check CUDA BEFORE loading PaddleOCR ── | |
| try: | |
| import paddle | |
| except ImportError as e: | |
| print(f"Error: paddle not installed. {e}", file=sys.stderr) | |
| sys.exit(1) | |
| compiled_with_cuda = False | |
| try: | |
| compiled_with_cuda = paddle.device.is_compiled_with_cuda() | |
| except Exception: | |
| pass | |
| # Log configuration and GPU check | |
| print(f"[OCR CONFIG] requested device: {args.device}") | |
| print(f"[OCR CONFIG] allow_cpu_fallback: {allow_cpu_fallback}") | |
| print(f"[OCR GPU CHECK] paddle compiled with cuda: {compiled_with_cuda}") | |
| try: | |
| paddle_device = paddle.device.get_device() | |
| except Exception: | |
| paddle_device = "unknown" | |
| print(f"[OCR GPU CHECK] paddle current device: {paddle_device}") | |
| # ── Step 3: FAIL-FAST gate — exit before any heavy model loading ── | |
| run_on_gpu = False | |
| fallback_reason = None | |
| if args.device == "cuda": | |
| if not compiled_with_cuda: | |
| fallback_reason = "Paddle is not compiled with CUDA." | |
| if not allow_cpu_fallback: | |
| print("[OCR ERROR] PaddleOCR GPU requested but Paddle is CPU-only. Failing early.", file=sys.stderr) | |
| print(f"Reason: {fallback_reason}", file=sys.stderr) | |
| sys.exit(3) | |
| run_on_gpu = False | |
| else: | |
| run_on_gpu = True | |
| elif args.device == "cpu": | |
| run_on_gpu = False | |
| else: # auto | |
| if compiled_with_cuda and prefer_gpu: | |
| run_on_gpu = True | |
| else: | |
| if not compiled_with_cuda: | |
| fallback_reason = "Paddle is not compiled with CUDA." | |
| else: | |
| fallback_reason = "GPU not preferred by configuration." | |
| if not allow_cpu_fallback: | |
| print("[OCR ERROR] PaddleOCR GPU requested but Paddle is CPU-only. Failing early.", file=sys.stderr) | |
| print(f"Reason: {fallback_reason}", file=sys.stderr) | |
| sys.exit(3) | |
| run_on_gpu = False | |
| # Log device selection | |
| if run_on_gpu: | |
| print("[OCR] selected device: gpu") | |
| else: | |
| print(f"[OCR WARNING] PaddleOCR GPU unavailable, falling back to CPU.") | |
| print(f"Reason: {fallback_reason or 'User requested CPU execution.'}") | |
| if run_on_gpu: | |
| try: | |
| paddle.set_device("gpu:0") | |
| except Exception as e: | |
| print(f"Warning: paddle.set_device('gpu:0') failed: {e}") | |
| # ── Step 4: NOW import PaddleOCR (heavy — downloads/loads models) ── | |
| print("Loading PaddleOCR model...") | |
| try: | |
| from paddleocr import PaddleOCR | |
| except ImportError as e: | |
| print(f"Error: paddleocr not installed. {e}", file=sys.stderr) | |
| sys.exit(1) | |
| # Inspect PaddleOCR constructor signature | |
| import inspect | |
| sig = inspect.signature(PaddleOCR) | |
| print(f"[DEBUG] PaddleOCR constructor signature: {sig}") | |
| ocr_params = { | |
| "use_angle_cls": True, | |
| "lang": 'ch', | |
| "show_log": False, | |
| "enable_mkldnn": False | |
| } | |
| if "device" in sig.parameters: | |
| ocr_params["device"] = "gpu" if run_on_gpu else "cpu" | |
| else: | |
| ocr_params["use_gpu"] = run_on_gpu | |
| if run_on_gpu: | |
| ocr_params["gpu_mem"] = args.gpu_mem_mb | |
| try: | |
| reader = PaddleOCR(**ocr_params) | |
| except Exception as e: | |
| if not allow_cpu_fallback and run_on_gpu: | |
| print("PaddleOCR GPU requested but unavailable.", file=sys.stderr) | |
| print(f"Reason: PaddleOCR GPU init failed (Error: {e}).", file=sys.stderr) | |
| sys.exit(3) | |
| print(f"Warning: Failed to load PaddleOCR on GPU. Falling back to CPU. Error: {e}") | |
| # Try CPU fallback | |
| if "device" in sig.parameters: | |
| ocr_params["device"] = "cpu" | |
| else: | |
| ocr_params["use_gpu"] = False | |
| if "gpu_mem" in ocr_params: | |
| del ocr_params["gpu_mem"] | |
| try: | |
| reader = PaddleOCR(**ocr_params) | |
| except Exception as ex: | |
| print(f"PaddleOCR initialization failed.\nReason: {ex}", file=sys.stderr) | |
| sys.exit(4) | |
| # ── OCR GPU Preflight Inference Check ── | |
| if run_on_gpu: | |
| print("[OCR GPU PREFLIGHT] running first inference test...") | |
| try: | |
| import numpy as np | |
| dummy_img = np.zeros((64, 64, 3), dtype=np.uint8) | |
| reader.ocr(dummy_img, cls=True) | |
| print("[OCR GPU PREFLIGHT] PASS") | |
| except Exception as e: | |
| err_msg = str(e) | |
| if "cudnn64_8.dll" in err_msg or "PreconditionNotMet" in err_msg or "dynamic library" in err_msg or "DLL load failed" in err_msg: | |
| print("[OCR GPU PREFLIGHT] FAIL: cudnn64_8.dll missing or incompatible", file=sys.stderr) | |
| print("[OCR GPU PREFLIGHT] FAIL: cudnn64_8.dll missing or incompatible") | |
| if not allow_cpu_fallback: | |
| print("PaddleOCR GPU preflight failed because cudnn64_8.dll is missing or incompatible.", file=sys.stderr) | |
| sys.exit(5) | |
| else: | |
| print(f"[OCR GPU PREFLIGHT] FAIL: {e}", file=sys.stderr) | |
| print(f"[OCR GPU PREFLIGHT] FAIL: {e}") | |
| if not allow_cpu_fallback: | |
| print(f"PaddleOCR GPU preflight failed: {e}", file=sys.stderr) | |
| sys.exit(5) | |
| # Handle CPU fallback if allowed | |
| if allow_cpu_fallback: | |
| print("Warning: OCR GPU preflight failed. Falling back to CPU.") | |
| if "device" in sig.parameters: | |
| ocr_params["device"] = "cpu" | |
| else: | |
| ocr_params["use_gpu"] = False | |
| if "gpu_mem" in ocr_params: | |
| del ocr_params["gpu_mem"] | |
| try: | |
| reader = PaddleOCR(**ocr_params) | |
| run_on_gpu = False | |
| except Exception as ex: | |
| print(f"PaddleOCR CPU fallback initialization failed.\nReason: {ex}", file=sys.stderr) | |
| sys.exit(4) | |
| if run_on_gpu: | |
| print("[OCR GPU QUALITY PREFLIGHT] running Chinese text probe...") | |
| try: | |
| quality_ok, quality_details = _run_ocr_quality_probe(reader) | |
| print("[OCR GPU QUALITY PREFLIGHT] " + json.dumps(quality_details, ensure_ascii=False)) | |
| except Exception as e: | |
| quality_ok = False | |
| quality_details = {"ok": False, "error": str(e)} | |
| print(f"[OCR GPU QUALITY PREFLIGHT] FAIL: {e}", file=sys.stderr) | |
| if quality_ok: | |
| print("[OCR GPU QUALITY PREFLIGHT] PASS") | |
| elif not allow_cpu_fallback: | |
| print("PaddleOCR GPU quality preflight failed. GPU OCR is returning invalid Chinese text.", file=sys.stderr) | |
| print(json.dumps(quality_details, ensure_ascii=False), file=sys.stderr) | |
| sys.exit(5) | |
| else: | |
| print("Warning: OCR GPU quality preflight failed. Falling back to CPU.") | |
| if "device" in sig.parameters: | |
| ocr_params["device"] = "cpu" | |
| else: | |
| ocr_params["use_gpu"] = False | |
| if "gpu_mem" in ocr_params: | |
| del ocr_params["gpu_mem"] | |
| try: | |
| reader = PaddleOCR(**ocr_params) | |
| cpu_ok, cpu_details = _run_ocr_quality_probe(reader) | |
| print("[OCR CPU QUALITY PREFLIGHT] " + json.dumps(cpu_details, ensure_ascii=False)) | |
| if not cpu_ok: | |
| print("PaddleOCR CPU quality preflight also failed.", file=sys.stderr) | |
| sys.exit(5) | |
| run_on_gpu = False | |
| print("[OCR CPU QUALITY PREFLIGHT] PASS") | |
| except SystemExit: | |
| raise | |
| except Exception as ex: | |
| print(f"PaddleOCR CPU fallback quality preflight failed.\nReason: {ex}", file=sys.stderr) | |
| sys.exit(5) | |
| cap = cv2.VideoCapture(str(video_path)) | |
| if not cap.isOpened(): | |
| print(f"Error: Cannot open video file {video_path}", file=sys.stderr) | |
| sys.exit(2) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| duration = total_frames / fps if fps > 0 else 0 | |
| vid_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| vid_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| print(f"Video resolution: {vid_w}x{vid_h} | FPS: {fps:.2f} | Duration: {duration:.2f}s") | |
| if args.region: | |
| try: | |
| rx, ry, rw, rh, orig_w, orig_h = map(int, args.region.split(',')) | |
| norm_x = rx / orig_w if orig_w > 0 else 0.0 | |
| norm_y = ry / orig_h if orig_h > 0 else 0.0 | |
| norm_w = rw / orig_w if orig_w > 0 else 0.0 | |
| norm_h = rh / orig_h if orig_h > 0 else 0.0 | |
| real_x = int(round(norm_x * vid_w)) | |
| real_y = int(round(norm_y * vid_h)) | |
| real_w = max(1, int(round(norm_w * vid_w))) | |
| real_h = max(1, int(round(norm_h * vid_h))) | |
| real_w = min(real_w, vid_w) | |
| real_h = min(real_h, vid_h) | |
| real_x = max(0, min(real_x, vid_w - real_w)) | |
| real_y = max(0, min(real_y, vid_h - real_h)) | |
| print(f"[OCR] [UI Cords] region={rx},{ry},{rw},{rh}@{orig_w}x{orig_h}") | |
| print(f"[OCR] [Normalized Cords] x={norm_x:.4f}, y={norm_y:.4f}, w={norm_w:.4f}, h={norm_h:.4f}") | |
| print(f"[OCR] [Native Frame Cords] x={real_x}, y={real_y}, w={real_w}, h={real_h} (frame {vid_w}x{vid_h})") | |
| except Exception as e: | |
| print(f"Error parsing region '{args.region}': {e}", file=sys.stderr) | |
| sys.exit(1) | |
| else: | |
| # Default to bottom 25% of the frame | |
| real_x = 0 | |
| real_y = int(vid_h * 0.75) | |
| real_w = vid_w | |
| real_h = int(vid_h * 0.25) | |
| print(f"OCR Region: default bottom-center (x={real_x}, y={real_y}, w={real_w}, h={real_h})") | |
| sample_interval_sec = args.sample_interval | |
| frame_interval = max(1, int(fps * sample_interval_sec)) | |
| print(f"Analyzing video (sampling every {sample_interval_sec}s)...") | |
| raw_detections = [] | |
| prev_gray = None | |
| last_text = "" | |
| last_raw_text = "" | |
| last_confidence = 0.0 | |
| last_reject_reason = None | |
| ocr_calls = 0 | |
| skipped_calls = 0 # Duplicate-frame skips | |
| ocr_time_sum = 0.0 | |
| frame_idx = 0 | |
| start_time = time.time() | |
| try: | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if frame_idx % frame_interval == 0: | |
| current_time_sec = frame_idx / fps | |
| crop = frame[real_y:real_y+real_h, real_x:real_x+real_w] | |
| if crop.size == 0: | |
| frame_idx += 1 | |
| continue | |
| gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) | |
| is_duplicate = False | |
| if prev_gray is not None: | |
| diff = cv2.absdiff(gray, prev_gray) | |
| mean_diff = diff.mean() | |
| if mean_diff < 2.0: | |
| is_duplicate = True | |
| skipped_calls += 1 | |
| if is_duplicate: | |
| text = last_text | |
| raw_text = last_raw_text | |
| avg_confidence = last_confidence | |
| reject_reason = last_reject_reason | |
| else: | |
| ocr_start = time.time() | |
| results = reader.ocr(gray, cls=True) | |
| ocr_time_sum += (time.time() - ocr_start) | |
| ocr_calls += 1 | |
| text_list = [] | |
| scores = [] | |
| if results and results[0]: | |
| for detected_text, score in _extract_ocr_lines(results, args.min_line_confidence): | |
| text_list.append(detected_text) | |
| scores.append(score) | |
| raw_text = " ".join(text_list).strip() | |
| raw_text = " ".join(raw_text.split()) | |
| avg_confidence = sum(scores) / len(scores) if scores else 0.0 | |
| reject_reason = _ocr_garbage_reason(raw_text, args.max_frame_text_chars) | |
| text = "" if reject_reason else raw_text | |
| last_text = text | |
| last_raw_text = raw_text | |
| last_confidence = avg_confidence | |
| last_reject_reason = reject_reason | |
| prev_gray = gray | |
| raw_detections.append((current_time_sec, text, raw_text, avg_confidence, reject_reason)) | |
| if total_frames > 0: | |
| pct = min(100, int((frame_idx / total_frames) * 100)) | |
| print(f"PROGRESS: {pct}%", flush=True) | |
| frame_idx += 1 | |
| except Exception as e: | |
| print(f"OCR processing failed during scanning: {e}", file=sys.stderr) | |
| traceback.print_exc() | |
| sys.exit(5) | |
| cap.release() | |
| print("PROGRESS: 100%", flush=True) | |
| elapsed = time.time() - start_time | |
| # Total frames skipped from OCR includes both the interval skips and duplicate skips | |
| skipped_frames_total = (frame_idx - ocr_calls) | |
| avg_ocr_time = (ocr_time_sum / ocr_calls) if ocr_calls > 0 else 0.0 | |
| print(f"Scanning completed in {elapsed:.1f}s.") | |
| print(f"OCR calls: {ocr_calls}") | |
| print(f"Skipped frames: {skipped_frames_total}") | |
| print(f"Duplicate-frame skips: {skipped_calls}") | |
| print(f"Average OCR time per call: {avg_ocr_time:.3f}s") | |
| accepted_frames = sum(1 for _t, txt, _raw, _conf, _reason in raw_detections if txt) | |
| raw_text_frames = sum(1 for _t, _txt, raw, _conf, _reason in raw_detections if raw) | |
| rejected_frames = sum(1 for _t, _txt, _raw, _conf, reason in raw_detections if reason) | |
| # Blocker 5: Save raw detections to ocr_raw.json | |
| raw_json_path = output_srt_path.parent / "ocr_raw.json" | |
| raw_payload = [ | |
| { | |
| "timestamp": t, | |
| "text": txt, | |
| "raw_text": raw_txt, | |
| "confidence": round(float(conf), 4), | |
| "rejected": bool(reason), | |
| "reject_reason": reason, | |
| } | |
| for t, txt, raw_txt, conf, reason in raw_detections | |
| ] | |
| try: | |
| _atomic_write_json(raw_json_path, raw_payload) | |
| print(f"Saved raw detections to {raw_json_path}") | |
| except Exception as e: | |
| print(f"Warning: Failed to save ocr_raw.json: {e}", file=sys.stderr) | |
| # Temporal Smoothing & Merge | |
| srt_blocks = [] | |
| block_id = 1 | |
| active_text = "" | |
| active_start = -1.0 | |
| active_end = -1.0 | |
| gap_threshold = 0.4 | |
| similarity_threshold = 0.7 | |
| min_duration = 0.3 | |
| try: | |
| for current_time, text, _raw_text, _confidence, _reject_reason in raw_detections: | |
| if len(text) < 1: | |
| text = "" | |
| if text: | |
| if not active_text: | |
| active_text = text | |
| active_start = current_time | |
| active_end = current_time + sample_interval_sec | |
| else: | |
| sim = get_similarity(active_text, text) | |
| forced_split = ( | |
| float(args.max_block_duration) > 0 | |
| and (current_time + sample_interval_sec - active_start) > float(args.max_block_duration) | |
| ) | |
| is_same_subtitle = ( | |
| sim >= similarity_threshold | |
| or (text in active_text and len(text) >= 2) | |
| or (active_text in text and len(active_text) >= 2) | |
| ) | |
| if is_same_subtitle and not forced_split: | |
| active_end = current_time + sample_interval_sec | |
| if len(text) > len(active_text): | |
| active_text = text | |
| else: | |
| dur = active_end - active_start | |
| if dur >= min_duration: | |
| srt_blocks.append({ | |
| "id": block_id, | |
| "start": active_start, | |
| "end": active_end, | |
| "text": active_text | |
| }) | |
| block_id += 1 | |
| active_text = text | |
| active_start = current_time | |
| active_end = current_time + sample_interval_sec | |
| else: | |
| if active_text: | |
| if current_time - active_end <= gap_threshold: | |
| continue | |
| else: | |
| dur = active_end - active_start | |
| if dur >= min_duration: | |
| srt_blocks.append({ | |
| "id": block_id, | |
| "start": active_start, | |
| "end": active_end, | |
| "text": active_text | |
| }) | |
| block_id += 1 | |
| active_text = "" | |
| active_start = -1.0 | |
| active_end = -1.0 | |
| if active_text: | |
| dur = active_end - active_start | |
| if dur >= min_duration: | |
| srt_blocks.append({ | |
| "id": block_id, | |
| "start": active_start, | |
| "end": active_end, | |
| "text": active_text | |
| }) | |
| except Exception as e: | |
| print(f"OCR processing failed during smoothing: {e}", file=sys.stderr) | |
| traceback.print_exc() | |
| sys.exit(5) | |
| # Blocker 5: Save srt_blocks to ocr_grouped_blocks.json | |
| grouped_json_path = output_srt_path.parent / "ocr_grouped_blocks.json" | |
| grouped_payload = [ | |
| {"id": b["id"], "start": b["start"], "end": b["end"], "text": b["text"]} | |
| for b in srt_blocks | |
| ] | |
| try: | |
| _atomic_write_json(grouped_json_path, grouped_payload) | |
| print(f"Saved grouped blocks to {grouped_json_path}") | |
| except Exception as e: | |
| print(f"Warning: Failed to save ocr_grouped_blocks.json: {e}", file=sys.stderr) | |
| summary_path = output_srt_path.parent / "ocr_summary.json" | |
| summary_payload = { | |
| "status": "OK" if srt_blocks else "EMPTY", | |
| "device": "gpu" if run_on_gpu else "cpu", | |
| "duration_sec": round(float(duration), 3), | |
| "sample_interval_sec": sample_interval_sec, | |
| "ocr_calls": ocr_calls, | |
| "raw_detection_frames": len(raw_detections), | |
| "raw_text_frames": raw_text_frames, | |
| "accepted_frames": accepted_frames, | |
| "rejected_frames": rejected_frames, | |
| "srt_blocks": len(srt_blocks), | |
| "region": { | |
| "x": real_x, | |
| "y": real_y, | |
| "w": real_w, | |
| "h": real_h, | |
| "video_w": vid_w, | |
| "video_h": vid_h, | |
| }, | |
| } | |
| try: | |
| _atomic_write_json(summary_path, summary_payload) | |
| print(f"Saved OCR summary to {summary_path}") | |
| except Exception as e: | |
| print(f"Warning: Failed to save ocr_summary.json: {e}", file=sys.stderr) | |
| # Output SRT file | |
| srt_content = [] | |
| for block in srt_blocks: | |
| srt_content.append(f"{block['id']}\n") | |
| srt_content.append(f"{format_srt_time(block['start'])} --> {format_srt_time(block['end'])}\n") | |
| srt_content.append(f"{block['text']}\n\n") | |
| _atomic_write_text(output_srt_path, "".join(srt_content)) | |
| if not srt_blocks: | |
| print( | |
| "[OCR EMPTY] No valid subtitle blocks were produced. " | |
| f"raw_text_frames={raw_text_frames}, accepted_frames={accepted_frames}, rejected_frames={rejected_frames}. " | |
| "Check OCR region, GPU OCR quality, and ocr_raw.json/ocr_summary.json.", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(6) | |
| print(f"OCR transcription completed successfully. Saved {len(srt_blocks)} blocks to {output_srt_path}") | |
| # Cleanup VRAM and exit | |
| try: | |
| del reader | |
| except NameError: | |
| pass | |
| import gc | |
| gc.collect() | |
| try: | |
| import torch | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| except Exception: | |
| pass | |
| sys.exit(0) | |
| if __name__ == "__main__": | |
| main() | |