lazylabel-gui 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. lazylabel/__init__.py +8 -8
  2. lazylabel/config/__init__.py +6 -6
  3. lazylabel/config/hotkeys.py +168 -168
  4. lazylabel/config/paths.py +40 -40
  5. lazylabel/config/settings.py +65 -65
  6. lazylabel/core/__init__.py +6 -6
  7. lazylabel/core/file_manager.py +105 -105
  8. lazylabel/core/model_manager.py +97 -94
  9. lazylabel/core/segment_manager.py +171 -140
  10. lazylabel/main.py +36 -22
  11. lazylabel/models/__init__.py +4 -4
  12. lazylabel/models/sam_model.py +195 -154
  13. lazylabel/ui/__init__.py +7 -7
  14. lazylabel/ui/control_panel.py +241 -220
  15. lazylabel/ui/editable_vertex.py +64 -51
  16. lazylabel/ui/hotkey_dialog.py +383 -383
  17. lazylabel/ui/hoverable_pixelmap_item.py +22 -22
  18. lazylabel/ui/hoverable_polygon_item.py +39 -39
  19. lazylabel/ui/main_window.py +1659 -1264
  20. lazylabel/ui/numeric_table_widget_item.py +9 -9
  21. lazylabel/ui/photo_viewer.py +54 -54
  22. lazylabel/ui/reorderable_class_table.py +61 -61
  23. lazylabel/ui/right_panel.py +315 -239
  24. lazylabel/ui/widgets/__init__.py +8 -7
  25. lazylabel/ui/widgets/adjustments_widget.py +108 -107
  26. lazylabel/ui/widgets/model_selection_widget.py +93 -93
  27. lazylabel/ui/widgets/settings_widget.py +105 -105
  28. lazylabel/ui/widgets/status_bar.py +109 -0
  29. lazylabel/utils/__init__.py +5 -5
  30. lazylabel/utils/custom_file_system_model.py +132 -132
  31. lazylabel/utils/utils.py +12 -12
  32. {lazylabel_gui-1.1.0.dist-info → lazylabel_gui-1.1.2.dist-info}/METADATA +197 -197
  33. lazylabel_gui-1.1.2.dist-info/RECORD +37 -0
  34. {lazylabel_gui-1.1.0.dist-info → lazylabel_gui-1.1.2.dist-info}/licenses/LICENSE +21 -21
  35. lazylabel_gui-1.1.0.dist-info/RECORD +0 -36
  36. {lazylabel_gui-1.1.0.dist-info → lazylabel_gui-1.1.2.dist-info}/WHEEL +0 -0
  37. {lazylabel_gui-1.1.0.dist-info → lazylabel_gui-1.1.2.dist-info}/entry_points.txt +0 -0
  38. {lazylabel_gui-1.1.0.dist-info → lazylabel_gui-1.1.2.dist-info}/top_level.txt +0 -0
