250 lines
9.5 KiB
Python
250 lines
9.5 KiB
Python
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)
|