DRIPPY4 / app /core /ocr_worker_cli.py
hoangtaiii's picture
Upload 92 files
16c3ac7 verified
Raw History Blame Contribute Delete
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()