lazylabel-gui 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

lazylabel/config/hotkeys.py

@@ -44,7 +44,7 @@ class HotkeyManager:
             ),
             HotkeyAction("fit_view", "Fit View", ".", category="Navigation"),
             # Modes
-            HotkeyAction("sam_mode", "Point Mode (SAM)", "1", category="Modes"),
+            HotkeyAction("sam_mode", "AI Mode (Points + Box)", "1", category="Modes"),
             HotkeyAction("polygon_mode", "Polygon Mode", "2", category="Modes"),
             HotkeyAction("bbox_mode", "Bounding Box Mode", "3", category="Modes"),
             HotkeyAction("selection_mode", "Selection Mode", "E", category="Modes"),
@@ -95,7 +95,7 @@ class HotkeyManager:
             # Mouse-related (cannot be reassigned)
             HotkeyAction(
                 "left_click",
-                "Add Positive Point / Select",
+                "AI: Point (click) / Box (drag) / Select",
                 "Left Click",
                 category="Mouse",
                 mouse_related=True,

lazylabel/core/model_manager.py

@@ -7,6 +7,18 @@ from ..config import Paths
 from ..models.sam_model import SamModel
 from ..utils.logger import logger
 
+# Optional SAM-2 support
+try:
+    from ..models.sam2_model import Sam2Model
+
+    SAM2_AVAILABLE = True
+except ImportError:
+    logger.info(
+        "SAM-2 not available. Install with: pip install git+https://github.com/facebookresearch/sam2.git"
+    )
+    Sam2Model = None
+    SAM2_AVAILABLE = False
+
 
 class ModelManager:
     """Manages SAM model loading and selection."""
@@ -42,7 +54,7 @@ class ModelManager:
         pth_files = []
         for root, _dirs, files in os.walk(folder_path):
             for file in files:
-                if file.lower().endswith(".pth"):
+                if file.lower().endswith(".pth") or file.lower().endswith(".pt"):
                     full_path = os.path.join(root, file)
                     rel_path = os.path.relpath(full_path, folder_path)
                     pth_files.append((rel_path, full_path))
@@ -52,13 +64,34 @@ class ModelManager:
     def detect_model_type(self, model_path: str) -> str:
         """Detect model type from filename."""
         filename = os.path.basename(model_path).lower()
-        if "vit_l" in filename or "large" in filename:
-            return "vit_l"
-        elif "vit_b" in filename or "base" in filename:
-            return "vit_b"
-        elif "vit_h" in filename or "huge" in filename:
-            return "vit_h"
-        return "vit_h"  # default
+
+        # Check if it's a SAM2 model
+        if self._is_sam2_model(model_path):
+            if "tiny" in filename or "_t" in filename:
+                return "sam2_tiny"
+            elif "small" in filename or "_s" in filename:
+                return "sam2_small"
+            elif "base_plus" in filename or "_b+" in filename:
+                return "sam2_base_plus"
+            elif "large" in filename or "_l" in filename:
+                return "sam2_large"
+            else:
+                return "sam2_large"  # default for SAM2
+        else:
+            # Original SAM model types
+            if "vit_l" in filename or "large" in filename:
+                return "vit_l"
+            elif "vit_b" in filename or "base" in filename:
+                return "vit_b"
+            elif "vit_h" in filename or "huge" in filename:
+                return "vit_h"
+            return "vit_h"  # default for SAM1
+
+    def _is_sam2_model(self, model_path: str) -> bool:
+        """Check if the model is a SAM2 model based on filename patterns."""
+        filename = os.path.basename(model_path).lower()
+        sam2_indicators = ["sam2", "sam2.1", "hiera", "_t.", "_s.", "_b+.", "_l."]
+        return any(indicator in filename for indicator in sam2_indicators)
 
     def load_custom_model(self, model_path: str) -> bool:
         """Load a custom model from path.
@@ -66,20 +99,62 @@ class ModelManager:
         Returns:
             True if successful, False otherwise
         """
-        if not self.sam_model:
-            return False
-
         if not os.path.exists(model_path):
             return False
 
         model_type = self.detect_model_type(model_path)
-        success = self.sam_model.load_custom_model(model_path, model_type)
 
-        if success and self.on_model_changed:
-            model_name = os.path.basename(model_path)
-            self.on_model_changed(f"Current: {model_name}")
+        try:
+            # Clear existing model from memory
+            if self.sam_model is not None:
+                del self.sam_model
+                import torch
+
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+            # Create appropriate model instance
+            if self._is_sam2_model(model_path):
+                if not SAM2_AVAILABLE:
+                    logger.warning(
+                        f"SAM-2 model detected but SAM-2 not installed: {model_path}"
+                    )
+                    logger.info(
+                        "Install SAM-2 with: pip install git+https://github.com/facebookresearch/sam2.git"
+                    )
+                    return False
+
+                logger.info(f"Loading SAM2 model: {model_type}")
+                self.sam_model = Sam2Model(model_path)
+            else:
+                logger.info(f"Loading SAM1 model: {model_type}")
+                # Convert SAM2 model types back to SAM1 types for compatibility
+                sam1_model_type = model_type
+                if model_type.startswith("sam2_"):
+                    type_mapping = {
+                        "sam2_tiny": "vit_b",
+                        "sam2_small": "vit_b",
+                        "sam2_base_plus": "vit_l",
+                        "sam2_large": "vit_h",
+                    }
+                    sam1_model_type = type_mapping.get(model_type, "vit_h")
+
+                # Create SAM1 model with custom path
+                self.sam_model = SamModel(
+                    model_type=sam1_model_type, custom_model_path=model_path
+                )
+
+            success = self.sam_model.is_loaded
+
+            if success and self.on_model_changed:
+                model_name = os.path.basename(model_path)
+                self.on_model_changed(f"Current: {model_name}")
+
+            return success
 
-        return success
+        except Exception as e:
+            logger.error(f"Failed to load custom model: {e}")
+            self.sam_model = None
+            return False
 
     def set_models_folder(self, folder_path: str) -> None:
         """Set the current models folder."""

lazylabel/models/sam2_model.py

@@ -0,0 +1,223 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from ..utils.logger import logger
+
+# SAM-2 specific imports - will fail gracefully if not available
+try:
+    from sam2.build_sam import build_sam2
+    from sam2.sam2_image_predictor import SAM2ImagePredictor
+except ImportError as e:
+    logger.error(f"SAM-2 dependencies not found: {e}")
+    logger.info(
+        "Install SAM-2 with: pip install git+https://github.com/facebookresearch/sam2.git"
+    )
+    raise ImportError("SAM-2 dependencies required for Sam2Model") from e
+
+
+class Sam2Model:
+    """SAM2 model wrapper that provides the same interface as SamModel."""
+
+    def __init__(self, model_path: str, config_path: str | None = None):
+        """Initialize SAM2 model.
+
+        Args:
+            model_path: Path to the SAM2 model checkpoint (.pt file)
+            config_path: Path to the config file (optional, will auto-detect if None)
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"SAM2: Detected device: {str(self.device).upper()}")
+
+        self.current_model_path = model_path
+        self.model = None
+        self.predictor = None
+        self.image = None
+        self.is_loaded = False
+
+        # Auto-detect config if not provided
+        if config_path is None:
+            config_path = self._auto_detect_config(model_path)
+
+        try:
+            logger.info(f"SAM2: Loading model from {model_path}...")
+            logger.info(f"SAM2: Using config: {config_path}")
+
+            # Build SAM2 model
+            self.model = build_sam2(config_path, model_path, device=self.device)
+
+            # Create predictor
+            self.predictor = SAM2ImagePredictor(self.model)
+
+            self.is_loaded = True
+            logger.info("SAM2: Model loaded successfully.")
+
+        except Exception as e:
+            logger.error(f"SAM2: Failed to load model: {e}")
+            logger.warning("SAM2: SAM2 functionality will be disabled.")
+            self.is_loaded = False
+
+    def _auto_detect_config(self, model_path: str) -> str:
+        """Auto-detect the appropriate config file based on model filename."""
+        filename = os.path.basename(model_path).lower()
+
+        # Get the sam2 package directory
+        try:
+            import sam2
+
+            sam2_dir = os.path.dirname(sam2.__file__)
+            configs_dir = os.path.join(sam2_dir, "configs")
+
+            # Map model types to config files
+            if "tiny" in filename or "_t" in filename:
+                config_file = (
+                    "sam2.1_hiera_t.yaml" if "2.1" in filename else "sam2_hiera_t.yaml"
+                )
+            elif "small" in filename or "_s" in filename:
+                config_file = (
+                    "sam2.1_hiera_s.yaml" if "2.1" in filename else "sam2_hiera_s.yaml"
+                )
+            elif "base_plus" in filename or "_b+" in filename:
+                config_file = (
+                    "sam2.1_hiera_b+.yaml"
+                    if "2.1" in filename
+                    else "sam2_hiera_b+.yaml"
+                )
+            elif "large" in filename or "_l" in filename:
+                config_file = (
+                    "sam2.1_hiera_l.yaml" if "2.1" in filename else "sam2_hiera_l.yaml"
+                )
+            else:
+                # Default to large model
+                config_file = "sam2.1_hiera_l.yaml"
+
+            # Check sam2.1 configs first, then fall back to sam2
+            if "2.1" in filename:
+                config_path = os.path.join(configs_dir, "sam2.1", config_file)
+            else:
+                config_path = os.path.join(
+                    configs_dir, "sam2", config_file.replace("2.1_", "")
+                )
+
+            if os.path.exists(config_path):
+                return config_path
+
+            # Fallback to default large config
+            fallback_config = os.path.join(configs_dir, "sam2.1", "sam2.1_hiera_l.yaml")
+            if os.path.exists(fallback_config):
+                return fallback_config
+
+            raise FileNotFoundError(f"No suitable config found for {filename}")
+
+        except Exception as e:
+            logger.error(f"SAM2: Failed to auto-detect config: {e}")
+            # Return a reasonable default path
+            return "sam2.1_hiera_l.yaml"
+
+    def set_image_from_path(self, image_path: str) -> bool:
+        """Set image for SAM2 model from file path."""
+        if not self.is_loaded:
+            return False
+        try:
+            self.image = cv2.imread(image_path)
+            self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
+            self.predictor.set_image(self.image)
+            return True
+        except Exception as e:
+            logger.error(f"SAM2: Error setting image from path: {e}")
+            return False
+
+    def set_image_from_array(self, image_array: np.ndarray) -> bool:
+        """Set image for SAM2 model from numpy array."""
+        if not self.is_loaded:
+            return False
+        try:
+            self.image = image_array
+            self.predictor.set_image(self.image)
+            return True
+        except Exception as e:
+            logger.error(f"SAM2: Error setting image from array: {e}")
+            return False
+
+    def predict(self, positive_points, negative_points):
+        """Generate predictions using SAM2."""
+        if not self.is_loaded or not positive_points:
+            return None
+
+        try:
+            points = np.array(positive_points + negative_points)
+            labels = np.array([1] * len(positive_points) + [0] * len(negative_points))
+
+            masks, scores, logits = self.predictor.predict(
+                point_coords=points,
+                point_labels=labels,
+                multimask_output=True,
+            )
+
+            # Return the mask with the highest score
+            best_mask_idx = np.argmax(scores)
+            return masks[best_mask_idx], scores[best_mask_idx], logits[best_mask_idx]
+
+        except Exception as e:
+            logger.error(f"SAM2: Error during prediction: {e}")
+            return None
+
+    def predict_from_box(self, box):
+        """Generate predictions from bounding box using SAM2."""
+        if not self.is_loaded:
+            return None
+
+        try:
+            masks, scores, logits = self.predictor.predict(
+                box=np.array(box),
+                multimask_output=True,
+            )
+
+            # Return the mask with the highest score
+            best_mask_idx = np.argmax(scores)
+            return masks[best_mask_idx], scores[best_mask_idx], logits[best_mask_idx]
+
+        except Exception as e:
+            logger.error(f"SAM2: Error during box prediction: {e}")
+            return None
+
+    def load_custom_model(
+        self, model_path: str, config_path: str | None = None
+    ) -> bool:
+        """Load a custom SAM2 model from the specified path."""
+        if not os.path.exists(model_path):
+            logger.warning(f"SAM2: Model file not found: {model_path}")
+            return False
+
+        logger.info(f"SAM2: Loading custom model from {model_path}...")
+        try:
+            # Clear existing model from memory
+            if hasattr(self, "model") and self.model is not None:
+                del self.model
+                del self.predictor
+                torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+            # Auto-detect config if not provided
+            if config_path is None:
+                config_path = self._auto_detect_config(model_path)
+
+            # Load new model
+            self.model = build_sam2(config_path, model_path, device=self.device)
+            self.predictor = SAM2ImagePredictor(self.model)
+            self.current_model_path = model_path
+            self.is_loaded = True
+
+            # Re-set image if one was previously loaded
+            if self.image is not None:
+                self.predictor.set_image(self.image)
+
+            logger.info("SAM2: Custom model loaded successfully.")
+            return True
+        except Exception as e:
+            logger.error(f"SAM2: Error loading custom model: {e}")
+            self.is_loaded = False
+            self.model = None
+            self.predictor = None
+            return False

lazylabel/models/sam_model.py

@@ -200,12 +200,34 @@ class SamModel:
             points = np.array(positive_points + negative_points)
             labels = np.array([1] * len(positive_points) + [0] * len(negative_points))
 
-            masks, _, _ = self.predictor.predict(
+            masks, scores, logits = self.predictor.predict(
                 point_coords=points,
                 point_labels=labels,
-                multimask_output=False,
+                multimask_output=True,
             )
-            return masks[0]
+
+            # Return the mask with the highest score (consistent with SAM2)
+            best_mask_idx = np.argmax(scores)
+            return masks[best_mask_idx], scores[best_mask_idx], logits[best_mask_idx]
         except Exception as e:
             logger.error(f"Error during prediction: {e}")
             return None
+
+    def predict_from_box(self, box):
+        """Generate predictions from bounding box using SAM."""
+        if not self.is_loaded:
+            return None
+
+        try:
+            masks, scores, logits = self.predictor.predict(
+                box=np.array(box),
+                multimask_output=True,
+            )
+
+            # Return the mask with the highest score
+            best_mask_idx = np.argmax(scores)
+            return masks[best_mask_idx], scores[best_mask_idx], logits[best_mask_idx]
+
+        except Exception as e:
+            logger.error(f"Error during box prediction: {e}")
+            return None