napari-tmidas 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

napari_tmidas/_crop_anything.py

@@ -9,32 +9,99 @@ The plugin supports both 2D (YX) and 3D (TYX/ZYX) data.
 import contextlib
 import os
 import sys
+from pathlib import Path
 
 import numpy as np
-import requests
-import torch
-from magicgui import magicgui
-from napari.layers import Labels
-from napari.viewer import Viewer
-from qtpy.QtCore import Qt
-from qtpy.QtWidgets import (
-    QCheckBox,
-    QFileDialog,
-    QHBoxLayout,
-    QHeaderView,
-    QLabel,
-    QMessageBox,
-    QPushButton,
-    QScrollArea,
-    QTableWidget,
-    QTableWidgetItem,
-    QVBoxLayout,
-    QWidget,
-)
-from skimage.io import imread
-from skimage.transform import resize
-from tifffile import imwrite
 
+# Lazy imports for optional heavy dependencies
+try:
+    import requests
+
+    _HAS_REQUESTS = True
+except ImportError:
+    requests = None
+    _HAS_REQUESTS = False
+
+try:
+    import torch
+
+    _HAS_TORCH = True
+except ImportError:
+    torch = None
+    _HAS_TORCH = False
+
+try:
+    from magicgui import magicgui
+
+    _HAS_MAGICGUI = True
+except ImportError:
+    # Create stub decorator
+    def magicgui(*args, **kwargs):
+        def decorator(func):
+            return func
+
+        if len(args) == 1 and callable(args[0]) and not kwargs:
+            return args[0]
+        return decorator
+
+    _HAS_MAGICGUI = False
+
+try:
+    from napari.layers import Labels
+    from napari.viewer import Viewer
+
+    _HAS_NAPARI = True
+except ImportError:
+    Labels = None
+    Viewer = None
+    _HAS_NAPARI = False
+
+try:
+    from qtpy.QtCore import Qt
+    from qtpy.QtWidgets import (
+        QCheckBox,
+        QHBoxLayout,
+        QHeaderView,
+        QLabel,
+        QMessageBox,
+        QPushButton,
+        QScrollArea,
+        QTableWidget,
+        QTableWidgetItem,
+        QVBoxLayout,
+        QWidget,
+    )
+
+    _HAS_QTPY = True
+except ImportError:
+    Qt = None
+    QCheckBox = QHBoxLayout = QHeaderView = QLabel = QMessageBox = None
+    QPushButton = QScrollArea = QTableWidget = QTableWidgetItem = None
+    QVBoxLayout = QWidget = None
+    _HAS_QTPY = False
+
+try:
+    from skimage.io import imread
+    from skimage.transform import resize
+
+    _HAS_SKIMAGE = True
+except ImportError:
+    imread = None
+    resize = None
+    _HAS_SKIMAGE = False
+
+try:
+    from tifffile import imwrite
+
+    _HAS_TIFFFILE = True
+except ImportError:
+    imwrite = None
+    _HAS_TIFFFILE = False
+
+from napari_tmidas._file_selector import (
+    load_image_file as load_any_image,
+)
+from napari_tmidas._ui_utils import add_browse_button_to_folder_field
 from napari_tmidas.processing_functions.sam2_mp4 import tif_to_mp4
 
 sam2_paths = [
@@ -98,6 +165,7 @@ class BatchCropAnything:
         self.image_layer = None
         self.label_layer = None
         self.label_table_widget = None
+        self.shapes_layer = None
 
         # State tracking
         self.selected_labels = set()
@@ -106,6 +174,9 @@ class BatchCropAnything:
         # Segmentation parameters
         self.sensitivity = 50  # Default sensitivity (0-100 scale)
 
+        # Prompt mode: 'point' or 'box'
+        self.prompt_mode = "point"
+
         # Initialize the SAM2 model
         self._initialize_sam2()
 
@@ -131,17 +202,45 @@ class BatchCropAnything:
 
         try:
             # import torch
+            print("DEBUG: Starting SAM2 initialization...")
 
             self.device = get_device()
+            print(f"DEBUG: Device set to {self.device}")
 
             # Download checkpoint if needed
             checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt"
             checkpoint_path = download_checkpoint(
                 checkpoint_url, "/opt/sam2/checkpoints/"
             )
+            print(f"DEBUG: Checkpoint path: {checkpoint_path}")
+
+            # Use relative config path for SAM2's Hydra config system
             model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+            print(f"DEBUG: Model config: {model_cfg}")
+
+            # Verify the actual config file exists in the SAM2 installation
+            sam2_base_path = None
+            for path in sam2_paths:
+                if path and os.path.exists(path):
+                    sam2_base_path = path
+                    break
+
+            if sam2_base_path is not None:
+                full_config_path = os.path.join(
+                    sam2_base_path, "sam2", model_cfg
+                )
+                if not os.path.exists(full_config_path):
+                    raise FileNotFoundError(
+                        f"SAM2 config file not found at: {full_config_path}"
+                    )
+                print(f"DEBUG: Verified config exists at: {full_config_path}")
+            else:
+                print(
+                    "DEBUG: Warning - could not verify config file exists, but proceeding with relative path"
+                )
 
             if self.use_3d:
+                print("DEBUG: Initializing SAM2 Video Predictor...")
                 from sam2.build_sam import build_sam2_video_predictor
 
                 self.predictor = build_sam2_video_predictor(
@@ -150,7 +249,9 @@ class BatchCropAnything:
                 self.viewer.status = (
                     f"Initialized SAM2 Video Predictor on {self.device}"
                 )
+                print("DEBUG: SAM2 Video Predictor initialized successfully")
             else:
+                print("DEBUG: Initializing SAM2 Image Predictor...")
                 from sam2.build_sam import build_sam2
                 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
@@ -160,6 +261,7 @@ class BatchCropAnything:
                 self.viewer.status = (
                     f"Initialized SAM2 Image Predictor on {self.device}"
                 )
+                print("DEBUG: SAM2 Image Predictor initialized successfully")
 
         except (
             ImportError,
@@ -167,37 +269,79 @@ class BatchCropAnything:
             ValueError,
             FileNotFoundError,
             requests.RequestException,
+            AttributeError,
+            ModuleNotFoundError,
         ) as e:
             import traceback
 
-            self.viewer.status = f"Error initializing SAM2: {str(e)}"
+            error_msg = f"SAM2 initialization failed: {str(e)}"
+            error_type = type(e).__name__
+            self.viewer.status = (
+                f"{error_msg} - Images will load without segmentation"
+            )
             self.predictor = None
+            print(f"DEBUG: SAM2 Error ({error_type}): {error_msg}")
+            print("DEBUG: Full traceback:")
             print(traceback.format_exc())
+            print(
+                "DEBUG: Note: Images will still load, but automatic segmentation will not be available."
+            )
+
+            # Provide specific guidance based on error type
+            if isinstance(e, FileNotFoundError):
+                print(
+                    "DEBUG: This appears to be a missing file issue. Check SAM2 installation and config paths."
+                )
+            elif isinstance(e, (ImportError, ModuleNotFoundError)):
+                print(
+                    "DEBUG: This appears to be a SAM2 import issue. Check SAM2 installation."
+                )
+            elif isinstance(e, RuntimeError):
+                print(
+                    "DEBUG: This appears to be a runtime issue, possibly GPU/CUDA related."
+                )
+            else:
+                print(f"DEBUG: Unexpected error type: {error_type}")
 
     def load_images(self, folder_path: str):
         """Load images from the specified folder path."""
+        print(f"DEBUG: Loading images from folder: {folder_path}")
         if not os.path.exists(folder_path):
             self.viewer.status = f"Folder not found: {folder_path}"
+            print(f"DEBUG: Folder does not exist: {folder_path}")
             return
 
         files = os.listdir(folder_path)
-        self.images = [
-            os.path.join(folder_path, file)
-            for file in files
-            if file.lower().endswith(".tif")
-            or file.lower().endswith(".tiff")
-            and "label" not in file.lower()
-            and "cropped" not in file.lower()
-            and "_labels_" not in file.lower()
-            and "_cropped_" not in file.lower()
-        ]
+        print(f"DEBUG: Found {len(files)} files in folder")
+        self.images = []
+        for file in files:
+            full = os.path.join(folder_path, file)
+            low = file.lower()
+            if (
+                low.endswith((".tif", ".tiff"))
+                or (os.path.isdir(full) and low.endswith(".zarr"))
+            ) and (
+                "label" not in low
+                and "_labels_" not in low
+                and "sam2"
+                not in low  # Exclude any SAM2-related files (including output from this tool)
+            ):
+                self.images.append(full)
+                print(f"DEBUG: Added image: {file}")
+            else:
+                print(
+                    f"DEBUG: Excluded file: {file} (reason: filtering criteria)"
+                )
 
         if not self.images:
             self.viewer.status = "No compatible images found in the folder."
+            print("DEBUG: No compatible images found")
             return
 
+        print(f"DEBUG: Total compatible images found: {len(self.images)}")
         self.viewer.status = f"Found {len(self.images)} .tif images."
         self.current_index = 0
+        print(f"DEBUG: About to load first image: {self.images[0]}")
         self._load_current_image()
 
     def next_image(self):
@@ -250,25 +394,69 @@ class BatchCropAnything:
 
     def _load_current_image(self):
         """Load the current image and generate segmentation."""
+        print("DEBUG: _load_current_image called")
         if not self.images:
             self.viewer.status = "No images to process."
-            return
-
-        if self.predictor is None:
-            self.viewer.status = (
-                "SAM2 model not initialized. Cannot segment images."
-            )
+            print("DEBUG: No images to process")
             return
 
         image_path = self.images[self.current_index]
-        self.viewer.status = f"Processing {os.path.basename(image_path)}"
+        print(f"DEBUG: Loading image at path: {image_path}")
+
+        if self.predictor is None:
+            self.viewer.status = f"Loading {os.path.basename(image_path)} (SAM2 model not initialized - no segmentation will be available)"
+            print("DEBUG: SAM2 predictor is None")
+        else:
+            self.viewer.status = f"Processing {os.path.basename(image_path)}"
+            print("DEBUG: SAM2 predictor is available")
 
         try:
+            print("DEBUG: About to clear viewer layers")
             # Clear existing layers
             self.viewer.layers.clear()
+            print("DEBUG: Viewer layers cleared")
 
+            print("DEBUG: About to load image file")
             # Load and process image
-            self.original_image = imread(image_path)
+            if image_path.lower().endswith(".zarr") or (
+                os.path.isdir(image_path)
+                and image_path.lower().endswith(".zarr")
+            ):
+                print("DEBUG: Loading Zarr file")
+                data = load_any_image(image_path)
+                # If multiple layers returned, take first image layer
+                if isinstance(data, list):
+                    img = None
+                    for entry in data:
+                        if isinstance(entry, tuple) and len(entry) == 3:
+                            d, _kwargs, layer_type = entry
+                            if layer_type == "image":
+                                img = d
+                                break
+                        elif isinstance(entry, tuple) and len(entry) == 2:
+                            d, _kwargs = entry
+                            img = d
+                            break
+                        else:
+                            img = entry
+                            break
+                    if img is None:
+                        raise ValueError("No image layer found in Zarr store")
+                else:
+                    img = data
+
+                # Compute dask arrays to numpy if needed
+                if hasattr(img, "compute"):
+                    img = img.compute()
+
+                self.original_image = img
+            else:
+                print("DEBUG: Loading TIFF file")
+                self.original_image = imread(image_path)
+
+            print(
+                f"DEBUG: Image loaded, shape: {self.original_image.shape}, dtype: {self.original_image.dtype}"
+            )
 
             # For 3D/4D data, determine dimensions
             if self.use_3d and len(self.original_image.shape) >= 3:
@@ -284,10 +472,12 @@ class BatchCropAnything:
 
                     if time_dim_idx == 0:  # TZYX format
                         # Keep as is, T is already the first dimension
+                        print("DEBUG: Adding 4D image (TZYX format) to viewer")
                         self.image_layer = self.viewer.add_image(
                             self.original_image,
                             name=f"Image ({os.path.basename(image_path)})",
                         )
+                        print(f"DEBUG: Added image layer: {self.image_layer}")
                         # Store time dimension info
                         self.time_dim_size = self.original_image.shape[0]
                         self.has_z_dim = True
@@ -309,19 +499,23 @@ class BatchCropAnything:
                             transposed_image  # Replace with transposed version
                         )
 
+                        print("DEBUG: Adding transposed 4D image to viewer")
                         self.image_layer = self.viewer.add_image(
                             self.original_image,
                             name=f"Image ({os.path.basename(image_path)})",
                         )
+                        print(f"DEBUG: Added image layer: {self.image_layer}")
                         # Store time dimension info
                         self.time_dim_size = self.original_image.shape[0]
                         self.has_z_dim = True
                     else:
                         # No time dimension found, treat as ZYX
+                        print("DEBUG: Adding 4D image (ZYX format) to viewer")
                         self.image_layer = self.viewer.add_image(
                             self.original_image,
                             name=f"Image ({os.path.basename(image_path)})",
                         )
+                        print(f"DEBUG: Added image layer: {self.image_layer}")
                         self.time_dim_size = 1
                         self.has_z_dim = True
                 elif (
@@ -330,30 +524,37 @@ class BatchCropAnything:
                     # Check if first dimension is likely time (> 4, < 400)
                     if 4 < self.original_image.shape[0] < 400:
                         # Likely TYX format
+                        print("DEBUG: Adding 3D image (TYX format) to viewer")
                         self.image_layer = self.viewer.add_image(
                             self.original_image,
                             name=f"Image ({os.path.basename(image_path)})",
                         )
+                        print(f"DEBUG: Added image layer: {self.image_layer}")
                         self.time_dim_size = self.original_image.shape[0]
                         self.has_z_dim = False
                     else:
                         # Likely ZYX format or another 3D format
+                        print("DEBUG: Adding 3D image (ZYX format) to viewer")
                         self.image_layer = self.viewer.add_image(
                             self.original_image,
                             name=f"Image ({os.path.basename(image_path)})",
                         )
+                        print(f"DEBUG: Added image layer: {self.image_layer}")
                         self.time_dim_size = 1
                         self.has_z_dim = True
                 else:
                     # Should not reach here with use_3d=True, but just in case
+                    print("DEBUG: Adding 3D image (fallback) to viewer")
                     self.image_layer = self.viewer.add_image(
                         self.original_image,
                         name=f"Image ({os.path.basename(image_path)})",
                     )
+                    print(f"DEBUG: Added image layer: {self.image_layer}")
                     self.time_dim_size = 1
                     self.has_z_dim = False
             else:
                 # Handle 2D data as before
+                print("DEBUG: Processing 2D image")
                 if self.original_image.dtype != np.uint8:
                     image_for_display = (
                         self.original_image
@@ -364,18 +565,42 @@ class BatchCropAnything:
                     image_for_display = self.original_image
 
                 # Add image to viewer
+                print("DEBUG: Adding 2D image to viewer")
                 self.image_layer = self.viewer.add_image(
                     image_for_display,
                     name=f"Image ({os.path.basename(image_path)})",
                 )
+                print(f"DEBUG: Added image layer: {self.image_layer}")
+
+            # Generate segmentation only if predictor is available
+            if self.predictor is not None:
+                print("DEBUG: About to generate segmentation")
+                self._generate_segmentation(self.original_image, image_path)
+                print("DEBUG: Segmentation generation completed")
+            else:
+                print("DEBUG: Creating empty segmentation (no predictor)")
+                # Create empty segmentation when predictor is not available
+                if self.use_3d:
+                    shape = self.original_image.shape
+                else:
+                    shape = self.original_image.shape[:2]
+
+                self.segmentation_result = np.zeros(shape, dtype=np.uint32)
+                self.label_layer = self.viewer.add_labels(
+                    self.segmentation_result,
+                    name="No Segmentation (SAM2 not available)",
+                )
+                print(f"DEBUG: Added empty label layer: {self.label_layer}")
 
-            # Generate segmentation
-            self._generate_segmentation(self.original_image, image_path)
+            print("DEBUG: _load_current_image completed successfully")
 
         except (FileNotFoundError, ValueError, TypeError, OSError) as e:
             import traceback
 
-            self.viewer.status = f"Error processing image: {str(e)}"
+            error_msg = f"Error processing image: {str(e)}"
+            self.viewer.status = error_msg
+            print(f"DEBUG: Exception in _load_current_image: {error_msg}")
+            print("DEBUG: Full traceback:")
             traceback.print_exc()
 
             # Create empty segmentation in case of error
@@ -392,6 +617,7 @@ class BatchCropAnything:
                 self.label_layer = self.viewer.add_labels(
                     self.segmentation_result, name="Error: No Segmentation"
                 )
+                print(f"DEBUG: Added error label layer: {self.label_layer}")
 
     def _generate_segmentation(self, image, image_path: str):
         """Generate segmentation for the current image using SAM2."""
@@ -447,7 +673,8 @@ class BatchCropAnything:
             traceback.print_exc()
 
     def _generate_2d_segmentation(self, confidence_threshold):
-        """Generate 2D segmentation using SAM2 Image Predictor."""
+        """Generate initial 2D segmentation - start with empty labels for interactive mode."""
+        device_type = "cuda" if self.device.type == "cuda" else "cpu"
         # Ensure image is in the correct format for SAM2
         image = self.current_image_for_segmentation
 
@@ -469,9 +696,7 @@ class BatchCropAnything:
                 (new_height, new_width),
                 anti_aliasing=True,
                 preserve_range=True,
-            ).astype(
-                np.float32
-            )  # Convert to float32
+            ).astype(np.float32)
 
             self.current_scale_factor = scale_factor
         else:
@@ -497,73 +722,54 @@ class BatchCropAnything:
         if resized_image.max() > 1.0:
             resized_image = resized_image / 255.0
 
-        # Set SAM2 prediction parameters based on sensitivity
-        with torch.inference_mode(), torch.autocast(
-            "cuda", dtype=torch.float32
-        ):
-            # Set the image in the predictor
-            self.predictor.set_image(resized_image)
+        # Store the prepared image for later use
+        self.prepared_sam2_image = resized_image
 
-            # Use automatic points generation with confidence threshold
-            masks, scores, _ = self.predictor.predict(
-                point_coords=None,
-                point_labels=None,
-                box=None,
-                multimask_output=True,
-            )
+        # Initialize empty segmentation result
+        self.segmentation_result = np.zeros(orig_shape, dtype=np.uint32)
+        self.label_info = {}
 
-            # Filter masks by confidence threshold
-            valid_masks = scores > confidence_threshold
-            masks = masks[valid_masks]
-            scores = scores[valid_masks]
-
-        # Convert masks to label image
-        labels = np.zeros(resized_image.shape[:2], dtype=np.uint32)
-        self.label_info = {}  # Reset label info
-
-        for i, mask in enumerate(masks):
-            label_id = i + 1  # Start label IDs from 1
-            labels[mask] = label_id
-
-            # Calculate label information
-            area = np.sum(mask)
-            y_indices, x_indices = np.where(mask)
-            center_y = np.mean(y_indices) if len(y_indices) > 0 else 0
-            center_x = np.mean(x_indices) if len(x_indices) > 0 else 0
-
-            # Store label info
-            self.label_info[label_id] = {
-                "area": area,
-                "center_y": center_y,
-                "center_x": center_x,
-                "score": float(scores[i]),
-            }
-
-        # Handle upscaling if needed
-        if self.current_scale_factor < 1.0:
-            labels = resize(
-                labels,
-                orig_shape,
-                order=0,  # Nearest neighbor interpolation
-                preserve_range=True,
-                anti_aliasing=False,
-            ).astype(np.uint32)
-
-        # Sort labels by area (largest first)
-        self.label_info = dict(
-            sorted(
-                self.label_info.items(),
-                key=lambda item: item[1]["area"],
-                reverse=True,
-            )
+        # Initialize tracking for interactive segmentation
+        self.current_points = []
+        self.current_labels = []
+        self.current_obj_id = 1
+        self.next_obj_id = 1
+
+        # Initialize object tracking dictionaries
+        self.obj_points = {}
+        self.obj_labels = {}
+
+        # Reset SAM2-specific tracking dictionaries for 2D mode
+        self.sam2_points_by_obj = {}
+        self.sam2_labels_by_obj = {}
+        self._sam2_next_obj_id = 1
+        print(
+            "DEBUG: Reset _sam2_next_obj_id to 1 in _generate_2d_segmentation"
         )
 
-        # Save segmentation result
-        self.segmentation_result = labels
+        # Set the image in the predictor for later use (2D mode only)
+        device_type = "cuda" if self.device.type == "cuda" else "cpu"
+        if hasattr(self.predictor, "set_image"):
+            with (
+                torch.inference_mode(),
+                torch.autocast(device_type, dtype=torch.float32),
+            ):
+                self.predictor.set_image(resized_image)
+        else:
+            print(
+                "DEBUG: Skipping set_image - predictor doesn't support it (likely VideoPredictor)"
+            )
 
         # Update the label layer
         self._update_label_layer()
 
+        # Show instructions
+        self.viewer.status = (
+            "2D Mode: Click on the image to add objects. Use Shift+click for negative points to refine. "
+            "Click existing objects to select them for cropping. "
+            "Note: For stacks, interactive segmentation only works in 2D view mode."
+        )
+
     def _generate_3d_segmentation(self, confidence_threshold, image_path):
         """
         Initialize 3D segmentation using SAM2 Video Predictor.
@@ -584,9 +790,7 @@ class BatchCropAnything:
             import tempfile
 
             temp_dir = tempfile.gettempdir()
-            mp4_path = os.path.join(
-                temp_dir, f"temp_volume_{os.path.basename(image_path)}.mp4"
-            )
+            mp4_path = None
 
             # If we need to save a modified version for MP4 conversion
             need_temp_tif = False
@@ -616,31 +820,72 @@ class BatchCropAnything:
                 imwrite(temp_tif_path, projected_volume)
                 need_temp_tif = True
 
-                # Convert the projected TIF to MP4
-                self.viewer.status = (
-                    "Converting projected 3D volume to MP4 format for SAM2..."
-                )
-                mp4_path = tif_to_mp4(temp_tif_path)
+                # Check if MP4 already exists
+                expected_mp4 = str(Path(temp_tif_path).with_suffix(".mp4"))
+                if os.path.exists(expected_mp4):
+                    self.viewer.status = (
+                        f"Using existing MP4: {os.path.basename(expected_mp4)}"
+                    )
+                    print(
+                        f"DEBUG: MP4 already exists, skipping conversion: {expected_mp4}"
+                    )
+                    mp4_path = expected_mp4
+                else:
+                    # Convert the projected TIF to MP4
+                    self.viewer.status = "Converting projected 3D volume to MP4 format for SAM2..."
+                    mp4_path = tif_to_mp4(temp_tif_path)
             else:
-                # Convert original volume to video format for SAM2
-                self.viewer.status = (
-                    "Converting 3D volume to MP4 format for SAM2..."
-                )
-                mp4_path = tif_to_mp4(image_path)
+                # Check if MP4 already exists for the original image
+                expected_mp4 = str(Path(image_path).with_suffix(".mp4"))
+                if os.path.exists(expected_mp4):
+                    self.viewer.status = (
+                        f"Using existing MP4: {os.path.basename(expected_mp4)}"
+                    )
+                    print(
+                        f"DEBUG: MP4 already exists, skipping conversion: {expected_mp4}"
+                    )
+                    mp4_path = expected_mp4
+                else:
+                    # Convert original volume to video format for SAM2
+                    self.viewer.status = (
+                        "Converting 3D volume to MP4 format for SAM2..."
+                    )
+                    mp4_path = tif_to_mp4(image_path)
 
             # Initialize SAM2 state with the video
             self.viewer.status = "Initializing SAM2 Video Predictor..."
-            with torch.inference_mode(), torch.autocast(
-                "cuda", dtype=torch.bfloat16
-            ):
-                self._sam2_state = self.predictor.init_state(mp4_path)
+            try:
+                device_type = "cuda" if self.device.type == "cuda" else "cpu"
+                with (
+                    torch.inference_mode(),
+                    torch.autocast(device_type, dtype=torch.float32),
+                ):
+                    self._sam2_state = self.predictor.init_state(mp4_path)
+            except (
+                RuntimeError,
+                ValueError,
+                TypeError,
+                torch.cuda.OutOfMemoryError,
+            ) as e:
+                self.viewer.status = (
+                    f"Error initializing SAM2 video predictor: {str(e)}"
+                )
+                print(f"SAM2 video predictor initialization failed: {e}")
+                return
 
             # Store needed state for 3D processing
             self._sam2_next_obj_id = 1
+            print(
+                "DEBUG: Reset _sam2_next_obj_id to 1 in _generate_3d_segmentation"
+            )
             self._sam2_prompts = (
                 {}
             )  # Store prompts for each object (points, labels, box)
 
+            # Reset SAM2-specific tracking dictionaries for 3D mode
+            self.sam2_points_by_obj = {}
+            self.sam2_labels_by_obj = {}
+
             # Update the label layer with empty segmentation
             self._update_label_layer()
 
@@ -648,8 +893,10 @@ class BatchCropAnything:
             if self.label_layer is not None and hasattr(
                 self.label_layer, "mouse_drag_callbacks"
             ):
+                # Safely remove all existing callbacks
                 for callback in list(self.label_layer.mouse_drag_callbacks):
-                    self.label_layer.mouse_drag_callbacks.remove(callback)
+                    with contextlib.suppress(ValueError):
+                        self.label_layer.mouse_drag_callbacks.remove(callback)
 
                 # Add 3D-specific click handler
                 self.label_layer.mouse_drag_callbacks.append(
@@ -673,8 +920,8 @@ class BatchCropAnything:
 
             # Show instructions
             self.viewer.status = (
-                "3D Mode active: Navigate to the first frame where object appears, then click. "
-                "Use Shift+click for negative points (to remove areas). "
+                "3D Mode active: IMPORTANT - Navigate to the FIRST SLICE where object appears (using slider), "
+                "then click on object in 2D view (not 3D view). Use Shift+click for negative points. "
                 "Segmentation will be propagated to all frames automatically."
             )
 
@@ -728,6 +975,9 @@ class BatchCropAnything:
                 # Create new object for positive points on background
                 ann_obj_id = self._sam2_next_obj_id
                 if point_label > 0 and label_id == 0:
+                    print(
+                        f"DEBUG: Incrementing _sam2_next_obj_id from {self._sam2_next_obj_id} to {self._sam2_next_obj_id + 1}"
+                    )
                     self._sam2_next_obj_id += 1
 
             # Find or create points layer for this object
@@ -915,8 +1165,10 @@ class BatchCropAnything:
             # Try to perform SAM2 propagation with error handling
             try:
                 # Use torch.inference_mode() and torch.autocast to ensure consistent dtypes
-                with torch.inference_mode(), torch.autocast(
-                    "cuda", dtype=torch.float32
+                device_type = "cuda" if self.device.type == "cuda" else "cpu"
+                with (
+                    torch.inference_mode(),
+                    torch.autocast(device_type, dtype=torch.float32),
                 ):
                     # Attempt to run SAM2 propagation - this will iterate through all frames
                     for (
@@ -1012,7 +1264,11 @@ class BatchCropAnything:
                 time.sleep(2)
                 for layer in list(self.viewer.layers):
                     if "Propagation Progress" in layer.name:
-                        self.viewer.layers.remove(layer)
+                        # Clean up callbacks before removing the layer to prevent cleanup issues
+                        if hasattr(layer, "mouse_drag_callbacks"):
+                            layer.mouse_drag_callbacks.clear()
+                        with contextlib.suppress(ValueError):
+                            self.viewer.layers.remove(layer)
 
             threading.Thread(target=remove_progress).start()
 
@@ -1035,6 +1291,7 @@ class BatchCropAnything:
         Given a 3D coordinate (x, y, z), run SAM2 video predictor to segment the object at that point,
         update the segmentation result and label layer.
         """
+        device_type = "cuda" if self.device.type == "cuda" else "cpu"
         if not hasattr(self, "_sam2_state") or self._sam2_state is None:
             self.viewer.status = "SAM2 3D state not initialized."
             return
@@ -1048,8 +1305,9 @@ class BatchCropAnything:
         point_coords = np.array([[x, y, z]])
         point_labels = np.array([1])  # 1 = foreground
 
-        with torch.inference_mode(), torch.autocast(
-            "cuda", dtype=torch.bfloat16
+        with (
+            torch.inference_mode(),
+            torch.autocast(device_type, dtype=torch.float32),
         ):
             masks, scores, _ = self.predictor.predict(
                 state=self._sam2_state,
@@ -1103,7 +1361,11 @@ class BatchCropAnything:
         # Remove existing label layer if it exists
         for layer in list(self.viewer.layers):
             if isinstance(layer, Labels) and "Segmentation" in layer.name:
-                self.viewer.layers.remove(layer)
+                # Clean up callbacks before removing the layer to prevent cleanup issues
+                if hasattr(layer, "mouse_drag_callbacks"):
+                    layer.mouse_drag_callbacks.clear()
+                with contextlib.suppress(ValueError):
+                    self.viewer.layers.remove(layer)
 
         # Add label layer to viewer
         self.label_layer = self.viewer.add_labels(
@@ -1112,10 +1374,36 @@ class BatchCropAnything:
             opacity=0.7,
         )
 
-        # Create points layer for interaction if it doesn't exist
+        # Connect click handler to the label layer for selection and deletion
+        if hasattr(self.label_layer, "mouse_drag_callbacks"):
+            # Clear existing callbacks to avoid duplicates
+            self.label_layer.mouse_drag_callbacks.clear()
+            # Add our click handler
+            self.label_layer.mouse_drag_callbacks.append(
+                self._on_label_clicked
+            )
+
+        # Create or update interaction layers based on mode
+        if self.prompt_mode == "point":
+            self._ensure_points_layer()
+            self._remove_shapes_layer()
+        else:  # box mode
+            self._ensure_shapes_layer()
+            self._remove_points_layer()
+
+        # Update status
+        n_labels = len(np.unique(self.segmentation_result)) - (
+            1 if 0 in np.unique(self.segmentation_result) else 0
+        )
+        self.viewer.status = f"Loaded image {self.current_index + 1}/{len(self.images)} - Found {n_labels} segments"
+
+    def _ensure_points_layer(self):
+        """Ensure points layer exists and is properly configured."""
         points_layer = None
         for layer in list(self.viewer.layers):
-            if "Points" in layer.name:
+            if (
+                "Points" in layer.name and "Object" not in layer.name
+            ):  # Main points layer
                 points_layer = layer
                 break
 
@@ -1131,141 +1419,193 @@ class BatchCropAnything:
                 opacity=0.8,
             )
 
-            with contextlib.suppress(AttributeError, ValueError):
-                points_layer.mouse_drag_callbacks.remove(
-                    self._on_points_clicked
-                )
+            # Connect points layer mouse click event
+            if hasattr(points_layer, "mouse_drag_callbacks"):
+                points_layer.mouse_drag_callbacks.clear()
                 points_layer.mouse_drag_callbacks.append(
                     self._on_points_clicked
                 )
 
-            # Connect points layer mouse click event
-            points_layer.mouse_drag_callbacks.append(self._on_points_clicked)
-
         # Make the points layer active to encourage interaction with it
         self.viewer.layers.selection.active = points_layer
 
-        # Update status
-        n_labels = len(np.unique(self.segmentation_result)) - (
-            1 if 0 in np.unique(self.segmentation_result) else 0
-        )
-        self.viewer.status = f"Loaded image {self.current_index + 1}/{len(self.images)} - Found {n_labels} segments"
+    def _ensure_shapes_layer(self):
+        """Ensure shapes layer exists and is properly configured."""
+        shapes_layer = None
+        for layer in list(self.viewer.layers):
+            if "Rectangles" in layer.name:
+                shapes_layer = layer
+                break
 
-    def _on_points_clicked(self, layer, event):
-        """Handle clicks on the points layer for adding/removing points."""
-        try:
-            # Only process clicks, not drags
-            if event.type != "mouse_press":
+        if shapes_layer is None:
+            # Initialize an empty shapes layer
+            shapes_layer = self.viewer.add_shapes(
+                None,
+                shape_type="rectangle",
+                edge_width=3,
+                edge_color="green",
+                face_color="transparent",
+                name="Rectangles (Draw to Segment)",
+            )
+
+        # Store reference
+        self.shapes_layer = shapes_layer
+
+        # Initialize processing flag to prevent re-entry
+        if not hasattr(self, "_processing_rectangle"):
+            self._processing_rectangle = False
+
+        # Always ensure the event is connected (disconnect old ones first to avoid duplicates)
+        # Remove any existing callbacks
+        with contextlib.suppress(Exception):
+            shapes_layer.events.data.disconnect()
+
+        # Connect shape added event
+        @shapes_layer.events.data.connect
+        def on_shape_added(event):
+            print(
+                f"DEBUG: Shape event triggered! Shapes: {len(shapes_layer.data)}, Processing: {self._processing_rectangle}"
+            )
+
+            # Ignore if we're already processing or if there are no shapes
+            if self._processing_rectangle:
+                print("DEBUG: Already processing a rectangle, ignoring event")
                 return
 
-            # Get coordinates of mouse click
-            coords = np.round(event.position).astype(int)
+            if len(shapes_layer.data) == 0:
+                print("DEBUG: No shapes present, ignoring event")
+                return
 
-            # Check if Shift is pressed for negative points
-            is_negative = "Shift" in event.modifiers
-            point_label = -1 if is_negative else 1
+            # Only process if we have exactly 1 shape (newly drawn)
+            if len(shapes_layer.data) == 1:
+                print("DEBUG: New shape detected, processing...")
+                # Set flag to prevent re-entry
+                self._processing_rectangle = True
+                try:
+                    # Get the shape
+                    self._on_rectangle_added(shapes_layer.data[-1])
+                finally:
+                    # Always reset flag
+                    self._processing_rectangle = False
+            else:
+                print(
+                    f"DEBUG: Multiple shapes present ({len(shapes_layer.data)}), skipping"
+                )
 
-            # Handle 2D vs 3D coordinates
-            if self.use_3d:
-                if len(coords) == 3:
-                    t, y, x = map(int, coords)
-                elif len(coords) == 2:
-                    t = int(self.viewer.dims.current_step[0])
-                    y, x = map(int, coords)
-                else:
-                    self.viewer.status = (
-                        f"Unexpected coordinate dimensions: {coords}"
-                    )
-                    return
+        # Make the shapes layer active
+        self.viewer.layers.selection.active = shapes_layer
 
-                # Add point to the layer immediately for visual feedback
-                new_point = np.array([[t, y, x]])
-                if len(layer.data) == 0:
-                    layer.data = new_point
-                else:
-                    layer.data = np.vstack([layer.data, new_point])
+    def _remove_points_layer(self):
+        """Remove points layer when not in point mode."""
+        for layer in list(self.viewer.layers):
+            if "Points" in layer.name and "Object" not in layer.name:
+                if hasattr(layer, "mouse_drag_callbacks"):
+                    layer.mouse_drag_callbacks.clear()
+                with contextlib.suppress(ValueError):
+                    self.viewer.layers.remove(layer)
 
-                # Update point colors
-                colors = layer.face_color
-                if isinstance(colors, list):
-                    colors.append("red" if is_negative else "green")
-                else:
-                    n_points = len(layer.data)
-                    colors = ["green"] * (n_points - 1)
-                    colors.append("red" if is_negative else "green")
-                layer.face_color = colors
+    def _remove_shapes_layer(self):
+        """Remove shapes layer when not in box mode."""
+        for layer in list(self.viewer.layers):
+            if "Rectangles" in layer.name:
+                with contextlib.suppress(ValueError):
+                    self.viewer.layers.remove(layer)
+        self.shapes_layer = None
 
-                # Get the object ID
-                # If clicking on existing segmentation with negative point
-                label_id = self.segmentation_result[t, y, x]
-                if is_negative and label_id > 0:
-                    obj_id = label_id
+    def _on_rectangle_added(self, rectangle_coords):
+        """Handle rectangle selection for segmentation."""
+        print("DEBUG: _on_rectangle_added called!")
+        device_type = "cuda" if self.device.type == "cuda" else "cpu"
+        try:
+            # Rectangle coords are in the form of a 4x2 or 4x3 array (corners)
+            # Convert to bounding box format [x_min, y_min, x_max, y_max]
+
+            # Debug info
+            print(f"DEBUG: Rectangle coords: {rectangle_coords}")
+            print(f"DEBUG: Rectangle coords shape: {rectangle_coords.shape}")
+            print(f"DEBUG: use_3d flag: {self.use_3d}")
+            print(
+                f"DEBUG: Has predictor: {hasattr(self, 'predictor') and self.predictor is not None}"
+            )
+            if hasattr(self, "predictor") and self.predictor is not None:
+                print(
+                    f"DEBUG: Predictor type: {type(self.predictor).__name__}"
+                )
+            else:
+                print("DEBUG: No predictor available!")
+                self.viewer.status = "Error: Predictor not initialized"
+                return
+
+            # Check if we're in 3D mode (use the flag, not coordinate shape)
+            # In 3D mode, even when drawing on a 2D slice, we get (4, 2) coords
+            # but we need to treat it as 3D with propagation
+            if (
+                self.use_3d
+                and len(rectangle_coords.shape) == 2
+                and rectangle_coords.shape[0] == 4
+            ):
+                print("DEBUG: Processing as 3D rectangle (will propagate)")
+
+                # Get current frame/slice
+                t = int(self.viewer.dims.current_step[0])
+                print(f"DEBUG: Current frame/slice: {t}")
+
+                # Get Y and X bounds from 2D coordinates
+                if rectangle_coords.shape[1] == 3:
+                    # If we somehow got 3D coords (T/Z, Y, X)
+                    y_coords = rectangle_coords[:, 1]
+                    x_coords = rectangle_coords[:, 2]
+                elif rectangle_coords.shape[1] == 2:
+                    # More common: 2D coords (Y, X) when drawing on a slice
+                    y_coords = rectangle_coords[:, 0]
+                    x_coords = rectangle_coords[:, 1]
                 else:
-                    # For new objects or negative on background
-                    if not hasattr(self, "_sam2_next_obj_id"):
-                        self._sam2_next_obj_id = 1
-                    obj_id = self._sam2_next_obj_id
-                    if point_label > 0 and label_id == 0:
-                        self._sam2_next_obj_id += 1
+                    print(
+                        f"DEBUG: Unexpected coordinate dimensions: {rectangle_coords.shape[1]}"
+                    )
+                    self.viewer.status = "Error: Unexpected rectangle format"
+                    return
 
-                # Store point information
-                if not hasattr(self, "points_data"):
-                    self.points_data = {}
-                    self.points_labels = {}
+                y_min, y_max = int(min(y_coords)), int(max(y_coords))
+                x_min, x_max = int(min(x_coords)), int(max(x_coords))
 
-                if obj_id not in self.points_data:
-                    self.points_data[obj_id] = []
-                    self.points_labels[obj_id] = []
+                box = np.array([x_min, y_min, x_max, y_max], dtype=np.float32)
+                print(f"DEBUG: Box coordinates: {box}")
 
-                self.points_data[obj_id].append(
-                    [x, y]
-                )  # Note: SAM2 expects [x,y] format
-                self.points_labels[obj_id].append(point_label)
+                # Use SAM2 with box prompt - use _sam2_next_obj_id for 3D mode
+                if not hasattr(self, "_sam2_next_obj_id"):
+                    self._sam2_next_obj_id = 1
+                obj_id = self._sam2_next_obj_id
+                self._sam2_next_obj_id += 1
+                print(
+                    f"DEBUG: Box mode - using object ID {obj_id}, next will be {self._sam2_next_obj_id}"
+                )
 
-                # Perform segmentation
+                # Store box for this object
+                if not hasattr(self, "obj_boxes"):
+                    self.obj_boxes = {}
+                self.obj_boxes[obj_id] = box
+
+                # Perform segmentation with 3D propagation
                 if (
                     hasattr(self, "_sam2_state")
                     and self._sam2_state is not None
                 ):
-                    # Prepare points
-                    points = np.array(
-                        self.points_data[obj_id], dtype=np.float32
-                    )
-                    labels = np.array(
-                        self.points_labels[obj_id], dtype=np.int32
+                    self.viewer.status = (
+                        f"Segmenting object {obj_id} with box at frame {t}..."
                     )
+                    print(f"DEBUG: Starting segmentation for object {obj_id}")
 
-                    # Create progress layer for visual feedback
-                    progress_layer = None
-                    for existing_layer in self.viewer.layers:
-                        if "Propagation Progress" in existing_layer.name:
-                            progress_layer = existing_layer
-                            break
-
-                    if progress_layer is None:
-                        progress_data = np.zeros_like(self.segmentation_result)
-                        progress_layer = self.viewer.add_image(
-                            progress_data,
-                            name="Propagation Progress",
-                            colormap="magma",
-                            opacity=0.5,
-                            visible=True,
-                        )
-
-                    # First update the current frame immediately
-                    self.viewer.status = f"Processing object at frame {t}..."
-
-                    # Run SAM2 on current frame
                     _, out_obj_ids, out_mask_logits = (
                         self.predictor.add_new_points_or_box(
                             inference_state=self._sam2_state,
                             frame_idx=t,
                             obj_id=obj_id,
-                            points=points,
-                            labels=labels,
+                            box=box,
                         )
                     )
 
+                    print("DEBUG: Segmentation complete, processing mask")
                     # Update current frame
                     mask = (out_mask_logits[0] > 0.0).cpu().numpy()
                     if mask.ndim > 2:
@@ -1283,21 +1623,380 @@ class BatchCropAnything:
                             anti_aliasing=False,
                         ).astype(bool)
 
-                    # Update segmentation for this frame
-                    if point_label < 0:
-                        # For negative points, only remove from this object
-                        self.segmentation_result[t][
-                            (self.segmentation_result[t] == obj_id) & mask
-                        ] = 0
-                    else:
-                        # For positive points, only replace background
-                        self.segmentation_result[t][
-                            mask & (self.segmentation_result[t] == 0)
-                        ] = obj_id
+                    # Update segmentation
+                    self.segmentation_result[t][
+                        mask & (self.segmentation_result[t] == 0)
+                    ] = obj_id
 
-                    # Update progress layer for this frame
-                    progress_data = progress_layer.data
-                    progress_data[t] = (
+                    print(f"DEBUG: Starting propagation for object {obj_id}")
+                    # Propagate to all frames
+                    self._propagate_mask_for_current_object(obj_id, t)
+
+                    # Update UI
+                    print("DEBUG: Updating label layer")
+                    self._update_label_layer()
+                    if (
+                        hasattr(self, "label_table_widget")
+                        and self.label_table_widget is not None
+                    ):
+                        self._populate_label_table(self.label_table_widget)
+
+                    self.viewer.status = (
+                        f"Segmented and propagated object {obj_id} from box"
+                    )
+                    print("DEBUG: Rectangle processing complete!")
+
+                    # Keep the rectangle visible after processing
+                    # Users can manually delete it if needed
+                    # if self.shapes_layer is not None:
+                    #     self.shapes_layer.data = []
+                else:
+                    print("DEBUG: _sam2_state not available")
+                    self.viewer.status = (
+                        "Error: 3D segmentation state not initialized"
+                    )
+
+            elif (
+                not self.use_3d
+                and len(rectangle_coords.shape) == 2
+                and rectangle_coords.shape[1] == 2
+            ):
+                # 2D case: rectangle_coords shape is (4, 2) for Y, X
+                if rectangle_coords.shape[0] == 4:
+                    # Get Y and X bounds
+                    y_coords = rectangle_coords[:, 0]
+                    x_coords = rectangle_coords[:, 1]
+                    y_min, y_max = int(min(y_coords)), int(max(y_coords))
+                    x_min, x_max = int(min(x_coords)), int(max(x_coords))
+
+                    box = np.array(
+                        [x_min, y_min, x_max, y_max], dtype=np.float32
+                    )
+
+                    # Use SAM2 with box prompt - use next_obj_id for 2D mode
+                    if not hasattr(self, "next_obj_id"):
+                        self.next_obj_id = 1
+                    obj_id = self.next_obj_id
+                    self.next_obj_id += 1
+                    print(
+                        f"DEBUG: 2D Box mode - using object ID {obj_id}, next will be {self.next_obj_id}"
+                    )
+
+                    # Store box for this object
+                    if not hasattr(self, "obj_boxes"):
+                        self.obj_boxes = {}
+                    self.obj_boxes[obj_id] = box
+
+                    # Perform segmentation
+                    if (
+                        hasattr(self, "predictor")
+                        and self.predictor is not None
+                    ):
+                        # Make sure image is loaded
+                        if self.current_image_for_segmentation is None:
+                            self.viewer.status = (
+                                "No image loaded for segmentation"
+                            )
+                            return
+
+                        # Prepare image for SAM2
+                        image = self.current_image_for_segmentation
+                        if len(image.shape) == 2:
+                            image = np.stack([image] * 3, axis=-1)
+                        elif len(image.shape) == 3 and image.shape[2] == 1:
+                            image = np.concatenate([image] * 3, axis=2)
+                        elif len(image.shape) == 3 and image.shape[2] > 3:
+                            image = image[:, :, :3]
+
+                        if image.dtype != np.uint8:
+                            image = (image / np.max(image) * 255).astype(
+                                np.uint8
+                            )
+
+                        # Set the image in the predictor (only for ImagePredictor, not VideoPredictor)
+                        if hasattr(self.predictor, "set_image"):
+                            self.predictor.set_image(image)
+                        else:
+                            self.viewer.status = "Error: Rectangle mode requires Image Predictor (2D mode)"
+                            return
+
+                        self.viewer.status = (
+                            f"Segmenting object {obj_id} with box..."
+                        )
+
+                        with (
+                            torch.inference_mode(),
+                            torch.autocast(device_type),
+                        ):
+                            masks, scores, _ = self.predictor.predict(
+                                box=box,
+                                multimask_output=False,
+                            )
+
+                            # Get the mask
+                            if len(masks) > 0:
+                                best_mask = masks[0]
+
+                                # Resize if needed
+                                if (
+                                    best_mask.shape
+                                    != self.segmentation_result.shape
+                                ):
+                                    from skimage.transform import resize
+
+                                    best_mask = resize(
+                                        best_mask.astype(float),
+                                        self.segmentation_result.shape,
+                                        order=0,
+                                        preserve_range=True,
+                                        anti_aliasing=False,
+                                    ).astype(bool)
+
+                                # Apply mask (only overwrite background)
+                                mask_condition = np.logical_and(
+                                    best_mask, (self.segmentation_result == 0)
+                                )
+                                self.segmentation_result[mask_condition] = (
+                                    obj_id
+                                )
+
+                                # Update label info
+                                area = np.sum(
+                                    self.segmentation_result == obj_id
+                                )
+                                y_indices, x_indices = np.where(
+                                    self.segmentation_result == obj_id
+                                )
+                                center_y = (
+                                    np.mean(y_indices)
+                                    if len(y_indices) > 0
+                                    else 0
+                                )
+                                center_x = (
+                                    np.mean(x_indices)
+                                    if len(x_indices) > 0
+                                    else 0
+                                )
+
+                                self.label_info[obj_id] = {
+                                    "area": area,
+                                    "center_y": center_y,
+                                    "center_x": center_x,
+                                    "score": float(scores[0]),
+                                }
+
+                                self.viewer.status = (
+                                    f"Segmented object {obj_id} from box"
+                                )
+                            else:
+                                self.viewer.status = "No valid mask produced"
+
+                        # Update the UI
+                        self._update_label_layer()
+                        if (
+                            hasattr(self, "label_table_widget")
+                            and self.label_table_widget is not None
+                        ):
+                            self._populate_label_table(self.label_table_widget)
+
+                        # Keep the rectangle visible after processing
+                        # Users can manually delete it if needed
+                        # if self.shapes_layer is not None:
+                        #     self.shapes_layer.data = []
+            else:
+                # Unexpected shape dimensions
+                print(
+                    f"DEBUG: Unexpected rectangle shape: {rectangle_coords.shape}"
+                )
+                self.viewer.status = f"Error: Unexpected rectangle dimensions {rectangle_coords.shape}. Expected (4,2) for 2D or (4,3) for 3D."
+
+        except (
+            IndexError,
+            KeyError,
+            ValueError,
+            RuntimeError,
+            TypeError,
+        ) as e:
+            import traceback
+
+            self.viewer.status = f"Error in rectangle handling: {str(e)}"
+            print("DEBUG: Exception in _on_rectangle_added:")
+            traceback.print_exc()
+
+    def _on_points_clicked(self, layer, event):
+        """Handle clicks on the points layer for adding/removing points."""
+        device_type = "cuda" if self.device.type == "cuda" else "cpu"
+        try:
+            # Only process clicks, not drags
+            if event.type != "mouse_press":
+                return
+
+            # Check if segmentation result exists
+            if self.segmentation_result is None:
+                self.viewer.status = (
+                    "Segmentation not ready. Please wait for image to load."
+                )
+                return
+
+            # Get coordinates of mouse click
+            coords = np.round(event.position).astype(int)
+
+            # Check if Shift is pressed for negative points
+            is_negative = "Shift" in event.modifiers
+            point_label = -1 if is_negative else 1
+
+            # Handle 2D vs 3D coordinates
+            if self.use_3d:
+                if len(coords) == 3:
+                    t, y, x = map(int, coords)
+                elif len(coords) == 2:
+                    t = int(self.viewer.dims.current_step[0])
+                    y, x = map(int, coords)
+                else:
+                    self.viewer.status = (
+                        f"Unexpected coordinate dimensions: {coords}"
+                    )
+                    return
+
+                # Add point to the layer immediately for visual feedback
+                new_point = np.array([[t, y, x]])
+                if len(layer.data) == 0:
+                    layer.data = new_point
+                else:
+                    layer.data = np.vstack([layer.data, new_point])
+
+                # Update point colors
+                colors = layer.face_color
+                if isinstance(colors, list):
+                    colors.append("red" if is_negative else "green")
+                else:
+                    n_points = len(layer.data)
+                    colors = ["green"] * (n_points - 1)
+                    colors.append("red" if is_negative else "green")
+                layer.face_color = colors
+
+                # Validate coordinates are within segmentation bounds
+                if (
+                    t < 0
+                    or t >= self.segmentation_result.shape[0]
+                    or y < 0
+                    or y >= self.segmentation_result.shape[1]
+                    or x < 0
+                    or x >= self.segmentation_result.shape[2]
+                ):
+                    self.viewer.status = (
+                        f"Click at ({t}, {y}, {x}) is out of bounds for "
+                        f"segmentation shape {self.segmentation_result.shape}. "
+                        f"Please click within the image bounds."
+                    )
+                    # Remove the invalid point that was just added
+                    if len(layer.data) > 0:
+                        layer.data = layer.data[:-1]
+                    return
+
+                # Get the object ID
+                # If clicking on existing segmentation with negative point
+                label_id = self.segmentation_result[t, y, x]
+                if is_negative and label_id > 0:
+                    obj_id = label_id
+                else:
+                    # For new objects or negative on background
+                    if not hasattr(self, "_sam2_next_obj_id"):
+                        self._sam2_next_obj_id = 1
+                    obj_id = self._sam2_next_obj_id
+                    if point_label > 0 and label_id == 0:
+                        self._sam2_next_obj_id += 1
+
+                # Store point information
+                if not hasattr(self, "points_data"):
+                    self.points_data = {}
+                    self.points_labels = {}
+
+                if obj_id not in self.points_data:
+                    self.points_data[obj_id] = []
+                    self.points_labels[obj_id] = []
+
+                self.points_data[obj_id].append(
+                    [x, y]
+                )  # Note: SAM2 expects [x,y] format
+                self.points_labels[obj_id].append(point_label)
+
+                # Perform segmentation
+                if (
+                    hasattr(self, "_sam2_state")
+                    and self._sam2_state is not None
+                ):
+                    # Prepare points
+                    points = np.array(
+                        self.points_data[obj_id], dtype=np.float32
+                    )
+                    labels = np.array(
+                        self.points_labels[obj_id], dtype=np.int32
+                    )
+
+                    # Create progress layer for visual feedback
+                    progress_layer = None
+                    for existing_layer in self.viewer.layers:
+                        if "Propagation Progress" in existing_layer.name:
+                            progress_layer = existing_layer
+                            break
+
+                    if progress_layer is None:
+                        progress_data = np.zeros_like(self.segmentation_result)
+                        progress_layer = self.viewer.add_image(
+                            progress_data,
+                            name="Propagation Progress",
+                            colormap="magma",
+                            opacity=0.5,
+                            visible=True,
+                        )
+
+                    # First update the current frame immediately
+                    self.viewer.status = f"Processing object at frame {t}..."
+
+                    # Run SAM2 on current frame
+                    _, out_obj_ids, out_mask_logits = (
+                        self.predictor.add_new_points_or_box(
+                            inference_state=self._sam2_state,
+                            frame_idx=t,
+                            obj_id=obj_id,
+                            points=points,
+                            labels=labels,
+                        )
+                    )
+
+                    # Update current frame
+                    mask = (out_mask_logits[0] > 0.0).cpu().numpy()
+                    if mask.ndim > 2:
+                        mask = mask.squeeze()
+
+                    # Resize if needed
+                    if mask.shape != self.segmentation_result[t].shape:
+                        from skimage.transform import resize
+
+                        mask = resize(
+                            mask.astype(float),
+                            self.segmentation_result[t].shape,
+                            order=0,
+                            preserve_range=True,
+                            anti_aliasing=False,
+                        ).astype(bool)
+
+                    # Update segmentation for this frame
+                    if point_label < 0:
+                        # For negative points, only remove from this object
+                        self.segmentation_result[t][
+                            (self.segmentation_result[t] == obj_id) & mask
+                        ] = 0
+                    else:
+                        # For positive points, only replace background
+                        self.segmentation_result[t][
+                            mask & (self.segmentation_result[t] == 0)
+                        ] = obj_id
+
+                    # Update progress layer for this frame
+                    progress_data = progress_layer.data
+                    progress_data[t] = (
                         mask.astype(float) * 0.5
                     )  # Highlight current frame
                     progress_layer.data = progress_data
@@ -1398,7 +2097,11 @@ class BatchCropAnything:
                         time.sleep(2)
                         for layer in list(self.viewer.layers):
                             if "Propagation Progress" in layer.name:
-                                self.viewer.layers.remove(layer)
+                                # Clean up callbacks before removing the layer to prevent cleanup issues
+                                if hasattr(layer, "mouse_drag_callbacks"):
+                                    layer.mouse_drag_callbacks.clear()
+                                with contextlib.suppress(ValueError):
+                                    self.viewer.layers.remove(layer)
 
                     threading.Thread(target=remove_progress).start()
 
@@ -1439,6 +2142,23 @@ class BatchCropAnything:
                     colors.append("red" if is_negative else "green")
                 layer.face_color = colors
 
+                # Validate coordinates are within segmentation bounds
+                if (
+                    y < 0
+                    or y >= self.segmentation_result.shape[0]
+                    or x < 0
+                    or x >= self.segmentation_result.shape[1]
+                ):
+                    self.viewer.status = (
+                        f"Click at ({y}, {x}) is out of bounds for "
+                        f"segmentation shape {self.segmentation_result.shape}. "
+                        f"Please click within the image bounds."
+                    )
+                    # Remove the invalid point that was just added
+                    if len(layer.data) > 0:
+                        layer.data = layer.data[:-1]
+                    return
+
                 # Get object ID
                 label_id = self.segmentation_result[y, x]
                 if is_negative and label_id > 0:
@@ -1483,8 +2203,14 @@ class BatchCropAnything:
                     if image.dtype != np.uint8:
                         image = (image / np.max(image) * 255).astype(np.uint8)
 
-                    # Set the image in the predictor
-                    self.predictor.set_image(image)
+                    # Set the image in the predictor (only for ImagePredictor, not VideoPredictor)
+                    if hasattr(self.predictor, "set_image"):
+                        self.predictor.set_image(image)
+                    else:
+                        self.viewer.status = (
+                            "Error: Point mode in 2D requires Image Predictor"
+                        )
+                        return
 
                     # Use only points for current object
                     points = np.array(
@@ -1494,7 +2220,7 @@ class BatchCropAnything:
 
                     self.viewer.status = f"Segmenting object {obj_id} with {len(points)} points..."
 
-                    with torch.inference_mode(), torch.autocast("cuda"):
+                    with torch.inference_mode(), torch.autocast(device_type):
                         masks, scores, _ = self.predictor.predict(
                             point_coords=points,
                             point_labels=labels,
@@ -1583,16 +2309,23 @@ class BatchCropAnything:
     def _on_label_clicked(self, layer, event):
         """Handle label selection and user prompts on mouse click."""
         try:
-            # Only process clicks, not drags
+            # Only process mouse press events
             if event.type != "mouse_press":
                 return
 
+            # Only handle left mouse button
+            if event.button != 1:
+                return
+
             # Get coordinates of mouse click
             coords = np.round(event.position).astype(int)
 
-            # Check if Shift is pressed (negative point)
+            # Check modifiers
             is_negative = "Shift" in event.modifiers
-            point_label = -1 if is_negative else 1
+            is_control = (
+                "Control" in event.modifiers or "Ctrl" in event.modifiers
+            )
+            # point_label = -1 if is_negative else 1
 
             # For 2D data
             if not self.use_3d:
@@ -1613,262 +2346,13 @@ class BatchCropAnything:
                 # Get the label ID at the clicked position
                 label_id = self.segmentation_result[y, x]
 
-                # Initialize a unique object ID for this click (if needed)
-                if not hasattr(self, "next_obj_id"):
-                    # Start with highest existing ID + 1
-                    if self.segmentation_result.max() > 0:
-                        self.next_obj_id = (
-                            int(self.segmentation_result.max()) + 1
-                        )
-                    else:
-                        self.next_obj_id = 1
-
-                # If clicking on background or using negative click, handle segmentation
-                if label_id == 0 or is_negative:
-                    # Find or create points layer for the current object we're working on
-                    current_obj_id = None
-
-                    # If negative point on existing label, use that label's ID
-                    if is_negative and label_id > 0:
-                        current_obj_id = label_id
-                    # For positive clicks on background, create a new object
-                    elif point_label > 0 and label_id == 0:
-                        current_obj_id = self.next_obj_id
-                        self.next_obj_id += 1
-                    # For negative on background, try to find most recent object
-                    elif point_label < 0 and label_id == 0:
-                        # Use most recently created object if available
-                        if hasattr(self, "obj_points") and self.obj_points:
-                            current_obj_id = max(self.obj_points.keys())
-                        else:
-                            self.viewer.status = "No existing object to modify with negative point"
-                            return
-
-                    if current_obj_id is None:
-                        self.viewer.status = (
-                            "Could not determine which object to modify"
-                        )
-                        return
-
-                    # Find or create points layer for this object
-                    points_layer = None
-                    for layer in list(self.viewer.layers):
-                        if f"Points for Object {current_obj_id}" in layer.name:
-                            points_layer = layer
-                            break
-
-                    # Initialize object tracking if needed
-                    if not hasattr(self, "obj_points"):
-                        self.obj_points = {}
-                        self.obj_labels = {}
-
-                    if current_obj_id not in self.obj_points:
-                        self.obj_points[current_obj_id] = []
-                        self.obj_labels[current_obj_id] = []
-
-                    # Create or update points layer for this object
-                    if points_layer is None:
-                        # First point for this object
-                        points_layer = self.viewer.add_points(
-                            np.array([[y, x]]),
-                            name=f"Points for Object {current_obj_id}",
-                            size=10,
-                            face_color=["green" if point_label > 0 else "red"],
-                            border_color="white",
-                            border_width=1,
-                            opacity=0.8,
-                        )
-                        with contextlib.suppress(AttributeError, ValueError):
-                            points_layer.mouse_drag_callbacks.remove(
-                                self._on_points_clicked
-                            )
-                            points_layer.mouse_drag_callbacks.append(
-                                self._on_points_clicked
-                            )
-
-                        self.obj_points[current_obj_id] = [[x, y]]
-                        self.obj_labels[current_obj_id] = [point_label]
-                    else:
-                        # Add point to existing layer
-                        current_points = points_layer.data
-                        current_colors = points_layer.face_color
-
-                        # Add new point
-                        new_points = np.vstack([current_points, [y, x]])
-                        new_color = "green" if point_label > 0 else "red"
-
-                        # Update points layer
-                        points_layer.data = new_points
-
-                        # Update colors
-                        if isinstance(current_colors, list):
-                            current_colors.append(new_color)
-                            points_layer.face_color = current_colors
-                        else:
-                            # If it's an array, create a list of colors
-                            colors = []
-                            for i in range(len(new_points)):
-                                if i < len(current_points):
-                                    colors.append(
-                                        "green" if point_label > 0 else "red"
-                                    )
-                                else:
-                                    colors.append(new_color)
-                            points_layer.face_color = colors
-
-                        # Update object tracking
-                        self.obj_points[current_obj_id].append([x, y])
-                        self.obj_labels[current_obj_id].append(point_label)
-
-                    # Now do the actual segmentation using SAM2
-                    if (
-                        hasattr(self, "predictor")
-                        and self.predictor is not None
-                    ):
-                        try:
-                            # Make sure image is loaded
-                            if self.current_image_for_segmentation is None:
-                                self.viewer.status = (
-                                    "No image loaded for segmentation"
-                                )
-                                return
-
-                            # Prepare image for SAM2
-                            image = self.current_image_for_segmentation
-                            if len(image.shape) == 2:
-                                image = np.stack([image] * 3, axis=-1)
-                            elif len(image.shape) == 3 and image.shape[2] == 1:
-                                image = np.concatenate([image] * 3, axis=2)
-                            elif len(image.shape) == 3 and image.shape[2] > 3:
-                                image = image[:, :, :3]
-
-                            if image.dtype != np.uint8:
-                                image = (image / np.max(image) * 255).astype(
-                                    np.uint8
-                                )
-
-                            # Set the image in the predictor
-                            self.predictor.set_image(image)
-
-                            # Only use the points for the current object being segmented
-                            points = np.array(
-                                self.obj_points[current_obj_id],
-                                dtype=np.float32,
-                            )
-                            labels = np.array(
-                                self.obj_labels[current_obj_id], dtype=np.int32
-                            )
-
-                            self.viewer.status = f"Segmenting object {current_obj_id} with {len(points)} points..."
-
-                            with torch.inference_mode(), torch.autocast(
-                                "cuda"
-                            ):
-                                masks, scores, _ = self.predictor.predict(
-                                    point_coords=points,
-                                    point_labels=labels,
-                                    multimask_output=True,
-                                )
-
-                                # Get best mask
-                                if len(masks) > 0:
-                                    best_mask = masks[0]
-
-                                    # Update segmentation result
-                                    if (
-                                        best_mask.shape
-                                        != self.segmentation_result.shape
-                                    ):
-                                        from skimage.transform import resize
-
-                                        best_mask = resize(
-                                            best_mask.astype(float),
-                                            self.segmentation_result.shape,
-                                            order=0,
-                                            preserve_range=True,
-                                            anti_aliasing=False,
-                                        ).astype(bool)
-
-                                    # CRITICAL FIX: For negative points, only remove from this object's mask
-                                    # For positive points, add to this object's mask without removing other objects
-                                    if point_label < 0:
-                                        # Remove only from current object's mask
-                                        self.segmentation_result[
-                                            (
-                                                self.segmentation_result
-                                                == current_obj_id
-                                            )
-                                            & best_mask
-                                        ] = 0
-                                    else:
-                                        # Add to current object's mask without affecting other objects
-                                        # Only overwrite background (value 0)
-                                        self.segmentation_result[
-                                            best_mask
-                                            & (self.segmentation_result == 0)
-                                        ] = current_obj_id
-
-                                    # Update label info
-                                    area = np.sum(
-                                        self.segmentation_result
-                                        == current_obj_id
-                                    )
-                                    y_indices, x_indices = np.where(
-                                        self.segmentation_result
-                                        == current_obj_id
-                                    )
-                                    center_y = (
-                                        np.mean(y_indices)
-                                        if len(y_indices) > 0
-                                        else 0
-                                    )
-                                    center_x = (
-                                        np.mean(x_indices)
-                                        if len(x_indices) > 0
-                                        else 0
-                                    )
-
-                                    self.label_info[current_obj_id] = {
-                                        "area": area,
-                                        "center_y": center_y,
-                                        "center_x": center_x,
-                                        "score": float(scores[0]),
-                                    }
-
-                                    self.viewer.status = (
-                                        f"Updated object {current_obj_id}"
-                                    )
-                                else:
-                                    self.viewer.status = (
-                                        "No valid mask produced"
-                                    )
-
-                            # Update the UI
-                            self._update_label_layer()
-                            if (
-                                hasattr(self, "label_table_widget")
-                                and self.label_table_widget is not None
-                            ):
-                                self._populate_label_table(
-                                    self.label_table_widget
-                                )
-
-                        except (
-                            IndexError,
-                            KeyError,
-                            ValueError,
-                            AttributeError,
-                            TypeError,
-                        ) as e:
-                            import traceback
-
-                            self.viewer.status = (
-                                f"Error in SAM2 processing: {str(e)}"
-                            )
-                            traceback.print_exc()
+                # Handle Ctrl+Click to clear a single label
+                if is_control and label_id > 0:
+                    self.clear_label_at_position(y, x)
+                    return
 
-                # If clicking on an existing label, toggle selection
-                elif label_id > 0:
+                # If clicking on an existing label (and not using modifiers), toggle selection
+                if label_id > 0 and not is_negative and not is_control:
                     # Toggle the label selection
                     if label_id in self.selected_labels:
                         self.selected_labels.remove(label_id)
@@ -1880,8 +2364,14 @@ class BatchCropAnything:
                     # Update table and preview
                     self._update_label_table()
                     self.preview_crop()
+                    return
+
+                # If clicking on background or using Shift (negative points), this should be handled by points layer
+                # Don't process these clicks here to avoid conflicts
+                if label_id == 0 or is_negative:
+                    return
 
-            # 3D case (handle differently)
+            # 3D case
             else:
                 if len(coords) == 3:
                     t, y, x = map(int, coords)
@@ -1910,12 +2400,13 @@ class BatchCropAnything:
                 # Get the label ID at the clicked position
                 label_id = self.segmentation_result[t, y, x]
 
-                # If background or shift is pressed, handle in _on_3d_label_clicked
-                if label_id == 0 or is_negative:
-                    # This will be handled by _on_3d_label_clicked already attached
-                    pass
-                # If clicking on an existing label, handle selection
-                elif label_id > 0:
+                # Handle Ctrl+Click to clear a single label
+                if is_control and label_id > 0:
+                    self.clear_label_at_position_3d(t, y, x)
+                    return
+
+                # If clicking on an existing label and not using negative points, handle selection
+                if label_id > 0 and not is_negative and not is_control:
                     # Toggle the label selection
                     if label_id in self.selected_labels:
                         self.selected_labels.remove(label_id)
@@ -1926,9 +2417,12 @@ class BatchCropAnything:
 
                     # Update table if it exists
                     self._update_label_table()
-
-                    # Update preview after selection changes
                     self.preview_crop()
+                    return
+
+                # For background clicks or negative points, let the 3D handler deal with it
+                if label_id == 0 or is_negative:
+                    return
 
         except (
             IndexError,
@@ -1942,12 +2436,74 @@ class BatchCropAnything:
             self.viewer.status = f"Error in click handling: {str(e)}"
             traceback.print_exc()
 
+    def _add_segmentation_point(self, x, y, event):
+        """Add a point for segmentation."""
+        is_negative = "Shift" in event.modifiers
+
+        # Initialize tracking if needed
+        if not hasattr(self, "current_points"):
+            self.current_points = []
+            self.current_labels = []
+            self.current_obj_id = 1
+
+        # Add point
+        self.current_points.append([x, y])
+        self.current_labels.append(0 if is_negative else 1)
+
+        # Run SAM2 prediction
+        if self.predictor is not None:
+            # Prepare image
+            image = self._prepare_image_for_sam2()
+
+            # Set the image in the predictor (only for ImagePredictor, not VideoPredictor)
+            if hasattr(self.predictor, "set_image"):
+                self.predictor.set_image(image)
+            else:
+                self.viewer.status = (
+                    "Error: This operation requires Image Predictor (2D mode)"
+                )
+                return
+
+            # Predict
+            device_type = "cuda" if self.device.type == "cuda" else "cpu"
+            with torch.inference_mode(), torch.autocast(device_type):
+                masks, scores, _ = self.predictor.predict(
+                    point_coords=np.array(
+                        self.current_points, dtype=np.float32
+                    ),
+                    point_labels=np.array(self.current_labels, dtype=np.int32),
+                    multimask_output=False,
+                )
+
+            # Update segmentation
+            if len(masks) > 0:
+                mask = masks[0] > 0.5
+                if self.current_scale_factor < 1.0:
+                    mask = resize(
+                        mask, self.segmentation_result.shape, order=0
+                    ).astype(bool)
+
+                # Update segmentation result
+                self.segmentation_result[mask] = self.current_obj_id
+
+                # Move to next object if adding positive point
+                if not is_negative:
+                    self.current_obj_id += 1
+                    self.current_points = []
+                    self.current_labels = []
+
+                self._update_label_layer()
+
     def _add_point_marker(self, coords, label_type):
         """Add a visible marker for where the user clicked."""
         # Remove previous point markers
         for layer in list(self.viewer.layers):
             if "Point Prompt" in layer.name:
-                self.viewer.layers.remove(layer)
+                # Clean up callbacks before removing the layer to prevent cleanup issues
+                if hasattr(layer, "mouse_drag_callbacks"):
+                    layer.mouse_drag_callbacks.clear()
+                with contextlib.suppress(ValueError):
+                    self.viewer.layers.remove(layer)
 
         # Create points layer
         color = (
@@ -2135,11 +2691,170 @@ class BatchCropAnything:
         self.viewer.status = f"Selected all {len(self.selected_labels)} labels"
 
     def clear_selection(self):
-        """Clear all selected labels."""
+        """Clear all labels from the segmentation.
+
+        This removes all segmented objects from the label layer, resets all tracking data,
+        and prepares the interface for new segmentations. Note: The method name is kept as
+        'clear_selection' for backwards compatibility, but it clears all labels, not just
+        the selection.
+        """
+        if self.segmentation_result is None:
+            self.viewer.status = "No segmentation available"
+            return
+
+        # Get all unique label IDs (excluding background 0)
+        unique_labels = np.unique(self.segmentation_result)
+        label_ids = [label for label in unique_labels if label > 0]
+
+        if len(label_ids) == 0:
+            self.viewer.status = "No labels to clear"
+            return
+
+        # Clear the entire segmentation result
+        self.segmentation_result[:] = 0
+
+        # Clear selected labels
         self.selected_labels = set()
+
+        # Clear label info
+        self.label_info = {}
+
+        # Remove any object-specific point layers
+        for layer in list(self.viewer.layers):
+            if "Points for Object" in layer.name:
+                # Clean up callbacks before removing the layer to prevent cleanup issues
+                if hasattr(layer, "mouse_drag_callbacks"):
+                    layer.mouse_drag_callbacks.clear()
+                with contextlib.suppress(ValueError):
+                    self.viewer.layers.remove(layer)
+
+        # Clean up object tracking data
+        if hasattr(self, "obj_points"):
+            self.obj_points = {}
+        if hasattr(self, "obj_labels"):
+            self.obj_labels = {}
+        if hasattr(self, "points_data"):
+            self.points_data = {}
+        if hasattr(self, "points_labels"):
+            self.points_labels = {}
+
+        # Reset object ID counters
+        if hasattr(self, "next_obj_id"):
+            self.next_obj_id = 1
+        if hasattr(self, "_sam2_next_obj_id"):
+            self._sam2_next_obj_id = 1
+
+        # Update UI
+        self._update_label_layer()
         self._update_label_table()
         self.preview_crop()
-        self.viewer.status = "Cleared all selections"
+
+        self.viewer.status = (
+            f"Cleared all {len(label_ids)} labels from segmentation"
+        )
+
+    def clear_label_at_position(self, y, x):
+        """Clear a single label at the specified 2D position."""
+        if self.segmentation_result is None:
+            self.viewer.status = "No segmentation available"
+            return
+
+        label_id = self.segmentation_result[y, x]
+        if label_id > 0:
+            # Remove all pixels with this label ID
+            self.segmentation_result[self.segmentation_result == label_id] = 0
+
+            # Remove from selected labels if it was selected
+            self.selected_labels.discard(label_id)
+
+            # Remove from label info
+            if label_id in self.label_info:
+                del self.label_info[label_id]
+
+            # Remove any object-specific point layers for this label
+            for layer in list(self.viewer.layers):
+                if f"Points for Object {label_id}" in layer.name:
+                    # Clean up callbacks before removing the layer to prevent cleanup issues
+                    if hasattr(layer, "mouse_drag_callbacks"):
+                        layer.mouse_drag_callbacks.clear()
+                    with contextlib.suppress(ValueError):
+                        self.viewer.layers.remove(layer)
+
+            # Clean up object tracking data
+            if hasattr(self, "obj_points") and label_id in self.obj_points:
+                del self.obj_points[label_id]
+            if hasattr(self, "obj_labels") and label_id in self.obj_labels:
+                del self.obj_labels[label_id]
+
+            # Update UI
+            self._update_label_layer()
+            self._update_label_table()
+            self.preview_crop()
+
+            self.viewer.status = f"Deleted label ID: {label_id}"
+        else:
+            self.viewer.status = "No label to delete at this position"
+
+    def clear_label_at_position_3d(self, t, y, x):
+        """Clear a single label at the specified 3D position."""
+        if self.segmentation_result is None:
+            self.viewer.status = "No segmentation available"
+            return
+
+        label_id = self.segmentation_result[t, y, x]
+        if label_id > 0:
+            # Remove all pixels with this label ID across all timeframes
+            self.segmentation_result[self.segmentation_result == label_id] = 0
+
+            # Remove from selected labels if it was selected
+            self.selected_labels.discard(label_id)
+
+            # Remove from label info
+            if label_id in self.label_info:
+                del self.label_info[label_id]
+
+            # Remove any object-specific point layers for this label
+            for layer in list(self.viewer.layers):
+                if f"Points for Object {label_id}" in layer.name:
+                    # Clean up callbacks before removing the layer to prevent cleanup issues
+                    if hasattr(layer, "mouse_drag_callbacks"):
+                        layer.mouse_drag_callbacks.clear()
+                    with contextlib.suppress(ValueError):
+                        self.viewer.layers.remove(layer)
+
+            # Clean up 3D object tracking data
+            if (
+                hasattr(self, "sam2_points_by_obj")
+                and label_id in self.sam2_points_by_obj
+            ):
+                del self.sam2_points_by_obj[label_id]
+            if (
+                hasattr(self, "sam2_labels_by_obj")
+                and label_id in self.sam2_labels_by_obj
+            ):
+                del self.sam2_labels_by_obj[label_id]
+            if hasattr(self, "points_data") and label_id in self.points_data:
+                del self.points_data[label_id]
+            if (
+                hasattr(self, "points_labels")
+                and label_id in self.points_labels
+            ):
+                del self.points_labels[label_id]
+
+            # Update UI
+            self._update_label_layer()
+            if (
+                hasattr(self, "label_table_widget")
+                and self.label_table_widget is not None
+            ):
+                self._populate_label_table(self.label_table_widget)
+            self.preview_crop()
+
+            self.viewer.status = (
+                f"Deleted label ID: {label_id} from all timeframes"
+            )
+        else:
+            self.viewer.status = "No label to delete at this position"
 
     def preview_crop(self, label_ids=None):
         """Preview the crop result with the selected label IDs."""
@@ -2159,7 +2874,11 @@ class BatchCropAnything:
                 # Remove previous preview if exists
                 for layer in list(self.viewer.layers):
                     if "Preview" in layer.name:
-                        self.viewer.layers.remove(layer)
+                        # Clean up callbacks before removing the layer to prevent cleanup issues
+                        if hasattr(layer, "mouse_drag_callbacks"):
+                            layer.mouse_drag_callbacks.clear()
+                        with contextlib.suppress(ValueError):
+                            self.viewer.layers.remove(layer)
 
                 # Make sure the segmentation layer is active again
                 if self.label_layer is not None:
@@ -2197,7 +2916,11 @@ class BatchCropAnything:
             # Remove previous preview if exists
             for layer in list(self.viewer.layers):
                 if "Preview" in layer.name:
-                    self.viewer.layers.remove(layer)
+                    # Clean up callbacks before removing the layer to prevent cleanup issues
+                    if hasattr(layer, "mouse_drag_callbacks"):
+                        layer.mouse_drag_callbacks.clear()
+                    with contextlib.suppress(ValueError):
+                        self.viewer.layers.remove(layer)
 
             # Add preview layer
             if label_ids:
@@ -2288,17 +3011,14 @@ class BatchCropAnything:
             # Save cropped image
             image_path = self.images[self.current_index]
             base_name, ext = os.path.splitext(image_path)
-            label_str = "_".join(
-                str(lid) for lid in sorted(self.selected_labels)
-            )
-            output_path = f"{base_name}_cropped_{label_str}.tif"
+            output_path = f"{base_name}_sam2_cropped.tif"
 
             # Save using tifffile with explicit parameters for best compatibility
             imwrite(output_path, cropped_image, compression="zlib")
             self.viewer.status = f"Saved cropped image to {output_path}"
 
             # Save the label image with exact same dimensions as original
-            label_output_path = f"{base_name}_labels_{label_str}.tif"
+            label_output_path = f"{base_name}_sam2_labels.tif"
             imwrite(label_output_path, label_image, compression="zlib")
             self.viewer.status += f"\nSaved label mask to {label_output_path}"
 
@@ -2312,6 +3032,27 @@ class BatchCropAnything:
             self.viewer.status = f"Error cropping image: {str(e)}"
             return False
 
+    def reset_sam2_state(self):
+        """Reset SAM2 predictor state for 2D segmentation."""
+        if not self.use_3d and hasattr(self, "prepared_sam2_image"):
+            # Re-set the image in the predictor (only for ImagePredictor)
+            device_type = "cuda" if self.device.type == "cuda" else "cpu"
+            try:
+                if hasattr(self.predictor, "set_image"):
+                    with (
+                        torch.inference_mode(),
+                        torch.autocast(device_type, dtype=torch.float32),
+                    ):
+                        self.predictor.set_image(self.prepared_sam2_image)
+                else:
+                    print(
+                        "DEBUG: reset_sam2_state - predictor doesn't have set_image method"
+                    )
+            except (RuntimeError, AssertionError, TypeError, ValueError) as e:
+                print(f"Error resetting SAM2 state: {e}")
+                # If there's an error, try to reinitialize
+                self._initialize_sam2()
+
 
 def create_crop_widget(processor):
     """Create the crop control widget."""
@@ -2322,27 +3063,70 @@ def create_crop_widget(processor):
 
     # Instructions
     dimension_type = "3D (TYX/ZYX)" if processor.use_3d else "2D (YX)"
-    instructions_label = QLabel(
-        f"<b>Processing {dimension_type} data</b><br><br>"
-        "To create/edit objects:<br>"
-        "1. <b>Click on the POINTS layer</b> to add positive points<br>"
-        "2. Use Shift+click for negative points to refine segmentation<br>"
-        "3. Click on existing objects in the Segmentation layer to select them<br>"
-        "4. Press 'Crop' to save the selected objects to disk"
-    )
+
+    if processor.use_3d:
+        instructions_text = (
+            f"<b>Processing {dimension_type} data</b><br><br>"
+            "<b>⚠️ IMPORTANT for 3D stacks:</b><br>"
+            "<ul>"
+            "<li><b>Navigate to the FIRST SLICE</b> where your object appears (use the time/Z slider)</li>"
+            "<li><b>Switch to 2D view</b> (click 2D icon in napari, NOT 3D view)</li>"
+            "<li><b>Point Mode:</b> Select Points layer and click on objects to segment them</li>"
+            "<li><b>Rectangle Mode:</b> Draw rectangles around objects to segment them</li>"
+            "<li>Segmentation will automatically propagate to all slices</li>"
+            "</ul><br>"
+            "<b>General Controls:</b><br>"
+            "<ul>"
+            "<li>Use <b>Shift+click</b> for negative points (remove areas from segmentation)</li>"
+            "<li>Click on existing objects in <b>Segmentation layer</b> to select for cropping</li>"
+            "<li>Press <b>CTRL+click</b> on labels in <b>Segmentation layer</b> to delete them</li>"
+            "<li>Press <b>'Crop'</b> to save selected objects to disk</li>"
+            "</ul>"
+        )
+    else:
+        instructions_text = (
+            f"<b>Processing {dimension_type} data</b><br><br>"
+            "<b>Point Mode:</b> Click on objects to segment them. Use Shift+click for negative points.<br>"
+            "<b>Rectangle Mode:</b> Draw rectangles around objects to segment them.<br><br>"
+            "<ul>"
+            "<li>Click on existing objects in <b>Segmentation layer</b> to select them for cropping</li>"
+            "<li>Press <b>CTRL+click</b> on labels in <b>Segmentation layer</b> to delete them</li>"
+            "<li>Press <b>'Crop'</b> to save selected objects to disk</li>"
+            "</ul>"
+        )
+
+    instructions_label = QLabel(instructions_text)
     instructions_label.setWordWrap(True)
     layout.addWidget(instructions_label)
 
-    # Add a button to ensure points layer is active
-    activate_button = QPushButton("Make Points Layer Active")
+    # Add mode selector
+    mode_layout = QHBoxLayout()
+    mode_label = QLabel("<b>Prompt Mode:</b>")
+    mode_layout.addWidget(mode_label)
+
+    point_mode_button = QPushButton("Points")
+    point_mode_button.setCheckable(True)
+    point_mode_button.setChecked(True)
+    mode_layout.addWidget(point_mode_button)
+
+    box_mode_button = QPushButton("Rectangle")
+    box_mode_button.setCheckable(True)
+    box_mode_button.setChecked(False)
+    mode_layout.addWidget(box_mode_button)
+
+    mode_layout.addStretch()
+    layout.addLayout(mode_layout)
+
+    # Add a button to ensure active layer is correct
+    activate_button = QPushButton("Make Prompt Layer Active")
     activate_button.clicked.connect(
-        lambda: processor._ensure_points_layer_active()
+        lambda: processor._ensure_active_prompt_layer()
     )
     layout.addWidget(activate_button)
 
-    # Add a "Clear Points" button to reset prompts
-    clear_points_button = QPushButton("Clear Points")
-    layout.addWidget(clear_points_button)
+    # Add a "Clear Prompts" button to reset prompts
+    clear_prompts_button = QPushButton("Clear Prompts")
+    layout.addWidget(clear_prompts_button)
 
     # Create label table
     label_table = processor.create_label_table(crop_widget)
@@ -2353,7 +3137,7 @@ def create_crop_widget(processor):
     # Selection buttons
     selection_layout = QHBoxLayout()
     select_all_button = QPushButton("Select All")
-    clear_selection_button = QPushButton("Clear Selection")
+    clear_selection_button = QPushButton("Clear All Labels")
     selection_layout.addWidget(select_all_button)
     selection_layout.addWidget(clear_selection_button)
     layout.addLayout(selection_layout)
@@ -2391,51 +3175,152 @@ def create_crop_widget(processor):
         # Create new table
         label_table = processor.create_label_table(crop_widget)
         label_table.setMinimumHeight(200)
-        layout.insertWidget(3, label_table)  # Insert after clear points button
+        layout.insertWidget(
+            3, label_table
+        )  # Insert after clear prompts button
         return label_table
 
-    # Add helper method to ensure points layer is active
-    def _ensure_points_layer_active():
-        points_layer = None
-        for layer in list(processor.viewer.layers):
-            if "Points" in layer.name:
-                points_layer = layer
-                break
+    # Add helper method to ensure active prompt layer is selected based on mode
+    def _ensure_active_prompt_layer():
+        if processor.prompt_mode == "point":
+            points_layer = None
+            for layer in list(processor.viewer.layers):
+                if "Points" in layer.name and "Object" not in layer.name:
+                    points_layer = layer
+                    break
 
-        if points_layer is not None:
-            processor.viewer.layers.selection.active = points_layer
-            status_label.setText(
-                "Points layer is now active - click to add points"
-            )
-        else:
-            status_label.setText(
-                "No points layer found. Please load an image first."
-            )
+            if points_layer is not None:
+                processor.viewer.layers.selection.active = points_layer
+                if processor.use_3d:
+                    status_label.setText(
+                        "Points layer active - Navigate to FIRST SLICE of object, ensure 2D view, then click"
+                    )
+                else:
+                    status_label.setText(
+                        "Points layer is now active - click to add points"
+                    )
+            else:
+                status_label.setText(
+                    "No points layer found. Please load an image first."
+                )
+        else:  # box mode
+            shapes_layer = None
+            for layer in list(processor.viewer.layers):
+                if "Rectangles" in layer.name:
+                    shapes_layer = layer
+                    break
+
+            if shapes_layer is not None:
+                processor.viewer.layers.selection.active = shapes_layer
+                status_label.setText(
+                    "Rectangles layer is now active - draw rectangles"
+                )
+            else:
+                status_label.setText(
+                    "No rectangles layer found. Please load an image first."
+                )
+
+    processor._ensure_active_prompt_layer = _ensure_active_prompt_layer
+
+    # Keep the old method for backward compatibility
+    processor._ensure_points_layer_active = _ensure_active_prompt_layer
 
-    processor._ensure_points_layer_active = _ensure_points_layer_active
+    def on_clear_prompts_clicked():
+        # Find and clear/remove prompt layers based on mode
+        main_points_layer = None
+        object_points_layers = []
+        shapes_layer = None
 
-    # Connect button signals
-    def on_clear_points_clicked():
-        # Remove all point layers
         for layer in list(processor.viewer.layers):
             if "Points" in layer.name:
+                if "Object" in layer.name:
+                    object_points_layers.append(layer)
+                else:
+                    main_points_layer = layer
+            elif "Rectangles" in layer.name:
+                shapes_layer = layer
+
+        # Remove object-specific point layers (these are created dynamically)
+        for layer in object_points_layers:
+            # Clean up callbacks before removing the layer to prevent cleanup issues
+            if hasattr(layer, "mouse_drag_callbacks"):
+                layer.mouse_drag_callbacks.clear()
+            with contextlib.suppress(ValueError):
                 processor.viewer.layers.remove(layer)
 
-        # Reset point tracking attributes
-        if hasattr(processor, "points_data"):
-            processor.points_data = {}
-            processor.points_labels = {}
+        # Clear shapes layer
+        if shapes_layer is not None:
+            shapes_layer.data = []
 
-        if hasattr(processor, "obj_points"):
-            processor.obj_points = {}
-            processor.obj_labels = {}
+        # Clear data from main points layer instead of removing it
+        if main_points_layer is not None:
+            # Clear the points data
+            main_points_layer.data = np.zeros(
+                (0, 2 if not processor.use_3d else 3)
+            )
+            main_points_layer.face_color = "green"
 
-        # Re-create empty points layer
-        processor._update_label_layer()
-        processor._ensure_points_layer_active()
+            # Ensure the click callback is still connected
+            if (
+                hasattr(main_points_layer, "mouse_drag_callbacks")
+                and processor._on_points_clicked
+                not in main_points_layer.mouse_drag_callbacks
+            ):
+                main_points_layer.mouse_drag_callbacks.append(
+                    processor._on_points_clicked
+                )
+
+        # Reset all tracking attributes for 2D
+        if not processor.use_3d:
+            # Reset current segmentation tracking
+            if hasattr(processor, "current_points"):
+                processor.current_points = []
+                processor.current_labels = []
+
+            # Reset object tracking
+            if hasattr(processor, "obj_points"):
+                processor.obj_points = {}
+                processor.obj_labels = {}
+
+            # Reset box tracking
+            if hasattr(processor, "obj_boxes"):
+                processor.obj_boxes = {}
+
+            # Reset object ID counters
+            if hasattr(processor, "current_obj_id"):
+                # Find the highest existing label ID
+                if processor.segmentation_result is not None:
+                    max_label = processor.segmentation_result.max()
+                    processor.current_obj_id = max(int(max_label) + 1, 1)
+                    processor.next_obj_id = processor.current_obj_id
+                else:
+                    processor.current_obj_id = 1
+                    processor.next_obj_id = 1
+
+            # Reset SAM2 predictor state
+            processor.reset_sam2_state()
+
+        # For 3D, reset video-specific tracking
+        else:
+            if hasattr(processor, "sam2_points_by_obj"):
+                processor.sam2_points_by_obj = {}
+                processor.sam2_labels_by_obj = {}
+
+            # Reset box tracking
+            if hasattr(processor, "obj_boxes"):
+                processor.obj_boxes = {}
+
+            if hasattr(processor, "points_data"):
+                processor.points_data = {}
+                processor.points_labels = {}
+
+            # Note: We don't reset _sam2_state for 3D as it needs to maintain video state
+
+        # Make the appropriate prompt layer active based on mode
+        _ensure_active_prompt_layer()
 
         status_label.setText(
-            "Cleared all points. Click on Points layer to add new points."
+            "Cleared all prompts. Ready to add new segmentation prompts."
         )
 
     def on_select_all_clicked():
@@ -2459,8 +3344,14 @@ def create_crop_widget(processor):
             )
 
     def on_next_clicked():
-        # Clear points before moving to next image
-        on_clear_points_clicked()
+        # Check if we can move to the next image before clearing prompts
+        if processor.current_index >= len(processor.images) - 1:
+            next_button.setEnabled(False)
+            status_label.setText("No more images. Processing complete.")
+            return
+
+        # Clear prompts before moving to next image
+        on_clear_prompts_clicked()
 
         if not processor.next_image():
             next_button.setEnabled(False)
@@ -2470,11 +3361,17 @@ def create_crop_widget(processor):
             status_label.setText(
                 f"Showing image {processor.current_index + 1}/{len(processor.images)}"
             )
-            processor._ensure_points_layer_active()
+            processor._ensure_active_prompt_layer()
 
     def on_prev_clicked():
-        # Clear points before moving to previous image
-        on_clear_points_clicked()
+        # Check if we can move to the previous image before clearing prompts
+        if processor.current_index <= 0:
+            prev_button.setEnabled(False)
+            status_label.setText("Already at the first image.")
+            return
+
+        # Clear prompts before moving to previous image
+        on_clear_prompts_clicked()
 
         if not processor.previous_image():
             prev_button.setEnabled(False)
@@ -2484,15 +3381,33 @@ def create_crop_widget(processor):
             status_label.setText(
                 f"Showing image {processor.current_index + 1}/{len(processor.images)}"
             )
-            processor._ensure_points_layer_active()
+            processor._ensure_active_prompt_layer()
+
+    def on_point_mode_clicked():
+        processor.prompt_mode = "point"
+        point_mode_button.setChecked(True)
+        box_mode_button.setChecked(False)
+        processor._update_label_layer()
+        status_label.setText("Point mode active - click on objects to segment")
 
-    clear_points_button.clicked.connect(on_clear_points_clicked)
+    def on_box_mode_clicked():
+        processor.prompt_mode = "box"
+        point_mode_button.setChecked(False)
+        box_mode_button.setChecked(True)
+        processor._update_label_layer()
+        status_label.setText(
+            "Rectangle mode active - draw rectangles around objects"
+        )
+
+    clear_prompts_button.clicked.connect(on_clear_prompts_clicked)
     select_all_button.clicked.connect(on_select_all_clicked)
     clear_selection_button.clicked.connect(on_clear_selection_clicked)
     crop_button.clicked.connect(on_crop_clicked)
     next_button.clicked.connect(on_next_clicked)
     prev_button.clicked.connect(on_prev_clicked)
-    activate_button.clicked.connect(_ensure_points_layer_active)
+    activate_button.clicked.connect(_ensure_active_prompt_layer)
+    point_mode_button.clicked.connect(on_point_mode_clicked)
+    box_mode_button.clicked.connect(on_box_mode_clicked)
 
     return crop_widget
 
@@ -2511,6 +3426,19 @@ def batch_crop_anything(
     viewer: Viewer = None,
 ):
     """MagicGUI widget for starting Batch Crop Anything using SAM2."""
+    # Check if torch is available
+    if not _HAS_TORCH:
+        QMessageBox.critical(
+            None,
+            "Missing Dependency",
+            "PyTorch not found. Batch Crop Anything requires PyTorch and SAM2.\n\n"
+            "To install the required dependencies, run:\n"
+            "pip install 'napari-tmidas[deep-learning]'\n\n"
+            "Then follow SAM2 installation instructions at:\n"
+            "https://github.com/MercaderLabAnatomy/napari-tmidas#installation",
+        )
+        return
+
     # Check if SAM2 is available
     try:
         import importlib.util
@@ -2521,15 +3449,15 @@ def batch_crop_anything(
                 None,
                 "Missing Dependency",
                 "SAM2 not found. Please follow installation instructions at:\n"
-                "https://github.com/MercaderLabAnatomy/napari-tmidas?tab=readme-ov-file#dependencies\n",
+                "https://github.com/MercaderLabAnatomy/napari-tmidas#installation\n",
             )
             return
     except ImportError:
         QMessageBox.critical(
             None,
             "Missing Dependency",
-            "SAM2 package cannot be imported. Please follow installation instructions at\n"
-            "https://github.com/MercaderLabAnatomy/napari-tmidas?tab=readme-ov-file#dependencies",
+            "SAM2 package cannot be imported. Please follow installation instructions at:\n"
+            "https://github.com/MercaderLabAnatomy/napari-tmidas#installation",
         )
         return
 
@@ -2557,24 +3485,7 @@ def batch_crop_anything_widget():
     # Create the magicgui widget
     widget = batch_crop_anything
 
-    # Create and add browse button for folder path
-    folder_browse_button = QPushButton("Browse...")
-
-    def on_folder_browse_clicked():
-        folder = QFileDialog.getExistingDirectory(
-            None,
-            "Select Folder",
-            os.path.expanduser("~"),
-            QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks,
-        )
-        if folder:
-            # Update the folder_path field
-            widget.folder_path.value = folder
-
-    folder_browse_button.clicked.connect(on_folder_browse_clicked)
-
-    # Insert the browse button next to the folder_path field
-    folder_layout = widget.folder_path.native.parent().layout()
-    folder_layout.addWidget(folder_browse_button)
+    # Add browse button using common utility
+    add_browse_button_to_folder_field(widget, "folder_path")
 
     return widget