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

lazylabel/models/__init__.py

@@ -0,0 +1,5 @@
+"""Model-related modules."""
+
+from .sam_model import SamModel
+
+__all__ = ['SamModel']

lazylabel/models/sam_model.py

@@ -0,0 +1,154 @@
+import os
+import cv2
+import numpy as np
+import torch
+import requests
+from tqdm import tqdm
+from segment_anything import sam_model_registry, SamPredictor
+
+
+def download_model(url, download_path):
+    """Downloads file with a progress bar."""
+    print(
+        f"SAM model not found. Downloading from Meta's GitHub repository to: {download_path}"
+    )
+    try:
+        response = requests.get(url, stream=True, timeout=30)
+        response.raise_for_status()
+        total_size_in_bytes = int(response.headers.get("content-length", 0))
+        block_size = 1024  # 1 Kibibyte
+
+        progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
+        with open(download_path, "wb") as file:
+            for data in response.iter_content(block_size):
+                progress_bar.update(len(data))
+                file.write(data)
+        progress_bar.close()
+
+        if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+            raise RuntimeError("Download incomplete - file size mismatch")
+            
+        print("Model download completed successfully.")
+        
+    except requests.exceptions.RequestException as e:
+        raise RuntimeError(f"Network error during download: {e}")
+    except Exception as e:
+        # Clean up partial download
+        if os.path.exists(download_path):
+            os.remove(download_path)
+        raise RuntimeError(f"Download failed: {e}")
+
+
+class SamModel:
+    def __init__(self, model_type="vit_h", model_filename="sam_vit_h_4b8939.pth", custom_model_path=None):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.current_model_type = model_type
+        self.current_model_path = custom_model_path
+        self.model = None
+        self.predictor = None
+        self.image = None
+        self.is_loaded = False
+        
+        try:
+            if custom_model_path and os.path.exists(custom_model_path):
+                # Use custom model path
+                model_path = custom_model_path
+                print(f"Loading custom SAM model from {model_path}...")
+            else:
+                # Use default model with download if needed - store in models folder
+                model_url = f"https://dl.fbaipublicfiles.com/segment_anything/{model_filename}"
+                
+                # Use models folder instead of cache folder
+                models_dir = os.path.dirname(__file__)  # Already in models directory
+                os.makedirs(models_dir, exist_ok=True)
+                model_path = os.path.join(models_dir, model_filename)
+                
+                # Also check the old cache location and move it if it exists
+                old_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "lazylabel")
+                old_model_path = os.path.join(old_cache_dir, model_filename)
+                
+                if os.path.exists(old_model_path) and not os.path.exists(model_path):
+                    print(f"Moving existing model from cache to models folder...")
+                    import shutil
+                    shutil.move(old_model_path, model_path)
+                elif not os.path.exists(model_path):
+                    # Download the model if it doesn't exist
+                    download_model(model_url, model_path)
+                
+                print(f"Loading default SAM model from {model_path}...")
+
+            self.model = sam_model_registry[model_type](checkpoint=model_path).to(
+                self.device
+            )
+            self.predictor = SamPredictor(self.model)
+            self.is_loaded = True
+            print("SAM model loaded successfully.")
+            
+        except Exception as e:
+            print(f"Failed to load SAM model: {e}")
+            print("SAM point functionality will be disabled.")
+            self.is_loaded = False
+    
+    def load_custom_model(self, model_path, model_type="vit_h"):
+        """Load a custom model from the specified path."""
+        if not os.path.exists(model_path):
+            print(f"Model file not found: {model_path}")
+            return False
+        
+        print(f"Loading custom SAM 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
+            
+            # Load new model
+            self.model = sam_model_registry[model_type](checkpoint=model_path).to(self.device)
+            self.predictor = SamPredictor(self.model)
+            self.current_model_type = model_type
+            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)
+                
+            print("Custom SAM model loaded successfully.")
+            return True
+        except Exception as e:
+            print(f"Error loading custom model: {e}")
+            self.is_loaded = False
+            self.model = None
+            self.predictor = None
+            return False
+
+    def set_image(self, image_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:
+            print(f"Error setting image: {e}")
+            return False
+
+    def predict(self, positive_points, negative_points):
+        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, _, _ = self.predictor.predict(
+                point_coords=points,
+                point_labels=labels,
+                multimask_output=False,
+            )
+            return masks[0]
+        except Exception as e:
+            print(f"Error during prediction: {e}")
+            return None

lazylabel/ui/__init__.py

@@ -0,0 +1,8 @@
+"""UI components for LazyLabel."""
+
+from .main_window import MainWindow
+from .control_panel import ControlPanel
+from .right_panel import RightPanel
+from .hotkey_dialog import HotkeyDialog
+
+__all__ = ['MainWindow', 'ControlPanel', 'RightPanel', 'HotkeyDialog']

lazylabel/ui/control_panel.py

@@ -0,0 +1,220 @@
+"""Left control panel with mode controls and settings."""
+
+from PyQt6.QtWidgets import (
+    QWidget, QVBoxLayout, QPushButton, QLabel, QFrame, QHBoxLayout,
+    QCheckBox, QSlider, QGroupBox, QComboBox
+)
+from PyQt6.QtCore import Qt, pyqtSignal
+
+from .widgets import ModelSelectionWidget, SettingsWidget, AdjustmentsWidget
+
+
+class ControlPanel(QWidget):
+    """Left control panel with mode controls and settings."""
+    
+    # Signals
+    sam_mode_requested = pyqtSignal()
+    polygon_mode_requested = pyqtSignal()
+    selection_mode_requested = pyqtSignal()
+    clear_points_requested = pyqtSignal()
+    fit_view_requested = pyqtSignal()
+    browse_models_requested = pyqtSignal()
+    refresh_models_requested = pyqtSignal()
+    model_selected = pyqtSignal(str)
+    annotation_size_changed = pyqtSignal(int)
+    pan_speed_changed = pyqtSignal(int)
+    join_threshold_changed = pyqtSignal(int)
+    hotkeys_requested = pyqtSignal()
+    
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.setFixedWidth(250)
+        self._setup_ui()
+        self._connect_signals()
+    
+    def _setup_ui(self):
+        """Setup the UI layout."""
+        layout = QVBoxLayout(self)
+        layout.setAlignment(Qt.AlignmentFlag.AlignTop)
+        
+        # Toggle button
+        toggle_layout = QHBoxLayout()
+        self.btn_toggle_visibility = QPushButton("< Hide")
+        self.btn_toggle_visibility.setToolTip("Hide this panel")
+        toggle_layout.addWidget(self.btn_toggle_visibility)
+        toggle_layout.addStretch()
+        layout.addLayout(toggle_layout)
+        
+        # Main controls widget
+        self.main_controls_widget = QWidget()
+        main_layout = QVBoxLayout(self.main_controls_widget)
+        main_layout.setContentsMargins(0, 0, 0, 0)
+        
+        # Mode label
+        self.mode_label = QLabel("Mode: Points")
+        font = self.mode_label.font()
+        font.setPointSize(14)
+        font.setBold(True)
+        self.mode_label.setFont(font)
+        main_layout.addWidget(self.mode_label)
+        
+        # Mode buttons
+        self._add_mode_buttons(main_layout)
+        
+        # Separator
+        main_layout.addSpacing(20)
+        main_layout.addWidget(self._create_separator())
+        main_layout.addSpacing(10)
+        
+        # Model selection
+        self.model_widget = ModelSelectionWidget()
+        main_layout.addWidget(self.model_widget)
+        
+        # Separator
+        main_layout.addSpacing(10)
+        main_layout.addWidget(self._create_separator())
+        main_layout.addSpacing(10)
+        
+        # Action buttons
+        self._add_action_buttons(main_layout)
+        main_layout.addSpacing(10)
+        
+        # Settings
+        self.settings_widget = SettingsWidget()
+        main_layout.addWidget(self.settings_widget)
+        
+        # Adjustments
+        self.adjustments_widget = AdjustmentsWidget()
+        main_layout.addWidget(self.adjustments_widget)
+        
+        main_layout.addStretch()
+        
+        # Status labels
+        self.notification_label = QLabel("")
+        font = self.notification_label.font()
+        font.setItalic(True)
+        self.notification_label.setFont(font)
+        self.notification_label.setStyleSheet("color: #ffa500;")
+        self.notification_label.setWordWrap(True)
+        main_layout.addWidget(self.notification_label)
+        
+        self.device_label = QLabel("Device: Unknown")
+        main_layout.addWidget(self.device_label)
+        
+        layout.addWidget(self.main_controls_widget)
+    
+    def _add_mode_buttons(self, layout):
+        """Add mode control buttons."""
+        self.btn_sam_mode = QPushButton("Point Mode (1)")
+        self.btn_sam_mode.setToolTip("Switch to Point Mode for AI segmentation (1)")
+        
+        self.btn_polygon_mode = QPushButton("Polygon Mode (2)")
+        self.btn_polygon_mode.setToolTip("Switch to Polygon Drawing Mode (2)")
+        
+        self.btn_selection_mode = QPushButton("Selection Mode (E)")
+        self.btn_selection_mode.setToolTip("Toggle segment selection (E)")
+        
+        layout.addWidget(self.btn_sam_mode)
+        layout.addWidget(self.btn_polygon_mode)
+        layout.addWidget(self.btn_selection_mode)
+    
+    def _add_action_buttons(self, layout):
+        """Add action buttons."""
+        self.btn_fit_view = QPushButton("Fit View (.)")
+        self.btn_fit_view.setToolTip("Reset image zoom and pan to fit the view (.)")
+        
+        self.btn_clear_points = QPushButton("Clear Clicks (C)")
+        self.btn_clear_points.setToolTip("Clear current temporary points/vertices (C)")
+        
+        self.btn_hotkeys = QPushButton("Hotkeys")
+        self.btn_hotkeys.setToolTip("Configure keyboard shortcuts")
+        
+        layout.addWidget(self.btn_fit_view)
+        layout.addWidget(self.btn_clear_points)
+        layout.addWidget(self.btn_hotkeys)
+    
+    def _create_separator(self):
+        """Create a horizontal separator line."""
+        line = QFrame()
+        line.setFrameShape(QFrame.Shape.HLine)
+        return line
+    
+    def _connect_signals(self):
+        """Connect internal signals."""
+        self.btn_sam_mode.clicked.connect(self.sam_mode_requested)
+        self.btn_polygon_mode.clicked.connect(self.polygon_mode_requested)
+        self.btn_selection_mode.clicked.connect(self.selection_mode_requested)
+        self.btn_clear_points.clicked.connect(self.clear_points_requested)
+        self.btn_fit_view.clicked.connect(self.fit_view_requested)
+        self.btn_hotkeys.clicked.connect(self.hotkeys_requested)
+        
+        # Model widget signals
+        self.model_widget.browse_requested.connect(self.browse_models_requested)
+        self.model_widget.refresh_requested.connect(self.refresh_models_requested)
+        self.model_widget.model_selected.connect(self.model_selected)
+        
+        # Settings widget signals
+        self.adjustments_widget.annotation_size_changed.connect(self.annotation_size_changed)
+        self.adjustments_widget.pan_speed_changed.connect(self.pan_speed_changed)
+        self.adjustments_widget.join_threshold_changed.connect(self.join_threshold_changed)
+    
+    def toggle_visibility(self):
+        """Toggle panel visibility."""
+        is_visible = self.main_controls_widget.isVisible()
+        self.main_controls_widget.setVisible(not is_visible)
+        if is_visible:
+            self.btn_toggle_visibility.setText("> Show")
+            self.setFixedWidth(self.btn_toggle_visibility.sizeHint().width() + 20)
+        else:
+            self.btn_toggle_visibility.setText("< Hide")
+            self.setFixedWidth(250)
+    
+    def show_notification(self, message: str, duration: int = 3000):
+        """Show a notification message."""
+        self.notification_label.setText(message)
+        # Note: Timer should be handled by the caller
+    
+    def clear_notification(self):
+        """Clear the notification message."""
+        self.notification_label.clear()
+    
+    def set_mode_text(self, mode: str):
+        """Set the mode label text."""
+        self.mode_label.setText(f"Mode: {mode.replace('_', ' ').title()}")
+    
+    def set_device_text(self, device: str):
+        """Set the device label text."""
+        self.device_label.setText(f"Device: {device.upper()}")
+    
+    # Delegate methods for sub-widgets
+    def populate_models(self, models):
+        """Populate the models combo box."""
+        self.model_widget.populate_models(models)
+    
+    def set_current_model(self, model_name):
+        """Set the current model display."""
+        self.model_widget.set_current_model(model_name)
+    
+    def get_settings(self):
+        """Get current settings from the settings widget."""
+        return self.settings_widget.get_settings()
+    
+    def set_settings(self, settings):
+        """Set settings in the settings widget."""
+        self.settings_widget.set_settings(settings)
+    
+    def get_annotation_size(self):
+        """Get current annotation size."""
+        return self.adjustments_widget.get_annotation_size()
+    
+    def set_annotation_size(self, value):
+        """Set annotation size."""
+        self.adjustments_widget.set_annotation_size(value)
+    
+    def set_sam_mode_enabled(self, enabled: bool):
+        """Enable or disable the SAM mode button."""
+        self.btn_sam_mode.setEnabled(enabled)
+        if not enabled:
+            self.btn_sam_mode.setToolTip("Point Mode (SAM model not available)")
+        else:
+            self.btn_sam_mode.setToolTip("Switch to Point Mode for AI segmentation (1)")

lazylabel/{editable_vertex.py → ui/editable_vertex.py}

@@ -17,13 +17,35 @@ class EditableVertexItem(QGraphicsEllipseItem):
         self.setBrush(QBrush(color))
 
         self.setPen(QPen(Qt.GlobalColor.transparent))
+
+        # Set flags for dragging - use the original working approach
         self.setFlag(QGraphicsItem.GraphicsItemFlag.ItemIsMovable)
         self.setFlag(QGraphicsItem.GraphicsItemFlag.ItemSendsGeometryChanges)
+        self.setFlag(QGraphicsItem.GraphicsItemFlag.ItemIsSelectable)
+
+        # Accept mouse events
+        self.setAcceptHoverEvents(True)
 
     def itemChange(self, change, value):
         if change == QGraphicsItem.GraphicsItemChange.ItemPositionHasChanged:
             new_pos = value
-            self.main_window.update_vertex_pos(
-                self.segment_index, self.vertex_index, new_pos
-            )
+            if hasattr(self.main_window, "update_vertex_pos"):
+                self.main_window.update_vertex_pos(
+                    self.segment_index, self.vertex_index, new_pos
+                )
         return super().itemChange(change, value)
+
+    def mousePressEvent(self, event):
+        """Handle mouse press events."""
+        super().mousePressEvent(event)
+        event.accept()
+
+    def mouseMoveEvent(self, event):
+        """Handle mouse move events."""
+        super().mouseMoveEvent(event)
+        event.accept()
+
+    def mouseReleaseEvent(self, event):
+        """Handle mouse release events."""
+        super().mouseReleaseEvent(event)
+        event.accept()