Add batch processing for pending OCR
This commit is contained in:
219
scripts/process_pending_ocr.py
Executable file
219
scripts/process_pending_ocr.py
Executable file
@ -0,0 +1,219 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
"""
|
||||
Process all movies waiting for OCR in the WCX SQLite index.
|
||||
|
||||
The script queries the movie table for records whose OCR status is "pending"
|
||||
and invokes check_ocr.py once for each movie ID.
|
||||
|
||||
check_ocr.py performs the actual OCR request, parsing, validation, and database
|
||||
update. This batch script only coordinates those individual executions.
|
||||
|
||||
Processing continues if one movie fails. A final summary reports how many
|
||||
movies completed successfully and how many failed or require manual review.
|
||||
|
||||
By default, all pending movies are processed. Use --limit to restrict the
|
||||
number of movies during testing.
|
||||
|
||||
Usage:
|
||||
process_pending_ocr.py
|
||||
process_pending_ocr.py --limit 1
|
||||
process_pending_ocr.py --database /path/to/wcx-test.db
|
||||
|
||||
Requirements:
|
||||
GOOGLE_VISION_API_KEY must be available in the environment.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
DEFAULT_DATABASE_FILE = Path(
|
||||
"/storage/disk1/WCX/database/wcx.db"
|
||||
)
|
||||
CHECK_OCR_SCRIPT = Path(
|
||||
"/storage/disk1/WCX/scripts/check_ocr.py"
|
||||
)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Process movies whose OCR status is pending."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--database",
|
||||
type=Path,
|
||||
default=DEFAULT_DATABASE_FILE,
|
||||
help=f"SQLite database; default: {DEFAULT_DATABASE_FILE}",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
type=int,
|
||||
help="Process at most N pending movies.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def get_pending_movies(
|
||||
database_file: Path,
|
||||
limit: int | None,
|
||||
) -> list[tuple[str, str]]:
|
||||
sql = """
|
||||
SELECT id, name
|
||||
FROM movie
|
||||
WHERE ocr_status = 'pending'
|
||||
ORDER BY published, id
|
||||
"""
|
||||
|
||||
parameters: tuple[int, ...] = ()
|
||||
|
||||
if limit is not None:
|
||||
sql += "\nLIMIT ?"
|
||||
parameters = (limit,)
|
||||
|
||||
with sqlite3.connect(database_file) as connection:
|
||||
rows = connection.execute(sql, parameters).fetchall()
|
||||
|
||||
return [
|
||||
(str(movie_id), str(name))
|
||||
for movie_id, name in rows
|
||||
]
|
||||
|
||||
|
||||
def get_ocr_status(
|
||||
database_file: Path,
|
||||
movie_id: str,
|
||||
) -> tuple[str | None, str | None]:
|
||||
with sqlite3.connect(database_file) as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT ocr_status, ocr_error
|
||||
FROM movie
|
||||
WHERE id = ?
|
||||
""",
|
||||
(movie_id,),
|
||||
).fetchone()
|
||||
|
||||
if row is None:
|
||||
return None, "Movie disappeared from the database."
|
||||
|
||||
return row[0], row[1]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_arguments()
|
||||
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
raise ValueError("--limit must be greater than zero.")
|
||||
|
||||
if not args.database.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"Database file does not exist: {args.database}"
|
||||
)
|
||||
|
||||
if not CHECK_OCR_SCRIPT.is_file():
|
||||
raise FileNotFoundError(
|
||||
f"OCR processing script does not exist: {CHECK_OCR_SCRIPT}"
|
||||
)
|
||||
|
||||
if not os.environ.get("GOOGLE_VISION_API_KEY"):
|
||||
raise ValueError(
|
||||
"GOOGLE_VISION_API_KEY is not set in the environment."
|
||||
)
|
||||
|
||||
pending_movies = get_pending_movies(
|
||||
database_file=args.database,
|
||||
limit=args.limit,
|
||||
)
|
||||
|
||||
if not pending_movies:
|
||||
print("No movies are waiting for OCR.")
|
||||
return
|
||||
|
||||
print(f"Movies waiting for OCR: {len(pending_movies)}")
|
||||
print()
|
||||
|
||||
completed = 0
|
||||
manual_review = 0
|
||||
failed = 0
|
||||
unexpected = 0
|
||||
|
||||
for position, (movie_id, name) in enumerate(
|
||||
pending_movies,
|
||||
start=1,
|
||||
):
|
||||
print(
|
||||
f"[{position}/{len(pending_movies)}] "
|
||||
f"{movie_id} - {name}"
|
||||
)
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
str(CHECK_OCR_SCRIPT),
|
||||
movie_id,
|
||||
"--database",
|
||||
str(args.database),
|
||||
],
|
||||
text=True,
|
||||
capture_output=True,
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
print(result.stdout.strip())
|
||||
|
||||
if result.stderr.strip():
|
||||
print(result.stderr.strip(), file=sys.stderr)
|
||||
|
||||
status, error_message = get_ocr_status(
|
||||
database_file=args.database,
|
||||
movie_id=movie_id,
|
||||
)
|
||||
|
||||
if status == "completed":
|
||||
completed += 1
|
||||
print("Result: completed")
|
||||
elif status == "manual_review":
|
||||
manual_review += 1
|
||||
print("Result: manual_review")
|
||||
elif status == "failed":
|
||||
failed += 1
|
||||
print("Result: failed")
|
||||
else:
|
||||
unexpected += 1
|
||||
print(f"Result: unexpected status {status!r}")
|
||||
|
||||
if error_message:
|
||||
print(f"Error: {error_message}")
|
||||
|
||||
print()
|
||||
|
||||
print("OCR batch summary")
|
||||
print("-----------------")
|
||||
print(f"Selected: {len(pending_movies)}")
|
||||
print(f"Completed: {completed}")
|
||||
print(f"Manual review: {manual_review}")
|
||||
print(f"Failed: {failed}")
|
||||
print(f"Unexpected: {unexpected}")
|
||||
|
||||
if failed or unexpected:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except (
|
||||
FileNotFoundError,
|
||||
ValueError,
|
||||
sqlite3.Error,
|
||||
) as error:
|
||||
print(f"Error: {error}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
Reference in New Issue
Block a user