diff --git a/marimapper/detector_process.py b/marimapper/detector_process.py index 7a9d0d1..e1df99a 100644 --- a/marimapper/detector_process.py +++ b/marimapper/detector_process.py @@ -1,4 +1,4 @@ -from multiprocessing import Process, Queue, Event +from multiprocessing import get_logger, Process, Queue, Event from marimapper.detector import ( show_image, set_cam_default, @@ -7,10 +7,8 @@ set_cam_dark, enable_and_find_led, ) -from marimapper.led import LED2D -from marimapper.utils import get_backend -from multiprocessing import get_logger +from marimapper.utils import get_backend logger = get_logger() @@ -27,10 +25,9 @@ def __init__( display: bool = True, ): super().__init__() - self._detection_request = Queue() # {led_id, view_id} - self._detection_result = Queue() # LED3D - self._detection_request.cancel_join_thread() - self._detection_result.cancel_join_thread() + self._input_queue = Queue() # {led_id, view_id} + self._input_queue.cancel_join_thread() + self._output_queues: list[Queue] = [] # LED3D self._led_count = Queue() self._led_count.cancel_join_thread() self._exit_event = Event() @@ -42,11 +39,14 @@ def __init__( self._led_backend_server = led_backend_server self._display = display - def detect(self, led_id: int, view_id: int): - self._detection_request.put((led_id, view_id)) + def get_input_queue(self) -> Queue: + return self._input_queue - def get_results(self) -> LED2D: - return self._detection_result.get() + def add_output_queue(self, queue: Queue): + self._output_queues.append(queue) + + def detect(self, led_id: int, view_id: int): + self._input_queue.put((led_id, view_id)) def get_led_count(self): return self._led_count.get() @@ -66,9 +66,9 @@ def run(self): while not self._exit_event.is_set(): - if not self._detection_request.empty(): + if not self._input_queue.empty(): set_cam_dark(cam, self._dark_exposure) - led_id, view_id = self._detection_request.get() + led_id, view_id = self._input_queue.get() result = enable_and_find_led( cam, led_backend, @@ -79,7 +79,8 @@ def run(self): self._display, ) - self._detection_result.put(result) + for queue in self._output_queues: + queue.put(result) else: set_cam_default(cam) if self._display: @@ -90,7 +91,9 @@ def run(self): set_cam_default(cam) # clear the queues, don't ask why. - while not self._detection_request.empty(): - self._detection_request.get() - while not self._detection_result.empty(): - self._detection_result.get() + while not self._input_queue.empty(): + self._input_queue.get() + + for queue in self._output_queues: + while not queue.empty(): + queue.get() diff --git a/marimapper/file_tools.py b/marimapper/file_tools.py index 32e9475..f7c7efe 100644 --- a/marimapper/file_tools.py +++ b/marimapper/file_tools.py @@ -1,6 +1,7 @@ import os from marimapper.led import Point2D, LED3D, LED2D import typing +from pathlib import Path def load_detections(filename: os.path, view_id) -> typing.Optional[list[LED2D]]: @@ -51,7 +52,18 @@ def get_all_2d_led_maps(directory: os.path) -> list[LED2D]: return points -def write_3d_leds_to_file(leds: list[LED3D], filename: str): +def write_2d_leds_to_file(leds: list[LED2D], filename: Path): + + lines = ["index,u,v"] + + for led in sorted(leds, key=lambda led_t: led_t.led_id): + lines.append(f"{led.led_id}," f"{led.point.u():f}," f"{led.point.v():f}") + + with open(filename, "w") as f: + f.write("\n".join(lines)) + + +def write_3d_leds_to_file(leds: list[LED3D], filename: Path): lines = ["index,x,y,z,xn,yn,zn,error"] @@ -69,15 +81,3 @@ def write_3d_leds_to_file(leds: list[LED3D], filename: str): with open(filename, "w") as f: f.write("\n".join(lines)) - - -def write_2d_leds_to_file(leds: list[LED2D], filename: str): - - lines = ["index,u,v"] - - for led in sorted(leds, key=lambda led_t: led_t.led_id): - - lines.append(f"{led.led_id}," f"{led.point.u():f}," f"{led.point.v():f}") - - with open(filename, "w") as f: - f.write("\n".join(lines)) diff --git a/marimapper/file_writer_process.py b/marimapper/file_writer_process.py new file mode 100644 index 0000000..a33c5db --- /dev/null +++ b/marimapper/file_writer_process.py @@ -0,0 +1,61 @@ +from multiprocessing import Process, Queue, Event +import time +from marimapper.file_tools import write_3d_leds_to_file, write_2d_leds_to_file +from pathlib import Path + + +class FileWriterProcess(Process): + + def __init__(self, base_path: Path): + super().__init__() + self._input_queue_2d = Queue() + self._input_queue_2d.cancel_join_thread() + self._input_queue_3d = Queue() + self._input_queue_3d.cancel_join_thread() + self._exit_event = Event() + self._base_path = base_path + + def stop(self): + self._exit_event.set() + + def get_2d_input_queue(self): + return self._input_queue_2d + + def get_3d_input_queue(self): + return self._input_queue_3d + + def get_new_filename(self) -> Path: + string_time = time.strftime("%Y%m%d-%H%M%S") + return self._base_path / f"led_map_2d_{string_time}.csv" + + def run(self): + views = {} + view_id_to_filename = {} + + requires_update = set() + + while not self._exit_event.is_set(): + if not self._input_queue_3d.empty(): + leds = self._input_queue_3d.get() + write_3d_leds_to_file(leds, Path("led3d.csv")) + + if not self._input_queue_2d.empty(): + led = self._input_queue_2d.get() + if led.view_id not in view_id_to_filename: + view_id_to_filename[led.view_id] = self.get_new_filename() + views[led.view_id] = [] + + views[led.view_id].append(led) + requires_update.add(led.view_id) + + else: + for view_id in requires_update: + write_2d_leds_to_file(views[view_id], view_id_to_filename[view_id]) + requires_update = [] + time.sleep(1) + + # clear the queues, don't ask why. + while not self._input_queue_2d.empty(): + self._input_queue_2d.get() + while not self._input_queue_3d.empty(): + self._input_queue_3d.get() diff --git a/marimapper/scanner.py b/marimapper/scanner.py index c2aee62..326edf9 100644 --- a/marimapper/scanner.py +++ b/marimapper/scanner.py @@ -3,15 +3,15 @@ from marimapper.sfm_process import SFM import os -import time from tqdm import tqdm from pathlib import Path from marimapper.detector_process import DetectorProcess -from multiprocessing import get_logger -from marimapper.file_tools import get_all_2d_led_maps, write_2d_leds_to_file +from multiprocessing import get_logger, Queue +from marimapper.file_tools import get_all_2d_led_maps from marimapper.utils import get_user_confirmation from marimapper.visualize_process import VisualiseProcess from marimapper.led import last_view +from marimapper.file_writer_process import FileWriterProcess logger = get_logger() @@ -20,7 +20,7 @@ class Scanner: def __init__(self, cli_args): logger.debug("initialising scanner") - self.output_dir = cli_args.dir + self.output_dir = Path(cli_args.dir) os.makedirs(self.output_dir, exist_ok=True) self.detector = DetectorProcess( @@ -33,17 +33,29 @@ def __init__(self, cli_args): self.sfm = SFM() - self.leds = get_all_2d_led_maps(Path(self.output_dir)) - for led in self.leds: + self.file_writer = FileWriterProcess(self.output_dir) + + leds = get_all_2d_led_maps(self.output_dir) + + for led in leds: self.sfm.add_detection(led) - self.current_view = last_view(self.leds) + 1 + self.current_view = last_view(leds) + 1 + + self.renderer3d = VisualiseProcess() - self.renderer3d = VisualiseProcess(input_queue=self.sfm.get_output_queue()) + self.detector_update_queue = Queue() + self.detector.add_output_queue(self.sfm.get_input_queue()) + self.detector.add_output_queue(self.detector_update_queue) + self.detector.add_output_queue(self.file_writer.get_2d_input_queue()) + + self.sfm.add_output_queue(self.renderer3d.get_input_queue()) + self.sfm.add_output_queue(self.file_writer.get_3d_input_queue()) self.sfm.start() self.renderer3d.start() self.detector.start() + self.file_writer.start() self.led_id_range = range( cli_args.start, min(cli_args.end, self.detector.get_led_count()) @@ -57,11 +69,12 @@ def close(self): self.detector.stop() self.sfm.stop() self.renderer3d.stop() + self.file_writer.stop() self.sfm.join() self.renderer3d.join() self.detector.join() - + self.file_writer.join() logger.debug("scanner closed") def mainloop(self): @@ -74,8 +87,6 @@ def mainloop(self): print("exiting") return - leds = [] - for led_id in self.led_id_range: self.detector.detect(led_id, self.current_view) @@ -86,18 +97,6 @@ def mainloop(self): total=self.led_id_range.stop, smoothing=0, ): - led = self.detector.get_results() - - if led.point is None: - continue - - leds.append(led) - - self.sfm.add_detection(led) + self.detector_update_queue.get() self.current_view += 1 - - # The filename is made out of the date, then the resolution of the camera - string_time = time.strftime("%Y%m%d-%H%M%S") - filepath = os.path.join(self.output_dir, f"led_map_2d_{string_time}.csv") - write_2d_leds_to_file(leds, filepath) diff --git a/marimapper/sfm_process.py b/marimapper/sfm_process.py index bdacc23..a830044 100644 --- a/marimapper/sfm_process.py +++ b/marimapper/sfm_process.py @@ -40,20 +40,19 @@ class SFM(Process): def __init__(self): super().__init__() - self._output_queue = Queue() - self._output_queue.cancel_join_thread() self._input_queue = Queue() self._input_queue.cancel_join_thread() + self._output_queues: list[Queue] = [] self._exit_event = Event() + def get_input_queue(self) -> Queue: + return self._input_queue + def add_detection(self, led: LED2D): self._input_queue.put(led) - def get_output_queue(self): - return self._output_queue - - def get_results(self): - return self._output_queue.get() + def add_output_queue(self, queue: Queue): + self._output_queues.append(queue) def stop(self): self._exit_event.set() @@ -88,11 +87,13 @@ def run(self): add_normals(leds_3d) - self._output_queue.put(leds_3d) + for queue in self._output_queues: + queue.put(leds_3d) update_required = False # clear the queues, don't ask why. while not self._input_queue.empty(): self._input_queue.get() - while not self._output_queue.empty(): - self._output_queue.get() + for queue in self._output_queues: + while not queue.empty(): + queue.get() diff --git a/marimapper/visualize_process.py b/marimapper/visualize_process.py index 830ab86..35df8be 100644 --- a/marimapper/visualize_process.py +++ b/marimapper/visualize_process.py @@ -1,7 +1,6 @@ import numpy as np import open3d -from multiprocessing import get_logger -from multiprocessing import Process, Event +from multiprocessing import get_logger, Process, Event, Queue from marimapper.led import LED3D, View, get_next, get_distance import time @@ -23,17 +22,21 @@ def get_all_views(leds: list[LED3D]) -> list[View]: class VisualiseProcess(Process): - def __init__(self, input_queue): + def __init__(self): logger.debug("Renderer3D initialising") super().__init__() self._vis = None - self._input_queue = input_queue + self._input_queue = Queue() + self._input_queue.cancel_join_thread() self._exit_event = Event() self.point_cloud = None self.line_set = None self.strip_set = None logger.debug("Renderer3D initialised") + def get_input_queue(self) -> Queue: + return self._input_queue + def stop(self): self._exit_event.set()