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