""" 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)")