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

lazylabel/core/file_manager.py

@@ -1,106 +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."""
+"""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

@@ -1,94 +1,97 @@
-"""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)
+"""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:
+            print(f"[8/20] Loading {model_type.upper()} model...")
+            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"[8/20] 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
+        )