diff --git a/src/draw/__main__.py b/src/draw/__main__.py index 1aeccd7..bd2f211 100644 --- a/src/draw/__main__.py +++ b/src/draw/__main__.py @@ -1,4 +1,36 @@ +import argparse + from draw.run import run +def parse_command_line(): + parser = argparse.ArgumentParser('Detecting and Recognizing a Wide Range of Yu-Gi-Oh! Cards', add_help=True) + + parser.add_argument('--model_size', default=0, type=int, + help="The size of the model to use, 0 for the standard one and 1 for the large one (default: 0).") + + parser.add_argument("--deck_list", default="", type=str, + help="[Optional] Path or link to the deck list.") + + parser.add_argument("--minimum_out_of_screen_time", default=25, type=int, + help="Minimum time before a card is allow to be displayed again in seconds (default: 25s).") + + parser.add_argument('--minimum_screen_time', default=6, type=int, + help="Minimum time a card is displayed in seconds (default: 6s).") + parser.add_argument('--confidence_threshold', default=5, type=int, + help="Confidence threshold a detection can have to be displayed in % (default: 5%).") + + parser.add_argument('--channel', default=1, type=int, + help="Channel of the display source in the plugin (default: 1).") + + return parser.parse_args() + if __name__ == '__main__': - run() + args = parse_command_line() + run( + model_size=args.model_size, + deck_list=args.deck_list, + minimum_out_of_screen_time=args.minimum_out_of_screen_time, + minimum_screen_time=args.minimum_screen_time, + confidence_threshold=args.confidence_threshold, + channel=args.channel + ) diff --git a/src/draw/run.py b/src/draw/run.py index 5c1fbaa..da3e747 100644 --- a/src/draw/run.py +++ b/src/draw/run.py @@ -34,8 +34,8 @@ logging.info("Python subprocess started") -OBS_SHM_NAME = "obs_shared_memory" -PYTHON_SHM_NAME = "python_shared_memory" +OBS_SHM_NAME = "/obs_shared_memory" +PYTHON_SHM_NAME = "/python_shared_memory" HEADER_FORMAT = "II" HEADER_SIZE = struct.calcsize(HEADER_FORMAT) @@ -48,12 +48,22 @@ class DrawSharedMemoryHandler: - def __init__(self, model_id="HichTala/draw2", deck_list="", minimum_out_of_screen_time=25, minimum_screen_time=6, confidence_threshold=5): - self.draw = Draw(model=model_id, deck_lists=get_deck_list(deck_list), confidence_threshold=confidence_threshold) + def __init__( + self, + model_id="HichTala/draw2", + deck_list="", + minimum_out_of_screen_time=25, + minimum_screen_time=6, + confidence_threshold=5, + channel=1 + ): + self.draw = Draw(model=model_id, deck_lists=get_deck_list(deck_list), confidence_threshold=confidence_threshold + ) self.obs_shm = None self.python_shm = None self.shm_array = None + self.channel = channel self.minimum_out_of_screen_time = minimum_out_of_screen_time self.minimum_screen_time = minimum_screen_time @@ -66,38 +76,30 @@ def __init__(self, model_id="HichTala/draw2", deck_list="", minimum_out_of_scree self.waiting = True - def __call__(self, address=None, ready_ptr=None): - continue_execution = True if address is None else address.contents.value - if ready_ptr is not None: - ready_ptr.contents.value = True - + def __call__(self): print(f"Waiting for OBS to start...") - while continue_execution: - continue_execution = True if address is None else address.contents.value + while True: try: if sys.platform == 'win32': - self.obs_shm = mmap.mmap(-1, SIZE, tagname=OBS_SHM_NAME, access=mmap.ACCESS_READ) + self.obs_shm = mmap.mmap(-1, SIZE, tagname=f"{OBS_SHM_NAME}_{self.channel}", access=mmap.ACCESS_READ) print("mapped size:", self.obs_shm.size()) else: - self.obs_shm = shared_memory.SharedMemory(name=OBS_SHM_NAME) + self.obs_shm = shared_memory.SharedMemory(name=f"{OBS_SHM_NAME}_{self.channel}") print("mapped size:", self.obs_shm.size) except (FileNotFoundError, ValueError, OSError): continue except KeyboardInterrupt: return break - if not continue_execution: - print("OBS stopped before it could start.") - return print(f"Shared memory found") if sys.platform == 'win32': obs_buf = memoryview(self.obs_shm) else: obs_buf = memoryview(self.obs_shm.buf) - while continue_execution: + + while True: try: - continue_execution = True if address is None else address.contents.value image = read_shared_frame(obs_buf, HEADER_SIZE, HEADER_FORMAT) except TypeError as e: print("type error: ", e) @@ -160,7 +162,7 @@ def send_image_to_obs(self, image): if sys.platform == 'win32' and (self.python_shm is None or total_size != self.python_shm.size()): if self.python_shm is not None: self.python_shm.close() - self.python_shm = mmap.mmap(-1, total_size, tagname=PYTHON_SHM_NAME, access=mmap.ACCESS_WRITE) + self.python_shm = mmap.mmap(-1, total_size, tagname=f"{PYTHON_SHM_NAME}_{self.channel}", access=mmap.ACCESS_WRITE) self.python_shm.seek(0) self.python_shm.write(b"\x00" * total_size) python_buf = memoryview(self.python_shm) @@ -170,7 +172,7 @@ def send_image_to_obs(self, image): if self.python_shm is not None: self.python_shm.close() self.python_shm.unlink() - self.python_shm = shared_memory.SharedMemory(name=PYTHON_SHM_NAME, create=True, size=total_size) + self.python_shm = shared_memory.SharedMemory(name=f"{PYTHON_SHM_NAME}_{self.channel}", create=True, size=total_size) self.shm_array = np.ndarray((height, width, channels), dtype=np.uint8, buffer=self.python_shm.buf[HEADER_SIZE:]) self.python_shm.buf[:HEADER_SIZE] = struct.pack(HEADER_FORMAT, width, height) @@ -190,45 +192,24 @@ def close_shared_memory(shm): def run( - stop_flag=None, - model_ready=None, - update_flag=None, model_size=0, deck_list="", minimum_out_of_screen_time=25, minimum_screen_time=6, - confidence_threshold=5 + confidence_threshold=5, + channel=1 ): print("Starting Draw2...") - if stop_flag is not None and model_ready is not None and update_flag is not None: - addr = ctypes.cast(ctypes.pythonapi.PyCapsule_GetPointer(stop_flag, b"stop_flag"), - ctypes.POINTER(ctypes.c_bool)) - ready_ptr = ctypes.cast(ctypes.pythonapi.PyCapsule_GetPointer(model_ready, b"model_ready"), - ctypes.POINTER(ctypes.c_bool)) - update_ptr = ctypes.cast(ctypes.pythonapi.PyCapsule_GetPointer(update_flag, b"update_flag"), - ctypes.POINTER(ctypes.c_bool)) - - update_ptr.contents.value = True - sh_memory_handler = DrawSharedMemoryHandler( - model_id=["HichTala/draw2", "HichTala/draw2-large"][model_size], - deck_list=deck_list, - minimum_out_of_screen_time=minimum_out_of_screen_time, - minimum_screen_time=minimum_screen_time, - confidence_threshold=confidence_threshold - ) - update_ptr.contents.value = False - - print("Running Draw2 with OBS shared memory") - sh_memory_handler(address=addr, ready_ptr=ready_ptr) - else: - sh_memory_handler = DrawSharedMemoryHandler( - deck_list=deck_list, - minimum_out_of_screen_time=minimum_out_of_screen_time, - minimum_screen_time=minimum_screen_time, - confidence_threshold=confidence_threshold - ) - print("Running Draw2 without OBS shared memory") - sh_memory_handler() + + sh_memory_handler = DrawSharedMemoryHandler( + model_id=["HichTala/draw2", "HichTala/draw2-large"][model_size], + deck_list=deck_list, + minimum_out_of_screen_time=minimum_out_of_screen_time, + minimum_screen_time=minimum_screen_time, + confidence_threshold=confidence_threshold, + channel=channel + ) + sh_memory_handler() close_shared_memory(sh_memory_handler.obs_shm) close_shared_memory(sh_memory_handler.python_shm)