242 lines
8.1 KiB
Python
242 lines
8.1 KiB
Python
"""OCR a still. Prefer Umi-OCR HTTP, then tesseract, then Windows.Media.Ocr.
|
|
|
|
Writes UTF-8 text. Optional: add a role=ocr row to media.json.
|
|
|
|
Umi-OCR must be running (HTTP :1224). Start Umi-OCR.exe once. Then:
|
|
|
|
python ocr_pass.py --in still.png --engine umi
|
|
|
|
Uses tbpu.parser=single_code (keep indent) and English. ShareX has no OCR CLI.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from .media_manifest import (
|
|
_from_wire_dict,
|
|
_manifest_to_wire,
|
|
atomic_write_json,
|
|
upsert_derived_item,
|
|
)
|
|
except ImportError: # pragma: no cover - script-mode import
|
|
from media_manifest import (
|
|
_from_wire_dict,
|
|
_manifest_to_wire,
|
|
atomic_write_json,
|
|
upsert_derived_item,
|
|
)
|
|
|
|
UMI_URL: str = "http://127.0.0.1:1224/api/ocr"
|
|
|
|
|
|
def find_umi() -> str | None:
|
|
"""Locate the ``umi-ocr`` / ``Umi-OCR`` binary on PATH (or ``None``)."""
|
|
return shutil.which("umi-ocr") or shutil.which("Umi-OCR")
|
|
|
|
|
|
def find_tesseract() -> str | None:
|
|
"""Locate the ``tesseract`` binary on PATH (or ``None``)."""
|
|
return shutil.which("tesseract")
|
|
|
|
|
|
def umi_http_up() -> bool:
|
|
"""Return True if Umi-OCR's HTTP server is reachable on 127.0.0.1:1224."""
|
|
try:
|
|
urllib.request.urlopen("http://127.0.0.1:1224/api/ocr/get_options", timeout=2)
|
|
return True
|
|
except (urllib.error.URLError, TimeoutError, OSError):
|
|
return False
|
|
|
|
|
|
def run_umi(src: Path, dest: Path) -> None:
|
|
"""OCR ``src`` via Umi-OCR's HTTP API and write UTF-8 text to ``dest``."""
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
if not umi_http_up():
|
|
exe = find_umi()
|
|
if exe:
|
|
subprocess.Popen([exe])
|
|
raise SystemExit(
|
|
"Umi-OCR HTTP not on 127.0.0.1:1224. Start Umi-OCR.exe and retry."
|
|
)
|
|
payload = {
|
|
"base64": base64.b64encode(src.read_bytes()).decode("ascii"),
|
|
"options": {
|
|
"ocr.language": "English",
|
|
"tbpu.parser": "single_code",
|
|
"ocr.maxSideLen": 999999,
|
|
"data.format": "text",
|
|
},
|
|
}
|
|
req = urllib.request.Request(
|
|
UMI_URL,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
print(f"umi-http {src.name}")
|
|
raw = urllib.request.urlopen(req, timeout=180).read()
|
|
res = json.loads(raw.decode("utf-8"))
|
|
code = res.get("code")
|
|
data = res.get("data")
|
|
if code == 100 and isinstance(data, str):
|
|
text = data
|
|
elif code == 101:
|
|
text = ""
|
|
elif isinstance(data, list):
|
|
text = "".join(
|
|
str(block.get("text", "")) + str(block.get("end", "")) for block in data
|
|
)
|
|
else:
|
|
raise SystemExit(f"umi-ocr failed code={code} data={data!r}")
|
|
dest.write_text(text.replace("\r\n", "\n"), encoding="utf-8")
|
|
|
|
|
|
def run_tesseract(src: Path, dest: Path) -> None:
|
|
"""OCR ``src`` via the local ``tesseract`` binary and write UTF-8 text to ``dest``."""
|
|
exe = find_tesseract()
|
|
exe = find_tesseract()
|
|
if not exe:
|
|
raise SystemExit("tesseract not on PATH. scoop install tesseract")
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
out_base = dest.with_suffix("")
|
|
cmd = [exe, str(src), str(out_base), "-l", "eng"]
|
|
print(" ".join(cmd))
|
|
subprocess.check_call(cmd)
|
|
txt = out_base.with_suffix(".txt")
|
|
if txt != dest:
|
|
dest.write_text(txt.read_text(encoding="utf-8", errors="replace"), encoding="utf-8")
|
|
|
|
|
|
def run_windows_ocr(src: Path, dest: Path) -> None:
|
|
"""OCR ``src`` via the Windows.Media.Ocr PowerShell shim and write UTF-8 text to ``dest``."""
|
|
ps = r"""
|
|
$ErrorActionPreference = 'Stop'
|
|
Add-Type -AssemblyName System.Runtime.WindowsRuntime
|
|
$null = [Windows.Storage.StorageFile, Windows.Storage, ContentType = WindowsRuntime]
|
|
$null = [Windows.Graphics.Imaging.BitmapDecoder, Windows.Graphics.Imaging, ContentType = WindowsRuntime]
|
|
$null = [Windows.Media.Ocr.OcrEngine, Windows.Media.Ocr, ContentType = WindowsRuntime]
|
|
function Await($op) {
|
|
$asTask = [System.WindowsRuntimeSystemExtensions].GetMethods() | Where-Object {
|
|
$_.Name -eq 'AsTask' -and $_.GetParameters().Count -eq 1
|
|
} | Select-Object -First 1
|
|
$task = $asTask.MakeGenericMethod($op.GetType().GenericTypeArguments[0]).Invoke($null, @($op))
|
|
$task.Wait()
|
|
$task.Result
|
|
}
|
|
$src = $args[0]
|
|
$dest = $args[1]
|
|
$file = Await ([Windows.Storage.StorageFile]::GetFileFromPathAsync($src))
|
|
$stream = Await ($file.OpenAsync([Windows.Storage.FileAccessMode]::Read))
|
|
$decoder = Await ([Windows.Graphics.Imaging.BitmapDecoder]::CreateAsync($stream))
|
|
$bitmap = Await ($decoder.GetSoftwareBitmapAsync())
|
|
$engine = [Windows.Media.Ocr.OcrEngine]::TryCreateFromUserProfileLanguages()
|
|
if ($null -eq $engine) { throw 'Windows OCR engine unavailable' }
|
|
$result = Await ($engine.RecognizeAsync($bitmap))
|
|
[System.IO.File]::WriteAllText($dest, $result.Text, [System.Text.UTF8Encoding]::new($false))
|
|
"""
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
cmd = [
|
|
"powershell",
|
|
"-NoProfile",
|
|
"-Command",
|
|
ps,
|
|
str(src.resolve()),
|
|
str(dest.resolve()),
|
|
]
|
|
print("windows.media.ocr", src.name)
|
|
subprocess.check_call(cmd)
|
|
|
|
|
|
def append_media_row(media_json: Path, src: Path, dest: Path) -> None:
|
|
"""Append or update the ``role=ocr`` row in ``media_json`` for the (src -> dest) file pair."""
|
|
data = json.loads(media_json.read_text(encoding="utf-8"))
|
|
if data.get("schema_version") != 2:
|
|
raise SystemExit(
|
|
"legacy media.json requires: tw.py migrate-media --id <id> --apply"
|
|
)
|
|
post_id = ""
|
|
media_id = ""
|
|
handle = ""
|
|
for item in data.get("items") or []:
|
|
if str(item.get("filename") or "") in {src.name, dest.name}:
|
|
post_id = str(item.get("post_id") or "")
|
|
media_id = str(item.get("media_id") or "")
|
|
handle = str(item.get("handle") or "")
|
|
break
|
|
if not post_id and "_" in src.stem:
|
|
parts = src.stem.split("_")
|
|
post_id = parts[0]
|
|
media_id = parts[1] if len(parts) > 1 else ""
|
|
asset_dir = media_json.parent
|
|
manifest = _from_wire_dict(data)
|
|
manifest, _ = upsert_derived_item(
|
|
manifest,
|
|
post_id=post_id,
|
|
media_id=media_id,
|
|
handle=handle,
|
|
role="ocr",
|
|
filename=dest.name,
|
|
asset_dir=asset_dir,
|
|
now=datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
)
|
|
wire = _manifest_to_wire(manifest)
|
|
wire["mirrors"] = list(data.get("mirrors") or [])
|
|
atomic_write_json(media_json, wire)
|
|
print(f"updated {media_json}")
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--in", dest="src", type=Path, required=True)
|
|
parser.add_argument("--out", dest="dest", type=Path, default=None)
|
|
parser.add_argument(
|
|
"--engine",
|
|
choices=("auto", "umi", "tesseract", "windows"),
|
|
default="auto",
|
|
)
|
|
parser.add_argument("--media-json", type=Path, default=None)
|
|
args = parser.parse_args(argv)
|
|
|
|
if not args.src.is_file():
|
|
raise SystemExit(f"missing {args.src}")
|
|
dest = args.dest
|
|
if dest is None:
|
|
dest = args.src.with_name(args.src.stem.replace("_orig", "") + "_ocr.txt")
|
|
if dest.name == args.src.stem + "_ocr.txt" and args.src.stem.endswith("_orig"):
|
|
dest = args.src.with_name(args.src.stem[: -len("_orig")] + "_ocr.txt")
|
|
|
|
engine = args.engine
|
|
if engine == "auto":
|
|
if find_umi():
|
|
engine = "umi"
|
|
elif find_tesseract():
|
|
engine = "tesseract"
|
|
else:
|
|
engine = "windows"
|
|
print(f"engine={engine}")
|
|
if engine == "umi":
|
|
run_umi(args.src, dest)
|
|
elif engine == "tesseract":
|
|
run_tesseract(args.src, dest)
|
|
else:
|
|
run_windows_ocr(args.src, dest)
|
|
|
|
text = dest.read_text(encoding="utf-8", errors="replace")
|
|
print(f"wrote {dest} ({len(text)} chars)")
|
|
if args.media_json:
|
|
append_media_row(args.media_json, args.src, dest)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|