140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from botty_next.capture.mss_backend import MssCaptureBackend, save_frame
|
|
from botty_next.capture.window import find_window_region
|
|
from botty_next.config import load_config
|
|
from botty_next.vision.fixtures import load_image
|
|
from botty_next.vision.ocr import run_tesseract_ocr, save_ocr_preprocess_debug
|
|
from botty_next.vision.template_matching import match_template, save_match_debug
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(prog="botty-next")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
config_parser = subparsers.add_parser("config")
|
|
config_subparsers = config_parser.add_subparsers(dest="config_command", required=True)
|
|
validate_parser = config_subparsers.add_parser("validate")
|
|
validate_parser.add_argument("-c", "--config", required=True, type=Path)
|
|
validate_parser.set_defaults(handler=validate_config)
|
|
|
|
detect_parser = subparsers.add_parser("detect")
|
|
detect_parser.add_argument("detector", choices=["template"], help="detector to run")
|
|
detect_parser.add_argument("--image", required=True, type=Path)
|
|
detect_parser.add_argument("--template", required=True, type=Path)
|
|
detect_parser.add_argument("--threshold", type=float, default=0.85)
|
|
detect_parser.add_argument("--debug-output", type=Path)
|
|
detect_parser.set_defaults(handler=detect)
|
|
|
|
capture_parser = subparsers.add_parser("capture")
|
|
capture_parser.add_argument("--output", required=True, type=Path)
|
|
capture_parser.add_argument("--window-title", type=str)
|
|
capture_parser.set_defaults(handler=capture)
|
|
|
|
ocr_parser = subparsers.add_parser("ocr")
|
|
ocr_parser.add_argument("--image", required=True, type=Path)
|
|
ocr_parser.add_argument("--lang", default="eng")
|
|
ocr_parser.add_argument("--psm", type=int, default=6)
|
|
ocr_parser.add_argument("--tesseract-cmd")
|
|
ocr_parser.add_argument("--debug-output", type=Path)
|
|
ocr_parser.set_defaults(handler=ocr)
|
|
|
|
return parser
|
|
|
|
|
|
def validate_config(args: argparse.Namespace) -> int:
|
|
config = load_config(args.config)
|
|
print(json.dumps(config.model_dump(mode="json"), indent=2))
|
|
return 0
|
|
|
|
|
|
def detect(args: argparse.Namespace) -> int:
|
|
image = load_image(args.image)
|
|
template = load_image(args.template)
|
|
result = match_template(image, template, threshold=args.threshold)
|
|
|
|
if args.debug_output:
|
|
save_match_debug(image, result, args.debug_output)
|
|
|
|
print(json.dumps(_result_to_dict(result), indent=2))
|
|
return 0 if result.passed else 1
|
|
|
|
|
|
def capture(args: argparse.Namespace) -> int:
|
|
region = find_window_region(args.window_title) if args.window_title else None
|
|
frame = MssCaptureBackend().grab(region)
|
|
output = save_frame(frame, args.output)
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"output": str(output),
|
|
"shape": tuple(map(int, frame.shape)),
|
|
"window": region.title if region else None,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
def ocr(args: argparse.Namespace) -> int:
|
|
image = load_image(args.image)
|
|
if args.debug_output:
|
|
save_ocr_preprocess_debug(image, args.debug_output)
|
|
|
|
try:
|
|
result = run_tesseract_ocr(
|
|
image,
|
|
lang=args.lang,
|
|
psm=args.psm,
|
|
tesseract_cmd=args.tesseract_cmd,
|
|
)
|
|
except RuntimeError as exc:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"error": str(exc),
|
|
"debug_output": str(args.debug_output) if args.debug_output else None,
|
|
},
|
|
indent=2,
|
|
)
|
|
)
|
|
return 2
|
|
|
|
print(json.dumps(_ocr_result_to_dict(result), indent=2))
|
|
return 0
|
|
|
|
|
|
def _result_to_dict(result) -> dict:
|
|
return {
|
|
"confidence": result.confidence,
|
|
"bbox": result.bbox,
|
|
"passed": result.passed,
|
|
"method": result.method,
|
|
"debug": result.debug,
|
|
}
|
|
|
|
|
|
def _ocr_result_to_dict(result) -> dict:
|
|
return {
|
|
"text": result.text,
|
|
"confidence": result.confidence,
|
|
"bbox": result.bbox,
|
|
"debug": result.debug,
|
|
}
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
return args.handler(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|