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

lazylabel/__init__.py

@@ -0,0 +1,9 @@
+"""LazyLabel - AI-assisted image segmentation tool."""
+
+__version__ = "1.0.9"
+__author__ = "Deniz N. Cakan"
+__email__ = "deniz.n.cakan@gmail.com"
+
+from .main import main
+
+__all__ = ['main']

lazylabel/config/__init__.py

@@ -0,0 +1,7 @@
+"""Configuration management for LazyLabel."""
+
+from .settings import Settings, DEFAULT_SETTINGS
+from .paths import Paths
+from .hotkeys import HotkeyManager, HotkeyAction
+
+__all__ = ['Settings', 'DEFAULT_SETTINGS', 'Paths', 'HotkeyManager', 'HotkeyAction']

lazylabel/config/hotkeys.py

@@ -0,0 +1,169 @@
+"""Hotkey management system."""
+
+import json
+import os
+from dataclasses import dataclass, asdict
+from typing import Dict, Optional, List, Tuple
+from PyQt6.QtCore import Qt
+from PyQt6.QtGui import QKeySequence
+
+
+@dataclass
+class HotkeyAction:
+    """Represents a hotkey action with primary and secondary keys."""
+    name: str
+    description: str
+    primary_key: str
+    secondary_key: Optional[str] = None
+    category: str = "General"
+    mouse_related: bool = False  # Cannot be reassigned if True
+
+
+class HotkeyManager:
+    """Manages application hotkeys with persistence."""
+    
+    def __init__(self, config_dir: str):
+        self.config_dir = config_dir
+        self.hotkeys_file = os.path.join(config_dir, "hotkeys.json")
+        self.actions: Dict[str, HotkeyAction] = {}
+        self._initialize_default_hotkeys()
+        self.load_hotkeys()
+    
+    def _initialize_default_hotkeys(self):
+        """Initialize default hotkey mappings."""
+        default_hotkeys = [
+            # Navigation
+            HotkeyAction("load_next_image", "Load Next Image", "Right", category="Navigation"),
+            HotkeyAction("load_previous_image", "Load Previous Image", "Left", category="Navigation"),
+            HotkeyAction("fit_view", "Fit View", "Period", category="Navigation"),
+            
+            # Modes
+            HotkeyAction("sam_mode", "Point Mode (SAM)", "1", category="Modes"),
+            HotkeyAction("polygon_mode", "Polygon Mode", "2", category="Modes"),
+            HotkeyAction("selection_mode", "Selection Mode", "E", category="Modes"),
+            HotkeyAction("pan_mode", "Pan Mode", "Q", category="Modes"),
+            HotkeyAction("edit_mode", "Edit Mode", "R", category="Modes"),
+            
+            # Actions
+            HotkeyAction("clear_points", "Clear Points/Vertices", "C", category="Actions"),
+            HotkeyAction("save_segment", "Save Current Segment", "Space", category="Actions"),
+            HotkeyAction("save_output", "Save Output", "Return", category="Actions"),
+            HotkeyAction("save_output_alt", "Save Output (Alt)", "Enter", category="Actions"),
+            HotkeyAction("undo", "Undo Last Action", "Ctrl+Z", category="Actions"),
+            HotkeyAction("escape", "Cancel/Clear Selection", "Escape", category="Actions"),
+            
+            # Segments
+            HotkeyAction("merge_segments", "Merge Selected Segments", "M", category="Segments"),
+            HotkeyAction("delete_segments", "Delete Selected Segments", "V", category="Segments"),
+            HotkeyAction("delete_segments_alt", "Delete Selected Segments (Alt)", "Backspace", category="Segments"),
+            HotkeyAction("select_all", "Select All Segments", "Ctrl+A", category="Segments"),
+            
+            # View
+            HotkeyAction("zoom_in", "Zoom In", "Ctrl+Plus", category="View"),
+            HotkeyAction("zoom_out", "Zoom Out", "Ctrl+Minus", category="View"),
+            
+            # Movement (WASD)
+            HotkeyAction("pan_up", "Pan Up", "W", category="Movement"),
+            HotkeyAction("pan_down", "Pan Down", "S", category="Movement"),
+            HotkeyAction("pan_left", "Pan Left", "A", category="Movement"),
+            HotkeyAction("pan_right", "Pan Right", "D", category="Movement"),
+            
+            # Mouse-related (cannot be reassigned)
+            HotkeyAction("left_click", "Add Positive Point / Select", "Left Click", 
+                        category="Mouse", mouse_related=True),
+            HotkeyAction("right_click", "Add Negative Point", "Right Click", 
+                        category="Mouse", mouse_related=True),
+            HotkeyAction("mouse_drag", "Drag/Pan", "Mouse Drag", 
+                        category="Mouse", mouse_related=True),
+        ]
+        
+        for action in default_hotkeys:
+            self.actions[action.name] = action
+    
+    def get_action(self, action_name: str) -> Optional[HotkeyAction]:
+        """Get hotkey action by name."""
+        return self.actions.get(action_name)
+    
+    def get_actions_by_category(self) -> Dict[str, List[HotkeyAction]]:
+        """Get actions grouped by category."""
+        categories = {}
+        for action in self.actions.values():
+            if action.category not in categories:
+                categories[action.category] = []
+            categories[action.category].append(action)
+        return categories
+    
+    def set_primary_key(self, action_name: str, key: str) -> bool:
+        """Set primary key for an action."""
+        if action_name in self.actions and not self.actions[action_name].mouse_related:
+            self.actions[action_name].primary_key = key
+            return True
+        return False
+    
+    def set_secondary_key(self, action_name: str, key: Optional[str]) -> bool:
+        """Set secondary key for an action."""
+        if action_name in self.actions and not self.actions[action_name].mouse_related:
+            self.actions[action_name].secondary_key = key
+            return True
+        return False
+    
+    def get_key_for_action(self, action_name: str) -> Tuple[Optional[str], Optional[str]]:
+        """Get primary and secondary keys for an action."""
+        action = self.actions.get(action_name)
+        if action:
+            return action.primary_key, action.secondary_key
+        return None, None
+    
+    def is_key_in_use(self, key: str, exclude_action: str = None) -> Optional[str]:
+        """Check if a key is already in use by another action."""
+        for name, action in self.actions.items():
+            if name == exclude_action:
+                continue
+            if action.primary_key == key or action.secondary_key == key:
+                return name
+        return None
+    
+    def reset_to_defaults(self):
+        """Reset all hotkeys to default values."""
+        self._initialize_default_hotkeys()
+    
+    def save_hotkeys(self):
+        """Save hotkeys to file."""
+        os.makedirs(self.config_dir, exist_ok=True)
+        
+        # Convert to serializable format
+        data = {}
+        for name, action in self.actions.items():
+            if not action.mouse_related:  # Don't save mouse-related actions
+                data[name] = {
+                    'primary_key': action.primary_key,
+                    'secondary_key': action.secondary_key
+                }
+        
+        with open(self.hotkeys_file, 'w') as f:
+            json.dump(data, f, indent=4)
+    
+    def load_hotkeys(self):
+        """Load hotkeys from file."""
+        if not os.path.exists(self.hotkeys_file):
+            return
+        
+        try:
+            with open(self.hotkeys_file, 'r') as f:
+                data = json.load(f)
+            
+            for name, keys in data.items():
+                if name in self.actions and not self.actions[name].mouse_related:
+                    self.actions[name].primary_key = keys.get('primary_key', '')
+                    self.actions[name].secondary_key = keys.get('secondary_key')
+        except (json.JSONDecodeError, KeyError, FileNotFoundError):
+            # If loading fails, keep defaults
+            pass
+    
+    def key_sequence_to_string(self, key_sequence: QKeySequence) -> str:
+        """Convert QKeySequence to string representation."""
+        return key_sequence.toString()
+    
+    def string_to_key_sequence(self, key_string: str) -> QKeySequence:
+        """Convert string to QKeySequence."""
+        return QKeySequence(key_string)

