ScanSlicer/scanslicer/controller_detection.py

import traceback

import cv2
from PySide6.QtCore import QObject, QRunnable, QThreadPool, Signal
from PySide6.QtGui import QImage
from PySide6.QtWidgets import QMessageBox


class DetectionTaskSignals(QObject):
    status = Signal(str)
    finished = Signal(object, object, object)
    failed = Signal(str)


class DetectionTask(QRunnable):
    def __init__(self, autoslicer):
        super().__init__()
        self.autoslicer = autoslicer
        self.signals = DetectionTaskSignals()

    def run(self):
        try:
            boxes, filter_preview_image = self.autoslicer.autodetect_slices(
                self.signals.status.emit,
                preview_filter_output=True,
            )
            normal_preview_image = cv2.cvtColor(self.autoslicer.image, cv2.COLOR_BGR2RGB)
            self.signals.finished.emit(boxes, filter_preview_image, normal_preview_image)
        except Exception as exc:
            details = f"{exc}\n\n{traceback.format_exc()}"
            self.signals.failed.emit(details)


class AutotuneTaskSignals(QObject):
    status = Signal(str)
    finished = Signal(object)
    failed = Signal(str)


class AutotuneTask(QRunnable):
    def __init__(self, autoslicer, profile_name):
        super().__init__()
        self.autoslicer = autoslicer
        self.profile_name = profile_name
        self.signals = AutotuneTaskSignals()

    def run(self):
        try:
            result = self.autoslicer.autotune_for_current_image(
                self.signals.status.emit,
                profile_name=self.profile_name,
            )
            self.signals.finished.emit(result)
        except Exception as exc:
            details = f"{exc}\n\n{traceback.format_exc()}"
            self.signals.failed.emit(details)


def rgb_array_to_qimage(image):
    h, w, c = image.shape
    bytes_per_line = c * w
    return QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888).copy()


