dudong-v2/workers/maturity_worker.py

"""
Maturity Worker Module

Worker thread for async multispectral maturity classification processing.
"""

from typing import Optional
import logging

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

from workers.base_worker import BaseWorker, WorkerSignals
from models.maturity_model import MaturityModel

logger = logging.getLogger(__name__)


class MaturityWorkerSignals(WorkerSignals):
    """
    Signals specific to maturity processing.
    
    Signals:
        result_ready: Emitted when prediction is complete
            (gradcam_image: QImage, class_name: str, confidence: float, 
             probabilities: dict)
    """
    result_ready = pyqtSignal(QImage, str, float, dict)


class MaturityWorker(BaseWorker):
    """
    Worker for processing multispectral TIFF files and predicting maturity.

    Runs MaturityModel inference in a background thread without blocking UI.
    Supports 8-band multispectral TIFF images.

    Attributes:
        tif_path: Path to multispectral TIFF file
        model: MaturityModel instance
        signals: MaturityWorkerSignals for emitting results
    """
    
    def __init__(self, tif_path: str, model: Optional[MaturityModel] = None):
        """
        Initialize the maturity worker.

        Args:
            tif_path: Path to multispectral TIFF file to process
            model: MaturityModel instance (if None, creates new one)
        """
        super().__init__()
        self.tif_path = tif_path
        self.model = model
        
        # Replace base signals with maturity-specific signals
        self.signals = MaturityWorkerSignals()
        
        # If no model provided, create and load one
        if self.model is None:
            from utils.config import get_device
            self.model = MaturityModel(device=get_device())
            self.model.load()
        
        logger.info(f"MaturityWorker created for: {tif_path}")
    
    def process(self):
        """
        Process the multispectral TIFF file and predict maturity.
        
        Emits result_ready signal with prediction results.
        """
        if self.is_cancelled():
            logger.info("MaturityWorker cancelled before processing")
            return
        
        # Update progress
        self.emit_progress(10, "Loading multispectral 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 maturity model")
        
        # Process TIFF
        self.emit_progress(50, "Preprocessing multispectral data...")
        
        result = self.model.predict(self.tif_path)
        
        if self.is_cancelled():
            logger.info("MaturityWorker cancelled during processing")
            return
        
        if not result['success']:
            raise RuntimeError(result['error'])
        
        self.emit_progress(70, "Generating Grad-CAM visualization...")
        
        # Generate Grad-CAM visualization
        gradcam_image, pred_name = self.model.run_gradcam(self.tif_path)
        
        if self.is_cancelled():
            logger.info("MaturityWorker cancelled during Grad-CAM generation")
            return
        
        if gradcam_image is None:
            # Create a placeholder pixmap for None Grad-CAM
            gradcam_image = QImage(256, 256, QImage.Format_RGB888)
            gradcam_image.fill(0x2c3e50)  # Dark gray background
        
        self.emit_progress(90, "Finalizing results...")
        
        # Emit results
        self.signals.result_ready.emit(
            gradcam_image,
            result['prediction'],
            result['confidence'],
            result['probabilities']
        )
        
        self.emit_progress(100, "Complete!")
        
        logger.info(f"MaturityWorker completed: {result['prediction']} ({result['confidence']:.2%})")