lazylabel/config/paths.py

@@ -0,0 +1,41 @@
+"""Path management for LazyLabel."""
+
+import os
+from pathlib import Path
+
+
+class Paths:
+    """Centralized path management."""
+    
+    def __init__(self):
+        self.app_dir = Path(__file__).parent.parent
+        self.models_dir = self.app_dir / "models"
+        self.config_dir = Path.home() / ".config" / "lazylabel"
+        self.cache_dir = Path.home() / ".cache" / "lazylabel"
+        
+        # Ensure directories exist
+        self.models_dir.mkdir(exist_ok=True)
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        
+    @property
+    def settings_file(self) -> Path:
+        """Path to settings file."""
+        return self.config_dir / "settings.json"
+    
+    @property
+    def demo_pictures_dir(self) -> Path:
+        """Path to demo pictures directory."""
+        return self.app_dir / "demo_pictures"
+    
+    @property
+    def logo_path(self) -> Path:
+        """Path to application logo."""
+        return self.demo_pictures_dir / "logo2.png"
+    
+    def get_model_path(self, filename: str) -> Path:
+        """Get path for a model file."""
+        return self.models_dir / filename
+    
+    def get_old_cache_model_path(self, filename: str) -> Path:
+        """Get path for model in old cache location."""
+        return self.cache_dir / filename

