dudong-v2/workers/locule_worker.py

"""
Locule Worker Module

Worker thread for async locule segmentation and counting.
"""

from typing import Optional
import logging

from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QImage

from workers.base_worker import BaseWorker, WorkerSignals
from models.locule_model import LoculeModel
from utils.config import get_device

logger = logging.getLogger(__name__)


class LoculeWorkerSignals(WorkerSignals):
    """
    Signals specific to locule counting.
    
    Signals:
        result_ready: Emitted when counting is complete
            (annotated_image: QImage, locule_count: int, detections: list)
    """
    result_ready = pyqtSignal(QImage, int, list)


class LoculeWorker(BaseWorker):
    """
    Worker for processing images and counting locules.
    
    Runs LoculeModel inference in a background thread without blocking UI.
    
    Attributes:
        image_path: Path to image file
        model: LoculeModel instance
        signals: LoculeWorkerSignals for emitting results
    """
    
    def __init__(self, image_path: str, model: Optional[LoculeModel] = None):
        """
        Initialize the locule worker.
        
        Args:
            image_path: Path to cross-section image to process
            model: LoculeModel instance (if None, creates new one)
        """
        super().__init__()
        self.image_path = image_path
        self.model = model
        
        # Replace base signals with locule-specific signals
        self.signals = LoculeWorkerSignals()
        
        # If no model provided, create and load one
        if self.model is None:
            device = get_device()
            self.model = LoculeModel(device=device)
            self.model.load()
        
        logger.info(f"LoculeWorker created for: {image_path}")
    
    def process(self):
        """
        Process the image and count locules.
        
        Emits result_ready signal with counting results.
        """
        if self.is_cancelled():
            logger.info("LoculeWorker cancelled before processing")
            return
        
        # Update progress
        self.emit_progress(10, "Loading image...")
        
        if not self.model.is_loaded:
            logger.warning("Model not loaded, loading now...")
            self.emit_progress(30, "Loading model...")
            if not self.model.load():
                raise RuntimeError("Failed to load locule model")
        
        # Process image
        self.emit_progress(50, "Segmenting locules...")
        
        result = self.model.predict(self.image_path)
        
        if self.is_cancelled():
            logger.info("LoculeWorker cancelled during processing")
            return
        
        if not result['success']:
            raise RuntimeError(result['error'])
        
        self.emit_progress(90, "Finalizing results...")
        
        # Make a copy of the QImage to ensure thread safety
        # QImage data can be garbage collected if not copied
        annotated_image_copy = result['annotated_image'].copy()
        
        # Emit results
        self.signals.result_ready.emit(
            annotated_image_copy,
            result['locule_count'],
            result['detections']
        )
        
        self.emit_progress(100, "Complete!")
        
        logger.info(f"LoculeWorker completed: {result['locule_count']} locules detected")