quality_tab.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859
  1. """
  2. Quality Classification Tab
  3. This tab handles comprehensive quality assessment with multiple camera views and analysis.
  4. """
  5. import os
  6. from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QGridLayout,
  7. QLabel, QPushButton, QFrame, QGroupBox, QTableWidget,
  8. QTableWidgetItem, QHeaderView)
  9. from PyQt5.QtCore import Qt, pyqtSignal
  10. from PyQt5.QtGui import QFont, QPixmap, QImage
  11. from resources.styles import COLORS, STYLES
  12. from ui.panels.quality_rgb_top_panel import QualityRGBTopPanel
  13. from ui.panels.quality_rgb_side_panel import QualityRGBSidePanel
  14. from ui.panels.quality_thermal_panel import QualityThermalPanel
  15. from ui.panels.quality_defects_panel import QualityDefectsPanel
  16. from ui.panels.quality_control_panel import QualityControlPanel
  17. from ui.panels.quality_results_panel import QualityResultsPanel
  18. from ui.panels.quality_history_panel import QualityHistoryPanel
  19. from ui.dialogs.image_preview_dialog import ImagePreviewDialog
  20. from models.locule_model import LoculeModel
  21. from models.defect_model import DefectModel
  22. from utils.config import get_device
  23. class QualityTab(QWidget):
  24. """
  25. Comprehensive quality assessment tab with multiple camera views and analysis.
  26. Signals:
  27. load_image_requested: Emitted when user wants to load an image file
  28. """
  29. load_image_requested = pyqtSignal()
  30. def __init__(self, parent=None):
  31. super().__init__(parent)
  32. self.locule_model = None
  33. self.defect_model = None
  34. self.last_processed_file = None
  35. self.init_ui()
  36. self._initialize_models()
  37. def init_ui(self):
  38. """Initialize the UI components with new 2x2+2 grid layout."""
  39. layout = QVBoxLayout(self)
  40. layout.setContentsMargins(10, 10, 10, 10)
  41. layout.setSpacing(10)
  42. # # Title
  43. # title = QLabel("🔍 Quality Assessment")
  44. # title.setFont(QFont("Arial", 16, QFont.Bold))
  45. # title.setAlignment(Qt.AlignCenter)
  46. # title.setStyleSheet(f"color: {COLORS['text_primary']}; margin: 10px;")
  47. # layout.addWidget(title)
  48. # Main content area with 2x2+2 grid layout
  49. content_widget = QWidget()
  50. content_layout = QGridLayout(content_widget)
  51. content_layout.setContentsMargins(0, 0, 0, 0)
  52. content_layout.setSpacing(8)
  53. # Create all panels
  54. self.rgb_top_panel = QualityRGBTopPanel()
  55. self.rgb_side_panel = QualityRGBSidePanel()
  56. self.thermal_panel = QualityThermalPanel()
  57. self.defects_panel = QualityDefectsPanel()
  58. self.quality_results_panel = QualityResultsPanel()
  59. self.control_panel = QualityControlPanel()
  60. self.history_panel = QualityHistoryPanel()
  61. # Add panels to grid layout
  62. # Row 0, Col 0: RGB Top View Panel
  63. content_layout.addWidget(self.rgb_top_panel, 0, 0)
  64. # Row 0, Col 1: RGB Side View Panel
  65. content_layout.addWidget(self.rgb_side_panel, 0, 1)
  66. # Row 1, Col 0: Thermal Analysis Panel
  67. content_layout.addWidget(self.thermal_panel, 1, 0)
  68. # Row 1, Col 1: Defect Detection Panel
  69. content_layout.addWidget(self.defects_panel, 1, 1)
  70. # Row 0-1, Col 2: Quality Control Panel (span 2 rows)
  71. content_layout.addWidget(self.control_panel, 0, 2, 2, 1)
  72. # Row 0-1, Col 3: Quality Results Panel (span 1 row)
  73. content_layout.addWidget(self.quality_results_panel, 0, 3)
  74. # Row 1, Col 3: Quality History Panel (span 1 row)
  75. content_layout.addWidget(self.history_panel, 1, 3)
  76. # Set column stretches: [2, 2, 1, 2] (control panel gets less space)
  77. content_layout.setColumnStretch(0, 2)
  78. content_layout.setColumnStretch(1, 2)
  79. content_layout.setColumnStretch(2, 1)
  80. content_layout.setColumnStretch(3, 2)
  81. layout.addWidget(content_widget, 1)
  82. # Status bar at bottom
  83. self.status_label = QLabel("Ready for quality assessment. Use Control Panel to analyze samples.")
  84. self.status_label.setAlignment(Qt.AlignCenter)
  85. self.status_label.setStyleSheet(f"color: {COLORS['text_secondary']}; font-size: 11px; padding: 5px;")
  86. layout.addWidget(self.status_label)
  87. # Connect control panel signals
  88. self._connect_signals()
  89. def _initialize_models(self):
  90. """Initialize the locule and defect models."""
  91. try:
  92. # Get device configuration
  93. device = get_device()
  94. print(f"Initializing quality models on device: {device}")
  95. # Initialize locule model
  96. try:
  97. self.locule_model = LoculeModel(device=device)
  98. if not self.locule_model.load():
  99. print("⚠ Warning: Could not load locule model. Check if 'locule.pt' exists in project root.")
  100. self.status_label.setText("Warning: Locule model not loaded. Check model file: locule.pt")
  101. self.locule_model = None
  102. else:
  103. print(f"✓ Locule model loaded successfully on {device}")
  104. except Exception as e:
  105. print(f"✗ Error initializing locule model: {e}")
  106. import traceback
  107. traceback.print_exc()
  108. self.locule_model = None
  109. # Initialize defect model
  110. try:
  111. self.defect_model = DefectModel(device=device)
  112. if not self.defect_model.load():
  113. print("⚠ Warning: Could not load defect model. Check if 'best.pt' exists in project root.")
  114. self.status_label.setText("Warning: Defect model not loaded. Check model file: best.pt")
  115. self.defect_model = None
  116. else:
  117. print(f"✓ Defect model loaded successfully on {device}")
  118. except Exception as e:
  119. print(f"✗ Error initializing defect model: {e}")
  120. import traceback
  121. traceback.print_exc()
  122. self.defect_model = None
  123. # Update status if both models loaded
  124. if self.locule_model and self.locule_model.is_loaded and self.defect_model and self.defect_model.is_loaded:
  125. self.status_label.setText("Ready for quality assessment. Click 'OPEN FILE' to analyze an image.")
  126. self.status_label.setStyleSheet(f"color: {COLORS['success']}; font-size: 11px;")
  127. except Exception as e:
  128. print(f"✗ Error initializing models: {e}")
  129. import traceback
  130. traceback.print_exc()
  131. self.locule_model = None
  132. self.defect_model = None
  133. self.status_label.setText(f"Error initializing models: {str(e)}")
  134. self.status_label.setStyleSheet(f"color: {COLORS.get('error', '#e74c3c')}; font-size: 11px;")
  135. def _connect_signals(self):
  136. """Connect signals between panels."""
  137. # Connect control panel signals
  138. self.control_panel.analyze_requested.connect(self._on_analyze_requested)
  139. self.control_panel.open_file_requested.connect(self._on_open_file_requested)
  140. self.control_panel.parameter_changed.connect(self._on_parameter_changed)
  141. self.control_panel.mode_changed.connect(self._on_mode_changed)
  142. # Connect defect panel signals
  143. self.defects_panel.annotated_image_requested.connect(self._on_annotated_image_requested)
  144. self.defects_panel.defect_image_requested.connect(self._on_defect_image_requested)
  145. def _on_analyze_requested(self):
  146. """Handle analyze button click from control panel."""
  147. # For now, simulate processing by updating all panels with sample data
  148. self._simulate_processing()
  149. def _on_parameter_changed(self, parameter_name, value):
  150. """Handle parameter changes from control panel."""
  151. # Update status to show parameter change
  152. self.status_label.setText(f"Parameter '{parameter_name}' changed to {value}")
  153. def _on_mode_changed(self, mode: str):
  154. """Handle mode change from control panel."""
  155. if mode == 'file':
  156. self.status_label.setText("File mode selected. Click 'OPEN FILE' to select an image.")
  157. else:
  158. self.status_label.setText("Live mode selected (Coming Soon).")
  159. def _on_open_file_requested(self, file_path: str):
  160. """Handle file selection from control panel."""
  161. if file_path and file_path.strip():
  162. self._process_image_file(file_path)
  163. else:
  164. self.status_label.setText("No file selected. Please select an image file.")
  165. self.status_label.setStyleSheet(f"color: {COLORS['warning']}; font-size: 11px;")
  166. def _on_annotated_image_requested(self):
  167. """Handle annotated image request from clicking on locule count."""
  168. # Show the most recent annotated image if available
  169. if hasattr(self, 'last_processed_file') and self.last_processed_file:
  170. self._show_annotated_image_dialog(self.last_processed_file)
  171. else:
  172. self.status_label.setText("No processed image available to display.")
  173. def _on_defect_image_requested(self):
  174. """Handle defect image request from clicking on defect analysis."""
  175. # Show the defect-annotated image if available
  176. if hasattr(self, 'last_processed_file') and self.last_processed_file:
  177. self._show_defect_annotated_image_dialog(self.last_processed_file)
  178. else:
  179. self.status_label.setText("No processed image available to display.")
  180. def _simulate_processing(self):
  181. """Simulate processing for demonstration purposes."""
  182. # Set control panel to analyzing state
  183. self.control_panel.set_analyzing(True)
  184. # Update status
  185. self.status_label.setText("Processing sample...")
  186. # Simulate processing delay (in real app, this would be handled by workers)
  187. from PyQt5.QtCore import QTimer
  188. QTimer.singleShot(2000, self._complete_processing)
  189. def _complete_processing(self):
  190. """Complete the processing simulation."""
  191. # Reset control panel
  192. self.control_panel.set_analyzing(False)
  193. # Add result to history
  194. from datetime import datetime
  195. current_time = datetime.now().strftime("%H:%M:%S")
  196. test_id = f"Q-{len(self.history_panel.test_history) + 1:04d}"
  197. # Sample result based on current parameters
  198. params = self.control_panel.get_parameters()
  199. grade = "B" # Default grade
  200. score = 78.5 # Default score
  201. defects = "2" # Default defect count
  202. self.history_panel.add_test_result(current_time, test_id, grade, f"{score:.1f}%", defects)
  203. # Update status
  204. self.status_label.setText("Processing complete. Results updated.")
  205. def _process_image_file(self, file_path: str):
  206. """Process the selected image file with quality models."""
  207. import os
  208. from pathlib import Path
  209. try:
  210. # Validate file exists
  211. if not os.path.exists(file_path):
  212. raise FileNotFoundError(f"File not found: {file_path}")
  213. # Validate it's an image file
  214. valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.JPG', '.JPEG', '.PNG', '.BMP']
  215. file_ext = Path(file_path).suffix
  216. if file_ext not in valid_extensions:
  217. raise ValueError(f"Invalid image file format: {file_ext}. Supported formats: {', '.join(valid_extensions)}")
  218. # Set loading state
  219. self.control_panel.set_analyzing(True)
  220. filename = os.path.basename(file_path)
  221. self.status_label.setText(f"Processing: {filename}...")
  222. self.status_label.setStyleSheet(f"color: {COLORS['warning']}; font-size: 11px;")
  223. # Load and display the image in RGB panels
  224. self._display_image_in_panels(file_path)
  225. # Process with models immediately (no artificial delay)
  226. self._analyze_image_with_models(file_path)
  227. except Exception as e:
  228. error_msg = f"Error processing file: {str(e)}"
  229. self.status_label.setText(error_msg)
  230. self.status_label.setStyleSheet(f"color: {COLORS.get('error', '#e74c3c')}; font-size: 11px;")
  231. self.control_panel.set_analyzing(False)
  232. print(f"Error in _process_image_file: {e}")
  233. import traceback
  234. traceback.print_exc()
  235. def _display_image_in_panels(self, file_path: str):
  236. """Display the loaded image in the RGB panels."""
  237. from PyQt5.QtGui import QPixmap
  238. # Load image as pixmap
  239. pixmap = QPixmap(file_path)
  240. if pixmap.isNull():
  241. raise ValueError(f"Could not load image: {file_path}")
  242. # Scale to fit panels
  243. scaled_pixmap = pixmap.scaled(
  244. 250, 180, Qt.KeepAspectRatio, Qt.SmoothTransformation
  245. )
  246. # Update RGB panels with the loaded image
  247. # Note: In a real implementation, you would process the image with models first
  248. # and then display the annotated results
  249. self.rgb_top_panel.set_image(file_path)
  250. self.rgb_side_panel.set_image(file_path)
  251. def _analyze_image_with_models(self, file_path: str):
  252. """Analyze the image using defect and locule models."""
  253. # Process immediately without artificial delay
  254. # The models will handle their own processing time
  255. self._complete_image_analysis(file_path)
  256. def _complete_image_analysis(self, file_path: str):
  257. """Complete the image analysis with both locule and defect model results."""
  258. # Store the file path for later use
  259. self.last_processed_file = file_path
  260. try:
  261. # Initialize results
  262. locule_result = None
  263. defect_result = None
  264. # Check if models are loaded
  265. if not self.locule_model or not self.locule_model.is_loaded:
  266. print("Warning: Locule model not loaded. Check model path: locule.pt")
  267. if not self.defect_model or not self.defect_model.is_loaded:
  268. print("Warning: Defect model not loaded. Check model path: best.pt")
  269. # Use actual locule model if available
  270. if self.locule_model and self.locule_model.is_loaded:
  271. try:
  272. print(f"Running locule analysis on: {file_path}")
  273. locule_result = self.locule_model.predict(file_path)
  274. if locule_result['success']:
  275. print(f"✓ Locule analysis complete: {locule_result['locule_count']} locules detected")
  276. else:
  277. print(f"✗ Locule model failed: {locule_result.get('error', 'Unknown error')}")
  278. except Exception as e:
  279. print(f"Exception during locule prediction: {e}")
  280. import traceback
  281. traceback.print_exc()
  282. else:
  283. print("Locule model not available - skipping locule analysis")
  284. # Use actual defect model if available
  285. if self.defect_model and self.defect_model.is_loaded:
  286. try:
  287. print(f"Running defect analysis on: {file_path}")
  288. defect_result = self.defect_model.predict(file_path)
  289. if defect_result['success']:
  290. print(f"✓ Defect analysis complete: {defect_result['total_detections']} detections, primary class: {defect_result.get('primary_class', 'N/A')}")
  291. else:
  292. print(f"✗ Defect model failed: {defect_result.get('error', 'Unknown error')}")
  293. except Exception as e:
  294. print(f"Exception during defect prediction: {e}")
  295. import traceback
  296. traceback.print_exc()
  297. else:
  298. print("Defect model not available - skipping defect analysis")
  299. # Update panels with results
  300. self._update_panels_with_results(locule_result, defect_result, file_path)
  301. except Exception as e:
  302. print(f"Error in image analysis: {e}")
  303. import traceback
  304. traceback.print_exc()
  305. self._fallback_analysis(file_path)
  306. # Reset control panel
  307. self.control_panel.set_analyzing(False)
  308. # Update status
  309. import os
  310. filename = os.path.basename(file_path)
  311. self.status_label.setText(f"Analysis complete: {filename}")
  312. self.status_label.setStyleSheet(f"color: {COLORS['success']}; font-size: 11px;")
  313. # Show both dialogs - locule and defect analysis (only if we have results)
  314. if (locule_result and locule_result.get('success')) or (defect_result and defect_result.get('success')):
  315. self._show_dual_analysis_dialogs(file_path)
  316. def _calculate_quality_grade(self, locule_count: int, has_defects: bool) -> tuple:
  317. """
  318. Calculate quality grade based on locule count and defect status.
  319. Rules:
  320. - 5 locules + no defects = Class A
  321. - Less than 5 locules + no defects = Class B
  322. - Less than 5 locules + defects = Class C
  323. Args:
  324. locule_count: Number of locules detected
  325. has_defects: True if defects are present, False otherwise
  326. Returns:
  327. tuple: (grade_letter, score_percentage)
  328. """
  329. if locule_count == 5 and not has_defects:
  330. return ("A", 95.0)
  331. elif locule_count < 5 and not has_defects:
  332. return ("B", 85.0)
  333. elif locule_count < 5 and has_defects:
  334. return ("C", 70.0)
  335. else:
  336. # Edge cases: 5 locules with defects, or more than 5 locules
  337. # Default to B for 5 locules with defects, C for more than 5 with defects
  338. if locule_count >= 5 and has_defects:
  339. return ("B", 80.0) # Still good locule count but has defects
  340. elif locule_count > 5 and not has_defects:
  341. return ("B", 85.0) # More than ideal but no defects
  342. else:
  343. # Fallback
  344. return ("B", 75.0)
  345. def _update_panels_with_results(self, locule_result, defect_result, file_path: str):
  346. """Update all panels with locule and defect model results."""
  347. # Initialize values
  348. locule_count = 0
  349. has_defects = False
  350. total_detections = 0
  351. primary_class = "No Defects"
  352. # Handle locule results
  353. if locule_result and locule_result['success']:
  354. locule_count = locule_result['locule_count']
  355. locule_detections = locule_result['detections']
  356. # Update the clickable locule count widget with actual AI results
  357. self.defects_panel.update_locule_count(locule_count)
  358. # Update the locule count metric in quality results
  359. self.quality_results_panel.update_locule_count(locule_count, 94.5)
  360. else:
  361. # Fallback locule count
  362. locule_count = 4
  363. self.defects_panel.update_locule_count(locule_count)
  364. self.quality_results_panel.update_locule_count(locule_count)
  365. # Handle defect results
  366. if defect_result and defect_result['success']:
  367. total_detections = defect_result['total_detections']
  368. primary_class = defect_result['primary_class']
  369. class_counts = defect_result['class_counts']
  370. defect_detections = defect_result['detections']
  371. # Determine if defects are present
  372. # "No Defects" class means no defects, anything else means defects exist
  373. has_defects = (primary_class != "No Defects")
  374. # Update defect detection panel with actual defect results
  375. self.defects_panel.update_defect_count(total_detections, primary_class)
  376. # Create defect analysis results for the panel
  377. defect_analysis_results = []
  378. # Add individual defect detections
  379. for detection in defect_detections:
  380. defect_analysis_results.append({
  381. 'type': detection['class_name'],
  382. 'confidence': detection['confidence'],
  383. 'color': '#f39c12',
  384. 'category': 'warning'
  385. })
  386. # Add shape analysis (placeholder for now)
  387. defect_analysis_results.append({
  388. 'type': 'Shape Analysis',
  389. 'result': 'Regular',
  390. 'symmetry': '91.2%',
  391. 'confidence': 94.1,
  392. 'color': '#27ae60',
  393. 'category': 'success'
  394. })
  395. # Update defect detection panel with all results
  396. self.defects_panel.update_defects(defect_analysis_results)
  397. else:
  398. # Fallback defect analysis - assume no defects if model failed
  399. has_defects = False
  400. total_detections = 0
  401. primary_class = "No Defects"
  402. self.defects_panel.update_defect_count(0, "No defects")
  403. self.defects_panel.update_defects([
  404. {
  405. 'type': 'Shape Analysis',
  406. 'result': 'Regular',
  407. 'symmetry': '91.2%',
  408. 'confidence': 94.1,
  409. 'color': '#27ae60',
  410. 'category': 'success'
  411. }
  412. ])
  413. # Calculate quality grade based on locule count and defect status
  414. grade, score = self._calculate_quality_grade(locule_count, has_defects)
  415. # Update quality results panel with calculated grade
  416. self.quality_results_panel.update_grade(grade, score)
  417. # Update history with calculated grade
  418. from datetime import datetime
  419. current_time = datetime.now().strftime("%H:%M:%S")
  420. self.history_panel.add_test_result(
  421. time=current_time,
  422. test_id=f"Q-{len(self.history_panel.test_history) + 1:04d}",
  423. grade=grade,
  424. score=f"{score:.1f}%",
  425. defects=str(total_detections)
  426. )
  427. print(f"Quality Grade Calculated: Grade {grade} (Score: {score:.1f}%) - Locules: {locule_count}, Defects: {'Yes' if has_defects else 'No'}")
  428. def _fallback_analysis(self, file_path: str):
  429. """Fallback analysis when model is not available."""
  430. # Update locule count widget with fallback data
  431. self.defects_panel.update_locule_count(4)
  432. # Update locule count in quality results panel
  433. self.quality_results_panel.update_locule_count(4)
  434. # Generate sample analysis results
  435. defects_list = [
  436. {
  437. 'type': 'Shape Analysis',
  438. 'result': 'Regular',
  439. 'symmetry': '91.2%',
  440. 'confidence': 94.1,
  441. 'color': '#27ae60',
  442. 'category': 'success'
  443. }
  444. ]
  445. # Update defect detection panel
  446. self.defects_panel.update_defects(defects_list)
  447. # Calculate grade for fallback (assume 4 locules, no defects = Grade B)
  448. grade, score = self._calculate_quality_grade(4, False)
  449. # Update quality results panel with calculated grade
  450. self.quality_results_panel.update_grade(grade, score)
  451. # Update history
  452. from datetime import datetime
  453. current_time = datetime.now().strftime("%H:%M:%S")
  454. self.history_panel.add_test_result(
  455. time=current_time,
  456. test_id=f"Q-{len(self.history_panel.test_history) + 1:04d}",
  457. grade=grade,
  458. score=f"{score:.1f}%",
  459. defects="0"
  460. )
  461. def _generate_annotated_image(self, original_path: str):
  462. """Generate an annotated image using the actual LoculeModel."""
  463. try:
  464. # Check if locule model is available
  465. if self.locule_model is None or not self.locule_model.is_loaded:
  466. print("Locule model not available, using fallback annotation")
  467. return self._generate_fallback_annotated_image(original_path)
  468. # Use actual locule model for prediction
  469. result = self.locule_model.predict(original_path)
  470. if result['success']:
  471. # Return the actual annotated QImage from the model
  472. return self._qimage_to_pixmap(result['annotated_image'])
  473. else:
  474. print(f"Locule model prediction failed: {result['error']}")
  475. return self._generate_fallback_annotated_image(original_path)
  476. except Exception as e:
  477. print(f"Error in locule model prediction: {e}")
  478. return self._generate_fallback_annotated_image(original_path)
  479. def _generate_fallback_annotated_image(self, original_path: str):
  480. """Fallback annotation method when model is not available."""
  481. try:
  482. from PyQt5.QtGui import QPixmap, QPainter, QColor, QPen, QBrush
  483. import math
  484. # Load original image
  485. original_pixmap = QPixmap(original_path)
  486. if original_pixmap.isNull():
  487. return None
  488. # Create a copy for annotation
  489. annotated_pixmap = original_pixmap.copy()
  490. # Create painter for annotations
  491. painter = QPainter(annotated_pixmap)
  492. painter.setRenderHint(QPainter.Antialiasing)
  493. # Get image dimensions
  494. width = annotated_pixmap.width()
  495. height = annotated_pixmap.height()
  496. # Draw defect markers (sample data for now)
  497. defect_positions = [
  498. (width * 0.3, height * 0.4, "Mechanical Damage"),
  499. (width * 0.7, height * 0.6, "Surface Blemish")
  500. ]
  501. for x, y, defect_type in defect_positions:
  502. # Draw defect circle
  503. painter.setBrush(QBrush(QColor("#f39c12")))
  504. painter.setPen(QPen(QColor("#e67e22"), 3))
  505. painter.drawEllipse(int(x - 15), int(y - 15), 30, 30)
  506. # Draw confidence text
  507. painter.setPen(QPen(QColor("white"), 2))
  508. painter.drawText(int(x - 20), int(y - 20), f"87%")
  509. # Draw locule counting visualization
  510. center_x, center_y = width // 2, height // 2
  511. fruit_radius = min(width, height) // 4
  512. # Draw locule segments
  513. locule_colors = ["#3498db", "#e74c3c", "#2ecc71", "#f39c12"]
  514. num_locules = 4
  515. angle_step = 360 / num_locules
  516. for i in range(num_locules):
  517. start_angle = int(i * angle_step)
  518. span_angle = int(angle_step)
  519. # Draw locule segment
  520. painter.setBrush(QBrush(QColor(locule_colors[i])))
  521. painter.setPen(QPen(QColor("#34495e"), 2))
  522. painter.drawPie(
  523. center_x - fruit_radius, center_y - fruit_radius,
  524. fruit_radius * 2, fruit_radius * 2,
  525. start_angle * 16, span_angle * 16
  526. )
  527. painter.end()
  528. return annotated_pixmap
  529. except Exception as e:
  530. print(f"Error generating fallback annotated image: {e}")
  531. return None
  532. def _qimage_to_pixmap(self, qimage):
  533. """Convert QImage to QPixmap."""
  534. from PyQt5.QtGui import QPixmap
  535. return QPixmap.fromImage(qimage)
  536. def _show_annotated_image_dialog(self, file_path: str):
  537. """Show dialog with annotated image results."""
  538. try:
  539. # Generate annotated image
  540. annotated_pixmap = self._generate_annotated_image(file_path)
  541. if annotated_pixmap:
  542. # Show in preview dialog
  543. dialog = ImagePreviewDialog(
  544. annotated_pixmap,
  545. title=f"Quality Analysis Results - {os.path.basename(file_path)}",
  546. parent=self
  547. )
  548. dialog.exec_()
  549. else:
  550. # Fallback to original image
  551. original_pixmap = QPixmap(file_path)
  552. if not original_pixmap.isNull():
  553. dialog = ImagePreviewDialog(
  554. original_pixmap,
  555. title=f"Original Image - {os.path.basename(file_path)}",
  556. parent=self
  557. )
  558. dialog.exec_()
  559. except Exception as e:
  560. print(f"Error showing image dialog: {e}")
  561. def _show_defect_annotated_image_dialog(self, file_path: str):
  562. """Show dialog with defect-annotated image results."""
  563. try:
  564. # Generate defect-annotated image
  565. annotated_pixmap = self._generate_defect_annotated_image(file_path)
  566. if annotated_pixmap:
  567. # Show in preview dialog
  568. dialog = ImagePreviewDialog(
  569. annotated_pixmap,
  570. title=f"Defect Detection Results - {os.path.basename(file_path)}",
  571. parent=self
  572. )
  573. dialog.exec_()
  574. else:
  575. # Fallback to original image
  576. original_pixmap = QPixmap(file_path)
  577. if not original_pixmap.isNull():
  578. dialog = ImagePreviewDialog(
  579. original_pixmap,
  580. title=f"Original Image - {os.path.basename(file_path)}",
  581. parent=self
  582. )
  583. dialog.exec_()
  584. except Exception as e:
  585. print(f"Error showing defect image dialog: {e}")
  586. def _generate_defect_annotated_image(self, original_path: str):
  587. """Generate a defect-annotated image using the actual DefectModel."""
  588. try:
  589. # Check if defect model is available
  590. if self.defect_model is None or not self.defect_model.is_loaded:
  591. print("Defect model not available, using fallback annotation")
  592. return self._generate_fallback_defect_annotated_image(original_path)
  593. # Use actual defect model for prediction
  594. result = self.defect_model.predict(original_path)
  595. if result['success']:
  596. # Return the actual annotated QImage from the model
  597. return self._qimage_to_pixmap(result['annotated_image'])
  598. else:
  599. print(f"Defect model prediction failed: {result['error']}")
  600. return self._generate_fallback_defect_annotated_image(original_path)
  601. except Exception as e:
  602. print(f"Error in defect model prediction: {e}")
  603. return self._generate_fallback_defect_annotated_image(original_path)
  604. def _generate_fallback_defect_annotated_image(self, original_path: str):
  605. """Fallback defect annotation method when model is not available."""
  606. try:
  607. from PyQt5.QtGui import QPixmap, QPainter, QColor, QPen, QBrush
  608. # Load original image
  609. original_pixmap = QPixmap(original_path)
  610. if original_pixmap.isNull():
  611. return None
  612. # Create a copy for annotation
  613. annotated_pixmap = original_pixmap.copy()
  614. # Create painter for annotations
  615. painter = QPainter(annotated_pixmap)
  616. painter.setRenderHint(QPainter.Antialiasing)
  617. # Get image dimensions
  618. width = annotated_pixmap.width()
  619. height = annotated_pixmap.height()
  620. # Draw sample defect markers
  621. defect_positions = [
  622. (width * 0.3, height * 0.4, "Minor Defect", "#f39c12"),
  623. (width * 0.7, height * 0.6, "Surface Blemish", "#e67e22")
  624. ]
  625. for x, y, defect_type, color_hex in defect_positions:
  626. # Draw defect bounding box
  627. color = QColor(color_hex)
  628. painter.setPen(QPen(color, 3))
  629. painter.setBrush(QBrush(QColor(0, 0, 0, 0))) # Transparent fill
  630. # Draw rectangle
  631. box_size = 40
  632. painter.drawRect(int(x - box_size/2), int(y - box_size/2), box_size, box_size)
  633. # Draw confidence text
  634. painter.setPen(QPen(QColor("white"), 2))
  635. painter.drawText(int(x - 20), int(y - 25), f"87%")
  636. # Draw label
  637. painter.drawText(int(x - 30), int(y + 30), defect_type)
  638. painter.end()
  639. return annotated_pixmap
  640. except Exception as e:
  641. print(f"Error generating fallback defect annotated image: {e}")
  642. return None
  643. def _show_dual_analysis_dialogs(self, file_path: str):
  644. """Show both locule and defect analysis dialogs."""
  645. try:
  646. # Show locule analysis dialog first
  647. self._show_annotated_image_dialog(file_path)
  648. # Then show defect analysis dialog
  649. self._show_defect_annotated_image_dialog(file_path)
  650. except Exception as e:
  651. print(f"Error showing dual analysis dialogs: {e}")
  652. def set_loading(self, is_loading: bool):
  653. """Set loading state for all panels."""
  654. # Update control panel analyzing state
  655. self.control_panel.set_analyzing(is_loading)
  656. if is_loading:
  657. self.status_label.setText("Processing sample...")
  658. self.status_label.setStyleSheet(f"color: {COLORS['warning']}; font-size: 11px;")
  659. else:
  660. self.status_label.setText("Ready for quality assessment.")
  661. self.status_label.setStyleSheet(f"color: {COLORS['success']}; font-size: 11px;")
  662. def update_results(self, annotated_image: QImage, primary_class: str,
  663. class_counts: dict, total_detections: int, file_path: str):
  664. """
  665. Update the tab with processing results.
  666. This method is kept for compatibility with existing workers.
  667. Args:
  668. annotated_image: QImage with bounding boxes drawn
  669. primary_class: Primary defect class detected
  670. class_counts: Dictionary of class counts
  671. total_detections: Total number of detections
  672. file_path: Path to the image file
  673. """
  674. # For now, use the new panel structure
  675. # In a real implementation, this would update the appropriate panels
  676. # with the processed image data
  677. # Update defect detection results
  678. defects_list = []
  679. for class_name, count in class_counts.items():
  680. defects_list.append({
  681. 'type': class_name,
  682. 'confidence': 85.0, # Placeholder confidence
  683. 'color': '#f39c12',
  684. 'category': 'warning'
  685. })
  686. self.defects_panel.update_defects(defects_list)
  687. # Update file path in status
  688. import os
  689. filename = os.path.basename(file_path)
  690. self.status_label.setText(f"Processed: {filename}")
  691. self.set_loading(False)
  692. def clear_results(self):
  693. """Clear all displayed results."""
  694. # Clear defect detection results
  695. self.defects_panel.clear_defects()
  696. # Clear history selection
  697. self.history_panel.history_table.clearSelection()
  698. # Reset status
  699. self.status_label.setText("Ready for quality assessment. Use Control Panel to analyze samples.")
  700. self.status_label.setStyleSheet(f"color: {COLORS['text_secondary']}; font-size: 11px;")
  701. def get_panel_status(self):
  702. """Get status information from all panels."""
  703. return {
  704. 'rgb_top': 'ONLINE',
  705. 'rgb_side': 'ONLINE',
  706. 'thermal': 'OFFLINE',
  707. 'defects': 'ONLINE',
  708. 'control': 'READY',
  709. 'history': 'READY'
  710. }