lazylabel/config/settings.py

@@ -0,0 +1,66 @@
+"""Application settings and configuration."""
+
+import os
+from dataclasses import dataclass, asdict
+from typing import Dict, Any
+import json
+
+
+@dataclass
+class Settings:
+    """Application settings with defaults."""
+    
+    # UI Settings
+    window_width: int = 1600
+    window_height: int = 900
+    left_panel_width: int = 250
+    right_panel_width: int = 350
+    
+    # Annotation Settings
+    point_radius: float = 0.3
+    line_thickness: float = 0.5
+    pan_multiplier: float = 1.0
+    polygon_join_threshold: int = 2
+    
+    # Model Settings
+    default_model_type: str = "vit_h"
+    default_model_filename: str = "sam_vit_h_4b8939.pth"
+    
+    # Save Settings
+    auto_save: bool = True
+    save_npz: bool = True
+    save_txt: bool = True
+    save_class_aliases: bool = False
+    yolo_use_alias: bool = True
+    
+    # UI State
+    annotation_size_multiplier: float = 1.0
+    
+    def save_to_file(self, filepath: str) -> None:
+        """Save settings to JSON file."""
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+        with open(filepath, 'w') as f:
+            json.dump(asdict(self), f, indent=4)
+    
+    @classmethod
+    def load_from_file(cls, filepath: str) -> 'Settings':
+        """Load settings from JSON file."""
+        if not os.path.exists(filepath):
+            return cls()
+        
+        try:
+            with open(filepath, 'r') as f:
+                data = json.load(f)
+            return cls(**data)
+        except (json.JSONDecodeError, TypeError):
+            return cls()
+    
+    def update(self, **kwargs) -> None:
+        """Update settings with new values."""
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+
+
+# Default settings instance
+DEFAULT_SETTINGS = Settings()

lazylabel/core/__init__.py

@@ -0,0 +1,7 @@
+"""Core business logic for LazyLabel."""
+
+from .segment_manager import SegmentManager
+from .model_manager import ModelManager
+from .file_manager import FileManager
+
+__all__ = ['SegmentManager', 'ModelManager', 'FileManager']

lazylabel/core/file_manager.py

@@ -0,0 +1,106 @@
+"""File management functionality."""
+
+import os
+import json
+import numpy as np
+import cv2
+from typing import List, Dict, Any, Optional, Tuple
+from pathlib import Path
+
+from .segment_manager import SegmentManager
+
+
+class FileManager:
+    """Manages file operations for saving and loading."""
+    
+    def __init__(self, segment_manager: SegmentManager):
+        self.segment_manager = segment_manager
+    
+    def save_npz(self, image_path: str, image_size: Tuple[int, int], class_order: List[int]) -> str:
+        """Save segments as NPZ file."""
+        final_mask_tensor = self.segment_manager.create_final_mask_tensor(image_size, class_order)
+        npz_path = os.path.splitext(image_path)[0] + ".npz"
+        np.savez_compressed(npz_path, mask=final_mask_tensor.astype(np.uint8))
+        return npz_path
+    
+    def save_yolo_txt(self, image_path: str, image_size: Tuple[int, int], 
+                      class_order: List[int], class_labels: List[str]) -> Optional[str]:
+        """Save segments as YOLO format TXT file."""
+        final_mask_tensor = self.segment_manager.create_final_mask_tensor(image_size, class_order)
+        output_path = os.path.splitext(image_path)[0] + ".txt"
+        h, w = image_size
+        
+        yolo_annotations = []
+        for channel in range(final_mask_tensor.shape[2]):
+            single_channel_image = final_mask_tensor[:, :, channel]
+            if not np.any(single_channel_image):
+                continue
+            
+            contours, _ = cv2.findContours(
+                single_channel_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+            )
+            
+            class_label = class_labels[channel]
+            for contour in contours:
+                x, y, width, height = cv2.boundingRect(contour)
+                center_x = (x + width / 2) / w
+                center_y = (y + height / 2) / h
+                normalized_width = width / w
+                normalized_height = height / h
+                yolo_entry = f"{class_label} {center_x} {center_y} {normalized_width} {normalized_height}"
+                yolo_annotations.append(yolo_entry)
+        
+        if not yolo_annotations:
+            return None
+        
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, "w") as file:
+            for annotation in yolo_annotations:
+                file.write(annotation + "\n")
+        
+        return output_path
+    
+    def save_class_aliases(self, image_path: str) -> str:
+        """Save class aliases as JSON file."""
+        aliases_path = os.path.splitext(image_path)[0] + ".json"
+        aliases_to_save = {str(k): v for k, v in self.segment_manager.class_aliases.items()}
+        with open(aliases_path, "w") as f:
+            json.dump(aliases_to_save, f, indent=4)
+        return aliases_path
+    
+    def load_class_aliases(self, image_path: str) -> None:
+        """Load class aliases from JSON file."""
+        json_path = os.path.splitext(image_path)[0] + ".json"
+        if os.path.exists(json_path):
+            try:
+                with open(json_path, "r") as f:
+                    loaded_aliases = json.load(f)
+                    self.segment_manager.class_aliases = {int(k): v for k, v in loaded_aliases.items()}
+            except (json.JSONDecodeError, ValueError) as e:
+                print(f"Error loading class aliases from {json_path}: {e}")
+                self.segment_manager.class_aliases.clear()
+    
+    def load_existing_mask(self, image_path: str) -> None:
+        """Load existing mask from NPZ file."""
+        npz_path = os.path.splitext(image_path)[0] + ".npz"
+        if os.path.exists(npz_path):
+            with np.load(npz_path) as data:
+                if "mask" in data:
+                    mask_data = data["mask"]
+                    if mask_data.ndim == 2:
+                        mask_data = np.expand_dims(mask_data, axis=-1)
+                    
+                    num_classes = mask_data.shape[2]
+                    for i in range(num_classes):
+                        class_mask = mask_data[:, :, i].astype(bool)
+                        if np.any(class_mask):
+                            self.segment_manager.add_segment({
+                                "mask": class_mask,
+                                "type": "Loaded",
+                                "vertices": None,
+                                "class_id": i,
+                            })
+    
+    def is_image_file(self, filepath: str) -> bool:
+        """Check if file is a supported image format."""
+        return filepath.lower().endswith((".png", ".jpg", ".jpeg", ".tiff", ".tif"))

