diff --git a/oxart/frontend/aqctl_orca_camera.py b/oxart/frontend/aqctl_orca_camera.py index 9fd84bb..f6f7a86 100644 --- a/oxart/frontend/aqctl_orca_camera.py +++ b/oxart/frontend/aqctl_orca_camera.py @@ -1,17 +1,37 @@ #!/usr/bin/env python3.5 import argparse +from textwrap import dedent +from typing import Dict, Any, List, Callable + +import toml import zmq import logging from sipyco.pc_rpc import simple_server_loop import sipyco.common_args as sca -from orca_camera.driver import OrcaFusion +from orca_camera.driver import OrcaFusion, LibInstance logger = logging.getLogger(__name__) def get_argparser(): - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser( + epilog=dedent("""\ + Example configuration file: + ``` + [system1] + serial = "S/N: 123456" + broadcast_images = true + roi = [884, 200, 956, 48] + speed = 2 + + [system2] + serial = "S/N: 456789" + ``` + + The serial must match the string provided by the device exactly. + """) + ) sca.simple_network_args(parser, 4000) sca.verbosity_args(parser) parser.add_argument("--broadcast-images", default=True, action="store_true") @@ -22,50 +42,119 @@ def get_argparser(): default=3, type=int, help="integer value - (0: QUIET, 1: STANDARD, 2: FAST)") + parser.add_argument( + "--config-file", + help="Path to TOML configuration file that allows setting serial, roi, " + "speed and broadcast_images for multiple cameras." + ) return parser -def create_zmq_server(bind="*", port=5555): - context = zmq.Context() - socket = context.socket(zmq.PUB) - socket.set_hwm(1) - socket.bind("tcp://{}:{}".format(bind, port)) - return socket - - -def main(): - args = get_argparser().parse_args() - sca.init_logger_from_args(args) - - logger.info("Initialising cameras...") - - roi = [int(x) for x in args.roi.split(",")] - assert (len(roi) == 4) - - dev = OrcaFusion() - dev.open(camera_index=0, framebuffer_len=1000) - dev.set_subarray(*roi) - dev.set_readout_speed(args.speed) +def lazy_zmq_server(bind="*", port=5555): + socket = None + + def create_zmq_server(): + nonlocal socket + if socket is None: + context = zmq.Context() + socket = context.socket(zmq.PUB) + socket.set_hwm(1) + socket.bind("tcp://{}:{}".format(bind, port)) + return socket + + return create_zmq_server + + +def init_dev( + name: str, + args: Dict[str, Any], + devs: List[OrcaFusion], + create_zmq_server: Callable[[], Any], +) -> OrcaFusion: + + if args["serial"] is None and len(devs) != 1: + raise ValueError( + f"Multiple cameras present but no serial set for {name}. Camera " + f"serials present: {', '.join(repr(dev.get_serial()) for dev in devs)}" + ) + elif args["serial"] is None and len(devs) == 1: + dev = devs[0] + else: + for idx, dev in enumerate(devs): + if dev.get_serial() == args["serial"]: + break + else: + raise ValueError( + f"Unable to find camera with serial number {args['serial']} for " + f"{name}. Camera serials present: " + f"{', '.join(repr(dev.get_serial()) for dev in devs)}" + ) + # Remove the matched device so that we don't reference the same device + # twice with different arguments + del devs[idx] + + if "roi" in args: + dev.set_subarray(*args["roi"]) + if "speed" in args: + dev.set_readout_speed(args["speed"]) dev.start_capture() - if args.broadcast_images: - socket = create_zmq_server(args.zmq_bind, args.zmq_port) + if args.get("broadcast_images", False): + socket = create_zmq_server() def frame_callback(im): # We send a multi-part message with first part the serial number # this allows the subscriber to filter out unwanted images - socket.send_string(str(0), flags=zmq.SNDMORE) + socket.send_string(dev.get_serial(), flags=zmq.SNDMORE) socket.send_pyobj(im.transpose()) dev.register_callback(frame_callback) + return dev + + +def main(): + args = get_argparser().parse_args() + sca.init_logger_from_args(args) + + logger.info("Initialising cameras...") + + if args.config_file: + config = toml.load(args.config_file) + else: + roi = [int(x) for x in args.roi.split(",")] + if len(roi) != 4: + raise ValueError("--roi should be 4 comma separated numbers") + + config = { + "camera": { + "serial": None, + "broadcast_images": args.broadcast_images, + "roi": roi, + "speed": args.speed, + } + } + + devs = [ + OrcaFusion(idx, framebuffer_len=1000) + for idx in range(LibInstance.get().num_dev) + ] + all_devs = devs.copy() + create_zmq_server = lazy_zmq_server(args.zmq_bind, args.zmq_port) + + servers = { + name: init_dev(name, args, devs, create_zmq_server) + for name, args in config.items() + } + try: - simple_server_loop({"camera": dev}, args.bind, args.port) + simple_server_loop(servers, args.bind, args.port) except KeyboardInterrupt: pass finally: logger.info("Shutting down camera ...") - dev.close() + for dev in all_devs: + dev.close() if __name__ == "__main__": diff --git a/setup.py b/setup.py index c5d1806..d8b6541 100644 --- a/setup.py +++ b/setup.py @@ -43,7 +43,12 @@ author="Oxford Ion Trap Quantum Computing Group", packages=find_packages(), entry_points={"console_scripts": scripts}, - install_requires=["sipyco", "appdirs"], + install_requires=[ + "sipyco", + "appdirs", + "toml", + "pyzmq", + ], # zip_safe=False apparently improves compatibility for namespace packages: # https://github.com/pypa/sample-namespace-packages/issues/6 zip_safe=False,