@@ -1,1264 +1,1659 @@
1
- """Main application window."""
2
-
3
- import os
4
- import numpy as np
5
- import cv2
6
- from PyQt6.QtWidgets import (
7
- QMainWindow,
8
- QWidget,
9
- QHBoxLayout,
10
- QFileDialog,
11
- QApplication,
12
- QGraphicsEllipseItem,
13
- QGraphicsLineItem,
14
- QGraphicsPolygonItem,
15
- QTableWidgetItem,
16
- QTableWidgetSelectionRange,
17
- QHeaderView,
18
- )
19
- from PyQt6.QtGui import (
20
- QIcon,
21
- QKeySequence,
22
- QShortcut,
23
- QPixmap,
24
- QColor,
25
- QPen,
26
- QBrush,
27
- QPolygonF,
28
- QImage,
29
- )
30
- from PyQt6.QtCore import Qt, QTimer, QModelIndex, QPointF
31
-
32
- from .control_panel import ControlPanel
33
- from .right_panel import RightPanel
34
- from .photo_viewer import PhotoViewer
35
- from .hoverable_polygon_item import HoverablePolygonItem
36
- from .hoverable_pixelmap_item import HoverablePixmapItem
37
- from .editable_vertex import EditableVertexItem
38
- from .numeric_table_widget_item import NumericTableWidgetItem
39
- from ..core import SegmentManager, ModelManager, FileManager
40
- from ..config import Settings, Paths, HotkeyManager
41
- from ..utils import CustomFileSystemModel, mask_to_pixmap
42
- from .hotkey_dialog import HotkeyDialog
43
-
44
-
45
- class MainWindow(QMainWindow):
46
- """Main application window."""
47
-
48
- def __init__(self):
49
- super().__init__()
50
-
51
- # Initialize configuration
52
- self.paths = Paths()
53
- self.settings = Settings.load_from_file(str(self.paths.settings_file))
54
- self.hotkey_manager = HotkeyManager(str(self.paths.config_dir))
55
-
56
- # Initialize managers
57
- self.segment_manager = SegmentManager()
58
- self.model_manager = ModelManager(self.paths)
59
- self.file_manager = FileManager(self.segment_manager)
60
-
61
- # Initialize UI state
62
- self.mode = "sam_points"
63
- self.previous_mode = "sam_points"
64
- self.current_image_path = None
65
- self.current_file_index = QModelIndex()
66
-
67
- # Annotation state
68
- self.point_radius = self.settings.point_radius
69
- self.line_thickness = self.settings.line_thickness
70
- self.pan_multiplier = self.settings.pan_multiplier
71
- self.polygon_join_threshold = self.settings.polygon_join_threshold
72
-
73
- # Drawing state
74
- self.point_items, self.positive_points, self.negative_points = [], [], []
75
- self.polygon_points, self.polygon_preview_items = [], []
76
- self.rubber_band_line = None
77
- self.preview_mask_item = None
78
- self.segments, self.segment_items, self.highlight_items = [], {}, []
79
- self.edit_handles = []
80
- self.is_dragging_polygon, self.drag_start_pos, self.drag_initial_vertices = (
81
- False,
82
- None,
83
- {},
84
- )
85
-
86
- self._setup_ui()
87
- self._setup_model()
88
- self._setup_connections()
89
- self._setup_shortcuts()
90
- self._load_settings()
91
-
92
- def _setup_ui(self):
93
- """Setup the user interface."""
94
- self.setWindowTitle("LazyLabel by DNC")
95
- self.setGeometry(
96
- 50, 50, self.settings.window_width, self.settings.window_height
97
- )
98
-
99
- # Set window icon
100
- if self.paths.logo_path.exists():
101
- self.setWindowIcon(QIcon(str(self.paths.logo_path)))
102
-
103
- # Create panels
104
- self.control_panel = ControlPanel()
105
- self.right_panel = RightPanel()
106
- self.viewer = PhotoViewer(self)
107
- self.viewer.setMouseTracking(True)
108
-
109
- # Setup file model
110
- self.file_model = CustomFileSystemModel()
111
- self.right_panel.setup_file_model(self.file_model)
112
-
113
- # Layout
114
- main_layout = QHBoxLayout()
115
- main_layout.addWidget(self.control_panel)
116
- main_layout.addWidget(self.viewer, 1)
117
- main_layout.addWidget(self.right_panel)
118
-
119
- central_widget = QWidget()
120
- central_widget.setLayout(main_layout)
121
- self.setCentralWidget(central_widget)
122
-
123
- def _setup_model(self):
124
- """Setup the SAM model."""
125
- sam_model = self.model_manager.initialize_default_model(
126
- self.settings.default_model_type
127
- )
128
-
129
- if sam_model and sam_model.is_loaded:
130
- self.control_panel.set_device_text(str(sam_model.device))
131
- self._enable_sam_functionality(True)
132
- else:
133
- self.control_panel.set_device_text("No model loaded")
134
- self._enable_sam_functionality(False)
135
- print("SAM model failed to load. Point mode will be disabled.")
136
-
137
- # Setup model change callback
138
- self.model_manager.on_model_changed = self.control_panel.set_current_model
139
-
140
- # Initialize models list
141
- models = self.model_manager.get_available_models(str(self.paths.models_dir))
142
- self.control_panel.populate_models(models)
143
-
144
- def _enable_sam_functionality(self, enabled: bool):
145
- """Enable or disable SAM point functionality."""
146
- self.control_panel.set_sam_mode_enabled(enabled)
147
- if not enabled and self.mode == "sam_points":
148
- # Switch to polygon mode if SAM is disabled and we're in SAM mode
149
- self.set_polygon_mode()
150
-
151
- def _setup_connections(self):
152
- """Setup signal connections."""
153
- # Control panel connections
154
- self.control_panel.sam_mode_requested.connect(self.set_sam_mode)
155
- self.control_panel.polygon_mode_requested.connect(self.set_polygon_mode)
156
- self.control_panel.selection_mode_requested.connect(self.toggle_selection_mode)
157
- self.control_panel.clear_points_requested.connect(self.clear_all_points)
158
- self.control_panel.fit_view_requested.connect(self.viewer.fitInView)
159
- self.control_panel.hotkeys_requested.connect(self._show_hotkey_dialog)
160
-
161
- # Model management
162
- self.control_panel.browse_models_requested.connect(self._browse_models_folder)
163
- self.control_panel.refresh_models_requested.connect(self._refresh_models_list)
164
- self.control_panel.model_selected.connect(self._load_selected_model)
165
-
166
- # Adjustments
167
- self.control_panel.annotation_size_changed.connect(self._set_annotation_size)
168
- self.control_panel.pan_speed_changed.connect(self._set_pan_speed)
169
- self.control_panel.join_threshold_changed.connect(self._set_join_threshold)
170
-
171
- # Right panel connections
172
- self.right_panel.open_folder_requested.connect(self._open_folder_dialog)
173
- self.right_panel.image_selected.connect(self._load_selected_image)
174
- self.right_panel.merge_selection_requested.connect(
175
- self._assign_selected_to_class
176
- )
177
- self.right_panel.delete_selection_requested.connect(
178
- self._delete_selected_segments
179
- )
180
- self.right_panel.segments_selection_changed.connect(
181
- self._highlight_selected_segments
182
- )
183
- self.right_panel.class_alias_changed.connect(self._handle_alias_change)
184
- self.right_panel.reassign_classes_requested.connect(self._reassign_class_ids)
185
- self.right_panel.class_filter_changed.connect(self._update_segment_table)
186
-
187
- # Panel visibility
188
- self.control_panel.btn_toggle_visibility.clicked.connect(
189
- self.control_panel.toggle_visibility
190
- )
191
- self.right_panel.btn_toggle_visibility.clicked.connect(
192
- self.right_panel.toggle_visibility
193
- )
194
-
195
- # Mouse events (will be implemented in a separate handler)
196
- self._setup_mouse_events()
197
-
198
- def _setup_shortcuts(self):
199
- """Setup keyboard shortcuts based on hotkey manager."""
200
- self.shortcuts = [] # Keep track of shortcuts for updating
201
- self._update_shortcuts()
202
-
203
- def _update_shortcuts(self):
204
- """Update shortcuts based on current hotkey configuration."""
205
- # Clear existing shortcuts
206
- for shortcut in self.shortcuts:
207
- shortcut.setParent(None)
208
- self.shortcuts.clear()
209
-
210
- # Map action names to callbacks
211
- action_callbacks = {
212
- "load_next_image": self._load_next_image,
213
- "load_previous_image": self._load_previous_image,
214
- "sam_mode": self.set_sam_mode,
215
- "polygon_mode": self.set_polygon_mode,
216
- "selection_mode": self.toggle_selection_mode,
217
- "pan_mode": self.toggle_pan_mode,
218
- "edit_mode": self.toggle_edit_mode,
219
- "clear_points": self.clear_all_points,
220
- "escape": self._handle_escape_press,
221
- "delete_segments": self._delete_selected_segments,
222
- "delete_segments_alt": self._delete_selected_segments,
223
- "merge_segments": self._handle_merge_press,
224
- "undo": self._undo_last_action,
225
- "select_all": lambda: self.right_panel.select_all_segments(),
226
- "save_segment": self._handle_space_press,
227
- "save_output": self._handle_enter_press,
228
- "save_output_alt": self._handle_enter_press,
229
- "fit_view": self.viewer.fitInView,
230
- "zoom_in": self._handle_zoom_in,
231
- "zoom_out": self._handle_zoom_out,
232
- "pan_up": lambda: self._handle_pan_key("up"),
233
- "pan_down": lambda: self._handle_pan_key("down"),
234
- "pan_left": lambda: self._handle_pan_key("left"),
235
- "pan_right": lambda: self._handle_pan_key("right"),
236
- }
237
-
238
- # Create shortcuts for each action
239
- for action_name, callback in action_callbacks.items():
240
- primary_key, secondary_key = self.hotkey_manager.get_key_for_action(
241
- action_name
242
- )
243
-
244
- # Create primary shortcut
245
- if primary_key:
246
- shortcut = QShortcut(QKeySequence(primary_key), self, callback)
247
- self.shortcuts.append(shortcut)
248
-
249
- # Create secondary shortcut
250
- if secondary_key:
251
- shortcut = QShortcut(QKeySequence(secondary_key), self, callback)
252
- self.shortcuts.append(shortcut)
253
-
254
- def _load_settings(self):
255
- """Load and apply settings."""
256
- self.control_panel.set_settings(self.settings.__dict__)
257
- self.control_panel.set_annotation_size(
258
- int(self.settings.annotation_size_multiplier * 10)
259
- )
260
- # Set initial mode based on model availability
261
- if self.model_manager.is_model_available():
262
- self.set_sam_mode()
263
- else:
264
- self.set_polygon_mode()
265
-
266
- def _setup_mouse_events(self):
267
- """Setup mouse event handling."""
268
- self._original_mouse_press = self.viewer.scene().mousePressEvent
269
- self._original_mouse_move = self.viewer.scene().mouseMoveEvent
270
- self._original_mouse_release = self.viewer.scene().mouseReleaseEvent
271
-
272
- self.viewer.scene().mousePressEvent = self._scene_mouse_press
273
- self.viewer.scene().mouseMoveEvent = self._scene_mouse_move
274
- self.viewer.scene().mouseReleaseEvent = self._scene_mouse_release
275
-
276
- # Mode management methods
277
- def set_sam_mode(self):
278
- """Set SAM points mode."""
279
- if not self.model_manager.is_model_available():
280
- print("Cannot enter SAM mode: No model available")
281
- return
282
- self._set_mode("sam_points")
283
-
284
- def set_polygon_mode(self):
285
- """Set polygon drawing mode."""
286
- self._set_mode("polygon")
287
-
288
- def toggle_selection_mode(self):
289
- """Toggle selection mode."""
290
- self._toggle_mode("selection")
291
-
292
- def toggle_pan_mode(self):
293
- """Toggle pan mode."""
294
- self._toggle_mode("pan")
295
-
296
- def toggle_edit_mode(self):
297
- """Toggle edit mode."""
298
- self._toggle_mode("edit")
299
-
300
- def _set_mode(self, mode_name, is_toggle=False):
301
- """Set the current mode."""
302
- if not is_toggle and self.mode not in ["selection", "edit"]:
303
- self.previous_mode = self.mode
304
-
305
- self.mode = mode_name
306
- self.control_panel.set_mode_text(mode_name)
307
- self.clear_all_points()
308
-
309
- # Set cursor and drag mode based on mode
310
- cursor_map = {
311
- "sam_points": Qt.CursorShape.CrossCursor,
312
- "polygon": Qt.CursorShape.CrossCursor,
313
- "selection": Qt.CursorShape.ArrowCursor,
314
- "edit": Qt.CursorShape.SizeAllCursor,
315
- "pan": Qt.CursorShape.OpenHandCursor,
316
- }
317
- self.viewer.set_cursor(cursor_map.get(self.mode, Qt.CursorShape.ArrowCursor))
318
-
319
- drag_mode = (
320
- self.viewer.DragMode.ScrollHandDrag
321
- if self.mode == "pan"
322
- else self.viewer.DragMode.NoDrag
323
- )
324
- self.viewer.setDragMode(drag_mode)
325
-
326
- # Update highlights and handles based on the new mode
327
- self._highlight_selected_segments()
328
- if mode_name == "edit":
329
- self._display_edit_handles()
330
- else:
331
- self._clear_edit_handles()
332
-
333
- def _toggle_mode(self, new_mode):
334
- """Toggle between modes."""
335
- if self.mode == new_mode:
336
- self._set_mode(self.previous_mode, is_toggle=True)
337
- else:
338
- if self.mode not in ["selection", "edit"]:
339
- self.previous_mode = self.mode
340
- self._set_mode(new_mode, is_toggle=True)
341
-
342
- # Model management methods
343
- def _browse_models_folder(self):
344
- """Browse for models folder."""
345
- folder_path = QFileDialog.getExistingDirectory(self, "Select Models Folder")
346
- if folder_path:
347
- self.model_manager.set_models_folder(folder_path)
348
- models = self.model_manager.get_available_models(folder_path)
349
- self.control_panel.populate_models(models)
350
- self.viewer.setFocus()
351
-
352
- def _refresh_models_list(self):
353
- """Refresh the models list."""
354
- folder = self.model_manager.get_models_folder()
355
- if folder and os.path.exists(folder):
356
- models = self.model_manager.get_available_models(folder)
357
- self.control_panel.populate_models(models)
358
- self._show_notification("Models list refreshed.")
359
- else:
360
- self._show_notification("No models folder selected.")
361
-
362
- def _load_selected_model(self, model_text):
363
- """Load the selected model."""
364
- if not model_text or model_text == "Default (vit_h)":
365
- self.control_panel.set_current_model("Current: Default SAM Model")
366
- return
367
-
368
- model_path = self.control_panel.model_widget.get_selected_model_path()
369
- if not model_path or not os.path.exists(model_path):
370
- self._show_notification("Selected model file not found.")
371
- return
372
-
373
- self.control_panel.set_current_model("Loading model...")
374
- QApplication.processEvents()
375
-
376
- try:
377
- success = self.model_manager.load_custom_model(model_path)
378
- if success:
379
- # Re-enable SAM functionality if model loaded successfully
380
- self._enable_sam_functionality(True)
381
- if self.model_manager.sam_model:
382
- self.control_panel.set_device_text(
383
- str(self.model_manager.sam_model.device)
384
- )
385
- else:
386
- self.control_panel.set_current_model("Current: Default SAM Model")
387
- self._show_notification("Failed to load selected model. Using default.")
388
- self.control_panel.model_widget.reset_to_default()
389
- self._enable_sam_functionality(False)
390
- except Exception as e:
391
- self.control_panel.set_current_model("Current: Default SAM Model")
392
- self._show_notification(f"Error loading model: {str(e)}")
393
- self.control_panel.model_widget.reset_to_default()
394
- self._enable_sam_functionality(False)
395
-
396
- # Adjustment methods
397
- def _set_annotation_size(self, value):
398
- """Set annotation size."""
399
- multiplier = value / 10.0
400
- self.point_radius = self.settings.point_radius * multiplier
401
- self.line_thickness = self.settings.line_thickness * multiplier
402
- self.settings.annotation_size_multiplier = multiplier
403
- # Update display (implementation would go here)
404
-
405
- def _set_pan_speed(self, value):
406
- """Set pan speed."""
407
- self.pan_multiplier = value / 10.0
408
- self.settings.pan_multiplier = self.pan_multiplier
409
-
410
- def _set_join_threshold(self, value):
411
- """Set polygon join threshold."""
412
- self.polygon_join_threshold = value
413
- self.settings.polygon_join_threshold = value
414
-
415
- # File management methods
416
- def _open_folder_dialog(self):
417
- """Open folder dialog for images."""
418
- folder_path = QFileDialog.getExistingDirectory(self, "Select Image Folder")
419
- if folder_path:
420
- self.right_panel.set_folder(folder_path, self.file_model)
421
- self.viewer.setFocus()
422
-
423
- def _load_selected_image(self, index):
424
- """Load the selected image."""
425
- if not index.isValid() or not self.file_model.isDir(index.parent()):
426
- return
427
-
428
- self.current_file_index = index
429
- path = self.file_model.filePath(index)
430
-
431
- if os.path.isfile(path) and self.file_manager.is_image_file(path):
432
- self.current_image_path = path
433
- pixmap = QPixmap(self.current_image_path)
434
- if not pixmap.isNull():
435
- self._reset_state()
436
- self.viewer.set_photo(pixmap)
437
- if self.model_manager.is_model_available():
438
- self.model_manager.sam_model.set_image(self.current_image_path)
439
- self.file_manager.load_class_aliases(self.current_image_path)
440
- self.file_manager.load_existing_mask(self.current_image_path)
441
- self.right_panel.file_tree.setCurrentIndex(index)
442
- self._update_all_lists()
443
- self.viewer.setFocus()
444
-
445
- def _load_next_image(self):
446
- """Load next image in the file list, with auto-save if enabled."""
447
- if not self.current_file_index.isValid():
448
- return
449
- # Auto-save if enabled
450
- if self.control_panel.get_settings().get("auto_save", True):
451
- self._save_output_to_npz()
452
- parent = self.current_file_index.parent()
453
- row = self.current_file_index.row()
454
- # Find next valid image file
455
- for next_row in range(row + 1, self.file_model.rowCount(parent)):
456
- next_index = self.file_model.index(next_row, 0, parent)
457
- path = self.file_model.filePath(next_index)
458
- if os.path.isfile(path) and self.file_manager.is_image_file(path):
459
- self._load_selected_image(next_index)
460
- return
461
-
462
- def _load_previous_image(self):
463
- """Load previous image in the file list, with auto-save if enabled."""
464
- if not self.current_file_index.isValid():
465
- return
466
- # Auto-save if enabled
467
- if self.control_panel.get_settings().get("auto_save", True):
468
- self._save_output_to_npz()
469
- parent = self.current_file_index.parent()
470
- row = self.current_file_index.row()
471
- # Find previous valid image file
472
- for prev_row in range(row - 1, -1, -1):
473
- prev_index = self.file_model.index(prev_row, 0, parent)
474
- path = self.file_model.filePath(prev_index)
475
- if os.path.isfile(path) and self.file_manager.is_image_file(path):
476
- self._load_selected_image(prev_index)
477
- return
478
-
479
- # Segment management methods
480
- def _assign_selected_to_class(self):
481
- """Assign selected segments to class."""
482
- selected_indices = self.right_panel.get_selected_segment_indices()
483
- self.segment_manager.assign_segments_to_class(selected_indices)
484
- self._update_all_lists()
485
-
486
- def _delete_selected_segments(self):
487
- """Delete selected segments and remove any highlight overlays."""
488
- # Remove highlight overlays before deleting segments
489
- if hasattr(self, "highlight_items"):
490
- for item in self.highlight_items:
491
- self.viewer.scene().removeItem(item)
492
- self.highlight_items = []
493
- selected_indices = self.right_panel.get_selected_segment_indices()
494
- self.segment_manager.delete_segments(selected_indices)
495
- self._update_all_lists()
496
-
497
- def _highlight_selected_segments(self):
498
- """Highlight selected segments. In edit mode, use a brighter hover-like effect."""
499
- # Remove previous highlight overlays
500
- if hasattr(self, "highlight_items"):
501
- for item in self.highlight_items:
502
- if item.scene():
503
- self.viewer.scene().removeItem(item)
504
- self.highlight_items = []
505
-
506
- selected_indices = self.right_panel.get_selected_segment_indices()
507
- if not selected_indices:
508
- return
509
-
510
- for i in selected_indices:
511
- seg = self.segment_manager.segments[i]
512
- base_color = self._get_color_for_class(seg.get("class_id"))
513
-
514
- if self.mode == "edit":
515
- # Use a brighter, hover-like highlight in edit mode
516
- highlight_brush = QBrush(
517
- QColor(base_color.red(), base_color.green(), base_color.blue(), 170)
518
- )
519
- else:
520
- # Use the standard yellow overlay for selection
521
- highlight_brush = QBrush(QColor(255, 255, 0, 180))
522
-
523
- if seg["type"] == "Polygon" and seg.get("vertices"):
524
- poly_item = QGraphicsPolygonItem(QPolygonF(seg["vertices"]))
525
- poly_item.setBrush(highlight_brush)
526
- poly_item.setPen(QPen(Qt.GlobalColor.transparent))
527
- poly_item.setZValue(99)
528
- self.viewer.scene().addItem(poly_item)
529
- self.highlight_items.append(poly_item)
530
- elif seg.get("mask") is not None:
531
- # For non-polygon types, we still use the mask-to-pixmap approach.
532
- # If in edit mode, we could consider skipping non-polygons.
533
- if self.mode != "edit":
534
- mask = seg.get("mask")
535
- pixmap = mask_to_pixmap(mask, (255, 255, 0), alpha=180)
536
- highlight_item = self.viewer.scene().addPixmap(pixmap)
537
- highlight_item.setZValue(100)
538
- self.highlight_items.append(highlight_item)
539
-
540
- def _handle_alias_change(self, class_id, alias):
541
- """Handle class alias change."""
542
- self.segment_manager.set_class_alias(class_id, alias)
543
- self._update_all_lists()
544
-
545
- def _reassign_class_ids(self):
546
- """Reassign class IDs."""
547
- new_order = self.right_panel.get_class_order()
548
- self.segment_manager.reassign_class_ids(new_order)
549
- self._update_all_lists()
550
-
551
- def _update_segment_table(self):
552
- """Update segment table."""
553
- table = self.right_panel.segment_table
554
- table.blockSignals(True)
555
- selected_indices = self.right_panel.get_selected_segment_indices()
556
- table.clearContents()
557
- table.setRowCount(0)
558
-
559
- # Get current filter
560
- filter_text = self.right_panel.class_filter_combo.currentText()
561
- show_all = filter_text == "All Classes"
562
- filter_class_id = -1
563
- if not show_all:
564
- try:
565
- # Parse format like "Alias: ID" or "Class ID"
566
- if ":" in filter_text:
567
- filter_class_id = int(filter_text.split(":")[-1].strip())
568
- else:
569
- filter_class_id = int(filter_text.split()[-1])
570
- except (ValueError, IndexError):
571
- show_all = True # If parsing fails, show all
572
-
573
- # Filter segments based on class filter
574
- display_segments = []
575
- for i, seg in enumerate(self.segment_manager.segments):
576
- seg_class_id = seg.get("class_id")
577
- should_include = show_all or seg_class_id == filter_class_id
578
- if should_include:
579
- display_segments.append((i, seg))
580
-
581
- table.setRowCount(len(display_segments))
582
-
583
- # Populate table rows
584
- for row, (original_index, seg) in enumerate(display_segments):
585
- class_id = seg.get("class_id")
586
- color = self._get_color_for_class(class_id)
587
- class_id_str = str(class_id) if class_id is not None else "N/A"
588
-
589
- alias_str = "N/A"
590
- if class_id is not None:
591
- alias_str = self.segment_manager.get_class_alias(class_id)
592
-
593
- # Create table items (1-based segment ID for display)
594
- index_item = NumericTableWidgetItem(str(original_index + 1))
595
- class_item = NumericTableWidgetItem(class_id_str)
596
- alias_item = QTableWidgetItem(alias_str)
597
-
598
- # Set items as non-editable
599
- index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
600
- class_item.setFlags(class_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
601
- alias_item.setFlags(alias_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
602
-
603
- # Store original index for selection tracking
604
- index_item.setData(Qt.ItemDataRole.UserRole, original_index)
605
-
606
- # Set items in table
607
- table.setItem(row, 0, index_item)
608
- table.setItem(row, 1, class_item)
609
- table.setItem(row, 2, alias_item)
610
-
611
- # Set background color based on class
612
- for col in range(table.columnCount()):
613
- if table.item(row, col):
614
- table.item(row, col).setBackground(QBrush(color))
615
-
616
- # Restore selection
617
- table.setSortingEnabled(False)
618
- for row in range(table.rowCount()):
619
- item = table.item(row, 0)
620
- if item and item.data(Qt.ItemDataRole.UserRole) in selected_indices:
621
- table.selectRow(row)
622
- table.setSortingEnabled(True)
623
-
624
- table.blockSignals(False)
625
- self.viewer.setFocus()
626
-
627
- def _update_all_lists(self):
628
- """Update all UI lists."""
629
- self._update_class_list()
630
- self._update_segment_table()
631
- self._update_class_filter()
632
- self._display_all_segments()
633
- if self.mode == "edit":
634
- self._display_edit_handles()
635
- else:
636
- self._clear_edit_handles()
637
-
638
- def _update_class_list(self):
639
- """Update the class list in the right panel."""
640
- class_table = self.right_panel.class_table
641
- class_table.blockSignals(True)
642
-
643
- # Get unique class IDs
644
- unique_class_ids = self.segment_manager.get_unique_class_ids()
645
-
646
- class_table.clearContents()
647
- class_table.setRowCount(len(unique_class_ids))
648
-
649
- for row, cid in enumerate(unique_class_ids):
650
- alias_item = QTableWidgetItem(self.segment_manager.get_class_alias(cid))
651
- id_item = QTableWidgetItem(str(cid))
652
- id_item.setFlags(id_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
653
-
654
- color = self._get_color_for_class(cid)
655
- alias_item.setBackground(QBrush(color))
656
- id_item.setBackground(QBrush(color))
657
-
658
- class_table.setItem(row, 0, alias_item)
659
- class_table.setItem(row, 1, id_item)
660
-
661
- class_table.blockSignals(False)
662
-
663
- def _update_class_filter(self):
664
- """Update the class filter combo box."""
665
- combo = self.right_panel.class_filter_combo
666
- current_text = combo.currentText()
667
-
668
- combo.blockSignals(True)
669
- combo.clear()
670
- combo.addItem("All Classes")
671
-
672
- # Add class options
673
- unique_class_ids = self.segment_manager.get_unique_class_ids()
674
- for class_id in unique_class_ids:
675
- alias = self.segment_manager.get_class_alias(class_id)
676
- display_text = f"{alias}: {class_id}" if alias else f"Class {class_id}"
677
- combo.addItem(display_text)
678
-
679
- # Restore selection if possible
680
- index = combo.findText(current_text)
681
- if index >= 0:
682
- combo.setCurrentIndex(index)
683
- else:
684
- combo.setCurrentIndex(0)
685
-
686
- combo.blockSignals(False)
687
-
688
- def _display_all_segments(self):
689
- """Display all segments on the viewer."""
690
- # Clear existing segment items
691
- for i, items in self.segment_items.items():
692
- for item in items:
693
- if item.scene():
694
- self.viewer.scene().removeItem(item)
695
- self.segment_items.clear()
696
- self._clear_edit_handles()
697
-
698
- # Display segments from segment manager
699
- for i, segment in enumerate(self.segment_manager.segments):
700
- self.segment_items[i] = []
701
- class_id = segment.get("class_id")
702
- base_color = self._get_color_for_class(class_id)
703
-
704
- if segment["type"] == "Polygon" and segment.get("vertices"):
705
- poly_item = HoverablePolygonItem(QPolygonF(segment["vertices"]))
706
- default_brush = QBrush(
707
- QColor(base_color.red(), base_color.green(), base_color.blue(), 70)
708
- )
709
- hover_brush = QBrush(
710
- QColor(base_color.red(), base_color.green(), base_color.blue(), 170)
711
- )
712
- poly_item.set_brushes(default_brush, hover_brush)
713
- poly_item.setPen(QPen(Qt.GlobalColor.transparent))
714
- self.viewer.scene().addItem(poly_item)
715
- self.segment_items[i].append(poly_item)
716
- elif segment.get("mask") is not None:
717
- default_pixmap = mask_to_pixmap(
718
- segment["mask"], base_color.getRgb()[:3], alpha=70
719
- )
720
- hover_pixmap = mask_to_pixmap(
721
- segment["mask"], base_color.getRgb()[:3], alpha=170
722
- )
723
- pixmap_item = HoverablePixmapItem()
724
- pixmap_item.set_pixmaps(default_pixmap, hover_pixmap)
725
- self.viewer.scene().addItem(pixmap_item)
726
- pixmap_item.setZValue(i + 1)
727
- self.segment_items[i].append(pixmap_item)
728
-
729
- # Event handlers
730
- def _handle_escape_press(self):
731
- """Handle escape key press."""
732
- self.right_panel.clear_selections()
733
- self.clear_all_points()
734
- self.viewer.setFocus()
735
-
736
- def _handle_space_press(self):
737
- """Handle space key press."""
738
- if self.mode == "polygon" and self.polygon_points:
739
- self._finalize_polygon()
740
- else:
741
- self._save_current_segment()
742
-
743
- def _handle_enter_press(self):
744
- """Handle enter key press."""
745
- if self.mode == "polygon" and self.polygon_points:
746
- self._finalize_polygon()
747
- else:
748
- self._save_output_to_npz()
749
-
750
- def _save_current_segment(self):
751
- """Save current SAM segment."""
752
- if (
753
- self.mode != "sam_points"
754
- or not hasattr(self, "preview_mask_item")
755
- or not self.preview_mask_item
756
- or not self.model_manager.is_model_available()
757
- ):
758
- return
759
-
760
- mask = self.model_manager.sam_model.predict(
761
- self.positive_points, self.negative_points
762
- )
763
- if mask is not None:
764
- self.segment_manager.add_segment(
765
- {
766
- "mask": mask,
767
- "type": "SAM",
768
- "vertices": None,
769
- }
770
- )
771
- self.clear_all_points()
772
- self._update_all_lists()
773
-
774
- def _finalize_polygon(self):
775
- """Finalize polygon drawing."""
776
- if len(self.polygon_points) < 3:
777
- return
778
-
779
- self.segment_manager.add_segment(
780
- {
781
- "vertices": list(self.polygon_points),
782
- "type": "Polygon",
783
- "mask": None,
784
- }
785
- )
786
- self.polygon_points.clear()
787
- self.clear_all_points()
788
- self._update_all_lists()
789
-
790
- def _save_output_to_npz(self):
791
- """Save output to NPZ and TXT files as enabled, and update file list tickboxes/highlight. If no segments, delete associated files."""
792
- if not self.current_image_path:
793
- self._show_notification("No image loaded.")
794
- return
795
-
796
- # If no segments, delete associated files
797
- if not self.segment_manager.segments:
798
- base, _ = os.path.splitext(self.current_image_path)
799
- deleted_files = []
800
- for ext in [".npz", ".txt", ".json"]:
801
- file_path = base + ext
802
- if os.path.exists(file_path):
803
- try:
804
- os.remove(file_path)
805
- deleted_files.append(file_path)
806
- self.file_model.update_cache_for_path(file_path)
807
- except Exception as e:
808
- self._show_notification(f"Error deleting {file_path}: {e}")
809
- if deleted_files:
810
- self._show_notification(
811
- f"Deleted: {', '.join(os.path.basename(f) for f in deleted_files)}"
812
- )
813
- else:
814
- self._show_notification("No segments to save.")
815
- return
816
-
817
- try:
818
- settings = self.control_panel.get_settings()
819
- npz_path = None
820
- txt_path = None
821
- if settings.get("save_npz", True):
822
- h, w = (
823
- self.viewer._pixmap_item.pixmap().height(),
824
- self.viewer._pixmap_item.pixmap().width(),
825
- )
826
- class_order = self.segment_manager.get_unique_class_ids()
827
- if class_order:
828
- npz_path = self.file_manager.save_npz(
829
- self.current_image_path, (h, w), class_order
830
- )
831
- self._show_notification(f"Saved: {os.path.basename(npz_path)}")
832
- else:
833
- self._show_notification("No classes defined for saving.")
834
- if settings.get("save_txt", True):
835
- h, w = (
836
- self.viewer._pixmap_item.pixmap().height(),
837
- self.viewer._pixmap_item.pixmap().width(),
838
- )
839
- class_order = self.segment_manager.get_unique_class_ids()
840
- if settings.get("yolo_use_alias", True):
841
- class_labels = [
842
- self.segment_manager.get_class_alias(cid) for cid in class_order
843
- ]
844
- else:
845
- class_labels = [str(cid) for cid in class_order]
846
- if class_order:
847
- txt_path = self.file_manager.save_yolo_txt(
848
- self.current_image_path, (h, w), class_order, class_labels
849
- )
850
- # Efficiently update file list tickboxes and highlight
851
- for path in [npz_path, txt_path]:
852
- if path:
853
- self.file_model.update_cache_for_path(path)
854
- self.file_model.set_highlighted_path(path)
855
- QTimer.singleShot(
856
- 1500,
857
- lambda p=path: (
858
- self.file_model.set_highlighted_path(None)
859
- if self.file_model.highlighted_path == p
860
- else None
861
- ),
862
- )
863
- except Exception as e:
864
- self._show_notification(f"Error saving: {str(e)}")
865
-
866
- def _handle_merge_press(self):
867
- """Handle merge key press."""
868
- self._assign_selected_to_class()
869
- self.right_panel.clear_selections()
870
-
871
- def _undo_last_action(self):
872
- """Undo last action."""
873
- # Implementation would go here
874
- pass
875
-
876
- def clear_all_points(self):
877
- """Clear all temporary points."""
878
- if hasattr(self, "rubber_band_line") and self.rubber_band_line:
879
- self.viewer.scene().removeItem(self.rubber_band_line)
880
- self.rubber_band_line = None
881
-
882
- self.positive_points.clear()
883
- self.negative_points.clear()
884
-
885
- for item in self.point_items:
886
- self.viewer.scene().removeItem(item)
887
- self.point_items.clear()
888
-
889
- self.polygon_points.clear()
890
- for item in self.polygon_preview_items:
891
- self.viewer.scene().removeItem(item)
892
- self.polygon_preview_items.clear()
893
-
894
- if hasattr(self, "preview_mask_item") and self.preview_mask_item:
895
- self.viewer.scene().removeItem(self.preview_mask_item)
896
- self.preview_mask_item = None
897
-
898
- def _show_notification(self, message, duration=3000):
899
- """Show notification message."""
900
- self.control_panel.show_notification(message)
901
- QTimer.singleShot(duration, self.control_panel.clear_notification)
902
-
903
- def _show_hotkey_dialog(self):
904
- """Show the hotkey configuration dialog."""
905
- dialog = HotkeyDialog(self.hotkey_manager, self)
906
- dialog.exec()
907
- # Update shortcuts after dialog closes
908
- self._update_shortcuts()
909
-
910
- def _handle_zoom_in(self):
911
- """Handle zoom in."""
912
- current_val = self.control_panel.get_annotation_size()
913
- self.control_panel.set_annotation_size(min(current_val + 1, 50))
914
-
915
- def _handle_zoom_out(self):
916
- """Handle zoom out."""
917
- current_val = self.control_panel.get_annotation_size()
918
- self.control_panel.set_annotation_size(max(current_val - 1, 1))
919
-
920
- def _handle_pan_key(self, direction):
921
- """Handle WASD pan keys."""
922
- if not hasattr(self, "viewer"):
923
- return
924
-
925
- amount = int(self.viewer.height() * 0.1 * self.pan_multiplier)
926
-
927
- if direction == "up":
928
- self.viewer.verticalScrollBar().setValue(
929
- self.viewer.verticalScrollBar().value() - amount
930
- )
931
- elif direction == "down":
932
- self.viewer.verticalScrollBar().setValue(
933
- self.viewer.verticalScrollBar().value() + amount
934
- )
935
- elif direction == "left":
936
- amount = int(self.viewer.width() * 0.1 * self.pan_multiplier)
937
- self.viewer.horizontalScrollBar().setValue(
938
- self.viewer.horizontalScrollBar().value() - amount
939
- )
940
- elif direction == "right":
941
- amount = int(self.viewer.width() * 0.1 * self.pan_multiplier)
942
- self.viewer.horizontalScrollBar().setValue(
943
- self.viewer.horizontalScrollBar().value() + amount
944
- )
945
-
946
- def closeEvent(self, event):
947
- """Handle application close."""
948
- # Save settings
949
- self.settings.save_to_file(str(self.paths.settings_file))
950
- super().closeEvent(event)
951
-
952
- def _reset_state(self):
953
- """Reset application state."""
954
- self.clear_all_points()
955
- self.segment_manager.clear()
956
- self._update_all_lists()
957
- items_to_remove = [
958
- item
959
- for item in self.viewer.scene().items()
960
- if item is not self.viewer._pixmap_item
961
- ]
962
- for item in items_to_remove:
963
- self.viewer.scene().removeItem(item)
964
- self.segment_items.clear()
965
- self.highlight_items.clear()
966
-
967
- def _scene_mouse_press(self, event):
968
- """Handle mouse press events in the scene."""
969
- # Map scene coordinates to the view so items() works correctly.
970
- view_pos = self.viewer.mapFromScene(event.scenePos())
971
- items_at_pos = self.viewer.items(view_pos)
972
- is_handle_click = any(
973
- isinstance(item, EditableVertexItem) for item in items_at_pos
974
- )
975
-
976
- # Allow vertex handles to process their own mouse events.
977
- if is_handle_click:
978
- self._original_mouse_press(event)
979
- return
980
-
981
- if self.mode == "edit" and event.button() == Qt.MouseButton.LeftButton:
982
- pos = event.scenePos()
983
- if self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint()):
984
- self.is_dragging_polygon = True
985
- self.drag_start_pos = pos
986
- selected_indices = self.right_panel.get_selected_segment_indices()
987
- self.drag_initial_vertices = {
988
- i: list(self.segment_manager.segments[i]["vertices"])
989
- for i in selected_indices
990
- if self.segment_manager.segments[i].get("type") == "Polygon"
991
- }
992
- event.accept()
993
- return
994
-
995
- # Call the original scene handler.
996
- self._original_mouse_press(event)
997
- # Skip further processing unless we're in selection mode.
998
- if event.isAccepted() and self.mode != "selection":
999
- return
1000
-
1001
- if self.is_dragging_polygon:
1002
- return
1003
-
1004
- pos = event.scenePos()
1005
- if (
1006
- self.viewer._pixmap_item.pixmap().isNull()
1007
- or not self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint())
1008
- ):
1009
- return
1010
-
1011
- if self.mode == "pan":
1012
- self.viewer.set_cursor(Qt.CursorShape.ClosedHandCursor)
1013
- elif self.mode == "sam_points":
1014
- if event.button() == Qt.MouseButton.LeftButton:
1015
- self._add_point(pos, positive=True)
1016
- self._update_segmentation()
1017
- elif event.button() == Qt.MouseButton.RightButton:
1018
- self._add_point(pos, positive=False)
1019
- self._update_segmentation()
1020
- elif self.mode == "polygon":
1021
- if event.button() == Qt.MouseButton.LeftButton:
1022
- self._handle_polygon_click(pos)
1023
- elif self.mode == "selection":
1024
- if event.button() == Qt.MouseButton.LeftButton:
1025
- self._handle_segment_selection_click(pos)
1026
-
1027
- def _scene_mouse_move(self, event):
1028
- """Handle mouse move events in the scene."""
1029
- if self.mode == "edit" and self.is_dragging_polygon:
1030
- delta = event.scenePos() - self.drag_start_pos
1031
- for i, initial_verts in self.drag_initial_vertices.items():
1032
- self.segment_manager.segments[i]["vertices"] = [
1033
- QPointF(v) + delta for v in initial_verts
1034
- ]
1035
- self._update_polygon_item(i)
1036
- self._display_edit_handles() # Redraw handles at new positions
1037
- self._highlight_selected_segments() # Redraw highlight at new position
1038
- event.accept()
1039
- return
1040
-
1041
- self._original_mouse_move(event)
1042
-
1043
- def _scene_mouse_release(self, event):
1044
- """Handle mouse release events in the scene."""
1045
- if self.mode == "edit" and self.is_dragging_polygon:
1046
- self.is_dragging_polygon = False
1047
- self.drag_initial_vertices.clear()
1048
- event.accept()
1049
- return
1050
-
1051
- if self.mode == "pan":
1052
- self.viewer.set_cursor(Qt.CursorShape.OpenHandCursor)
1053
- self._original_mouse_release(event)
1054
-
1055
- def _add_point(self, pos, positive):
1056
- """Add a point for SAM segmentation."""
1057
- point_list = self.positive_points if positive else self.negative_points
1058
- point_list.append([int(pos.x()), int(pos.y())])
1059
-
1060
- point_color = (
1061
- QColor(Qt.GlobalColor.green) if positive else QColor(Qt.GlobalColor.red)
1062
- )
1063
- point_color.setAlpha(150)
1064
- point_diameter = self.point_radius * 2
1065
-
1066
- point_item = QGraphicsEllipseItem(
1067
- pos.x() - self.point_radius,
1068
- pos.y() - self.point_radius,
1069
- point_diameter,
1070
- point_diameter,
1071
- )
1072
- point_item.setBrush(QBrush(point_color))
1073
- point_item.setPen(QPen(Qt.GlobalColor.transparent))
1074
- self.viewer.scene().addItem(point_item)
1075
- self.point_items.append(point_item)
1076
-
1077
- def _update_segmentation(self):
1078
- """Update SAM segmentation preview."""
1079
- if hasattr(self, "preview_mask_item") and self.preview_mask_item:
1080
- self.viewer.scene().removeItem(self.preview_mask_item)
1081
- if not self.positive_points or not self.model_manager.is_model_available():
1082
- return
1083
-
1084
- mask = self.model_manager.sam_model.predict(
1085
- self.positive_points, self.negative_points
1086
- )
1087
- if mask is not None:
1088
- pixmap = mask_to_pixmap(mask, (255, 255, 0))
1089
- self.preview_mask_item = self.viewer.scene().addPixmap(pixmap)
1090
- self.preview_mask_item.setZValue(50)
1091
-
1092
- def _handle_polygon_click(self, pos):
1093
- """Handle polygon drawing clicks."""
1094
- # Check if clicking near the first point to close polygon
1095
- if self.polygon_points and len(self.polygon_points) > 2:
1096
- first_point = self.polygon_points[0]
1097
- distance_squared = (pos.x() - first_point.x()) ** 2 + (
1098
- pos.y() - first_point.y()
1099
- ) ** 2
1100
- if distance_squared < self.polygon_join_threshold**2:
1101
- self._finalize_polygon()
1102
- return
1103
-
1104
- # Add new point to polygon
1105
- self.polygon_points.append(pos)
1106
-
1107
- # Create visual point
1108
- point_diameter = self.point_radius * 2
1109
- point_color = QColor(Qt.GlobalColor.blue)
1110
- point_color.setAlpha(150)
1111
- dot = QGraphicsEllipseItem(
1112
- pos.x() - self.point_radius,
1113
- pos.y() - self.point_radius,
1114
- point_diameter,
1115
- point_diameter,
1116
- )
1117
- dot.setBrush(QBrush(point_color))
1118
- dot.setPen(QPen(Qt.GlobalColor.transparent))
1119
- self.viewer.scene().addItem(dot)
1120
- self.polygon_preview_items.append(dot)
1121
-
1122
- # Update polygon preview
1123
- self._draw_polygon_preview()
1124
-
1125
- def _draw_polygon_preview(self):
1126
- """Draw polygon preview lines and fill."""
1127
- # Remove old preview lines and polygons (keep dots)
1128
- for item in self.polygon_preview_items[:]:
1129
- if not isinstance(item, QGraphicsEllipseItem):
1130
- if item.scene():
1131
- self.viewer.scene().removeItem(item)
1132
- self.polygon_preview_items.remove(item)
1133
-
1134
- if len(self.polygon_points) > 2:
1135
- # Create preview polygon fill
1136
- preview_poly = QGraphicsPolygonItem(QPolygonF(self.polygon_points))
1137
- preview_poly.setBrush(QBrush(QColor(0, 255, 255, 100)))
1138
- preview_poly.setPen(QPen(Qt.GlobalColor.transparent))
1139
- self.viewer.scene().addItem(preview_poly)
1140
- self.polygon_preview_items.append(preview_poly)
1141
-
1142
- if len(self.polygon_points) > 1:
1143
- # Create preview lines between points
1144
- line_color = QColor(Qt.GlobalColor.cyan)
1145
- line_color.setAlpha(150)
1146
- for i in range(len(self.polygon_points) - 1):
1147
- line = QGraphicsLineItem(
1148
- self.polygon_points[i].x(),
1149
- self.polygon_points[i].y(),
1150
- self.polygon_points[i + 1].x(),
1151
- self.polygon_points[i + 1].y(),
1152
- )
1153
- line.setPen(QPen(line_color, self.line_thickness))
1154
- self.viewer.scene().addItem(line)
1155
- self.polygon_preview_items.append(line)
1156
-
1157
- def _handle_segment_selection_click(self, pos):
1158
- """Handle segment selection clicks (toggle behavior)."""
1159
- x, y = int(pos.x()), int(pos.y())
1160
- for i in range(len(self.segment_manager.segments) - 1, -1, -1):
1161
- seg = self.segment_manager.segments[i]
1162
- # Determine mask for hit-testing
1163
- if seg["type"] == "Polygon" and seg.get("vertices"):
1164
- # Rasterize polygon
1165
- if self.viewer._pixmap_item.pixmap().isNull():
1166
- continue
1167
- h = self.viewer._pixmap_item.pixmap().height()
1168
- w = self.viewer._pixmap_item.pixmap().width()
1169
- points_np = np.array(
1170
- [[p.x(), p.y()] for p in seg["vertices"]], dtype=np.int32
1171
- )
1172
- # Ensure points are within bounds
1173
- points_np = np.clip(points_np, 0, [w - 1, h - 1])
1174
- mask = np.zeros((h, w), dtype=np.uint8)
1175
- cv2.fillPoly(mask, [points_np], 1)
1176
- mask = mask.astype(bool)
1177
- else:
1178
- mask = seg.get("mask")
1179
- if (
1180
- mask is not None
1181
- and y < mask.shape[0]
1182
- and x < mask.shape[1]
1183
- and mask[y, x]
1184
- ):
1185
- # Find the corresponding row in the segment table and toggle selection
1186
- table = self.right_panel.segment_table
1187
- for j in range(table.rowCount()):
1188
- item = table.item(j, 0)
1189
- if item and item.data(Qt.ItemDataRole.UserRole) == i:
1190
- # Toggle selection for this row using the original working method
1191
- is_selected = table.item(j, 0).isSelected()
1192
- range_to_select = QTableWidgetSelectionRange(
1193
- j, 0, j, table.columnCount() - 1
1194
- )
1195
- table.setRangeSelected(range_to_select, not is_selected)
1196
- self._highlight_selected_segments()
1197
- return
1198
- self.viewer.setFocus()
1199
-
1200
- def _get_color_for_class(self, class_id):
1201
- """Get color for a class ID."""
1202
- if class_id is None:
1203
- return QColor.fromHsv(0, 0, 128)
1204
- hue = int((class_id * 222.4922359) % 360)
1205
- color = QColor.fromHsv(hue, 220, 220)
1206
- if not color.isValid():
1207
- return QColor(Qt.GlobalColor.white)
1208
- return color
1209
-
1210
- def _display_edit_handles(self):
1211
- """Display draggable vertex handles for selected polygons in edit mode."""
1212
- self._clear_edit_handles()
1213
- if self.mode != "edit":
1214
- return
1215
- selected_indices = self.right_panel.get_selected_segment_indices()
1216
- handle_radius = self.point_radius
1217
- handle_diam = handle_radius * 2
1218
- for seg_idx in selected_indices:
1219
- seg = self.segment_manager.segments[seg_idx]
1220
- if seg["type"] == "Polygon" and seg.get("vertices"):
1221
- for v_idx, pt in enumerate(seg["vertices"]):
1222
- handle = EditableVertexItem(
1223
- self,
1224
- seg_idx,
1225
- v_idx,
1226
- -handle_radius,
1227
- -handle_radius,
1228
- handle_diam,
1229
- handle_diam,
1230
- )
1231
- handle.setPos(pt) # Use setPos to handle zoom correctly
1232
- handle.setZValue(200) # Ensure handles are on top
1233
- # Make sure the handle can receive mouse events
1234
- handle.setAcceptHoverEvents(True)
1235
- self.viewer.scene().addItem(handle)
1236
- self.edit_handles.append(handle)
1237
-
1238
- def _clear_edit_handles(self):
1239
- """Remove all editable vertex handles from the scene."""
1240
- if hasattr(self, "edit_handles"):
1241
- for h in self.edit_handles:
1242
- if h.scene():
1243
- self.viewer.scene().removeItem(h)
1244
- self.edit_handles = []
1245
-
1246
- def update_vertex_pos(self, segment_index, vertex_index, new_pos):
1247
- """Update the position of a vertex in a polygon segment."""
1248
- seg = self.segment_manager.segments[segment_index]
1249
- if seg.get("type") == "Polygon":
1250
- seg["vertices"][
1251
- vertex_index
1252
- ] = new_pos # new_pos is already the correct scene coordinate
1253
- self._update_polygon_item(segment_index)
1254
- self._highlight_selected_segments() # Keep the highlight in sync with the new shape
1255
-
1256
- def _update_polygon_item(self, segment_index):
1257
- """Efficiently update the visual polygon item for a given segment."""
1258
- items = self.segment_items.get(segment_index, [])
1259
- for item in items:
1260
- if isinstance(item, HoverablePolygonItem):
1261
- item.setPolygon(
1262
- QPolygonF(self.segment_manager.segments[segment_index]["vertices"])
1263
- )
1264
- return
1
+ """Main application window."""
2
+
3
+ import os
4
+ import numpy as np
5
+ import cv2
6
+ from PyQt6.QtWidgets import (
7
+ QMainWindow,
8
+ QWidget,
9
+ QHBoxLayout,
10
+ QVBoxLayout,
11
+ QFileDialog,
12
+ QApplication,
13
+ QGraphicsEllipseItem,
14
+ QGraphicsLineItem,
15
+ QGraphicsPolygonItem,
16
+ QTableWidgetItem,
17
+ QTableWidgetSelectionRange,
18
+ QHeaderView,
19
+ QSplitter,
20
+ QDialog,
21
+ )
22
+ from PyQt6.QtGui import (
23
+ QIcon,
24
+ QKeySequence,
25
+ QShortcut,
26
+ QPixmap,
27
+ QColor,
28
+ QPen,
29
+ QBrush,
30
+ QPolygonF,
31
+ QImage,
32
+ )
33
+ from PyQt6.QtCore import Qt, QTimer, QModelIndex, QPointF, pyqtSignal
34
+
35
+ from .control_panel import ControlPanel
36
+ from .right_panel import RightPanel
37
+ from .photo_viewer import PhotoViewer
38
+ from .hoverable_polygon_item import HoverablePolygonItem
39
+ from .hoverable_pixelmap_item import HoverablePixmapItem
40
+ from .editable_vertex import EditableVertexItem
41
+ from .numeric_table_widget_item import NumericTableWidgetItem
42
+ from ..core import SegmentManager, ModelManager, FileManager
43
+ from ..config import Settings, Paths, HotkeyManager
44
+ from ..utils import CustomFileSystemModel, mask_to_pixmap
45
+ from .hotkey_dialog import HotkeyDialog
46
+ from .widgets import StatusBar
47
+
48
+
49
+ class PanelPopoutWindow(QDialog):
50
+ """Pop-out window for draggable panels."""
51
+
52
+ panel_closed = pyqtSignal(QWidget) # Signal emitted when panel window is closed
53
+
54
+ def __init__(self, panel_widget, title="Panel", parent=None):
55
+ super().__init__(parent)
56
+ self.panel_widget = panel_widget
57
+ self.setWindowTitle(title)
58
+ self.setWindowFlags(Qt.WindowType.Window) # Allow moving to other monitors
59
+
60
+ # Make window resizable
61
+ self.setMinimumSize(200, 300)
62
+ self.resize(400, 600)
63
+
64
+ # Set up layout
65
+ layout = QVBoxLayout(self)
66
+ layout.setContentsMargins(5, 5, 5, 5)
67
+ layout.addWidget(panel_widget)
68
+
69
+ # Store original parent for restoration
70
+ self.original_parent = parent
71
+
72
+ def closeEvent(self, event):
73
+ """Handle window close - emit signal to return panel to main window."""
74
+ self.panel_closed.emit(self.panel_widget)
75
+ super().closeEvent(event)
76
+
77
+
78
+ class MainWindow(QMainWindow):
79
+ """Main application window."""
80
+
81
+ def __init__(self):
82
+ super().__init__()
83
+
84
+ print("[3/20] Starting LazyLabel...")
85
+ print("[4/20] Loading configuration and settings...")
86
+
87
+ # Initialize configuration
88
+ self.paths = Paths()
89
+ self.settings = Settings.load_from_file(str(self.paths.settings_file))
90
+ self.hotkey_manager = HotkeyManager(str(self.paths.config_dir))
91
+
92
+ print("[5/20] Initializing core managers...")
93
+
94
+ # Initialize managers
95
+ self.segment_manager = SegmentManager()
96
+ self.model_manager = ModelManager(self.paths)
97
+ self.file_manager = FileManager(self.segment_manager)
98
+
99
+ print("[6/20] Setting up user interface...")
100
+
101
+ # Initialize UI state
102
+ self.mode = "sam_points"
103
+ self.previous_mode = "sam_points"
104
+ self.current_image_path = None
105
+ self.current_file_index = QModelIndex()
106
+
107
+ # Panel pop-out state
108
+ self.left_panel_popout = None
109
+ self.right_panel_popout = None
110
+
111
+ # Annotation state
112
+ self.point_radius = self.settings.point_radius
113
+ self.line_thickness = self.settings.line_thickness
114
+ self.pan_multiplier = self.settings.pan_multiplier
115
+ self.polygon_join_threshold = self.settings.polygon_join_threshold
116
+
117
+ # Drawing state
118
+ self.point_items, self.positive_points, self.negative_points = [], [], []
119
+ self.polygon_points, self.polygon_preview_items = [], []
120
+ self.rubber_band_line = None
121
+ self.preview_mask_item = None
122
+ self.segments, self.segment_items, self.highlight_items = [], {}, []
123
+ self.edit_handles = []
124
+ self.is_dragging_polygon, self.drag_start_pos, self.drag_initial_vertices = (
125
+ False,
126
+ None,
127
+ {},
128
+ )
129
+ self.action_history = []
130
+
131
+ # Update state flags to prevent recursion
132
+ self._updating_lists = False
133
+
134
+ self._setup_ui()
135
+ self._setup_model()
136
+
137
+ print("[17/20] Connecting UI signals and shortcuts...")
138
+ self._setup_connections()
139
+ self._setup_shortcuts()
140
+ self._load_settings()
141
+
142
+ print("[18/20] LazyLabel initialization complete!")
143
+
144
+ def _setup_ui(self):
145
+ """Setup the user interface."""
146
+ self.setWindowTitle("LazyLabel by DNC")
147
+ self.setGeometry(
148
+ 50, 50, self.settings.window_width, self.settings.window_height
149
+ )
150
+
151
+ # Set window icon
152
+ if self.paths.logo_path.exists():
153
+ self.setWindowIcon(QIcon(str(self.paths.logo_path)))
154
+
155
+ # Create panels
156
+ self.control_panel = ControlPanel()
157
+ self.right_panel = RightPanel()
158
+ self.viewer = PhotoViewer(self)
159
+ self.viewer.setMouseTracking(True)
160
+
161
+ # Setup file model
162
+ self.file_model = CustomFileSystemModel()
163
+ self.right_panel.setup_file_model(self.file_model)
164
+
165
+ # Create status bar
166
+ self.status_bar = StatusBar()
167
+ self.setStatusBar(self.status_bar)
168
+
169
+ # Create horizontal splitter for main panels
170
+ self.main_splitter = QSplitter(Qt.Orientation.Horizontal)
171
+ self.main_splitter.addWidget(self.control_panel)
172
+ self.main_splitter.addWidget(self.viewer)
173
+ self.main_splitter.addWidget(self.right_panel)
174
+
175
+ # Set minimum sizes for panels to prevent shrinking below preferred width
176
+ self.control_panel.setMinimumWidth(self.control_panel.preferred_width)
177
+ self.right_panel.setMinimumWidth(self.right_panel.preferred_width)
178
+
179
+ # Set splitter sizes - give most space to viewer
180
+ self.main_splitter.setSizes([250, 800, 350])
181
+ self.main_splitter.setStretchFactor(0, 0) # Control panel doesn't stretch
182
+ self.main_splitter.setStretchFactor(1, 1) # Viewer stretches
183
+ self.main_splitter.setStretchFactor(2, 0) # Right panel doesn't stretch
184
+
185
+ # Set splitter child sizes policy
186
+ self.main_splitter.setChildrenCollapsible(True)
187
+
188
+ # Connect splitter signals for intelligent expand/collapse
189
+ self.main_splitter.splitterMoved.connect(self._handle_splitter_moved)
190
+
191
+ # Main vertical layout to accommodate status bar
192
+ main_layout = QVBoxLayout()
193
+ main_layout.setContentsMargins(0, 0, 0, 0)
194
+ main_layout.setSpacing(0)
195
+ main_layout.addWidget(self.main_splitter, 1)
196
+
197
+ central_widget = QWidget()
198
+ central_widget.setLayout(main_layout)
199
+ self.setCentralWidget(central_widget)
200
+
201
+ def _setup_model(self):
202
+ """Setup the SAM model."""
203
+ print("[7/20] Initializing SAM model (this may take a moment)...")
204
+
205
+ sam_model = self.model_manager.initialize_default_model(
206
+ self.settings.default_model_type
207
+ )
208
+
209
+ if sam_model and sam_model.is_loaded:
210
+ device_text = str(sam_model.device).upper()
211
+ print(f"[14/20] SAM model loaded successfully on {device_text}")
212
+ self.status_bar.set_permanent_message(f"Device: {device_text}")
213
+ self._enable_sam_functionality(True)
214
+ elif sam_model is None:
215
+ print(
216
+ "[14/20] SAM model initialization failed. Point mode will be disabled."
217
+ )
218
+ self.status_bar.set_permanent_message("Model initialization failed")
219
+ self._enable_sam_functionality(False)
220
+ else:
221
+ print("[14/20] SAM model failed to load. Point mode will be disabled.")
222
+ self.status_bar.set_permanent_message("Model loading failed")
223
+ self._enable_sam_functionality(False)
224
+
225
+ print("[15/20] Scanning available models...")
226
+
227
+ # Setup model change callback
228
+ self.model_manager.on_model_changed = self.control_panel.set_current_model
229
+
230
+ # Initialize models list
231
+ models = self.model_manager.get_available_models(str(self.paths.models_dir))
232
+ self.control_panel.populate_models(models)
233
+
234
+ if models:
235
+ print(f"[16/20] Found {len(models)} model(s) in models directory")
236
+
237
+ def _enable_sam_functionality(self, enabled: bool):
238
+ """Enable or disable SAM point functionality."""
239
+ self.control_panel.set_sam_mode_enabled(enabled)
240
+ if not enabled and self.mode == "sam_points":
241
+ # Switch to polygon mode if SAM is disabled and we're in SAM mode
242
+ self.set_polygon_mode()
243
+
244
+ def _setup_connections(self):
245
+ """Setup signal connections."""
246
+ # Control panel connections
247
+ self.control_panel.sam_mode_requested.connect(self.set_sam_mode)
248
+ self.control_panel.polygon_mode_requested.connect(self.set_polygon_mode)
249
+ self.control_panel.selection_mode_requested.connect(self.toggle_selection_mode)
250
+ self.control_panel.clear_points_requested.connect(self.clear_all_points)
251
+ self.control_panel.fit_view_requested.connect(self.viewer.fitInView)
252
+ self.control_panel.hotkeys_requested.connect(self._show_hotkey_dialog)
253
+
254
+ # Model management
255
+ self.control_panel.browse_models_requested.connect(self._browse_models_folder)
256
+ self.control_panel.refresh_models_requested.connect(self._refresh_models_list)
257
+ self.control_panel.model_selected.connect(self._load_selected_model)
258
+
259
+ # Adjustments
260
+ self.control_panel.annotation_size_changed.connect(self._set_annotation_size)
261
+ self.control_panel.pan_speed_changed.connect(self._set_pan_speed)
262
+ self.control_panel.join_threshold_changed.connect(self._set_join_threshold)
263
+
264
+ # Right panel connections
265
+ self.right_panel.open_folder_requested.connect(self._open_folder_dialog)
266
+ self.right_panel.image_selected.connect(self._load_selected_image)
267
+ self.right_panel.merge_selection_requested.connect(
268
+ self._assign_selected_to_class
269
+ )
270
+ self.right_panel.delete_selection_requested.connect(
271
+ self._delete_selected_segments
272
+ )
273
+ self.right_panel.segments_selection_changed.connect(
274
+ self._highlight_selected_segments
275
+ )
276
+ self.right_panel.class_alias_changed.connect(self._handle_alias_change)
277
+ self.right_panel.reassign_classes_requested.connect(self._reassign_class_ids)
278
+ self.right_panel.class_filter_changed.connect(self._update_segment_table)
279
+ self.right_panel.class_toggled.connect(self._handle_class_toggle)
280
+
281
+ # Panel pop-out functionality
282
+ self.control_panel.pop_out_requested.connect(self._pop_out_left_panel)
283
+ self.right_panel.pop_out_requested.connect(self._pop_out_right_panel)
284
+
285
+ # Mouse events (will be implemented in a separate handler)
286
+ self._setup_mouse_events()
287
+
288
+ def _setup_shortcuts(self):
289
+ """Setup keyboard shortcuts based on hotkey manager."""
290
+ self.shortcuts = [] # Keep track of shortcuts for updating
291
+ self._update_shortcuts()
292
+
293
+ def _update_shortcuts(self):
294
+ """Update shortcuts based on current hotkey configuration."""
295
+ # Clear existing shortcuts
296
+ for shortcut in self.shortcuts:
297
+ shortcut.setParent(None)
298
+ self.shortcuts.clear()
299
+
300
+ # Map action names to callbacks
301
+ action_callbacks = {
302
+ "load_next_image": self._load_next_image,
303
+ "load_previous_image": self._load_previous_image,
304
+ "sam_mode": self.set_sam_mode,
305
+ "polygon_mode": self.set_polygon_mode,
306
+ "selection_mode": self.toggle_selection_mode,
307
+ "pan_mode": self.toggle_pan_mode,
308
+ "edit_mode": self.toggle_edit_mode,
309
+ "clear_points": self.clear_all_points,
310
+ "escape": self._handle_escape_press,
311
+ "delete_segments": self._delete_selected_segments,
312
+ "delete_segments_alt": self._delete_selected_segments,
313
+ "merge_segments": self._handle_merge_press,
314
+ "undo": self._undo_last_action,
315
+ "select_all": lambda: self.right_panel.select_all_segments(),
316
+ "save_segment": self._handle_space_press,
317
+ "save_output": self._handle_enter_press,
318
+ "save_output_alt": self._handle_enter_press,
319
+ "fit_view": self.viewer.fitInView,
320
+ "zoom_in": self._handle_zoom_in,
321
+ "zoom_out": self._handle_zoom_out,
322
+ "pan_up": lambda: self._handle_pan_key("up"),
323
+ "pan_down": lambda: self._handle_pan_key("down"),
324
+ "pan_left": lambda: self._handle_pan_key("left"),
325
+ "pan_right": lambda: self._handle_pan_key("right"),
326
+ }
327
+
328
+ # Create shortcuts for each action
329
+ for action_name, callback in action_callbacks.items():
330
+ primary_key, secondary_key = self.hotkey_manager.get_key_for_action(
331
+ action_name
332
+ )
333
+
334
+ # Create primary shortcut
335
+ if primary_key:
336
+ shortcut = QShortcut(QKeySequence(primary_key), self, callback)
337
+ self.shortcuts.append(shortcut)
338
+
339
+ # Create secondary shortcut
340
+ if secondary_key:
341
+ shortcut = QShortcut(QKeySequence(secondary_key), self, callback)
342
+ self.shortcuts.append(shortcut)
343
+
344
+ def _load_settings(self):
345
+ """Load and apply settings."""
346
+ self.control_panel.set_settings(self.settings.__dict__)
347
+ self.control_panel.set_annotation_size(
348
+ int(self.settings.annotation_size_multiplier * 10)
349
+ )
350
+ self.control_panel.set_join_threshold(self.settings.polygon_join_threshold)
351
+ # Set initial mode based on model availability
352
+ if self.model_manager.is_model_available():
353
+ self.set_sam_mode()
354
+ else:
355
+ self.set_polygon_mode()
356
+
357
+ def _setup_mouse_events(self):
358
+ """Setup mouse event handling."""
359
+ self._original_mouse_press = self.viewer.scene().mousePressEvent
360
+ self._original_mouse_move = self.viewer.scene().mouseMoveEvent
361
+ self._original_mouse_release = self.viewer.scene().mouseReleaseEvent
362
+
363
+ self.viewer.scene().mousePressEvent = self._scene_mouse_press
364
+ self.viewer.scene().mouseMoveEvent = self._scene_mouse_move
365
+ self.viewer.scene().mouseReleaseEvent = self._scene_mouse_release
366
+
367
+ # Mode management methods
368
+ def set_sam_mode(self):
369
+ """Set SAM points mode."""
370
+ if not self.model_manager.is_model_available():
371
+ print("Cannot enter SAM mode: No model available")
372
+ return
373
+ self._set_mode("sam_points")
374
+
375
+ def set_polygon_mode(self):
376
+ """Set polygon drawing mode."""
377
+ self._set_mode("polygon")
378
+
379
+ def toggle_selection_mode(self):
380
+ """Toggle selection mode."""
381
+ self._toggle_mode("selection")
382
+
383
+ def toggle_pan_mode(self):
384
+ """Toggle pan mode."""
385
+ self._toggle_mode("pan")
386
+
387
+ def toggle_edit_mode(self):
388
+ """Toggle edit mode."""
389
+ self._toggle_mode("edit")
390
+
391
+ def _set_mode(self, mode_name, is_toggle=False):
392
+ """Set the current mode."""
393
+ if not is_toggle and self.mode not in ["selection", "edit"]:
394
+ self.previous_mode = self.mode
395
+
396
+ self.mode = mode_name
397
+ self.control_panel.set_mode_text(mode_name)
398
+ self.clear_all_points()
399
+
400
+ # Set cursor and drag mode based on mode
401
+ cursor_map = {
402
+ "sam_points": Qt.CursorShape.CrossCursor,
403
+ "polygon": Qt.CursorShape.CrossCursor,
404
+ "selection": Qt.CursorShape.ArrowCursor,
405
+ "edit": Qt.CursorShape.SizeAllCursor,
406
+ "pan": Qt.CursorShape.OpenHandCursor,
407
+ }
408
+ self.viewer.set_cursor(cursor_map.get(self.mode, Qt.CursorShape.ArrowCursor))
409
+
410
+ drag_mode = (
411
+ self.viewer.DragMode.ScrollHandDrag
412
+ if self.mode == "pan"
413
+ else self.viewer.DragMode.NoDrag
414
+ )
415
+ self.viewer.setDragMode(drag_mode)
416
+
417
+ # Update highlights and handles based on the new mode
418
+ self._highlight_selected_segments()
419
+ if mode_name == "edit":
420
+ self._display_edit_handles()
421
+ else:
422
+ self._clear_edit_handles()
423
+
424
+ def _toggle_mode(self, new_mode):
425
+ """Toggle between modes."""
426
+ if self.mode == new_mode:
427
+ self._set_mode(self.previous_mode, is_toggle=True)
428
+ else:
429
+ if self.mode not in ["selection", "edit"]:
430
+ self.previous_mode = self.mode
431
+ self._set_mode(new_mode, is_toggle=True)
432
+
433
+ # Model management methods
434
+ def _browse_models_folder(self):
435
+ """Browse for models folder."""
436
+ folder_path = QFileDialog.getExistingDirectory(self, "Select Models Folder")
437
+ if folder_path:
438
+ self.model_manager.set_models_folder(folder_path)
439
+ models = self.model_manager.get_available_models(folder_path)
440
+ self.control_panel.populate_models(models)
441
+ self.viewer.setFocus()
442
+
443
+ def _refresh_models_list(self):
444
+ """Refresh the models list."""
445
+ folder = self.model_manager.get_models_folder()
446
+ if folder and os.path.exists(folder):
447
+ models = self.model_manager.get_available_models(folder)
448
+ self.control_panel.populate_models(models)
449
+ self._show_success_notification("Models list refreshed.")
450
+ else:
451
+ self._show_warning_notification("No models folder selected.")
452
+
453
+ def _load_selected_model(self, model_text):
454
+ """Load the selected model."""
455
+ if not model_text or model_text == "Default (vit_h)":
456
+ self.control_panel.set_current_model("Current: Default SAM Model")
457
+ return
458
+
459
+ model_path = self.control_panel.model_widget.get_selected_model_path()
460
+ if not model_path or not os.path.exists(model_path):
461
+ self._show_error_notification("Selected model file not found.")
462
+ return
463
+
464
+ self.control_panel.set_current_model("Loading model...")
465
+ QApplication.processEvents()
466
+
467
+ try:
468
+ success = self.model_manager.load_custom_model(model_path)
469
+ if success:
470
+ # Re-enable SAM functionality if model loaded successfully
471
+ self._enable_sam_functionality(True)
472
+ if self.model_manager.sam_model:
473
+ device_text = str(self.model_manager.sam_model.device).upper()
474
+ self.status_bar.set_permanent_message(f"Device: {device_text}")
475
+ else:
476
+ self.control_panel.set_current_model("Current: Default SAM Model")
477
+ self._show_error_notification(
478
+ "Failed to load selected model. Using default."
479
+ )
480
+ self.control_panel.model_widget.reset_to_default()
481
+ self._enable_sam_functionality(False)
482
+ except Exception as e:
483
+ self.control_panel.set_current_model("Current: Default SAM Model")
484
+ self._show_error_notification(f"Error loading model: {str(e)}")
485
+ self.control_panel.model_widget.reset_to_default()
486
+ self._enable_sam_functionality(False)
487
+
488
+ # Adjustment methods
489
+ def _set_annotation_size(self, value):
490
+ """Set annotation size."""
491
+ multiplier = value / 10.0
492
+ self.point_radius = self.settings.point_radius * multiplier
493
+ self.line_thickness = self.settings.line_thickness * multiplier
494
+ self.settings.annotation_size_multiplier = multiplier
495
+ # Update display (implementation would go here)
496
+
497
+ def _set_pan_speed(self, value):
498
+ """Set pan speed."""
499
+ self.pan_multiplier = value / 10.0
500
+ self.settings.pan_multiplier = self.pan_multiplier
501
+
502
+ def _set_join_threshold(self, value):
503
+ """Set polygon join threshold."""
504
+ self.polygon_join_threshold = value
505
+ self.settings.polygon_join_threshold = value
506
+
507
+ # File management methods
508
+ def _open_folder_dialog(self):
509
+ """Open folder dialog for images."""
510
+ folder_path = QFileDialog.getExistingDirectory(self, "Select Image Folder")
511
+ if folder_path:
512
+ self.right_panel.set_folder(folder_path, self.file_model)
513
+ self.viewer.setFocus()
514
+
515
+ def _load_selected_image(self, index):
516
+ """Load the selected image."""
517
+ if not index.isValid() or not self.file_model.isDir(index.parent()):
518
+ return
519
+
520
+ self.current_file_index = index
521
+ path = self.file_model.filePath(index)
522
+
523
+ if os.path.isfile(path) and self.file_manager.is_image_file(path):
524
+ self.current_image_path = path
525
+ pixmap = QPixmap(self.current_image_path)
526
+ if not pixmap.isNull():
527
+ self._reset_state()
528
+ self.viewer.set_photo(pixmap)
529
+ if self.model_manager.is_model_available():
530
+ self.model_manager.sam_model.set_image(self.current_image_path)
531
+ self.file_manager.load_class_aliases(self.current_image_path)
532
+ self.file_manager.load_existing_mask(self.current_image_path)
533
+ self.right_panel.file_tree.setCurrentIndex(index)
534
+ self._update_all_lists()
535
+ self.viewer.setFocus()
536
+
537
+ def _load_next_image(self):
538
+ """Load next image in the file list, with auto-save if enabled."""
539
+ if not self.current_file_index.isValid():
540
+ return
541
+ # Auto-save if enabled
542
+ if self.control_panel.get_settings().get("auto_save", True):
543
+ self._save_output_to_npz()
544
+ parent = self.current_file_index.parent()
545
+ row = self.current_file_index.row()
546
+ # Find next valid image file
547
+ for next_row in range(row + 1, self.file_model.rowCount(parent)):
548
+ next_index = self.file_model.index(next_row, 0, parent)
549
+ path = self.file_model.filePath(next_index)
550
+ if os.path.isfile(path) and self.file_manager.is_image_file(path):
551
+ self._load_selected_image(next_index)
552
+ return
553
+
554
+ def _load_previous_image(self):
555
+ """Load previous image in the file list, with auto-save if enabled."""
556
+ if not self.current_file_index.isValid():
557
+ return
558
+ # Auto-save if enabled
559
+ if self.control_panel.get_settings().get("auto_save", True):
560
+ self._save_output_to_npz()
561
+ parent = self.current_file_index.parent()
562
+ row = self.current_file_index.row()
563
+ # Find previous valid image file
564
+ for prev_row in range(row - 1, -1, -1):
565
+ prev_index = self.file_model.index(prev_row, 0, parent)
566
+ path = self.file_model.filePath(prev_index)
567
+ if os.path.isfile(path) and self.file_manager.is_image_file(path):
568
+ self._load_selected_image(prev_index)
569
+ return
570
+
571
+ # Segment management methods
572
+ def _assign_selected_to_class(self):
573
+ """Assign selected segments to class."""
574
+ selected_indices = self.right_panel.get_selected_segment_indices()
575
+ self.segment_manager.assign_segments_to_class(selected_indices)
576
+ self._update_all_lists()
577
+
578
+ def _delete_selected_segments(self):
579
+ """Delete selected segments and remove any highlight overlays."""
580
+ # Remove highlight overlays before deleting segments
581
+ if hasattr(self, "highlight_items"):
582
+ for item in self.highlight_items:
583
+ self.viewer.scene().removeItem(item)
584
+ self.highlight_items = []
585
+ selected_indices = self.right_panel.get_selected_segment_indices()
586
+ self.segment_manager.delete_segments(selected_indices)
587
+ self._update_all_lists()
588
+
589
+ def _highlight_selected_segments(self):
590
+ """Highlight selected segments. In edit mode, use a brighter hover-like effect."""
591
+ # Remove previous highlight overlays
592
+ if hasattr(self, "highlight_items"):
593
+ for item in self.highlight_items:
594
+ if item.scene():
595
+ self.viewer.scene().removeItem(item)
596
+ self.highlight_items = []
597
+
598
+ selected_indices = self.right_panel.get_selected_segment_indices()
599
+ if not selected_indices:
600
+ return
601
+
602
+ for i in selected_indices:
603
+ seg = self.segment_manager.segments[i]
604
+ base_color = self._get_color_for_class(seg.get("class_id"))
605
+
606
+ if self.mode == "edit":
607
+ # Use a brighter, hover-like highlight in edit mode
608
+ highlight_brush = QBrush(
609
+ QColor(base_color.red(), base_color.green(), base_color.blue(), 170)
610
+ )
611
+ else:
612
+ # Use the standard yellow overlay for selection
613
+ highlight_brush = QBrush(QColor(255, 255, 0, 180))
614
+
615
+ if seg["type"] == "Polygon" and seg.get("vertices"):
616
+ poly_item = QGraphicsPolygonItem(QPolygonF(seg["vertices"]))
617
+ poly_item.setBrush(highlight_brush)
618
+ poly_item.setPen(QPen(Qt.GlobalColor.transparent))
619
+ poly_item.setZValue(99)
620
+ self.viewer.scene().addItem(poly_item)
621
+ self.highlight_items.append(poly_item)
622
+ elif seg.get("mask") is not None:
623
+ # For non-polygon types, we still use the mask-to-pixmap approach.
624
+ # If in edit mode, we could consider skipping non-polygons.
625
+ if self.mode != "edit":
626
+ mask = seg.get("mask")
627
+ pixmap = mask_to_pixmap(mask, (255, 255, 0), alpha=180)
628
+ highlight_item = self.viewer.scene().addPixmap(pixmap)
629
+ highlight_item.setZValue(100)
630
+ self.highlight_items.append(highlight_item)
631
+
632
+ def _handle_alias_change(self, class_id, alias):
633
+ """Handle class alias change."""
634
+ if self._updating_lists:
635
+ return # Prevent recursion
636
+ self.segment_manager.set_class_alias(class_id, alias)
637
+ self._update_all_lists()
638
+
639
+ def _reassign_class_ids(self):
640
+ """Reassign class IDs."""
641
+ new_order = self.right_panel.get_class_order()
642
+ self.segment_manager.reassign_class_ids(new_order)
643
+ self._update_all_lists()
644
+
645
+ def _update_segment_table(self):
646
+ """Update segment table."""
647
+ table = self.right_panel.segment_table
648
+ table.blockSignals(True)
649
+ selected_indices = self.right_panel.get_selected_segment_indices()
650
+ table.clearContents()
651
+ table.setRowCount(0)
652
+
653
+ # Get current filter
654
+ filter_text = self.right_panel.class_filter_combo.currentText()
655
+ show_all = filter_text == "All Classes"
656
+ filter_class_id = -1
657
+ if not show_all:
658
+ try:
659
+ # Parse format like "Alias: ID" or "Class ID"
660
+ if ":" in filter_text:
661
+ filter_class_id = int(filter_text.split(":")[-1].strip())
662
+ else:
663
+ filter_class_id = int(filter_text.split()[-1])
664
+ except (ValueError, IndexError):
665
+ show_all = True # If parsing fails, show all
666
+
667
+ # Filter segments based on class filter
668
+ display_segments = []
669
+ for i, seg in enumerate(self.segment_manager.segments):
670
+ seg_class_id = seg.get("class_id")
671
+ should_include = show_all or seg_class_id == filter_class_id
672
+ if should_include:
673
+ display_segments.append((i, seg))
674
+
675
+ table.setRowCount(len(display_segments))
676
+
677
+ # Populate table rows
678
+ for row, (original_index, seg) in enumerate(display_segments):
679
+ class_id = seg.get("class_id")
680
+ color = self._get_color_for_class(class_id)
681
+ class_id_str = str(class_id) if class_id is not None else "N/A"
682
+
683
+ alias_str = "N/A"
684
+ if class_id is not None:
685
+ alias_str = self.segment_manager.get_class_alias(class_id)
686
+
687
+ # Create table items (1-based segment ID for display)
688
+ index_item = NumericTableWidgetItem(str(original_index + 1))
689
+ class_item = NumericTableWidgetItem(class_id_str)
690
+ alias_item = QTableWidgetItem(alias_str)
691
+
692
+ # Set items as non-editable
693
+ index_item.setFlags(index_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
694
+ class_item.setFlags(class_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
695
+ alias_item.setFlags(alias_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
696
+
697
+ # Store original index for selection tracking
698
+ index_item.setData(Qt.ItemDataRole.UserRole, original_index)
699
+
700
+ # Set items in table
701
+ table.setItem(row, 0, index_item)
702
+ table.setItem(row, 1, class_item)
703
+ table.setItem(row, 2, alias_item)
704
+
705
+ # Set background color based on class
706
+ for col in range(table.columnCount()):
707
+ if table.item(row, col):
708
+ table.item(row, col).setBackground(QBrush(color))
709
+
710
+ # Restore selection
711
+ table.setSortingEnabled(False)
712
+ for row in range(table.rowCount()):
713
+ item = table.item(row, 0)
714
+ if item and item.data(Qt.ItemDataRole.UserRole) in selected_indices:
715
+ table.selectRow(row)
716
+ table.setSortingEnabled(True)
717
+
718
+ table.blockSignals(False)
719
+ self.viewer.setFocus()
720
+
721
+ # Update active class display
722
+ active_class = self.segment_manager.get_active_class()
723
+ self.right_panel.update_active_class_display(active_class)
724
+
725
+ def _update_all_lists(self):
726
+ """Update all UI lists."""
727
+ if self._updating_lists:
728
+ return # Prevent recursion
729
+
730
+ self._updating_lists = True
731
+ try:
732
+ self._update_class_list()
733
+ self._update_segment_table()
734
+ self._update_class_filter()
735
+ self._display_all_segments()
736
+ if self.mode == "edit":
737
+ self._display_edit_handles()
738
+ else:
739
+ self._clear_edit_handles()
740
+ finally:
741
+ self._updating_lists = False
742
+
743
+ def _update_class_list(self):
744
+ """Update the class list in the right panel."""
745
+ class_table = self.right_panel.class_table
746
+ class_table.blockSignals(True)
747
+
748
+ # Get unique class IDs
749
+ unique_class_ids = self.segment_manager.get_unique_class_ids()
750
+
751
+ class_table.clearContents()
752
+ class_table.setRowCount(len(unique_class_ids))
753
+
754
+ for row, cid in enumerate(unique_class_ids):
755
+ alias_item = QTableWidgetItem(self.segment_manager.get_class_alias(cid))
756
+ id_item = QTableWidgetItem(str(cid))
757
+ id_item.setFlags(id_item.flags() & ~Qt.ItemFlag.ItemIsEditable)
758
+
759
+ color = self._get_color_for_class(cid)
760
+ alias_item.setBackground(QBrush(color))
761
+ id_item.setBackground(QBrush(color))
762
+
763
+ class_table.setItem(row, 0, alias_item)
764
+ class_table.setItem(row, 1, id_item)
765
+
766
+ # Update active class display BEFORE re-enabling signals
767
+ active_class = self.segment_manager.get_active_class()
768
+ self.right_panel.update_active_class_display(active_class)
769
+
770
+ class_table.blockSignals(False)
771
+
772
+ def _update_class_filter(self):
773
+ """Update the class filter combo box."""
774
+ combo = self.right_panel.class_filter_combo
775
+ current_text = combo.currentText()
776
+
777
+ combo.blockSignals(True)
778
+ combo.clear()
779
+ combo.addItem("All Classes")
780
+
781
+ # Add class options
782
+ unique_class_ids = self.segment_manager.get_unique_class_ids()
783
+ for class_id in unique_class_ids:
784
+ alias = self.segment_manager.get_class_alias(class_id)
785
+ display_text = f"{alias}: {class_id}" if alias else f"Class {class_id}"
786
+ combo.addItem(display_text)
787
+
788
+ # Restore selection if possible
789
+ index = combo.findText(current_text)
790
+ if index >= 0:
791
+ combo.setCurrentIndex(index)
792
+ else:
793
+ combo.setCurrentIndex(0)
794
+
795
+ combo.blockSignals(False)
796
+
797
+ def _display_all_segments(self):
798
+ """Display all segments on the viewer."""
799
+ # Clear existing segment items
800
+ for i, items in self.segment_items.items():
801
+ for item in items:
802
+ if item.scene():
803
+ self.viewer.scene().removeItem(item)
804
+ self.segment_items.clear()
805
+ self._clear_edit_handles()
806
+
807
+ # Display segments from segment manager
808
+ for i, segment in enumerate(self.segment_manager.segments):
809
+ self.segment_items[i] = []
810
+ class_id = segment.get("class_id")
811
+ base_color = self._get_color_for_class(class_id)
812
+
813
+ if segment["type"] == "Polygon" and segment.get("vertices"):
814
+ poly_item = HoverablePolygonItem(QPolygonF(segment["vertices"]))
815
+ default_brush = QBrush(
816
+ QColor(base_color.red(), base_color.green(), base_color.blue(), 70)
817
+ )
818
+ hover_brush = QBrush(
819
+ QColor(base_color.red(), base_color.green(), base_color.blue(), 170)
820
+ )
821
+ poly_item.set_brushes(default_brush, hover_brush)
822
+ poly_item.setPen(QPen(Qt.GlobalColor.transparent))
823
+ self.viewer.scene().addItem(poly_item)
824
+ self.segment_items[i].append(poly_item)
825
+ elif segment.get("mask") is not None:
826
+ default_pixmap = mask_to_pixmap(
827
+ segment["mask"], base_color.getRgb()[:3], alpha=70
828
+ )
829
+ hover_pixmap = mask_to_pixmap(
830
+ segment["mask"], base_color.getRgb()[:3], alpha=170
831
+ )
832
+ pixmap_item = HoverablePixmapItem()
833
+ pixmap_item.set_pixmaps(default_pixmap, hover_pixmap)
834
+ self.viewer.scene().addItem(pixmap_item)
835
+ pixmap_item.setZValue(i + 1)
836
+ self.segment_items[i].append(pixmap_item)
837
+
838
+ # Event handlers
839
+ def _handle_escape_press(self):
840
+ """Handle escape key press."""
841
+ self.right_panel.clear_selections()
842
+ self.clear_all_points()
843
+ self.viewer.setFocus()
844
+
845
+ def _handle_space_press(self):
846
+ """Handle space key press."""
847
+ if self.mode == "polygon" and self.polygon_points:
848
+ self._finalize_polygon()
849
+ else:
850
+ self._save_current_segment()
851
+
852
+ def _handle_enter_press(self):
853
+ """Handle enter key press."""
854
+ if self.mode == "polygon" and self.polygon_points:
855
+ self._finalize_polygon()
856
+ else:
857
+ self._save_output_to_npz()
858
+
859
+ def _save_current_segment(self):
860
+ """Save current SAM segment."""
861
+ if (
862
+ self.mode != "sam_points"
863
+ or not hasattr(self, "preview_mask_item")
864
+ or not self.preview_mask_item
865
+ or not self.model_manager.is_model_available()
866
+ ):
867
+ return
868
+
869
+ mask = self.model_manager.sam_model.predict(
870
+ self.positive_points, self.negative_points
871
+ )
872
+ if mask is not None:
873
+ new_segment = {
874
+ "mask": mask,
875
+ "type": "SAM",
876
+ "vertices": None,
877
+ }
878
+ self.segment_manager.add_segment(new_segment)
879
+ # Record the action for undo
880
+ self.action_history.append(
881
+ {
882
+ "type": "add_segment",
883
+ "segment_index": len(self.segment_manager.segments) - 1,
884
+ }
885
+ )
886
+ self.clear_all_points()
887
+ self._update_all_lists()
888
+
889
+ def _finalize_polygon(self):
890
+ """Finalize polygon drawing."""
891
+ if len(self.polygon_points) < 3:
892
+ return
893
+
894
+ new_segment = {
895
+ "vertices": list(self.polygon_points),
896
+ "type": "Polygon",
897
+ "mask": None,
898
+ }
899
+ self.segment_manager.add_segment(new_segment)
900
+ # Record the action for undo
901
+ self.action_history.append(
902
+ {
903
+ "type": "add_segment",
904
+ "segment_index": len(self.segment_manager.segments) - 1,
905
+ }
906
+ )
907
+
908
+ self.polygon_points.clear()
909
+ self.clear_all_points()
910
+ self._update_all_lists()
911
+
912
+ def _save_output_to_npz(self):
913
+ """Save output to NPZ and TXT files as enabled, and update file list tickboxes/highlight. If no segments, delete associated files."""
914
+ if not self.current_image_path:
915
+ self._show_warning_notification("No image loaded.")
916
+ return
917
+
918
+ # If no segments, delete associated files
919
+ if not self.segment_manager.segments:
920
+ base, _ = os.path.splitext(self.current_image_path)
921
+ deleted_files = []
922
+ for ext in [".npz", ".txt", ".json"]:
923
+ file_path = base + ext
924
+ if os.path.exists(file_path):
925
+ try:
926
+ os.remove(file_path)
927
+ deleted_files.append(file_path)
928
+ self.file_model.update_cache_for_path(file_path)
929
+ except Exception as e:
930
+ self._show_error_notification(
931
+ f"Error deleting {file_path}: {e}"
932
+ )
933
+ if deleted_files:
934
+ self._show_notification(
935
+ f"Deleted: {', '.join(os.path.basename(f) for f in deleted_files)}"
936
+ )
937
+ else:
938
+ self._show_warning_notification("No segments to save.")
939
+ return
940
+
941
+ try:
942
+ settings = self.control_panel.get_settings()
943
+ npz_path = None
944
+ txt_path = None
945
+ if settings.get("save_npz", True):
946
+ h, w = (
947
+ self.viewer._pixmap_item.pixmap().height(),
948
+ self.viewer._pixmap_item.pixmap().width(),
949
+ )
950
+ class_order = self.segment_manager.get_unique_class_ids()
951
+ if class_order:
952
+ npz_path = self.file_manager.save_npz(
953
+ self.current_image_path, (h, w), class_order
954
+ )
955
+ self._show_success_notification(
956
+ f"Saved: {os.path.basename(npz_path)}"
957
+ )
958
+ else:
959
+ self._show_warning_notification("No classes defined for saving.")
960
+ if settings.get("save_txt", True):
961
+ h, w = (
962
+ self.viewer._pixmap_item.pixmap().height(),
963
+ self.viewer._pixmap_item.pixmap().width(),
964
+ )
965
+ class_order = self.segment_manager.get_unique_class_ids()
966
+ if settings.get("yolo_use_alias", True):
967
+ class_labels = [
968
+ self.segment_manager.get_class_alias(cid) for cid in class_order
969
+ ]
970
+ else:
971
+ class_labels = [str(cid) for cid in class_order]
972
+ if class_order:
973
+ txt_path = self.file_manager.save_yolo_txt(
974
+ self.current_image_path, (h, w), class_order, class_labels
975
+ )
976
+ # Efficiently update file list tickboxes and highlight
977
+ for path in [npz_path, txt_path]:
978
+ if path:
979
+ self.file_model.update_cache_for_path(path)
980
+ self.file_model.set_highlighted_path(path)
981
+ QTimer.singleShot(
982
+ 1500,
983
+ lambda p=path: (
984
+ self.file_model.set_highlighted_path(None)
985
+ if self.file_model.highlighted_path == p
986
+ else None
987
+ ),
988
+ )
989
+ except Exception as e:
990
+ self._show_error_notification(f"Error saving: {str(e)}")
991
+
992
+ def _handle_merge_press(self):
993
+ """Handle merge key press."""
994
+ self._assign_selected_to_class()
995
+ self.right_panel.clear_selections()
996
+
997
+ def _undo_last_action(self):
998
+ """Undo the last action recorded in the history."""
999
+ if not self.action_history:
1000
+ self._show_notification("Nothing to undo.")
1001
+ return
1002
+
1003
+ last_action = self.action_history.pop()
1004
+ action_type = last_action.get("type")
1005
+
1006
+ if action_type == "add_segment":
1007
+ segment_index = last_action.get("segment_index")
1008
+ if segment_index is not None and 0 <= segment_index < len(
1009
+ self.segment_manager.segments
1010
+ ):
1011
+ # Remove the segment that was added
1012
+ self.segment_manager.delete_segments([segment_index])
1013
+ self.right_panel.clear_selections() # Clear selection to prevent phantom highlights
1014
+ self._update_all_lists()
1015
+ self._show_notification("Undid: Add Segment")
1016
+ elif action_type == "add_point":
1017
+ point_type = last_action.get("point_type")
1018
+ point_item = last_action.get("point_item")
1019
+ point_list = self.positive_points if point_type == "positive" else self.negative_points
1020
+ if point_list:
1021
+ point_list.pop()
1022
+ if point_item in self.point_items:
1023
+ self.point_items.remove(point_item)
1024
+ self.viewer.scene().removeItem(point_item)
1025
+ self._update_segmentation()
1026
+ self._show_notification("Undid: Add Point")
1027
+ elif action_type == "add_polygon_point":
1028
+ dot_item = last_action.get("dot_item")
1029
+ if self.polygon_points:
1030
+ self.polygon_points.pop()
1031
+ if dot_item in self.polygon_preview_items:
1032
+ self.polygon_preview_items.remove(dot_item)
1033
+ self.viewer.scene().removeItem(dot_item)
1034
+ self._draw_polygon_preview()
1035
+ self._show_notification("Undid: Add Polygon Point")
1036
+ elif action_type == "move_polygon":
1037
+ initial_vertices = last_action.get("initial_vertices")
1038
+ for i, vertices in initial_vertices.items():
1039
+ self.segment_manager.segments[i]["vertices"] = vertices
1040
+ self._update_polygon_item(i)
1041
+ self._display_edit_handles()
1042
+ self._highlight_selected_segments()
1043
+ self._show_notification("Undid: Move Polygon")
1044
+ elif action_type == "move_vertex":
1045
+ segment_index = last_action.get("segment_index")
1046
+ vertex_index = last_action.get("vertex_index")
1047
+ old_pos = last_action.get("old_pos")
1048
+ self.segment_manager.segments[segment_index]["vertices"][vertex_index] = old_pos
1049
+ self._update_polygon_item(segment_index)
1050
+ self._display_edit_handles()
1051
+ self._highlight_selected_segments()
1052
+ self._show_notification("Undid: Move Vertex")
1053
+
1054
+ # Add more undo logic for other action types here in the future
1055
+ else:
1056
+ self._show_warning_notification(f"Undo for action '{action_type}' not implemented.")
1057
+
1058
+ def clear_all_points(self):
1059
+ """Clear all temporary points."""
1060
+ if hasattr(self, "rubber_band_line") and self.rubber_band_line:
1061
+ self.viewer.scene().removeItem(self.rubber_band_line)
1062
+ self.rubber_band_line = None
1063
+
1064
+ self.positive_points.clear()
1065
+ self.negative_points.clear()
1066
+
1067
+ for item in self.point_items:
1068
+ self.viewer.scene().removeItem(item)
1069
+ self.point_items.clear()
1070
+
1071
+ self.polygon_points.clear()
1072
+ for item in self.polygon_preview_items:
1073
+ self.viewer.scene().removeItem(item)
1074
+ self.polygon_preview_items.clear()
1075
+
1076
+ if hasattr(self, "preview_mask_item") and self.preview_mask_item:
1077
+ self.viewer.scene().removeItem(self.preview_mask_item)
1078
+ self.preview_mask_item = None
1079
+
1080
+ def _show_notification(self, message, duration=3000):
1081
+ """Show notification message."""
1082
+ self.status_bar.show_message(message, duration)
1083
+
1084
+ def _show_error_notification(self, message, duration=8000):
1085
+ """Show error notification message."""
1086
+ self.status_bar.show_error_message(message, duration)
1087
+
1088
+ def _show_success_notification(self, message, duration=3000):
1089
+ """Show success notification message."""
1090
+ self.status_bar.show_success_message(message, duration)
1091
+
1092
+ def _show_warning_notification(self, message, duration=5000):
1093
+ """Show warning notification message."""
1094
+ self.status_bar.show_warning_message(message, duration)
1095
+
1096
+ def _show_hotkey_dialog(self):
1097
+ """Show the hotkey configuration dialog."""
1098
+ dialog = HotkeyDialog(self.hotkey_manager, self)
1099
+ dialog.exec()
1100
+ # Update shortcuts after dialog closes
1101
+ self._update_shortcuts()
1102
+
1103
+ def _handle_zoom_in(self):
1104
+ """Handle zoom in."""
1105
+ current_val = self.control_panel.get_annotation_size()
1106
+ self.control_panel.set_annotation_size(min(current_val + 1, 50))
1107
+
1108
+ def _handle_zoom_out(self):
1109
+ """Handle zoom out."""
1110
+ current_val = self.control_panel.get_annotation_size()
1111
+ self.control_panel.set_annotation_size(max(current_val - 1, 1))
1112
+
1113
+ def _handle_pan_key(self, direction):
1114
+ """Handle WASD pan keys."""
1115
+ if not hasattr(self, "viewer"):
1116
+ return
1117
+
1118
+ amount = int(self.viewer.height() * 0.1 * self.pan_multiplier)
1119
+
1120
+ if direction == "up":
1121
+ self.viewer.verticalScrollBar().setValue(
1122
+ self.viewer.verticalScrollBar().value() - amount
1123
+ )
1124
+ elif direction == "down":
1125
+ self.viewer.verticalScrollBar().setValue(
1126
+ self.viewer.verticalScrollBar().value() + amount
1127
+ )
1128
+ elif direction == "left":
1129
+ amount = int(self.viewer.width() * 0.1 * self.pan_multiplier)
1130
+ self.viewer.horizontalScrollBar().setValue(
1131
+ self.viewer.horizontalScrollBar().value() - amount
1132
+ )
1133
+ elif direction == "right":
1134
+ amount = int(self.viewer.width() * 0.1 * self.pan_multiplier)
1135
+ self.viewer.horizontalScrollBar().setValue(
1136
+ self.viewer.horizontalScrollBar().value() + amount
1137
+ )
1138
+
1139
+ def closeEvent(self, event):
1140
+ """Handle application close."""
1141
+ # Close any popped-out panels first
1142
+ if self.left_panel_popout is not None:
1143
+ self.left_panel_popout.close()
1144
+ if self.right_panel_popout is not None:
1145
+ self.right_panel_popout.close()
1146
+
1147
+ # Save settings
1148
+ self.settings.save_to_file(str(self.paths.settings_file))
1149
+ super().closeEvent(event)
1150
+
1151
+ def _reset_state(self):
1152
+ """Reset application state."""
1153
+ self.clear_all_points()
1154
+ self.segment_manager.clear()
1155
+ self._update_all_lists()
1156
+ items_to_remove = [
1157
+ item
1158
+ for item in self.viewer.scene().items()
1159
+ if item is not self.viewer._pixmap_item
1160
+ ]
1161
+ for item in items_to_remove:
1162
+ self.viewer.scene().removeItem(item)
1163
+ self.segment_items.clear()
1164
+ self.highlight_items.clear()
1165
+
1166
+ def _scene_mouse_press(self, event):
1167
+ """Handle mouse press events in the scene."""
1168
+ # Map scene coordinates to the view so items() works correctly.
1169
+ view_pos = self.viewer.mapFromScene(event.scenePos())
1170
+ items_at_pos = self.viewer.items(view_pos)
1171
+ is_handle_click = any(
1172
+ isinstance(item, EditableVertexItem) for item in items_at_pos
1173
+ )
1174
+
1175
+ # Allow vertex handles to process their own mouse events.
1176
+ if is_handle_click:
1177
+ self._original_mouse_press(event)
1178
+ return
1179
+
1180
+ if self.mode == "edit" and event.button() == Qt.MouseButton.LeftButton:
1181
+ pos = event.scenePos()
1182
+ if self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint()):
1183
+ self.is_dragging_polygon = True
1184
+ self.drag_start_pos = pos
1185
+ selected_indices = self.right_panel.get_selected_segment_indices()
1186
+ self.drag_initial_vertices = {
1187
+ i: [QPointF(p) for p in self.segment_manager.segments[i]["vertices"]]
1188
+ for i in selected_indices
1189
+ if self.segment_manager.segments[i].get("type") == "Polygon"
1190
+ }
1191
+ event.accept()
1192
+ return
1193
+
1194
+ # Call the original scene handler.
1195
+ self._original_mouse_press(event)
1196
+ # Skip further processing unless we're in selection mode.
1197
+ if event.isAccepted() and self.mode != "selection":
1198
+ return
1199
+
1200
+ if self.is_dragging_polygon:
1201
+ return
1202
+
1203
+ pos = event.scenePos()
1204
+ if (
1205
+ self.viewer._pixmap_item.pixmap().isNull()
1206
+ or not self.viewer._pixmap_item.pixmap().rect().contains(pos.toPoint())
1207
+ ):
1208
+ return
1209
+
1210
+ if self.mode == "pan":
1211
+ self.viewer.set_cursor(Qt.CursorShape.ClosedHandCursor)
1212
+ elif self.mode == "sam_points":
1213
+ if event.button() == Qt.MouseButton.LeftButton:
1214
+ self._add_point(pos, positive=True)
1215
+ self._update_segmentation()
1216
+ elif event.button() == Qt.MouseButton.RightButton:
1217
+ self._add_point(pos, positive=False)
1218
+ self._update_segmentation()
1219
+ elif self.mode == "polygon":
1220
+ if event.button() == Qt.MouseButton.LeftButton:
1221
+ self._handle_polygon_click(pos)
1222
+ elif self.mode == "selection":
1223
+ if event.button() == Qt.MouseButton.LeftButton:
1224
+ self._handle_segment_selection_click(pos)
1225
+
1226
+ def _scene_mouse_move(self, event):
1227
+ """Handle mouse move events in the scene."""
1228
+ if self.mode == "edit" and self.is_dragging_polygon:
1229
+ delta = event.scenePos() - self.drag_start_pos
1230
+ for i, initial_verts in self.drag_initial_vertices.items():
1231
+ self.segment_manager.segments[i]["vertices"] = [
1232
+ QPointF(v) + delta for v in initial_verts
1233
+ ]
1234
+ self._update_polygon_item(i)
1235
+ self._display_edit_handles() # Redraw handles at new positions
1236
+ self._highlight_selected_segments() # Redraw highlight at new position
1237
+ event.accept()
1238
+ return
1239
+
1240
+ self._original_mouse_move(event)
1241
+
1242
+ def _scene_mouse_release(self, event):
1243
+ """Handle mouse release events in the scene."""
1244
+ if self.mode == "edit" and self.is_dragging_polygon:
1245
+ # Record the action for undo
1246
+ final_vertices = {
1247
+ i: list(self.segment_manager.segments[i]["vertices"])
1248
+ for i in self.drag_initial_vertices.keys()
1249
+ }
1250
+ self.action_history.append(
1251
+ {
1252
+ "type": "move_polygon",
1253
+ "initial_vertices": {k: list(v) for k, v in self.drag_initial_vertices.items()},
1254
+ "final_vertices": final_vertices,
1255
+ }
1256
+ )
1257
+ self.is_dragging_polygon = False
1258
+ self.drag_initial_vertices.clear()
1259
+ event.accept()
1260
+ return
1261
+
1262
+ if self.mode == "pan":
1263
+ self.viewer.set_cursor(Qt.CursorShape.OpenHandCursor)
1264
+ self._original_mouse_release(event)
1265
+
1266
+ def _add_point(self, pos, positive):
1267
+ """Add a point for SAM segmentation."""
1268
+ point_list = self.positive_points if positive else self.negative_points
1269
+ point_list.append([int(pos.x()), int(pos.y())])
1270
+
1271
+ point_color = (
1272
+ QColor(Qt.GlobalColor.green) if positive else QColor(Qt.GlobalColor.red)
1273
+ )
1274
+ point_color.setAlpha(150)
1275
+ point_diameter = self.point_radius * 2
1276
+
1277
+ point_item = QGraphicsEllipseItem(
1278
+ pos.x() - self.point_radius,
1279
+ pos.y() - self.point_radius,
1280
+ point_diameter,
1281
+ point_diameter,
1282
+ )
1283
+ point_item.setBrush(QBrush(point_color))
1284
+ point_item.setPen(QPen(Qt.GlobalColor.transparent))
1285
+ self.viewer.scene().addItem(point_item)
1286
+ self.point_items.append(point_item)
1287
+
1288
+ # Record the action for undo
1289
+ self.action_history.append(
1290
+ {
1291
+ "type": "add_point",
1292
+ "point_type": "positive" if positive else "negative",
1293
+ "point_coords": [int(pos.x()), int(pos.y())],
1294
+ "point_item": point_item,
1295
+ }
1296
+ )
1297
+
1298
+ def _update_segmentation(self):
1299
+ """Update SAM segmentation preview."""
1300
+ if hasattr(self, "preview_mask_item") and self.preview_mask_item:
1301
+ self.viewer.scene().removeItem(self.preview_mask_item)
1302
+ if not self.positive_points or not self.model_manager.is_model_available():
1303
+ return
1304
+
1305
+ mask = self.model_manager.sam_model.predict(
1306
+ self.positive_points, self.negative_points
1307
+ )
1308
+ if mask is not None:
1309
+ pixmap = mask_to_pixmap(mask, (255, 255, 0))
1310
+ self.preview_mask_item = self.viewer.scene().addPixmap(pixmap)
1311
+ self.preview_mask_item.setZValue(50)
1312
+
1313
+ def _handle_polygon_click(self, pos):
1314
+ """Handle polygon drawing clicks."""
1315
+ # Check if clicking near the first point to close polygon
1316
+ if self.polygon_points and len(self.polygon_points) > 2:
1317
+ first_point = self.polygon_points[0]
1318
+ distance_squared = (pos.x() - first_point.x()) ** 2 + (
1319
+ pos.y() - first_point.y()
1320
+ ) ** 2
1321
+ if distance_squared < self.polygon_join_threshold**2:
1322
+ self._finalize_polygon()
1323
+ return
1324
+
1325
+ # Add new point to polygon
1326
+ self.polygon_points.append(pos)
1327
+
1328
+ # Create visual point
1329
+ point_diameter = self.point_radius * 2
1330
+ point_color = QColor(Qt.GlobalColor.blue)
1331
+ point_color.setAlpha(150)
1332
+ dot = QGraphicsEllipseItem(
1333
+ pos.x() - self.point_radius,
1334
+ pos.y() - self.point_radius,
1335
+ point_diameter,
1336
+ point_diameter,
1337
+ )
1338
+ dot.setBrush(QBrush(point_color))
1339
+ dot.setPen(QPen(Qt.GlobalColor.transparent))
1340
+ self.viewer.scene().addItem(dot)
1341
+ self.polygon_preview_items.append(dot)
1342
+
1343
+ # Update polygon preview
1344
+ self._draw_polygon_preview()
1345
+
1346
+ # Record the action for undo
1347
+ self.action_history.append(
1348
+ {
1349
+ "type": "add_polygon_point",
1350
+ "point_coords": pos,
1351
+ "dot_item": dot,
1352
+ }
1353
+ )
1354
+
1355
+ def _draw_polygon_preview(self):
1356
+ """Draw polygon preview lines and fill."""
1357
+ # Remove old preview lines and polygons (keep dots)
1358
+ for item in self.polygon_preview_items[:]:
1359
+ if not isinstance(item, QGraphicsEllipseItem):
1360
+ if item.scene():
1361
+ self.viewer.scene().removeItem(item)
1362
+ self.polygon_preview_items.remove(item)
1363
+
1364
+ if len(self.polygon_points) > 2:
1365
+ # Create preview polygon fill
1366
+ preview_poly = QGraphicsPolygonItem(QPolygonF(self.polygon_points))
1367
+ preview_poly.setBrush(QBrush(QColor(0, 255, 255, 100)))
1368
+ preview_poly.setPen(QPen(Qt.GlobalColor.transparent))
1369
+ self.viewer.scene().addItem(preview_poly)
1370
+ self.polygon_preview_items.append(preview_poly)
1371
+
1372
+ if len(self.polygon_points) > 1:
1373
+ # Create preview lines between points
1374
+ line_color = QColor(Qt.GlobalColor.cyan)
1375
+ line_color.setAlpha(150)
1376
+ for i in range(len(self.polygon_points) - 1):
1377
+ line = QGraphicsLineItem(
1378
+ self.polygon_points[i].x(),
1379
+ self.polygon_points[i].y(),
1380
+ self.polygon_points[i + 1].x(),
1381
+ self.polygon_points[i + 1].y(),
1382
+ )
1383
+ line.setPen(QPen(line_color, self.line_thickness))
1384
+ self.viewer.scene().addItem(line)
1385
+ self.polygon_preview_items.append(line)
1386
+
1387
+ def _handle_segment_selection_click(self, pos):
1388
+ """Handle segment selection clicks (toggle behavior)."""
1389
+ x, y = int(pos.x()), int(pos.y())
1390
+ for i in range(len(self.segment_manager.segments) - 1, -1, -1):
1391
+ seg = self.segment_manager.segments[i]
1392
+ # Determine mask for hit-testing
1393
+ if seg["type"] == "Polygon" and seg.get("vertices"):
1394
+ # Rasterize polygon
1395
+ if self.viewer._pixmap_item.pixmap().isNull():
1396
+ continue
1397
+ h = self.viewer._pixmap_item.pixmap().height()
1398
+ w = self.viewer._pixmap_item.pixmap().width()
1399
+ points_np = np.array(
1400
+ [[p.x(), p.y()] for p in seg["vertices"]], dtype=np.int32
1401
+ )
1402
+ # Ensure points are within bounds
1403
+ points_np = np.clip(points_np, 0, [w - 1, h - 1])
1404
+ mask = np.zeros((h, w), dtype=np.uint8)
1405
+ cv2.fillPoly(mask, [points_np], 1)
1406
+ mask = mask.astype(bool)
1407
+ else:
1408
+ mask = seg.get("mask")
1409
+ if (
1410
+ mask is not None
1411
+ and y < mask.shape[0]
1412
+ and x < mask.shape[1]
1413
+ and mask[y, x]
1414
+ ):
1415
+ # Find the corresponding row in the segment table and toggle selection
1416
+ table = self.right_panel.segment_table
1417
+ for j in range(table.rowCount()):
1418
+ item = table.item(j, 0)
1419
+ if item and item.data(Qt.ItemDataRole.UserRole) == i:
1420
+ # Toggle selection for this row using the original working method
1421
+ is_selected = table.item(j, 0).isSelected()
1422
+ range_to_select = QTableWidgetSelectionRange(
1423
+ j, 0, j, table.columnCount() - 1
1424
+ )
1425
+ table.setRangeSelected(range_to_select, not is_selected)
1426
+ self._highlight_selected_segments()
1427
+ return
1428
+ self.viewer.setFocus()
1429
+
1430
+ def _get_color_for_class(self, class_id):
1431
+ """Get color for a class ID."""
1432
+ if class_id is None:
1433
+ return QColor.fromHsv(0, 0, 128)
1434
+ hue = int((class_id * 222.4922359) % 360)
1435
+ color = QColor.fromHsv(hue, 220, 220)
1436
+ if not color.isValid():
1437
+ return QColor(Qt.GlobalColor.white)
1438
+ return color
1439
+
1440
+ def _display_edit_handles(self):
1441
+ """Display draggable vertex handles for selected polygons in edit mode."""
1442
+ self._clear_edit_handles()
1443
+ if self.mode != "edit":
1444
+ return
1445
+ selected_indices = self.right_panel.get_selected_segment_indices()
1446
+ handle_radius = self.point_radius
1447
+ handle_diam = handle_radius * 2
1448
+ for seg_idx in selected_indices:
1449
+ seg = self.segment_manager.segments[seg_idx]
1450
+ if seg["type"] == "Polygon" and seg.get("vertices"):
1451
+ for v_idx, pt in enumerate(seg["vertices"]):
1452
+ handle = EditableVertexItem(
1453
+ self,
1454
+ seg_idx,
1455
+ v_idx,
1456
+ -handle_radius,
1457
+ -handle_radius,
1458
+ handle_diam,
1459
+ handle_diam,
1460
+ )
1461
+ handle.setPos(pt) # Use setPos to handle zoom correctly
1462
+ handle.setZValue(200) # Ensure handles are on top
1463
+ # Make sure the handle can receive mouse events
1464
+ handle.setAcceptHoverEvents(True)
1465
+ self.viewer.scene().addItem(handle)
1466
+ self.edit_handles.append(handle)
1467
+
1468
+ def _clear_edit_handles(self):
1469
+ """Remove all editable vertex handles from the scene."""
1470
+ if hasattr(self, "edit_handles"):
1471
+ for h in self.edit_handles:
1472
+ if h.scene():
1473
+ self.viewer.scene().removeItem(h)
1474
+ self.edit_handles = []
1475
+
1476
+ def update_vertex_pos(self, segment_index, vertex_index, new_pos, record_undo=True):
1477
+ """Update the position of a vertex in a polygon segment."""
1478
+ seg = self.segment_manager.segments[segment_index]
1479
+ if seg.get("type") == "Polygon":
1480
+ old_pos = seg["vertices"][vertex_index]
1481
+ if record_undo:
1482
+ self.action_history.append(
1483
+ {
1484
+ "type": "move_vertex",
1485
+ "segment_index": segment_index,
1486
+ "vertex_index": vertex_index,
1487
+ "old_pos": old_pos,
1488
+ "new_pos": new_pos,
1489
+ }
1490
+ )
1491
+ seg["vertices"][
1492
+ vertex_index
1493
+ ] = new_pos # new_pos is already the correct scene coordinate
1494
+ self._update_polygon_item(segment_index)
1495
+ self._highlight_selected_segments() # Keep the highlight in sync with the new shape
1496
+
1497
+ def _update_polygon_item(self, segment_index):
1498
+ """Efficiently update the visual polygon item for a given segment."""
1499
+ items = self.segment_items.get(segment_index, [])
1500
+ for item in items:
1501
+ if isinstance(item, HoverablePolygonItem):
1502
+ item.setPolygon(
1503
+ QPolygonF(self.segment_manager.segments[segment_index]["vertices"])
1504
+ )
1505
+ return
1506
+
1507
+ def _handle_class_toggle(self, class_id):
1508
+ """Handle class toggle."""
1509
+ is_active = self.segment_manager.toggle_active_class(class_id)
1510
+
1511
+ if is_active:
1512
+ self._show_notification(f"Class {class_id} activated for new segments")
1513
+ # Update visual display
1514
+ self.right_panel.update_active_class_display(class_id)
1515
+ else:
1516
+ self._show_notification(
1517
+ "No active class - new segments will create new classes"
1518
+ )
1519
+ # Update visual display to clear active class
1520
+ self.right_panel.update_active_class_display(None)
1521
+
1522
+ def _pop_out_left_panel(self):
1523
+ """Pop out the left control panel into a separate window."""
1524
+ if self.left_panel_popout is not None:
1525
+ # Panel is already popped out, return it to main window
1526
+ self._return_left_panel(self.control_panel)
1527
+ return
1528
+
1529
+ # Remove panel from main splitter
1530
+ self.control_panel.setParent(None)
1531
+
1532
+ # Create pop-out window
1533
+ self.left_panel_popout = PanelPopoutWindow(
1534
+ self.control_panel, "Control Panel", self
1535
+ )
1536
+ self.left_panel_popout.panel_closed.connect(self._return_left_panel)
1537
+ self.left_panel_popout.show()
1538
+
1539
+ # Update panel's pop-out button
1540
+ self.control_panel.set_popout_mode(True)
1541
+
1542
+ # Make pop-out window resizable
1543
+ self.left_panel_popout.setMinimumSize(200, 400)
1544
+ self.left_panel_popout.resize(self.control_panel.preferred_width + 20, 600)
1545
+
1546
+ def _pop_out_right_panel(self):
1547
+ """Pop out the right panel into a separate window."""
1548
+ if self.right_panel_popout is not None:
1549
+ # Panel is already popped out, return it to main window
1550
+ self._return_right_panel(self.right_panel)
1551
+ return
1552
+
1553
+ # Remove panel from main splitter
1554
+ self.right_panel.setParent(None)
1555
+
1556
+ # Create pop-out window
1557
+ self.right_panel_popout = PanelPopoutWindow(
1558
+ self.right_panel, "File Explorer & Segments", self
1559
+ )
1560
+ self.right_panel_popout.panel_closed.connect(self._return_right_panel)
1561
+ self.right_panel_popout.show()
1562
+
1563
+ # Update panel's pop-out button
1564
+ self.right_panel.set_popout_mode(True)
1565
+
1566
+ # Make pop-out window resizable
1567
+ self.right_panel_popout.setMinimumSize(250, 400)
1568
+ self.right_panel_popout.resize(self.right_panel.preferred_width + 20, 600)
1569
+
1570
+ def _return_left_panel(self, panel_widget):
1571
+ """Return the left panel to the main window."""
1572
+ if self.left_panel_popout is not None:
1573
+ # Close the pop-out window
1574
+ self.left_panel_popout.close()
1575
+
1576
+ # Return panel to main splitter
1577
+ self.main_splitter.insertWidget(0, self.control_panel)
1578
+ self.left_panel_popout = None
1579
+
1580
+ # Update panel's pop-out button
1581
+ self.control_panel.set_popout_mode(False)
1582
+
1583
+ # Restore splitter sizes
1584
+ self.main_splitter.setSizes([250, 800, 350])
1585
+
1586
+ def _handle_splitter_moved(self, pos, index):
1587
+ """Handle splitter movement for intelligent expand/collapse behavior."""
1588
+ sizes = self.main_splitter.sizes()
1589
+
1590
+ # Left panel (index 0) - expand/collapse logic
1591
+ if index == 1: # Splitter between left panel and viewer
1592
+ left_size = sizes[0]
1593
+ # Only snap to collapsed if user drags very close to collapse
1594
+ if left_size < 50: # Collapsed threshold
1595
+ # Panel is being collapsed, snap to collapsed state
1596
+ new_sizes = [0] + sizes[1:]
1597
+ new_sizes[1] = new_sizes[1] + left_size # Give space back to viewer
1598
+ self.main_splitter.setSizes(new_sizes)
1599
+ # Temporarily override minimum width to allow collapsing
1600
+ self.control_panel.setMinimumWidth(0)
1601
+
1602
+ # Right panel (index 2) - expand/collapse logic
1603
+ elif index == 2: # Splitter between viewer and right panel
1604
+ right_size = sizes[2]
1605
+ # Only snap to collapsed if user drags very close to collapse
1606
+ if right_size < 50: # Collapsed threshold
1607
+ # Panel is being collapsed, snap to collapsed state
1608
+ new_sizes = sizes[:-1] + [0]
1609
+ new_sizes[1] = new_sizes[1] + right_size # Give space back to viewer
1610
+ self.main_splitter.setSizes(new_sizes)
1611
+ # Temporarily override minimum width to allow collapsing
1612
+ self.right_panel.setMinimumWidth(0)
1613
+
1614
+ def _expand_left_panel(self):
1615
+ """Expand the left panel to its preferred width."""
1616
+ sizes = self.main_splitter.sizes()
1617
+ if sizes[0] < 50: # Only expand if currently collapsed
1618
+ # Restore minimum width first
1619
+ self.control_panel.setMinimumWidth(self.control_panel.preferred_width)
1620
+
1621
+ space_needed = self.control_panel.preferred_width
1622
+ viewer_width = sizes[1] - space_needed
1623
+ if viewer_width > 400: # Ensure viewer has minimum space
1624
+ new_sizes = [self.control_panel.preferred_width, viewer_width] + sizes[
1625
+ 2:
1626
+ ]
1627
+ self.main_splitter.setSizes(new_sizes)
1628
+
1629
+ def _expand_right_panel(self):
1630
+ """Expand the right panel to its preferred width."""
1631
+ sizes = self.main_splitter.sizes()
1632
+ if sizes[2] < 50: # Only expand if currently collapsed
1633
+ # Restore minimum width first
1634
+ self.right_panel.setMinimumWidth(self.right_panel.preferred_width)
1635
+
1636
+ space_needed = self.right_panel.preferred_width
1637
+ viewer_width = sizes[1] - space_needed
1638
+ if viewer_width > 400: # Ensure viewer has minimum space
1639
+ new_sizes = sizes[:-1] + [
1640
+ viewer_width,
1641
+ self.right_panel.preferred_width,
1642
+ ]
1643
+ self.main_splitter.setSizes(new_sizes)
1644
+
1645
+ def _return_right_panel(self, panel_widget):
1646
+ """Return the right panel to the main window."""
1647
+ if self.right_panel_popout is not None:
1648
+ # Close the pop-out window
1649
+ self.right_panel_popout.close()
1650
+
1651
+ # Return panel to main splitter
1652
+ self.main_splitter.addWidget(self.right_panel)
1653
+ self.right_panel_popout = None
1654
+
1655
+ # Update panel's pop-out button
1656
+ self.right_panel.set_popout_mode(False)
1657
+
1658
+ # Restore splitter sizes
1659
+ self.main_splitter.setSizes([250, 800, 350])