import argparse import os import sys from .config import ( DEFAULT_POINTSIZE, DEFAULT_PORTAL_RESTORE_TOKEN, DEFAULT_REFERENCE_CROP, DEFAULT_REFERENCE_SIZE, DEFAULT_TEMPLATE_SET, ) from .ocr import ( read_queue_number, resolve_crop, templates_from_dataset, write_template_set, ) from .portal import PortalCaptureError from .reader import read_image_once, read_portal_window_once from .watcher import run_watch def parse_crop(value): parts = value.replace(",", " ").split() if len(parts) != 4: raise argparse.ArgumentTypeError("crop must be four numbers: x,y,width,height") try: return tuple(int(part) for part in parts) except ValueError as exc: raise argparse.ArgumentTypeError("crop values must be integers") from exc def parse_size(value): parts = value.lower().replace(",", "x").split("x") if len(parts) != 2: raise argparse.ArgumentTypeError("size must be WIDTHxHEIGHT") try: return int(parts[0]), int(parts[1]) except ValueError as exc: raise argparse.ArgumentTypeError("size values must be integers") from exc def image_paths(directory): extensions = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} return [ os.path.join(directory, name) for name in sorted(os.listdir(directory)) if os.path.splitext(name)[1].lower() in extensions ] def run_dataset(args): paths = image_paths(args.dataset) if not paths: print("no dataset images found") return 2 status = 0 passed = 0 failed = 0 for path in paths: crop = resolve_crop(path, args.crop, args.reference_crop, args.reference_size, args.scale_mode, args.cropped) if args.show_crop and crop: print(f"{path}\tcrop={crop[0]},{crop[1]},{crop[2]},{crop[3]}") number, details = read_queue_number(path, crop, args.template_set, args.font, args.pointsize, args.cropped) if not number: status = 2 failed += 1 suffix = "" if args.expect_filenames: expected = os.path.splitext(os.path.basename(path))[0] suffix = f"\texpected={expected}\tFAIL" print(f"{path}\tno digits found{suffix}") continue suffix = "" if args.expect_filenames: expected = os.path.splitext(os.path.basename(path))[0] if number == expected: passed += 1 suffix = f"\texpected={expected}\tPASS" else: failed += 1 status = 1 suffix = f"\texpected={expected}\tFAIL" print(f"{path}\t{number}{suffix}") if args.debug: for digit, bbox, area, hamming, iou in details: print(f" digit={digit} bbox={bbox} area={area} hamming={hamming:.3f} iou={iou:.3f}") if args.expect_filenames: print(f"summary: {passed} passed, {failed} failed") return status def build_template_set(args): paths = image_paths(args.build_template_set) if not paths: print("no template source images found") return 2 templates = templates_from_dataset(paths, args.reference_crop, args.reference_size, args.scale_mode, args.cropped) write_template_set(args.template_output, templates, paths) print(args.template_output) return 0 def read_single_image(args): result = read_image_once(args, args.image) if args.show_crop and result.crop: print(f"crop={result.crop[0]},{result.crop[1]},{result.crop[2]},{result.crop[3]}") if not result.number: print("no digits found") return 2 print(result.number) if args.debug: for digit, bbox, area, hamming, iou in result.details: print(f"digit={digit} bbox={bbox} area={area} hamming={hamming:.3f} iou={iou:.3f}") return 0 def read_portal_window(args): try: result = read_portal_window_once(args) except PortalCaptureError as exc: print(str(exc), file=sys.stderr) return 1 if args.show_crop and result.crop: print(f"crop={result.crop[0]},{result.crop[1]},{result.crop[2]},{result.crop[3]}") if not result.number: print("no digits found") return 2 print(result.number) if args.debug: portal_info = result.portal_info or {} print(f"portal-window node_id={portal_info.get('node_id', '')} streams={len(portal_info.get('streams', []))}") if portal_info.get("restore_token") and args.portal_restore_token: print(f"portal-restore-token {args.portal_restore_token}") for digit, bbox, area, hamming, iou in result.details: print(f"digit={digit} bbox={bbox} area={area} hamming={hamming:.3f} iou={iou:.3f}") return 0 def main(): parser = argparse.ArgumentParser(description="Read Arma Reforger queue number from a fixed screen crop.") parser.add_argument("--image", help="Read from an existing image.") parser.add_argument("--dataset", help="Read every image in a dataset directory.") parser.add_argument("--portal-window", action="store_true", help="Capture a user-selected Wayland window.") parser.add_argument("--watch", action="store_true", help="Keep reading and notify when the queue position is low.") parser.add_argument( "--watch-interval", type=float, default=15.0, help="Seconds between watch-mode polls.", ) parser.add_argument( "--alert-threshold", type=int, default=5, help="In watch mode, notify for detected queue positions at or below this number.", ) parser.add_argument( "--notify-command", default="notify-send", help="Notification command used by watch mode.", ) parser.add_argument( "--portal-restore-token", default=DEFAULT_PORTAL_RESTORE_TOKEN, help="Path used to store the portal restore token for later runs.", ) parser.add_argument( "--portal-reselect", action="store_true", help="Ignore the saved portal restore token and show the portal picker again.", ) parser.add_argument( "--portal-timeout", type=float, default=60.0, help="Seconds to wait for portal selection and the first PipeWire frame.", ) parser.add_argument("--cropped", action="store_true", help="Treat input image(s) as already cropped to the queue number.") parser.add_argument("--crop", type=parse_crop, help="Absolute crop as x,y,width,height.") parser.add_argument( "--reference-crop", type=parse_crop, default=DEFAULT_REFERENCE_CROP, help="Reference crop as x,y,width,height, scaled to the input image size when --crop is omitted.", ) parser.add_argument( "--reference-size", type=parse_size, default=DEFAULT_REFERENCE_SIZE, help="Reference screen size for --reference-crop scaling.", ) parser.add_argument( "--scale-mode", choices=("width", "independent"), default="width", help="How to scale --reference-crop. Width mode preserves a 16:9 UI layout and is the default.", ) parser.add_argument( "--template-set", default=DEFAULT_TEMPLATE_SET, help="Digit template JSON used by default recognition.", ) parser.add_argument("--font", help="Render synthetic digit templates from an ImageMagick font instead of using --template-set.") parser.add_argument("--pointsize", type=int, default=DEFAULT_POINTSIZE, help="Point size used for --font templates.") parser.add_argument("--debug", action="store_true", help="Print component and match details.") parser.add_argument("--show-crop", action="store_true", help="Print the resolved crop used for each full-size image.") parser.add_argument("--save-input", help="Save the live captured portal input image for debugging.") parser.add_argument( "--expect-filenames", action="store_true", help="In dataset mode, compare detected values to each image filename stem.", ) parser.add_argument("--build-template-set", metavar="DATASET", help=argparse.SUPPRESS) parser.add_argument("--template-output", default=DEFAULT_TEMPLATE_SET, help=argparse.SUPPRESS) args = parser.parse_args() if args.portal_reselect and not args.portal_window: parser.error("--portal-reselect can only be used with --portal-window") if args.save_input and not args.portal_window: parser.error("--save-input can only be used with --portal-window") if args.portal_timeout <= 0: parser.error("--portal-timeout must be positive") if args.watch_interval <= 0: parser.error("--watch-interval must be positive") if args.alert_threshold <= 1: parser.error("--alert-threshold must be greater than 1") if args.watch and not args.notify_command.strip(): parser.error("--notify-command must not be empty") selected_modes = [args.image, args.dataset, args.build_template_set, args.portal_window] if sum(bool(value) for value in selected_modes) != 1: parser.error("choose exactly one of --image, --dataset, --portal-window, or --build-template-set") if args.watch and not (args.image or args.portal_window): parser.error("--watch can only be used with --image or --portal-window") if args.build_template_set: return build_template_set(args) if args.dataset: return run_dataset(args) if args.watch: return run_watch(args) if args.image: return read_single_image(args) return read_portal_window(args)