lazylabel-gui 1.1.9__py3-none-any.whl → 1.2.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.
- lazylabel/config/hotkeys.py +2 -2
- lazylabel/core/model_manager.py +91 -16
- lazylabel/models/sam2_model.py +223 -0
- lazylabel/models/sam_model.py +25 -3
- lazylabel/ui/control_panel.py +37 -2
- lazylabel/ui/main_window.py +438 -33
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/METADATA +56 -96
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/RECORD +12 -11
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/WHEEL +0 -0
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/entry_points.txt +0 -0
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/licenses/LICENSE +0 -0
- {lazylabel_gui-1.1.9.dist-info → lazylabel_gui-1.2.1.dist-info}/top_level.txt +0 -0
lazylabel/ui/main_window.py
CHANGED
@@ -245,6 +245,11 @@ class MainWindow(QMainWindow):
|
|
245
245
|
self.rubber_band_line = None
|
246
246
|
self.rubber_band_rect = None # New attribute for bounding box
|
247
247
|
self.preview_mask_item = None
|
248
|
+
|
249
|
+
# AI mode state
|
250
|
+
self.ai_click_start_pos = None
|
251
|
+
self.ai_click_time = 0
|
252
|
+
self.ai_rubber_band_rect = None
|
248
253
|
self.segments, self.segment_items, self.highlight_items = [], {}, []
|
249
254
|
self.edit_handles = []
|
250
255
|
self.is_dragging_polygon, self.drag_start_pos, self.drag_initial_vertices = (
|
@@ -292,6 +297,10 @@ class MainWindow(QMainWindow):
|
|
292
297
|
self.sam_embedding_cache = {} # Cache SAM embeddings by content hash
|
293
298
|
self.current_sam_hash = None # Hash of currently loaded SAM image
|
294
299
|
|
300
|
+
# Add bounding box preview state
|
301
|
+
self.ai_bbox_preview_mask = None
|
302
|
+
self.ai_bbox_preview_rect = None
|
303
|
+
|
295
304
|
self._setup_ui()
|
296
305
|
self._setup_model()
|
297
306
|
self._setup_connections()
|
@@ -391,8 +400,8 @@ class MainWindow(QMainWindow):
|
|
391
400
|
def _enable_sam_functionality(self, enabled: bool):
|
392
401
|
"""Enable or disable SAM point functionality."""
|
393
402
|
self.control_panel.set_sam_mode_enabled(enabled)
|
394
|
-
if not enabled and self.mode
|
395
|
-
# Switch to polygon mode if SAM is disabled and we're in SAM mode
|
403
|
+
if not enabled and self.mode in ["sam_points", "ai"]:
|
404
|
+
# Switch to polygon mode if SAM is disabled and we're in SAM/AI mode
|
396
405
|
self.set_polygon_mode()
|
397
406
|
|
398
407
|
def _setup_connections(self):
|
@@ -559,12 +568,12 @@ class MainWindow(QMainWindow):
|
|
559
568
|
|
560
569
|
# Mode management methods
|
561
570
|
def set_sam_mode(self):
|
562
|
-
"""Set mode to SAM points."""
|
571
|
+
"""Set mode to AI (combines SAM points and bounding box)."""
|
563
572
|
if not self.model_manager.is_model_available():
|
564
|
-
logger.warning("Cannot enter
|
573
|
+
logger.warning("Cannot enter AI mode: No model available")
|
565
574
|
return
|
566
|
-
self._set_mode("
|
567
|
-
# Ensure SAM model is updated when entering
|
575
|
+
self._set_mode("ai")
|
576
|
+
# Ensure SAM model is updated when entering AI mode (lazy update)
|
568
577
|
self._ensure_sam_updated()
|
569
578
|
|
570
579
|
def set_polygon_mode(self):
|
@@ -625,6 +634,7 @@ class MainWindow(QMainWindow):
|
|
625
634
|
# Set cursor and drag mode based on mode
|
626
635
|
cursor_map = {
|
627
636
|
"sam_points": Qt.CursorShape.CrossCursor,
|
637
|
+
"ai": Qt.CursorShape.CrossCursor,
|
628
638
|
"polygon": Qt.CursorShape.CrossCursor,
|
629
639
|
"bbox": Qt.CursorShape.CrossCursor,
|
630
640
|
"selection": Qt.CursorShape.ArrowCursor,
|
@@ -690,6 +700,9 @@ class MainWindow(QMainWindow):
|
|
690
700
|
self.control_panel.set_current_model("Loading model...")
|
691
701
|
QApplication.processEvents()
|
692
702
|
|
703
|
+
# CRITICAL FIX: Reset SAM state before switching models
|
704
|
+
self._reset_sam_state_for_model_switch()
|
705
|
+
|
693
706
|
try:
|
694
707
|
success = self.model_manager.load_custom_model(model_path)
|
695
708
|
if success:
|
@@ -698,6 +711,9 @@ class MainWindow(QMainWindow):
|
|
698
711
|
if self.model_manager.sam_model:
|
699
712
|
device_text = str(self.model_manager.sam_model.device).upper()
|
700
713
|
self.status_bar.set_permanent_message(f"Device: {device_text}")
|
714
|
+
|
715
|
+
# Mark SAM as dirty to force update with new model
|
716
|
+
self._mark_sam_dirty()
|
701
717
|
else:
|
702
718
|
self.control_panel.set_current_model("Current: Default SAM Model")
|
703
719
|
self._show_error_notification(
|
@@ -773,9 +789,15 @@ class MainWindow(QMainWindow):
|
|
773
789
|
# Update the main window's settings object with the latest from the widget
|
774
790
|
self.settings.update(**self.control_panel.settings_widget.get_settings())
|
775
791
|
|
776
|
-
#
|
792
|
+
# When operate on view setting changes, we need to force SAM model to update
|
793
|
+
# with proper scale factor recalculation via the worker thread
|
777
794
|
if self.current_image_path:
|
778
|
-
|
795
|
+
# Mark SAM as dirty and reset scale factor to force proper recalculation
|
796
|
+
self.sam_is_dirty = True
|
797
|
+
self.sam_scale_factor = 1.0 # Reset to default
|
798
|
+
self.current_sam_hash = None # Invalidate cache
|
799
|
+
# Use the worker thread to properly calculate scale factor
|
800
|
+
self._ensure_sam_updated()
|
779
801
|
|
780
802
|
def _handle_image_adjustment_changed(self):
|
781
803
|
"""Handle changes in image adjustments (brightness, contrast, gamma)."""
|
@@ -1171,6 +1193,20 @@ class MainWindow(QMainWindow):
|
|
1171
1193
|
"""Handle escape key press."""
|
1172
1194
|
self.right_panel.clear_selections()
|
1173
1195
|
self.clear_all_points()
|
1196
|
+
|
1197
|
+
# Clear bounding box preview state if active
|
1198
|
+
if (
|
1199
|
+
hasattr(self, "ai_bbox_preview_mask")
|
1200
|
+
and self.ai_bbox_preview_mask is not None
|
1201
|
+
):
|
1202
|
+
self.ai_bbox_preview_mask = None
|
1203
|
+
self.ai_bbox_preview_rect = None
|
1204
|
+
|
1205
|
+
# Clear preview
|
1206
|
+
if hasattr(self, "preview_mask_item") and self.preview_mask_item:
|
1207
|
+
self.viewer.scene().removeItem(self.preview_mask_item)
|
1208
|
+
self.preview_mask_item = None
|
1209
|
+
|
1174
1210
|
self.viewer.setFocus()
|
1175
1211
|
|
1176
1212
|
def _handle_space_press(self):
|
@@ -1190,17 +1226,68 @@ class MainWindow(QMainWindow):
|
|
1190
1226
|
def _save_current_segment(self):
|
1191
1227
|
"""Save current SAM segment with fragment threshold filtering."""
|
1192
1228
|
if (
|
1193
|
-
self.mode
|
1194
|
-
or not hasattr(self, "preview_mask_item")
|
1195
|
-
or not self.preview_mask_item
|
1229
|
+
self.mode not in ["sam_points", "ai"]
|
1196
1230
|
or not self.model_manager.is_model_available()
|
1197
1231
|
):
|
1198
1232
|
return
|
1199
1233
|
|
1200
|
-
|
1234
|
+
# Check if we have a bounding box preview to save
|
1235
|
+
if (
|
1236
|
+
hasattr(self, "ai_bbox_preview_mask")
|
1237
|
+
and self.ai_bbox_preview_mask is not None
|
1238
|
+
):
|
1239
|
+
# Save bounding box preview
|
1240
|
+
mask = self.ai_bbox_preview_mask
|
1241
|
+
|
1242
|
+
# Apply fragment threshold filtering if enabled
|
1243
|
+
filtered_mask = self._apply_fragment_threshold(mask)
|
1244
|
+
if filtered_mask is not None:
|
1245
|
+
new_segment = {
|
1246
|
+
"mask": filtered_mask,
|
1247
|
+
"type": "SAM",
|
1248
|
+
"vertices": None,
|
1249
|
+
}
|
1250
|
+
self.segment_manager.add_segment(new_segment)
|
1251
|
+
# Record the action for undo
|
1252
|
+
self.action_history.append(
|
1253
|
+
{
|
1254
|
+
"type": "add_segment",
|
1255
|
+
"segment_index": len(self.segment_manager.segments) - 1,
|
1256
|
+
}
|
1257
|
+
)
|
1258
|
+
# Clear redo history when a new action is performed
|
1259
|
+
self.redo_history.clear()
|
1260
|
+
self._update_all_lists()
|
1261
|
+
self._show_success_notification("AI bounding box segmentation saved!")
|
1262
|
+
else:
|
1263
|
+
self._show_warning_notification(
|
1264
|
+
"All segments filtered out by fragment threshold"
|
1265
|
+
)
|
1266
|
+
|
1267
|
+
# Clear bounding box preview state
|
1268
|
+
self.ai_bbox_preview_mask = None
|
1269
|
+
self.ai_bbox_preview_rect = None
|
1270
|
+
|
1271
|
+
# Clear preview
|
1272
|
+
if hasattr(self, "preview_mask_item") and self.preview_mask_item:
|
1273
|
+
self.viewer.scene().removeItem(self.preview_mask_item)
|
1274
|
+
self.preview_mask_item = None
|
1275
|
+
return
|
1276
|
+
|
1277
|
+
# Handle point-based predictions (existing behavior)
|
1278
|
+
if not hasattr(self, "preview_mask_item") or not self.preview_mask_item:
|
1279
|
+
return
|
1280
|
+
|
1281
|
+
result = self.model_manager.sam_model.predict(
|
1201
1282
|
self.positive_points, self.negative_points
|
1202
1283
|
)
|
1203
|
-
if
|
1284
|
+
if result is not None:
|
1285
|
+
mask, scores, logits = result
|
1286
|
+
|
1287
|
+
# Ensure mask is boolean (SAM models can return float masks)
|
1288
|
+
if mask.dtype != bool:
|
1289
|
+
mask = mask > 0.5 # Convert float mask to boolean
|
1290
|
+
|
1204
1291
|
# COORDINATE TRANSFORMATION FIX: Scale mask back up to display size if needed
|
1205
1292
|
if (
|
1206
1293
|
self.sam_scale_factor != 1.0
|
@@ -1702,6 +1789,14 @@ class MainWindow(QMainWindow):
|
|
1702
1789
|
self.crop_start_pos = None
|
1703
1790
|
self.current_crop_coords = None
|
1704
1791
|
|
1792
|
+
# Reset AI mode state
|
1793
|
+
self.ai_click_start_pos = None
|
1794
|
+
self.ai_click_time = 0
|
1795
|
+
if hasattr(self, "ai_rubber_band_rect") and self.ai_rubber_band_rect:
|
1796
|
+
if self.ai_rubber_band_rect.scene():
|
1797
|
+
self.viewer.scene().removeItem(self.ai_rubber_band_rect)
|
1798
|
+
self.ai_rubber_band_rect = None
|
1799
|
+
|
1705
1800
|
items_to_remove = [
|
1706
1801
|
item
|
1707
1802
|
for item in self.viewer.scene().items()
|
@@ -1714,6 +1809,10 @@ class MainWindow(QMainWindow):
|
|
1714
1809
|
self.action_history.clear()
|
1715
1810
|
self.redo_history.clear()
|
1716
1811
|
|
1812
|
+
# Add bounding box preview state
|
1813
|
+
self.ai_bbox_preview_mask = None
|
1814
|
+
self.ai_bbox_preview_rect = None
|
1815
|
+
|
1717
1816
|
def _scene_mouse_press(self, event):
|
1718
1817
|
"""Handle mouse press events in the scene."""
|
1719
1818
|
# Map scene coordinates to the view so items() works correctly.
|
@@ -1763,10 +1862,19 @@ class MainWindow(QMainWindow):
|
|
1763
1862
|
elif self.mode == "sam_points":
|
1764
1863
|
if event.button() == Qt.MouseButton.LeftButton:
|
1765
1864
|
self._add_point(pos, positive=True)
|
1766
|
-
self._update_segmentation()
|
1767
1865
|
elif event.button() == Qt.MouseButton.RightButton:
|
1768
1866
|
self._add_point(pos, positive=False)
|
1769
|
-
|
1867
|
+
elif self.mode == "ai":
|
1868
|
+
if event.button() == Qt.MouseButton.LeftButton:
|
1869
|
+
# AI mode: single click adds point, drag creates bounding box
|
1870
|
+
self.ai_click_start_pos = pos
|
1871
|
+
self.ai_click_time = (
|
1872
|
+
event.timestamp() if hasattr(event, "timestamp") else 0
|
1873
|
+
)
|
1874
|
+
# We'll determine if it's a click or drag in mouse_release
|
1875
|
+
elif event.button() == Qt.MouseButton.RightButton:
|
1876
|
+
# Right-click adds negative point in AI mode
|
1877
|
+
self._add_point(pos, positive=False, update_segmentation=True)
|
1770
1878
|
elif self.mode == "polygon":
|
1771
1879
|
if event.button() == Qt.MouseButton.LeftButton:
|
1772
1880
|
self._handle_polygon_click(pos)
|
@@ -1816,6 +1924,40 @@ class MainWindow(QMainWindow):
|
|
1816
1924
|
event.accept()
|
1817
1925
|
return
|
1818
1926
|
|
1927
|
+
if (
|
1928
|
+
self.mode == "ai"
|
1929
|
+
and hasattr(self, "ai_click_start_pos")
|
1930
|
+
and self.ai_click_start_pos
|
1931
|
+
):
|
1932
|
+
current_pos = event.scenePos()
|
1933
|
+
# Check if we've moved enough to consider this a drag
|
1934
|
+
drag_distance = (
|
1935
|
+
(current_pos.x() - self.ai_click_start_pos.x()) ** 2
|
1936
|
+
+ (current_pos.y() - self.ai_click_start_pos.y()) ** 2
|
1937
|
+
) ** 0.5
|
1938
|
+
|
1939
|
+
if drag_distance > 5: # Minimum drag distance
|
1940
|
+
# Create rubber band if not exists
|
1941
|
+
if (
|
1942
|
+
not hasattr(self, "ai_rubber_band_rect")
|
1943
|
+
or not self.ai_rubber_band_rect
|
1944
|
+
):
|
1945
|
+
self.ai_rubber_band_rect = QGraphicsRectItem()
|
1946
|
+
self.ai_rubber_band_rect.setPen(
|
1947
|
+
QPen(
|
1948
|
+
Qt.GlobalColor.cyan,
|
1949
|
+
self.line_thickness,
|
1950
|
+
Qt.PenStyle.DashLine,
|
1951
|
+
)
|
1952
|
+
)
|
1953
|
+
self.viewer.scene().addItem(self.ai_rubber_band_rect)
|
1954
|
+
|
1955
|
+
# Update rubber band
|
1956
|
+
rect = QRectF(self.ai_click_start_pos, current_pos).normalized()
|
1957
|
+
self.ai_rubber_band_rect.setRect(rect)
|
1958
|
+
event.accept()
|
1959
|
+
return
|
1960
|
+
|
1819
1961
|
if self.mode == "crop" and self.crop_rect_item and self.crop_start_pos:
|
1820
1962
|
current_pos = event.scenePos()
|
1821
1963
|
rect = QRectF(self.crop_start_pos, current_pos).normalized()
|
@@ -1852,6 +1994,42 @@ class MainWindow(QMainWindow):
|
|
1852
1994
|
|
1853
1995
|
if self.mode == "pan":
|
1854
1996
|
self.viewer.set_cursor(Qt.CursorShape.OpenHandCursor)
|
1997
|
+
elif (
|
1998
|
+
self.mode == "ai"
|
1999
|
+
and hasattr(self, "ai_click_start_pos")
|
2000
|
+
and self.ai_click_start_pos
|
2001
|
+
):
|
2002
|
+
current_pos = event.scenePos()
|
2003
|
+
# Calculate drag distance
|
2004
|
+
drag_distance = (
|
2005
|
+
(current_pos.x() - self.ai_click_start_pos.x()) ** 2
|
2006
|
+
+ (current_pos.y() - self.ai_click_start_pos.y()) ** 2
|
2007
|
+
) ** 0.5
|
2008
|
+
|
2009
|
+
if (
|
2010
|
+
hasattr(self, "ai_rubber_band_rect")
|
2011
|
+
and self.ai_rubber_band_rect
|
2012
|
+
and drag_distance > 5
|
2013
|
+
):
|
2014
|
+
# This was a drag - use SAM bounding box prediction
|
2015
|
+
rect = self.ai_rubber_band_rect.rect()
|
2016
|
+
self.viewer.scene().removeItem(self.ai_rubber_band_rect)
|
2017
|
+
self.ai_rubber_band_rect = None
|
2018
|
+
self.ai_click_start_pos = None
|
2019
|
+
|
2020
|
+
if rect.width() > 10 and rect.height() > 10: # Minimum box size
|
2021
|
+
self._handle_ai_bounding_box(rect)
|
2022
|
+
else:
|
2023
|
+
# This was a click - add positive point
|
2024
|
+
self.ai_click_start_pos = None
|
2025
|
+
if hasattr(self, "ai_rubber_band_rect") and self.ai_rubber_band_rect:
|
2026
|
+
self.viewer.scene().removeItem(self.ai_rubber_band_rect)
|
2027
|
+
self.ai_rubber_band_rect = None
|
2028
|
+
|
2029
|
+
self._add_point(current_pos, positive=True, update_segmentation=True)
|
2030
|
+
|
2031
|
+
event.accept()
|
2032
|
+
return
|
1855
2033
|
elif self.mode == "bbox" and self.rubber_band_rect:
|
1856
2034
|
self.viewer.scene().removeItem(self.rubber_band_rect)
|
1857
2035
|
rect = self.rubber_band_rect.rect()
|
@@ -1909,14 +2087,88 @@ class MainWindow(QMainWindow):
|
|
1909
2087
|
|
1910
2088
|
self._original_mouse_release(event)
|
1911
2089
|
|
1912
|
-
def
|
2090
|
+
def _handle_ai_bounding_box(self, rect):
|
2091
|
+
"""Handle AI mode bounding box by using SAM's predict_from_box to create a preview."""
|
2092
|
+
if not self.model_manager.is_model_available():
|
2093
|
+
self._show_warning_notification("AI model not available", 2000)
|
2094
|
+
return
|
2095
|
+
|
2096
|
+
# Quick check - if currently updating, skip but don't block future attempts
|
2097
|
+
if self.sam_is_updating:
|
2098
|
+
self._show_warning_notification(
|
2099
|
+
"AI model is updating, please wait...", 2000
|
2100
|
+
)
|
2101
|
+
return
|
2102
|
+
|
2103
|
+
# Convert QRectF to SAM box format [x1, y1, x2, y2]
|
2104
|
+
# COORDINATE TRANSFORMATION FIX: Use proper coordinate mapping based on operate_on_view setting
|
2105
|
+
from PyQt6.QtCore import QPointF
|
2106
|
+
|
2107
|
+
top_left = QPointF(rect.left(), rect.top())
|
2108
|
+
bottom_right = QPointF(rect.right(), rect.bottom())
|
2109
|
+
|
2110
|
+
sam_x1, sam_y1 = self._transform_display_coords_to_sam_coords(top_left)
|
2111
|
+
sam_x2, sam_y2 = self._transform_display_coords_to_sam_coords(bottom_right)
|
2112
|
+
|
2113
|
+
box = [sam_x1, sam_y1, sam_x2, sam_y2]
|
2114
|
+
|
2115
|
+
try:
|
2116
|
+
result = self.model_manager.sam_model.predict_from_box(box)
|
2117
|
+
if result is not None:
|
2118
|
+
mask, scores, logits = result
|
2119
|
+
|
2120
|
+
# Ensure mask is boolean (SAM models can return float masks)
|
2121
|
+
if mask.dtype != bool:
|
2122
|
+
mask = mask > 0.5 # Convert float mask to boolean
|
2123
|
+
|
2124
|
+
# COORDINATE TRANSFORMATION FIX: Scale mask back up to display size if needed
|
2125
|
+
if (
|
2126
|
+
self.sam_scale_factor != 1.0
|
2127
|
+
and self.viewer._pixmap_item
|
2128
|
+
and not self.viewer._pixmap_item.pixmap().isNull()
|
2129
|
+
):
|
2130
|
+
# Get original image dimensions
|
2131
|
+
original_height = self.viewer._pixmap_item.pixmap().height()
|
2132
|
+
original_width = self.viewer._pixmap_item.pixmap().width()
|
2133
|
+
|
2134
|
+
# Resize mask back to original dimensions for saving
|
2135
|
+
mask_resized = cv2.resize(
|
2136
|
+
mask.astype(np.uint8),
|
2137
|
+
(original_width, original_height),
|
2138
|
+
interpolation=cv2.INTER_NEAREST,
|
2139
|
+
).astype(bool)
|
2140
|
+
mask = mask_resized
|
2141
|
+
|
2142
|
+
# Store the preview mask and rect for later confirmation
|
2143
|
+
self.ai_bbox_preview_mask = mask
|
2144
|
+
self.ai_bbox_preview_rect = rect
|
2145
|
+
|
2146
|
+
# Clear any existing preview
|
2147
|
+
if hasattr(self, "preview_mask_item") and self.preview_mask_item:
|
2148
|
+
self.viewer.scene().removeItem(self.preview_mask_item)
|
2149
|
+
|
2150
|
+
# Show preview with yellow color
|
2151
|
+
pixmap = mask_to_pixmap(mask, (255, 255, 0))
|
2152
|
+
self.preview_mask_item = self.viewer.scene().addPixmap(pixmap)
|
2153
|
+
self.preview_mask_item.setZValue(50)
|
2154
|
+
|
2155
|
+
self._show_success_notification(
|
2156
|
+
"AI bounding box preview ready - press Space to confirm!"
|
2157
|
+
)
|
2158
|
+
else:
|
2159
|
+
self._show_warning_notification("No prediction result from AI model")
|
2160
|
+
except Exception as e:
|
2161
|
+
logger.error(f"Error during AI bounding box prediction: {e}")
|
2162
|
+
self._show_error_notification("AI prediction failed")
|
2163
|
+
|
2164
|
+
def _add_point(self, pos, positive, update_segmentation=True):
|
1913
2165
|
"""Add a point for SAM segmentation."""
|
1914
2166
|
# RACE CONDITION FIX: Block clicks during SAM updates
|
1915
2167
|
if self.sam_is_updating:
|
1916
2168
|
self._show_warning_notification(
|
1917
2169
|
"AI model is updating, please wait...", 2000
|
1918
2170
|
)
|
1919
|
-
return
|
2171
|
+
return False
|
1920
2172
|
|
1921
2173
|
# Ensure SAM is updated before using it
|
1922
2174
|
self._ensure_sam_updated()
|
@@ -1926,11 +2178,10 @@ class MainWindow(QMainWindow):
|
|
1926
2178
|
self._show_warning_notification(
|
1927
2179
|
"AI model is updating, please wait...", 2000
|
1928
2180
|
)
|
1929
|
-
return
|
2181
|
+
return False
|
1930
2182
|
|
1931
|
-
# COORDINATE TRANSFORMATION FIX:
|
1932
|
-
sam_x =
|
1933
|
-
sam_y = int(pos.y() * self.sam_scale_factor)
|
2183
|
+
# COORDINATE TRANSFORMATION FIX: Use proper coordinate mapping based on operate_on_view setting
|
2184
|
+
sam_x, sam_y = self._transform_display_coords_to_sam_coords(pos)
|
1934
2185
|
|
1935
2186
|
point_list = self.positive_points if positive else self.negative_points
|
1936
2187
|
point_list.append([sam_x, sam_y])
|
@@ -1965,6 +2216,12 @@ class MainWindow(QMainWindow):
|
|
1965
2216
|
# Clear redo history when a new action is performed
|
1966
2217
|
self.redo_history.clear()
|
1967
2218
|
|
2219
|
+
# Update segmentation if requested and not currently updating
|
2220
|
+
if update_segmentation and not self.sam_is_updating:
|
2221
|
+
self._update_segmentation()
|
2222
|
+
|
2223
|
+
return True
|
2224
|
+
|
1968
2225
|
def _update_segmentation(self):
|
1969
2226
|
"""Update SAM segmentation preview."""
|
1970
2227
|
if hasattr(self, "preview_mask_item") and self.preview_mask_item:
|
@@ -1972,10 +2229,16 @@ class MainWindow(QMainWindow):
|
|
1972
2229
|
if not self.positive_points or not self.model_manager.is_model_available():
|
1973
2230
|
return
|
1974
2231
|
|
1975
|
-
|
2232
|
+
result = self.model_manager.sam_model.predict(
|
1976
2233
|
self.positive_points, self.negative_points
|
1977
2234
|
)
|
1978
|
-
if
|
2235
|
+
if result is not None:
|
2236
|
+
mask, scores, logits = result
|
2237
|
+
|
2238
|
+
# Ensure mask is boolean (SAM models can return float masks)
|
2239
|
+
if mask.dtype != bool:
|
2240
|
+
mask = mask > 0.5 # Convert float mask to boolean
|
2241
|
+
|
1979
2242
|
# COORDINATE TRANSFORMATION FIX: Scale mask back up to display size if needed
|
1980
2243
|
if (
|
1981
2244
|
self.sam_scale_factor != 1.0
|
@@ -2356,14 +2619,6 @@ class MainWindow(QMainWindow):
|
|
2356
2619
|
self.main_splitter.setSizes([250, 800, 350])
|
2357
2620
|
|
2358
2621
|
# Additional methods for new features
|
2359
|
-
def _handle_settings_changed(self):
|
2360
|
-
"""Handle changes in settings, e.g., 'Operate On View'."""
|
2361
|
-
# Update the main window's settings object with the latest from the widget
|
2362
|
-
self.settings.update(**self.control_panel.settings_widget.get_settings())
|
2363
|
-
|
2364
|
-
# Re-load the current image to apply the new 'Operate On View' setting
|
2365
|
-
if self.current_image_path:
|
2366
|
-
self._update_sam_model_image()
|
2367
2622
|
|
2368
2623
|
def _handle_channel_threshold_changed(self):
|
2369
2624
|
"""Handle changes in channel thresholding - optimized to avoid unnecessary work."""
|
@@ -2428,10 +2683,21 @@ class MainWindow(QMainWindow):
|
|
2428
2683
|
self.sam_is_dirty = False
|
2429
2684
|
return
|
2430
2685
|
|
2431
|
-
#
|
2686
|
+
# IMPROVED: More robust worker thread cleanup
|
2432
2687
|
if self.sam_worker_thread and self.sam_worker_thread.isRunning():
|
2433
2688
|
self.sam_worker_thread.stop()
|
2434
|
-
self.sam_worker_thread.
|
2689
|
+
self.sam_worker_thread.terminate()
|
2690
|
+
# Wait longer for proper cleanup
|
2691
|
+
self.sam_worker_thread.wait(5000) # Wait up to 5 seconds
|
2692
|
+
if self.sam_worker_thread.isRunning():
|
2693
|
+
# Force kill if still running
|
2694
|
+
self.sam_worker_thread.quit()
|
2695
|
+
self.sam_worker_thread.wait(2000)
|
2696
|
+
|
2697
|
+
# Clean up old worker thread
|
2698
|
+
if self.sam_worker_thread:
|
2699
|
+
self.sam_worker_thread.deleteLater()
|
2700
|
+
self.sam_worker_thread = None
|
2435
2701
|
|
2436
2702
|
# Show status message
|
2437
2703
|
if hasattr(self, "status_bar"):
|
@@ -2741,6 +3007,9 @@ class MainWindow(QMainWindow):
|
|
2741
3007
|
# Update the FFT threshold widget
|
2742
3008
|
self.control_panel.update_fft_threshold_for_image(image_array)
|
2743
3009
|
|
3010
|
+
# Auto-collapse FFT threshold panel if image is not black and white
|
3011
|
+
self.control_panel.auto_collapse_fft_threshold_for_image(image_array)
|
3012
|
+
|
2744
3013
|
# Border crop methods
|
2745
3014
|
def _start_crop_drawing(self):
|
2746
3015
|
"""Start crop drawing mode."""
|
@@ -3083,3 +3352,139 @@ class MainWindow(QMainWindow):
|
|
3083
3352
|
"""Update SAM model image after debounce delay."""
|
3084
3353
|
# This is called after the user stops interacting with sliders
|
3085
3354
|
self._update_sam_model_image()
|
3355
|
+
|
3356
|
+
def _reset_sam_state_for_model_switch(self):
|
3357
|
+
"""Reset SAM state completely when switching models to prevent worker thread conflicts."""
|
3358
|
+
|
3359
|
+
# CRITICAL: Force terminate any running SAM worker thread
|
3360
|
+
if self.sam_worker_thread and self.sam_worker_thread.isRunning():
|
3361
|
+
self.sam_worker_thread.stop()
|
3362
|
+
self.sam_worker_thread.terminate()
|
3363
|
+
self.sam_worker_thread.wait(3000) # Wait up to 3 seconds
|
3364
|
+
if self.sam_worker_thread.isRunning():
|
3365
|
+
# Force kill if still running
|
3366
|
+
self.sam_worker_thread.quit()
|
3367
|
+
self.sam_worker_thread.wait(1000)
|
3368
|
+
|
3369
|
+
# Clean up worker thread reference
|
3370
|
+
if self.sam_worker_thread:
|
3371
|
+
self.sam_worker_thread.deleteLater()
|
3372
|
+
self.sam_worker_thread = None
|
3373
|
+
|
3374
|
+
# Reset SAM update flags
|
3375
|
+
self.sam_is_updating = False
|
3376
|
+
self.sam_is_dirty = True # Force update with new model
|
3377
|
+
self.current_sam_hash = None # Invalidate cache
|
3378
|
+
self.sam_scale_factor = 1.0
|
3379
|
+
|
3380
|
+
# Clear all points but preserve segments
|
3381
|
+
self.clear_all_points()
|
3382
|
+
# Note: Segments are preserved when switching models
|
3383
|
+
self._update_all_lists()
|
3384
|
+
|
3385
|
+
# Clear preview items
|
3386
|
+
if hasattr(self, "preview_mask_item") and self.preview_mask_item:
|
3387
|
+
if self.preview_mask_item.scene():
|
3388
|
+
self.viewer.scene().removeItem(self.preview_mask_item)
|
3389
|
+
self.preview_mask_item = None
|
3390
|
+
|
3391
|
+
# Clean up crop visuals
|
3392
|
+
self._remove_crop_visual()
|
3393
|
+
self._remove_crop_hover_overlay()
|
3394
|
+
self._remove_crop_hover_effect()
|
3395
|
+
|
3396
|
+
# Reset crop state
|
3397
|
+
self.crop_mode = False
|
3398
|
+
self.crop_start_pos = None
|
3399
|
+
self.current_crop_coords = None
|
3400
|
+
|
3401
|
+
# Reset AI mode state
|
3402
|
+
self.ai_click_start_pos = None
|
3403
|
+
self.ai_click_time = 0
|
3404
|
+
if hasattr(self, "ai_rubber_band_rect") and self.ai_rubber_band_rect:
|
3405
|
+
if self.ai_rubber_band_rect.scene():
|
3406
|
+
self.viewer.scene().removeItem(self.ai_rubber_band_rect)
|
3407
|
+
self.ai_rubber_band_rect = None
|
3408
|
+
|
3409
|
+
# Clear all graphics items except the main image
|
3410
|
+
items_to_remove = [
|
3411
|
+
item
|
3412
|
+
for item in self.viewer.scene().items()
|
3413
|
+
if item is not self.viewer._pixmap_item
|
3414
|
+
]
|
3415
|
+
for item in items_to_remove:
|
3416
|
+
self.viewer.scene().removeItem(item)
|
3417
|
+
|
3418
|
+
# Reset all collections
|
3419
|
+
self.segment_items.clear()
|
3420
|
+
self.highlight_items.clear()
|
3421
|
+
self.action_history.clear()
|
3422
|
+
self.redo_history.clear()
|
3423
|
+
|
3424
|
+
# Reset bounding box preview state
|
3425
|
+
self.ai_bbox_preview_mask = None
|
3426
|
+
self.ai_bbox_preview_rect = None
|
3427
|
+
|
3428
|
+
# Clear status bar messages
|
3429
|
+
if hasattr(self, "status_bar"):
|
3430
|
+
self.status_bar.clear_message()
|
3431
|
+
|
3432
|
+
# Redisplay segments after model switch to restore visual representation
|
3433
|
+
self._display_all_segments()
|
3434
|
+
|
3435
|
+
def _transform_display_coords_to_sam_coords(self, pos):
|
3436
|
+
"""Transform display coordinates to SAM model coordinates.
|
3437
|
+
|
3438
|
+
When 'operate on view' is ON: SAM processes the displayed image
|
3439
|
+
When 'operate on view' is OFF: SAM processes the original image
|
3440
|
+
"""
|
3441
|
+
if self.settings.operate_on_view:
|
3442
|
+
# Simple case: SAM processes the same image the user sees
|
3443
|
+
sam_x = int(pos.x() * self.sam_scale_factor)
|
3444
|
+
sam_y = int(pos.y() * self.sam_scale_factor)
|
3445
|
+
else:
|
3446
|
+
# Complex case: Map display coordinates to original image coordinates
|
3447
|
+
# then scale for SAM processing
|
3448
|
+
|
3449
|
+
# Get displayed image dimensions (may include adjustments)
|
3450
|
+
if (
|
3451
|
+
not self.viewer._pixmap_item
|
3452
|
+
or self.viewer._pixmap_item.pixmap().isNull()
|
3453
|
+
):
|
3454
|
+
# Fallback: use simple scaling
|
3455
|
+
sam_x = int(pos.x() * self.sam_scale_factor)
|
3456
|
+
sam_y = int(pos.y() * self.sam_scale_factor)
|
3457
|
+
else:
|
3458
|
+
display_width = self.viewer._pixmap_item.pixmap().width()
|
3459
|
+
display_height = self.viewer._pixmap_item.pixmap().height()
|
3460
|
+
|
3461
|
+
# Get original image dimensions
|
3462
|
+
if not self.current_image_path:
|
3463
|
+
# Fallback: use simple scaling
|
3464
|
+
sam_x = int(pos.x() * self.sam_scale_factor)
|
3465
|
+
sam_y = int(pos.y() * self.sam_scale_factor)
|
3466
|
+
else:
|
3467
|
+
# Load original image to get true dimensions
|
3468
|
+
original_pixmap = QPixmap(self.current_image_path)
|
3469
|
+
if original_pixmap.isNull():
|
3470
|
+
# Fallback: use simple scaling
|
3471
|
+
sam_x = int(pos.x() * self.sam_scale_factor)
|
3472
|
+
sam_y = int(pos.y() * self.sam_scale_factor)
|
3473
|
+
else:
|
3474
|
+
original_width = original_pixmap.width()
|
3475
|
+
original_height = original_pixmap.height()
|
3476
|
+
|
3477
|
+
# Map display coordinates to original image coordinates
|
3478
|
+
if display_width > 0 and display_height > 0:
|
3479
|
+
original_x = pos.x() * (original_width / display_width)
|
3480
|
+
original_y = pos.y() * (original_height / display_height)
|
3481
|
+
|
3482
|
+
# Apply SAM scale factor to original coordinates
|
3483
|
+
sam_x = int(original_x * self.sam_scale_factor)
|
3484
|
+
sam_y = int(original_y * self.sam_scale_factor)
|
3485
|
+
else:
|
3486
|
+
# Fallback: use simple scaling
|
3487
|
+
sam_x = int(pos.x() * self.sam_scale_factor)
|
3488
|
+
sam_y = int(pos.y() * self.sam_scale_factor)
|
3489
|
+
|
3490
|
+
return sam_x, sam_y
|