import argparse import os import tempfile from .config import ( DEFAULT_POINTSIZE, DEFAULT_PORTAL_RESTORE_TOKEN, DEFAULT_REFERENCE_CROP, DEFAULT_REFERENCE_SIZE, DEFAULT_TEMPLATE_SET, ) from .magick import run_bytes from .ocr import ( read_queue_number, resolve_crop, templates_from_dataset, write_template_set, ) from .portal import capture_portal_window 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): crop = resolve_crop(args.image, args.crop, args.reference_crop, args.reference_size, args.scale_mode, args.cropped) if args.show_crop and crop: print(f"crop={crop[0]},{crop[1]},{crop[2]},{crop[3]}") number, details = read_queue_number(args.image, crop, args.template_set, args.font, args.pointsize, args.cropped) if not number: print("no digits found") return 2 print(number) 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}") return 0 def read_portal_window(args): fd, image_path = tempfile.mkstemp(prefix="reforger-queue-", suffix=".png") os.close(fd) try: portal_info = capture_portal_window( image_path, args.portal_restore_token, args.portal_reselect, args.portal_timeout, ) if args.save_input: run_bytes(["magick", image_path, "+repage", args.save_input]) crop = resolve_crop(image_path, args.crop, args.reference_crop, args.reference_size, args.scale_mode, args.cropped) if args.show_crop and crop: print(f"crop={crop[0]},{crop[1]},{crop[2]},{crop[3]}") number, details = read_queue_number(image_path, crop, args.template_set, args.font, args.pointsize, args.cropped) finally: try: os.unlink(image_path) except OSError: pass if not number: print("no digits found") return 2 print(number) if args.debug: 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 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( "--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") 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.build_template_set: return build_template_set(args) if args.dataset: return run_dataset(args) if args.image: return read_single_image(args) return read_portal_window(args)