lazylabel-gui 1.0.9__py3-none-any.whl → 1.1.1__py3-none-any.whl

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