80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Capture the QEMU VGA console and print its text with RapidOCR."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from rapidocr import RapidOCR
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
DEFAULT_IMAGE = ROOT / ".qemu" / "qemu-screen.png"
|
|
|
|
|
|
def capture() -> Path:
|
|
subprocess.run(
|
|
["ferro-vm", "screenshot"],
|
|
check=True,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
if not DEFAULT_IMAGE.is_file():
|
|
raise RuntimeError(f"QEMU screenshot was not created: {DEFAULT_IMAGE}")
|
|
return DEFAULT_IMAGE
|
|
|
|
|
|
def recognize(image: Path, min_score: float) -> list[dict[str, object]]:
|
|
logging.disable(logging.INFO)
|
|
result = RapidOCR()(image)
|
|
rows: list[dict[str, object]] = []
|
|
if result.txts is None or result.boxes is None or result.scores is None:
|
|
return rows
|
|
for box, text, score in zip(result.boxes, result.txts, result.scores):
|
|
if float(score) < min_score:
|
|
continue
|
|
rows.append(
|
|
{
|
|
"x": int(min(point[0] for point in box)),
|
|
"y": int(min(point[1] for point in box)),
|
|
"text": text,
|
|
"score": round(float(score), 5),
|
|
}
|
|
)
|
|
rows.sort(key=lambda row: (int(row["y"]), int(row["x"])))
|
|
return rows
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--image", type=Path, help="OCR an existing image instead of capturing QEMU")
|
|
parser.add_argument("--min-score", type=float, default=0.5)
|
|
parser.add_argument("--json", action="store_true", help="emit locations and scores as JSON")
|
|
parser.add_argument("-o", "--output", type=Path, help="also save output as UTF-8 text")
|
|
args = parser.parse_args()
|
|
|
|
image = args.image.resolve() if args.image else capture()
|
|
rows = recognize(image, args.min_score)
|
|
if args.json:
|
|
rendered = json.dumps({"image": str(image), "lines": rows}, ensure_ascii=False, indent=2)
|
|
else:
|
|
rendered = "\n".join(str(row["text"]) for row in rows)
|
|
if args.output:
|
|
args.output.write_text(rendered + ("\n" if rendered else ""), encoding="utf-8")
|
|
if rendered:
|
|
print(rendered)
|
|
return 0 if rows else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except (OSError, RuntimeError, subprocess.CalledProcessError) as exc:
|
|
print(f"qemu-ocr: {exc}", file=sys.stderr)
|
|
raise SystemExit(2)
|