diff --git a/owocr/run.py b/owocr/run.py index 388f033..478748d 100644 --- a/owocr/run.py +++ b/owocr/run.py @@ -1138,50 +1138,6 @@ class ScreenshotThread(threading.Thread): self.windows_window_tracker_instance.join() -class SecondPassThread: - def __init__(self): - self.input_queue = queue.Queue() - self.output_queue = queue.Queue() - self.ocr_thread = None - self.running = False - - def __del__(self): - self.stop() - - def start(self): - if self.ocr_thread is None or not self.ocr_thread.is_alive(): - self.running = True - self.ocr_thread = threading.Thread(target=self._process_ocr, daemon=True) - self.ocr_thread.start() - - def stop(self): - self.running = False - if self.ocr_thread and self.ocr_thread.is_alive(): - self.ocr_thread.join() - - def _process_ocr(self): - while self.running and not terminated: - try: - img, engine_instance = self.input_queue.get(timeout=0.1) - - start_time = time.time() - res, result_data = engine_instance(img) - end_time = time.time() - - self.output_queue.put((res, result_data, end_time - start_time)) - except queue.Empty: - continue - - def submit_task(self, img, engine_instance): - self.input_queue.put((img, engine_instance)) - - def get_result(self): - try: - return self.output_queue.get_nowait() - except queue.Empty: - return None - - class AutopauseTimer: def __init__(self, timeout): self.timeout = timeout @@ -1213,14 +1169,59 @@ class AutopauseTimer: pause_handler(True) +class SecondPassThread: + def __init__(self): + self.input_queue = queue.Queue() + self.output_queue = queue.Queue() + self.ocr_thread = None + self.running = False + + def __del__(self): + self.stop() + + def start(self): + if self.ocr_thread is None or not self.ocr_thread.is_alive(): + self.running = True + self.ocr_thread = threading.Thread(target=self._process_ocr, daemon=True) + self.ocr_thread.start() + + def stop(self): + self.running = False + if self.ocr_thread and self.ocr_thread.is_alive(): + self.ocr_thread.join() + while not self.input_queue.empty(): + self.input_queue.get() + while not self.output_queue.empty(): + self.output_queue.get() + + def _process_ocr(self): + while self.running and not terminated: + try: + img, engine_instance = self.input_queue.get(timeout=0.1) + + start_time = time.time() + res, result_data = engine_instance(img) + end_time = time.time() + + self.output_queue.put((engine_instance.readable_name, res, result_data, end_time - start_time)) + except queue.Empty: + continue + + def submit_task(self, img, engine_instance): + self.input_queue.put((img, engine_instance)) + + def get_result(self): + try: + return self.output_queue.get_nowait() + except queue.Empty: + return None + + class OutputResult: def __init__(self): self.filtering = TextFiltering() self.second_pass_thread = SecondPassThread() - def __del__(self): - self.second_pass_thread.stop() - def _post_process(self, text, strip_spaces): is_cj_text = self.filtering.cj_regex.search(''.join(text)) line_separator = '' if strip_spaces else ' ' @@ -1249,45 +1250,48 @@ class OutputResult: engine_color = config.get_general('engine_color') engine_instance = engine_instances[engine_index] two_pass_processing_active = False + result_data = None - if filter_text and engine_index_2 != -1 and engine_index_2 != engine_index: - self.second_pass_thread.start() - engine_instance_2 = engine_instances[engine_index_2] - start_time = time.time() - res2, result_data_2 = engine_instance_2(img_or_path) - end_time = time.time() - - if not res2: - logger.opt(ansi=True).warning(f'<{engine_color}>{engine_instance_2.readable_name} reported an error after {end_time - start_time:0.03f}s: {result_data_2}') - else: + if filter_text: + if engine_index_2 != -1 and engine_index_2 != engine_index and engine_instance.threading_support: two_pass_processing_active = True - changed_lines_count, changed_regions_image = self.filtering._find_changed_lines(img_or_path, result_data_2) + self.second_pass_thread.start() + engine_instance_2 = engine_instances[engine_index_2] + start_time = time.time() + res2, result_data_2 = engine_instance_2(img_or_path) + end_time = time.time() - if changed_lines_count: - logger.opt(ansi=True).info(f"<{engine_color}>{engine_instance_2.readable_name} found {changed_lines_count} changed line(s) in {end_time - start_time:0.03f}s, re-OCRing with <{engine_color}>{engine_instance.readable_name}") + if not res2: + logger.opt(ansi=True).warning(f'<{engine_color}>{engine_instance_2.readable_name} reported an error after {end_time - start_time:0.03f}s: {result_data_2}') + else: + changed_lines_count, changed_regions_image = self.filtering._find_changed_lines(img_or_path, result_data_2) - if output_format != 'json': - if changed_regions_image: - img_or_path = changed_regions_image + if changed_lines_count: + logger.opt(ansi=True).info(f"<{engine_color}>{engine_instance_2.readable_name} found {changed_lines_count} changed line(s) in {end_time - start_time:0.03f}s, re-OCRing with <{engine_color}>{engine_instance.readable_name}") + + if output_format != 'json': + if changed_regions_image: + img_or_path = changed_regions_image - if engine_instance.threading_support: self.second_pass_thread.submit_task(img_or_path, engine_instance) - else: - self.second_pass_thread.stop() - second_pass_result = self.second_pass_thread.get_result() - if second_pass_result: - res, result_data, processing_time = second_pass_result - two_pass_processing_active = True - elif two_pass_processing_active and engine_instance.threading_support: - return - else: + second_pass_result = self.second_pass_thread.get_result() + if second_pass_result: + engine_name, res, result_data, processing_time = second_pass_result + else: + return + else: + self.second_pass_thread.stop() + + if not result_data: start_time = time.time() res, result_data = engine_instance(img_or_path) end_time = time.time() + processing_time = end_time - start_time + engine_name = engine_instance.readable_name if not res: - logger.opt(ansi=True).warning(f'<{engine_color}>{engine_instance.readable_name} reported an error after {end_time - start_time:0.03f}s: {result_data}') + logger.opt(ansi=True).warning(f'<{engine_color}>{engine_name} reported an error after {processing_time:0.03f}s: {result_data}') return verbosity = config.get_general('verbosity') @@ -1327,7 +1331,7 @@ class OutputResult: else: log_message_terminal = ': ' + (log_message if len(log_message) <= verbosity else log_message[:verbosity] + '[...]') - logger.opt(ansi=True).info(f'Text recognized in {end_time - start_time:0.03f}s using <{engine_color}>{engine_instance.readable_name}{log_message_terminal}') + logger.opt(ansi=True).info(f'Text recognized in {processing_time:0.03f}s using <{engine_color}>{engine_name}{log_message_terminal}') if notify and config.get_general('notifications'): notifier.send(title='owocr', message='Text recognized: ' + log_message, urgency=get_notification_urgency())