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