lazylabel/core/model_manager.py

@@ -0,0 +1,94 @@
+"""Model management functionality."""
+
+import os
+import glob
+from typing import List, Tuple, Optional, Callable
+from pathlib import Path
+
+from ..models.sam_model import SamModel
+from ..config import Paths
+
+
+class ModelManager:
+    """Manages SAM model loading and selection."""
+    
+    def __init__(self, paths: Paths):
+        self.paths = paths
+        self.sam_model: Optional[SamModel] = None
+        self.current_models_folder: Optional[str] = None
+        self.on_model_changed: Optional[Callable[[str], None]] = None
+        
+    def initialize_default_model(self, model_type: str = "vit_h") -> Optional[SamModel]:
+        """Initialize the default SAM model.
+        
+        Returns:
+            SamModel instance if successful, None if failed
+        """
+        try:
+            self.sam_model = SamModel(model_type=model_type)
+            self.current_models_folder = str(self.paths.models_dir)
+            return self.sam_model
+        except Exception as e:
+            print(f"Failed to initialize default model: {e}")
+            self.sam_model = None
+            return None
+    
+    def get_available_models(self, folder_path: str) -> List[Tuple[str, str]]:
+        """Get list of available .pth models in folder.
+        
+        Returns:
+            List of (display_name, full_path) tuples
+        """
+        pth_files = []
+        for root, dirs, files in os.walk(folder_path):
+            for file in files:
+                if file.lower().endswith('.pth'):
+                    full_path = os.path.join(root, file)
+                    rel_path = os.path.relpath(full_path, folder_path)
+                    pth_files.append((rel_path, full_path))
+        
+        return sorted(pth_files, key=lambda x: x[0])
+    
+    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
+    
+    def load_custom_model(self, model_path: str) -> bool:
+        """Load a custom model from path.
+        
+        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}")
+            
+        return success
+    
+    def set_models_folder(self, folder_path: str) -> None:
+        """Set the current models folder."""
+        self.current_models_folder = folder_path
+    
+    def get_models_folder(self) -> Optional[str]:
+        """Get the current models folder."""
+        return self.current_models_folder
+    
+    def is_model_available(self) -> bool:
+        """Check if a SAM model is loaded and available."""
+        return self.sam_model is not None and getattr(self.sam_model, 'is_loaded', False)

lazylabel/core/segment_manager.py

@@ -0,0 +1,140 @@
+"""Segment management functionality."""
+
+import numpy as np
+import cv2
+from typing import List, Dict, Any, Optional, Tuple
+from PyQt6.QtCore import QPointF
+
+
+class SegmentManager:
+    """Manages image segments and classes."""
+    
+    def __init__(self):
+        self.segments: List[Dict[str, Any]] = []
+        self.class_aliases: Dict[int, str] = {}
+        self.next_class_id: int = 0
+        
+    def clear(self) -> None:
+        """Clear all segments and reset state."""
+        self.segments.clear()
+        self.class_aliases.clear()
+        self.next_class_id = 0
+    
+    def add_segment(self, segment_data: Dict[str, Any]) -> None:
+        """Add a new segment."""
+        if 'class_id' not in segment_data:
+            segment_data['class_id'] = self.next_class_id
+        self.segments.append(segment_data)
+        self._update_next_class_id()
+    
+    def delete_segments(self, indices: List[int]) -> None:
+        """Delete segments by indices."""
+        for i in sorted(indices, reverse=True):
+            if 0 <= i < len(self.segments):
+                del self.segments[i]
+        self._update_next_class_id()
+    
+    def assign_segments_to_class(self, indices: List[int]) -> None:
+        """Assign selected segments to a class."""
+        if not indices:
+            return
+            
+        existing_class_ids = [
+            self.segments[i]["class_id"]
+            for i in indices
+            if i < len(self.segments) and self.segments[i].get("class_id") is not None
+        ]
+        
+        if existing_class_ids:
+            target_class_id = min(existing_class_ids)
+        else:
+            target_class_id = self.next_class_id
+        
+        for i in indices:
+            if i < len(self.segments):
+                self.segments[i]["class_id"] = target_class_id
+        
+        self._update_next_class_id()
+    
+    def get_unique_class_ids(self) -> List[int]:
+        """Get sorted list of unique class IDs."""
+        return sorted(list({
+            seg.get("class_id")
+            for seg in self.segments
+            if seg.get("class_id") is not None
+        }))
+    
+    def rasterize_polygon(self, vertices: List[QPointF], image_size: Tuple[int, int]) -> Optional[np.ndarray]:
+        """Convert polygon vertices to binary mask."""
+        if not vertices:
+            return None
+            
+        h, w = image_size
+        points_np = np.array([[p.x(), p.y()] for p in vertices], dtype=np.int32)
+        mask = np.zeros((h, w), dtype=np.uint8)
+        cv2.fillPoly(mask, [points_np], 1)
+        return mask.astype(bool)
+    
+    def create_final_mask_tensor(self, image_size: Tuple[int, int], class_order: List[int]) -> np.ndarray:
+        """Create final mask tensor for saving."""
+        h, w = image_size
+        id_map = {old_id: new_id for new_id, old_id in enumerate(class_order)}
+        num_final_classes = len(class_order)
+        final_mask_tensor = np.zeros((h, w, num_final_classes), dtype=np.uint8)
+        
+        for seg in self.segments:
+            class_id = seg.get("class_id")
+            if class_id not in id_map:
+                continue
+                
+            new_channel_idx = id_map[class_id]
+            
+            if seg["type"] == "Polygon":
+                mask = self.rasterize_polygon(seg["vertices"], image_size)
+            else:
+                mask = seg.get("mask")
+                
+            if mask is not None:
+                final_mask_tensor[:, :, new_channel_idx] = np.logical_or(
+                    final_mask_tensor[:, :, new_channel_idx], mask
+                )
+        
+        return final_mask_tensor
+    
+    def reassign_class_ids(self, new_order: List[int]) -> None:
+        """Reassign class IDs based on new order."""
+        id_map = {old_id: new_id for new_id, old_id in enumerate(new_order)}
+        
+        for seg in self.segments:
+            old_id = seg.get("class_id")
+            if old_id in id_map:
+                seg["class_id"] = id_map[old_id]
+        
+        # Update aliases
+        new_aliases = {
+            id_map[old_id]: self.class_aliases.get(old_id, str(old_id))
+            for old_id in new_order
+            if old_id in self.class_aliases
+        }
+        self.class_aliases = new_aliases
+        self._update_next_class_id()
+    
+    def set_class_alias(self, class_id: int, alias: str) -> None:
+        """Set alias for a class."""
+        self.class_aliases[class_id] = alias
+    
+    def get_class_alias(self, class_id: int) -> str:
+        """Get alias for a class."""
+        return self.class_aliases.get(class_id, str(class_id))
+    
+    def _update_next_class_id(self) -> None:
+        """Update the next available class ID."""
+        all_ids = {
+            seg.get("class_id")
+            for seg in self.segments
+            if seg.get("class_id") is not None
+        }
+        if not all_ids:
+            self.next_class_id = 0
+        else:
+            self.next_class_id = max(all_ids) + 1