class ScanSlicerDetectionMixin:
    def _get_worker_pool(self):
        pool = getattr(self, "_worker_pool", None)
        if pool is None:
            pool = QThreadPool(self)
            pool.setMaxThreadCount(1)
            pool.setExpiryTimeout(-1)
            self._worker_pool = pool
        return pool

    def invalidate_preview_cache(self):
        self.last_detected_boxes = None
        self.last_preview_normal_image = None
        self.last_preview_filter_image = None

    def has_preview_cache(self):
        return (
            self.last_detected_boxes is not None
            and self.last_preview_normal_image is not None
            and self.last_preview_filter_image is not None
        )

    def render_cached_preview(self, new_image=False):
        if not self.has_preview_cache():
            return False
        image = self.last_preview_filter_image if self.preview_filter_enabled else self.last_preview_normal_image
        self.slicing_canvas.set_image(rgb_array_to_qimage(image), new_image)
        return True

    def update_preview(self, new_image=False, run_detection=True):
        if self.processing_busy:
            return
        if not self.autoslicer.image_loaded():
            self.update_statusbar("No image loaded.")
            return

        if not run_detection:
            if self.render_cached_preview(new_image=new_image):
                return
            run_detection = True

        if run_detection:
            self._start_detection_worker(new_image=bool(new_image))

    def _start_detection_worker(self, new_image=False):
        if getattr(self, "_is_closing", False):
            return
        if getattr(self, "_detection_running", False):
            return

        self._detection_task_id += 1
        task_id = int(self._detection_task_id)
        self.log_detection_config()
        self._pending_detection_new_image = bool(new_image)
        self._active_detection_task_id = task_id
        self._detection_running = True
        self.set_processing_busy(True, "Processing image...")

        task = DetectionTask(self.autoslicer)
        task.signals.status.connect(lambda text, tid=task_id: self._on_detection_status(tid, text))
        task.signals.finished.connect(
            lambda bbxs, filter_preview_image, normal_preview_image, tid=task_id: self._on_detection_finished(
                tid, bbxs, filter_preview_image, normal_preview_image
            )
        )
        task.signals.failed.connect(lambda error_text, tid=task_id: self._on_detection_failed(tid, error_text))
        self._detection_task = task

        self._get_worker_pool().start(task)

    def _is_active_detection_task(self, task_id):
        return (
            task_id is not None
            and int(task_id) == int(getattr(self, "_active_detection_task_id", -1))
            and bool(getattr(self, "_detection_running", False))
        )

    def _on_detection_status(self, task_id, text):
        if getattr(self, "_is_closing", False):
            return
        if not self._is_active_detection_task(task_id):
            return
        self.update_statusbar(text)

    def _on_detection_finished(self, task_id, bbxs, filter_preview_image, normal_preview_image):
        if not self._is_active_detection_task(task_id):
            return
        try:
            if getattr(self, "_is_closing", False):
                return
            self.last_detected_boxes = bbxs
            self.last_preview_normal_image = normal_preview_image
            self.last_preview_filter_image = filter_preview_image

            image_to_show = (
                self.last_preview_filter_image if self.preview_filter_enabled else self.last_preview_normal_image
            )
            self.slicing_canvas.set_image(
                rgb_array_to_qimage(image_to_show),
                getattr(self, "_pending_detection_new_image", False),
            )
            self.slicing_canvas.update_bboxes(bbxs)
            self.refresh_active_box_toolbar()
            if len(bbxs) == 0:
                self.update_statusbar("No pictures detected. Try adjusting parameters.")
            else:
                self.update_statusbar(f"Detected {len(bbxs)} picture(s).")
        finally:
            self._detection_running = False
            self._detection_task = None
            self._active_detection_task_id = None
            self._pending_detection_new_image = False
            if not getattr(self, "_is_closing", False):
                self.set_processing_busy(False)

    def _on_detection_failed(self, task_id, error_text):
        if not self._is_active_detection_task(task_id):
            return
        self._detection_running = False
        self._detection_task = None
        self._active_detection_task_id = None
        self._pending_detection_new_image = False
        if not getattr(self, "_is_closing", False):
            QMessageBox.critical(self, "Detection error", error_text)
            self.update_statusbar("Detection failed.")
            self.set_processing_busy(False)

    def _start_autotune_worker(self):
        if getattr(self, "_is_closing", False):
            return
        if getattr(self, "_autotune_running", False):
            return

        profile_name = self.autoslicer.normalize_autotune_profile(
            getattr(self, "autotune_profile", self.autoslicer.default_autotune_profile_name())
        )
        self.autotune_profile = profile_name

        self._autotune_task_id += 1
        task_id = int(self._autotune_task_id)
        self._active_autotune_task_id = task_id
        self._autotune_running = True
        self.set_processing_busy(True, f"Auto-tuning parameters ({profile_name})...")

        task = AutotuneTask(self.autoslicer, profile_name)
        task.signals.status.connect(lambda text, tid=task_id: self._on_autotune_status(tid, text))
        task.signals.finished.connect(lambda result, tid=task_id: self._on_autotune_finished(tid, result))
        task.signals.failed.connect(lambda error_text, tid=task_id: self._on_autotune_failed(tid, error_text))
        self._autotune_task = task

        self._get_worker_pool().start(task)

    def _is_active_autotune_task(self, task_id):
        return (
            task_id is not None
            and int(task_id) == int(getattr(self, "_active_autotune_task_id", -1))
            and bool(getattr(self, "_autotune_running", False))
        )

    def _on_autotune_status(self, task_id, text):
        if getattr(self, "_is_closing", False):
            return
        if not self._is_active_autotune_task(task_id):
            return
        self.update_statusbar(text)

    def _on_autotune_finished(self, task_id, result):
        if not self._is_active_autotune_task(task_id):
            return
        self._autotune_running = False
        self._autotune_task = None
        self._active_autotune_task_id = None
        if getattr(self, "_is_closing", False):
            return
        self.set_processing_busy(False)
        self.sync_param_controls_from_model()
        self.set_preset_dropdown_custom()
        profile_name = str(result.get("profile", getattr(self, "autotune_profile", "Balanced")))
        tested = int(result.get("candidates", 0))
        total = int(result.get("candidates_total", tested))
        refined = int(result.get("refined", 0))
        candidate_text = f"{tested}/{total}" if total > tested else str(tested)
        refine_text = f", refined={refined}" if refined > 0 else ""
        self.append_status_log(
            "Auto-tune result: "
            + f"profile={profile_name}, "
            + f"candidates={candidate_text}, "
            + f"refined={refined}, "
            + f"best_count={result['count']}, "
            + f"best_score={result['score']:.1f}"
        )
        self.update_statusbar(
            f"Auto-tune ({profile_name}) applied best candidate "
            + f"({result['count']} boxes from {candidate_text} candidates{refine_text})."
        )
        self.update_preview(new_image=False, run_detection=True)

    def _on_autotune_failed(self, task_id, error_text):
        if not self._is_active_autotune_task(task_id):
            return
        self._autotune_running = False
        self._autotune_task = None
        self._active_autotune_task_id = None
        if not getattr(self, "_is_closing", False):
            QMessageBox.critical(self, "Auto-tune failed", error_text)
            self.update_statusbar("Auto-tune failed.")
            self.set_processing_busy(False)