Files
WCX/scripts/check_ocr.py

361 lines
8.6 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Run and validate OCR processing for one movie in the SQLite index.
The script accepts a WCX movie ID, retrieves the movie name and thumbnail URL
from the selected database, and coordinates the complete OCR workflow:
1. Run ocr.sh for the thumbnail URL.
2. Pass the raw OCR text to parse_ocr.py.
3. Compare the OCR name with the movie name stored in the database.
4. Store validated metadata and OCR processing information.
When processing succeeds and the names match, the script updates:
nationality
shoot_location
shoot_date
ocr_raw_text
ocr_status = completed
ocr_processed_at
modified_at
If OCR execution fails, the movie is marked as failed.
If the OCR text cannot be parsed, required fields are missing, or the OCR name
does not match the database name, the raw OCR text is retained and the movie
is marked as manual_review. In these cases, the parsed metadata fields are not
updated.
Usage:
check_ocr.py <movie-id>
check_ocr.py <movie-id> --database /path/to/wcx.db
Example:
check_ocr.py susana-melo_6707
"""
import argparse
import sqlite3
import subprocess
import sys
from pathlib import Path
DEFAULT_DATABASE_FILE = Path(
"/storage/disk1/WCX/database/wcx.db"
)
OCR_SCRIPT = Path(
"/storage/disk1/WCX/scripts/ocr.sh"
)
PARSER_SCRIPT = Path(
"/storage/disk1/WCX/scripts/parse_ocr.py"
)
def parse_arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Run OCR processing for one WCX movie."
)
parser.add_argument(
"movie_id",
help="Movie ID to process.",
)
parser.add_argument(
"--database",
type=Path,
default=DEFAULT_DATABASE_FILE,
help=f"SQLite database; default: {DEFAULT_DATABASE_FILE}",
)
return parser.parse_args()
def parse_key_value_output(output: str) -> dict[str, str]:
result: dict[str, str] = {}
for line in output.splitlines():
if "=" not in line:
continue
key, value = line.split("=", 1)
result[key.strip()] = value.strip()
return result
def normalize_name(value: str) -> str:
return " ".join(value.casefold().split())
def set_ocr_failed(
connection: sqlite3.Connection,
movie_id: str,
error_message: str,
raw_text: str | None = None,
) -> None:
connection.execute(
"""
UPDATE movie
SET
ocr_status = 'failed',
ocr_raw_text = ?,
ocr_error = ?,
ocr_processed_at = CURRENT_TIMESTAMP,
modified_at = CURRENT_TIMESTAMP
WHERE id = ?
""",
(
raw_text,
error_message,
movie_id,
),
)
def set_manual_review(
connection: sqlite3.Connection,
movie_id: str,
raw_text: str,
error_message: str,
) -> None:
connection.execute(
"""
UPDATE movie
SET
ocr_status = 'manual_review',
ocr_raw_text = ?,
ocr_error = ?,
ocr_processed_at = CURRENT_TIMESTAMP,
modified_at = CURRENT_TIMESTAMP
WHERE id = ?
""",
(
raw_text,
error_message,
movie_id,
),
)
def set_ocr_completed(
connection: sqlite3.Connection,
movie_id: str,
raw_text: str,
nationality: str,
shoot_location: str,
shoot_date: str,
) -> None:
connection.execute(
"""
UPDATE movie
SET
nationality = ?,
shoot_location = ?,
shoot_date = ?,
ocr_status = 'completed',
ocr_raw_text = ?,
ocr_error = NULL,
ocr_processed_at = CURRENT_TIMESTAMP,
modified_at = CURRENT_TIMESTAMP
WHERE id = ?
""",
(
nationality,
shoot_location,
shoot_date,
raw_text,
movie_id,
),
)
def main() -> None:
args = parse_arguments()
movie_id = args.movie_id
database_file = args.database
if not database_file.is_file():
raise FileNotFoundError(
f"Database file does not exist: {database_file}"
)
if not OCR_SCRIPT.is_file():
raise FileNotFoundError(
f"OCR script does not exist: {OCR_SCRIPT}"
)
if not PARSER_SCRIPT.is_file():
raise FileNotFoundError(
f"OCR parser does not exist: {PARSER_SCRIPT}"
)
with sqlite3.connect(database_file) as connection:
row = connection.execute(
"""
SELECT name, thumbnail
FROM movie
WHERE id = ?
""",
(movie_id,),
).fetchone()
if row is None:
print(
f"Movie does not exist in database: {movie_id}",
file=sys.stderr,
)
sys.exit(1)
database_name, thumbnail = row
if not thumbnail:
error_message = "Movie has no thumbnail."
set_ocr_failed(
connection,
movie_id,
error_message,
)
print(error_message, file=sys.stderr)
sys.exit(1)
ocr_result = subprocess.run(
[str(OCR_SCRIPT), thumbnail],
capture_output=True,
text=True,
)
if ocr_result.returncode != 0:
error_message = (
ocr_result.stderr.strip()
or "OCR execution failed."
)
set_ocr_failed(
connection,
movie_id,
error_message,
ocr_result.stdout.strip() or None,
)
print("OCR execution failed:", file=sys.stderr)
print(error_message, file=sys.stderr)
sys.exit(1)
raw_text = ocr_result.stdout.strip()
parser_result = subprocess.run(
[str(PARSER_SCRIPT)],
input=raw_text,
capture_output=True,
text=True,
)
if parser_result.returncode != 0:
error_message = (
parser_result.stderr.strip()
or "OCR output parsing failed."
)
set_manual_review(
connection,
movie_id,
raw_text,
error_message,
)
print(
"OCR output parsing failed:",
file=sys.stderr,
)
print(error_message, file=sys.stderr)
sys.exit(1)
parsed = parse_key_value_output(parser_result.stdout)
required_fields = (
"ocr_name",
"nationality",
"shoot_location",
"shoot_date",
)
missing_fields = [
field
for field in required_fields
if not parsed.get(field)
]
if missing_fields:
error_message = (
"Parser output is missing fields: "
+ ", ".join(missing_fields)
)
set_manual_review(
connection,
movie_id,
raw_text,
error_message,
)
print(error_message, file=sys.stderr)
sys.exit(1)
ocr_name = parsed["ocr_name"]
names_match = (
normalize_name(database_name)
== normalize_name(ocr_name)
)
if not names_match:
error_message = (
"OCR name does not match database name: "
f"{ocr_name!r} != {database_name!r}"
)
set_manual_review(
connection,
movie_id,
raw_text,
error_message,
)
print(error_message, file=sys.stderr)
sys.exit(1)
set_ocr_completed(
connection=connection,
movie_id=movie_id,
raw_text=raw_text,
nationality=parsed["nationality"],
shoot_location=parsed["shoot_location"],
shoot_date=parsed["shoot_date"],
)
print(f"id={movie_id}")
print(f"database={database_file}")
print(f"database_name={database_name}")
print(f"ocr_name={ocr_name}")
print("name_match=yes")
print(f"nationality={parsed['nationality']}")
print(f"shoot_location={parsed['shoot_location']}")
print(f"shoot_date={parsed['shoot_date']}")
print("ocr_status=completed")
if __name__ == "__main__":
try:
main()
except (
FileNotFoundError,
sqlite3.Error,
) as error:
print(f"Error: {error}", file=sys.stderr)
sys.exit(1)