dudong-v2/workers/defect_worker.py

"""
Defect Worker Module

Worker thread for async defect detection and quality assessment.
"""

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.defect_model import DefectModel
from utils.config import get_device

logger = logging.getLogger(__name__)


class DefectWorkerSignals(WorkerSignals):
    """
    Signals specific to defect detection.
    
    Signals:
        result_ready: Emitted when detection is complete
            (annotated_image: QImage, primary_class: str, 
             class_counts: dict, total_detections: int)
    """
    result_ready = pyqtSignal(QImage, str, dict, int)


class DefectWorker(BaseWorker):
    """
    Worker for processing images and detecting defects.
    
    Runs DefectModel inference in a background thread without blocking UI.
    
    Attributes:
        image_path: Path to image file
        model: DefectModel instance
        signals: DefectWorkerSignals for emitting results
    """
    
    def __init__(self, image_path: str, model: Optional[DefectModel] = None):
        """
        Initialize the defect worker.
        
        Args:
            image_path: Path to image file to process
            model: DefectModel instance (if None, creates new one)
        """
        super().__init__()
        self.image_path = image_path
        self.model = model
        
        # Replace base signals with defect-specific signals
        self.signals = DefectWorkerSignals()
        
        # If no model provided, create and load one
        if self.model is None:
            device = get_device()
            self.model = DefectModel(device=device)
            self.model.load()
        
        logger.info(f"DefectWorker created for: {image_path}")
    
    def process(self):
        """
        Process the image and detect defects.
        
        Emits result_ready signal with detection results.
        """
        if self.is_cancelled():
            logger.info("DefectWorker 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 defect model")
        
        # Process image
        self.emit_progress(50, "Detecting defects...")
        
        result = self.model.predict(self.image_path)
        
        if self.is_cancelled():
            logger.info("DefectWorker 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['primary_class'],
            result['class_counts'],
            result['total_detections']
        )
        
        self.emit_progress(100, "Complete!")
        
        logger.info(f"DefectWorker completed: {result['primary_class']} ({result['total_detections']} detections)")