argus-cv 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

argus/__init__.py

@@ -1,3 +1,3 @@
 """Argus - Vision AI dataset toolkit."""
 
-__version__ = "1.2.0"
+__version__ = "1.3.0"

argus/cli.py

@@ -11,8 +11,8 @@ from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.table import Table
 
-from argus.core import COCODataset, Dataset, YOLODataset
-from argus.core.base import TaskType
+from argus.core import COCODataset, Dataset, MaskDataset, YOLODataset
+from argus.core.base import DatasetFormat, TaskType
 from argus.core.split import (
     is_coco_unsplit,
     parse_ratio,
@@ -128,12 +128,19 @@ def stats(
     dataset = _detect_dataset(dataset_path)
     if not dataset:
         console.print(
-            f"[red]Error: No YOLO or COCO dataset found at {dataset_path}[/red]\n"
+            f"[red]Error: No dataset found at {dataset_path}[/red]\n"
             "[yellow]Ensure the path points to a dataset root containing "
-            "data.yaml (YOLO) or annotations/ folder (COCO).[/yellow]"
+            "data.yaml (YOLO), annotations/ folder (COCO), or "
+            "images/ + masks/ directories (Mask).[/yellow]"
         )
         raise typer.Exit(1)
 
+    # Handle mask datasets with pixel statistics
+    if dataset.format == DatasetFormat.MASK:
+        assert isinstance(dataset, MaskDataset)
+        _show_mask_stats(dataset, dataset_path)
+        return
+
     # Get instance counts with progress indicator
     with Progress(
         SpinnerColumn(),
@@ -212,10 +219,100 @@ def stats(
             else:
                 image_parts.append(f"{split}: {img_total}")
 
-    console.print(f"\n[green]Dataset: {dataset.format.value.upper()} | "
-                  f"Task: {dataset.task.value} | "
-                  f"Classes: {len(sorted_classes)} | "
-                  f"Total instances: {grand_total}[/green]")
+    console.print(
+        f"\n[green]Dataset: {dataset.format.value.upper()} | "
+        f"Task: {dataset.task.value} | "
+        f"Classes: {len(sorted_classes)} | "
+        f"Total instances: {grand_total}[/green]"
+    )
+
+    if image_parts:
+        console.print(f"[blue]Images: {' | '.join(image_parts)}[/blue]")
+
+
+def _show_mask_stats(dataset: MaskDataset, dataset_path: Path) -> None:
+    """Show statistics for mask datasets with pixel-level information.
+
+    Args:
+        dataset: The MaskDataset instance.
+        dataset_path: Path to the dataset root.
+    """
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console,
+        transient=True,
+    ) as progress:
+        progress.add_task("Analyzing mask dataset...", total=None)
+        pixel_counts = dataset.get_pixel_counts()
+        image_presence = dataset.get_image_class_presence()
+        image_counts = dataset.get_image_counts()
+
+    # Get class mapping
+    class_mapping = dataset.get_class_mapping()
+
+    # Calculate total non-ignored pixels
+    total_pixels = sum(
+        count
+        for class_id, count in pixel_counts.items()
+        if class_id != dataset.ignore_index
+    )
+    ignored_pixels = pixel_counts.get(dataset.ignore_index, 0)
+
+    # Calculate total images
+    total_images = sum(ic["total"] for ic in image_counts.values())
+
+    # Create table
+    splits_str = ", ".join(dataset.splits) if dataset.splits else "unsplit"
+    title = f"Class Statistics: {dataset_path.name} ({splits_str})"
+    table = Table(title=title)
+    table.add_column("Class", style="cyan")
+    table.add_column("Total Pixels", justify="right", style="green")
+    table.add_column("% Coverage", justify="right", style="magenta")
+    table.add_column("Images With", justify="right", style="yellow")
+
+    # Sort classes by class_id
+    sorted_class_ids = sorted(class_mapping.keys())
+
+    for class_id in sorted_class_ids:
+        class_name = class_mapping[class_id]
+        pixels = pixel_counts.get(class_id, 0)
+        presence = image_presence.get(class_id, 0)
+
+        # Calculate coverage percentage
+        coverage = (pixels / total_pixels * 100) if total_pixels > 0 else 0.0
+
+        table.add_row(
+            class_name,
+            f"{pixels:,}",
+            f"{coverage:.1f}%",
+            str(presence),
+        )
+
+    # Add ignored row if there are ignored pixels
+    if ignored_pixels > 0:
+        table.add_section()
+        table.add_row(
+            "[dim](ignored)[/dim]",
+            f"[dim]{ignored_pixels:,}[/dim]",
+            "[dim]-[/dim]",
+            f"[dim]{total_images}[/dim]",
+        )
+
+    console.print(table)
+
+    # Summary line
+    console.print(f"\n[green]Dataset: {dataset_path}[/green]")
+    console.print(
+        f"[green]Format: {dataset.format.value.upper()} | "
+        f"Task: {dataset.task.value}[/green]"
+    )
+
+    # Image counts per split
+    image_parts = []
+    for split in dataset.splits if dataset.splits else ["unsplit"]:
+        if split in image_counts:
+            image_parts.append(f"{split}: {image_counts[split]['total']}")
 
     if image_parts:
         console.print(f"[blue]Images: {' | '.join(image_parts)}[/blue]")
@@ -247,6 +344,16 @@ def view(
             help="Maximum classes to show in grid (classification only).",
         ),
     ] = None,
+    opacity: Annotated[
+        float,
+        typer.Option(
+            "--opacity",
+            "-o",
+            help="Mask overlay opacity (0.0-1.0, mask datasets only).",
+            min=0.0,
+            max=1.0,
+        ),
+    ] = 0.5,
 ) -> None:
     """View annotated images in a dataset.
 
@@ -295,6 +402,41 @@ def view(
     # Generate consistent colors for each class
     class_colors = _generate_class_colors(dataset.class_names)
 
+    # Handle mask datasets with overlay viewer
+    if dataset.format == DatasetFormat.MASK:
+        assert isinstance(dataset, MaskDataset)
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            console=console,
+            transient=True,
+        ) as progress:
+            progress.add_task("Loading images...", total=None)
+            image_paths = dataset.get_image_paths(split)
+
+        if not image_paths:
+            console.print("[yellow]No images found in the dataset.[/yellow]")
+            return
+
+        console.print(
+            f"[green]Found {len(image_paths)} images. "
+            f"Opening mask viewer...[/green]\n"
+            "[dim]Controls: \u2190 / \u2192 or P / N to navigate, "
+            "Mouse wheel to zoom, Drag to pan, R to reset, T to toggle overlay, "
+            "Q / ESC to quit[/dim]"
+        )
+
+        viewer = _MaskViewer(
+            image_paths=image_paths,
+            dataset=dataset,
+            class_colors=class_colors,
+            window_name=f"Argus Mask Viewer - {dataset_path.name}",
+            opacity=opacity,
+        )
+        viewer.run()
+        console.print("[green]Viewer closed.[/green]")
+        return
+
     # Handle classification datasets with grid viewer
     if dataset.task == TaskType.CLASSIFICATION:
         # Use first split if specified, otherwise let get_images_by_class handle it
@@ -599,12 +741,16 @@ class _ImageViewer:
             info_text += " [Annotations: OFF]"
 
         cv2.putText(
-            display, info_text, (10, 30),
-            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2
+            display,
+            info_text,
+            (10, 30),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.7,
+            (255, 255, 255),
+            2,
         )
         cv2.putText(
-            display, info_text, (10, 30),
-            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1
+            display, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1
         )
 
         return display
@@ -733,9 +879,11 @@ class _ClassificationGridViewer:
         self.current_index = 0
 
         # Calculate max images across all classes
-        self.max_images = max(
-            len(imgs) for imgs in self.images_by_class.values()
-        ) if self.images_by_class else 0
+        self.max_images = (
+            max(len(imgs) for imgs in self.images_by_class.values())
+            if self.images_by_class
+            else 0
+        )
 
         # Calculate grid layout
         self.cols, self.rows = self._calculate_grid_layout()
@@ -883,6 +1031,267 @@ class _ClassificationGridViewer:
         cv2.destroyAllWindows()
 
 
+class _MaskViewer:
+    """Interactive viewer for semantic mask datasets with colored overlay."""
+
+    def __init__(
+        self,
+        image_paths: list[Path],
+        dataset: MaskDataset,
+        class_colors: dict[str, tuple[int, int, int]],
+        window_name: str,
+        opacity: float = 0.5,
+    ):
+        self.image_paths = image_paths
+        self.dataset = dataset
+        self.class_colors = class_colors
+        self.window_name = window_name
+        self.opacity = opacity
+
+        self.current_idx = 0
+        self.zoom = 1.0
+        self.pan_x = 0.0
+        self.pan_y = 0.0
+
+        # Mouse state for panning
+        self.dragging = False
+        self.drag_start_x = 0
+        self.drag_start_y = 0
+        self.pan_start_x = 0.0
+        self.pan_start_y = 0.0
+
+        # Current image cache
+        self.current_img: np.ndarray | None = None
+        self.overlay_img: np.ndarray | None = None
+
+        # Overlay visibility toggle
+        self.show_overlay = True
+
+        # Build class_id to color mapping
+        self._id_to_color: dict[int, tuple[int, int, int]] = {}
+        class_mapping = dataset.get_class_mapping()
+        for class_id, class_name in class_mapping.items():
+            if class_name in class_colors:
+                self._id_to_color[class_id] = class_colors[class_name]
+
+    def _load_current_image(self) -> bool:
+        """Load current image and create mask overlay."""
+        image_path = self.image_paths[self.current_idx]
+
+        img = cv2.imread(str(image_path))
+        if img is None:
+            return False
+
+        mask = self.dataset.load_mask(image_path)
+        if mask is None:
+            console.print(f"[yellow]Warning: No mask for {image_path}[/yellow]")
+            self.current_img = img
+            self.overlay_img = img.copy()
+            return True
+
+        # Validate dimensions
+        if img.shape[:2] != mask.shape[:2]:
+            console.print(
+                f"[red]Error: Dimension mismatch for {image_path.name}: "
+                f"image={img.shape[:2]}, mask={mask.shape[:2]}[/red]"
+            )
+            return False
+
+        self.current_img = img
+        self.overlay_img = self._create_overlay(img, mask)
+        return True
+
+    def _create_overlay(self, img: np.ndarray, mask: np.ndarray) -> np.ndarray:
+        """Create colored overlay from mask.
+
+        Args:
+            img: Original image (BGR).
+            mask: Grayscale mask with class IDs.
+
+        Returns:
+            Image with colored mask overlay.
+        """
+        # Create colored mask
+        h, w = mask.shape
+        colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
+
+        for class_id, color in self._id_to_color.items():
+            colored_mask[mask == class_id] = color
+
+        # Blend with original image
+        # Ignore pixels are fully transparent (not blended)
+        ignore_mask = mask == self.dataset.ignore_index
+        alpha = np.ones((h, w, 1), dtype=np.float32) * self.opacity
+        alpha[ignore_mask] = 0.0
+
+        # Blend: result = img * (1 - alpha) + colored_mask * alpha
+        blended = (
+            img.astype(np.float32) * (1 - alpha)
+            + colored_mask.astype(np.float32) * alpha
+        )
+        return blended.astype(np.uint8)
+
+    def _get_display_image(self) -> np.ndarray:
+        """Get the image transformed for current zoom/pan."""
+        if self.overlay_img is None:
+            return np.zeros((480, 640, 3), dtype=np.uint8)
+
+        if self.show_overlay:
+            img = self.overlay_img
+        elif self.current_img is not None:
+            img = self.current_img
+        else:
+            img = self.overlay_img
+
+        h, w = img.shape[:2]
+
+        if self.zoom == 1.0 and self.pan_x == 0.0 and self.pan_y == 0.0:
+            display = img.copy()
+        else:
+            # Calculate the visible region
+            view_w = int(w / self.zoom)
+            view_h = int(h / self.zoom)
+
+            # Center point with pan offset
+            cx = w / 2 + self.pan_x
+            cy = h / 2 + self.pan_y
+
+            # Calculate crop bounds
+            x1 = int(max(0, cx - view_w / 2))
+            y1 = int(max(0, cy - view_h / 2))
+            x2 = int(min(w, x1 + view_w))
+            y2 = int(min(h, y1 + view_h))
+
+            # Adjust if we hit boundaries
+            if x2 - x1 < view_w:
+                x1 = max(0, x2 - view_w)
+            if y2 - y1 < view_h:
+                y1 = max(0, y2 - view_h)
+
+            # Crop and resize
+            cropped = img[y1:y2, x1:x2]
+            display = cv2.resize(cropped, (w, h), interpolation=cv2.INTER_LINEAR)
+
+        # Add info overlay
+        image_path = self.image_paths[self.current_idx]
+        idx = self.current_idx + 1
+        total = len(self.image_paths)
+        info_text = f"[{idx}/{total}] {image_path.name}"
+        if self.zoom > 1.0:
+            info_text += f" (Zoom: {self.zoom:.1f}x)"
+        if not self.show_overlay:
+            info_text += " [Overlay: OFF]"
+
+        cv2.putText(
+            display,
+            info_text,
+            (10, 30),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.7,
+            (255, 255, 255),
+            2,
+        )
+        cv2.putText(
+            display, info_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 1
+        )
+
+        return display
+
+    def _mouse_callback(
+        self, event: int, x: int, y: int, flags: int, param: None
+    ) -> None:
+        """Handle mouse events for zoom and pan."""
+        if event == cv2.EVENT_MOUSEWHEEL:
+            # Zoom in/out
+            if flags > 0:
+                self.zoom = min(10.0, self.zoom * 1.2)
+            else:
+                self.zoom = max(1.0, self.zoom / 1.2)
+
+            # Reset pan if zoomed out to 1x
+            if self.zoom == 1.0:
+                self.pan_x = 0.0
+                self.pan_y = 0.0
+
+        elif event == cv2.EVENT_LBUTTONDOWN:
+            self.dragging = True
+            self.drag_start_x = x
+            self.drag_start_y = y
+            self.pan_start_x = self.pan_x
+            self.pan_start_y = self.pan_y
+
+        elif event == cv2.EVENT_MOUSEMOVE and self.dragging:
+            if self.zoom > 1.0 and self.overlay_img is not None:
+                h, w = self.overlay_img.shape[:2]
+                # Calculate pan delta (inverted for natural feel)
+                dx = (self.drag_start_x - x) / self.zoom
+                dy = (self.drag_start_y - y) / self.zoom
+
+                # Update pan with limits
+                max_pan_x = w * (1 - 1 / self.zoom) / 2
+                max_pan_y = h * (1 - 1 / self.zoom) / 2
+
+                self.pan_x = max(-max_pan_x, min(max_pan_x, self.pan_start_x + dx))
+                self.pan_y = max(-max_pan_y, min(max_pan_y, self.pan_start_y + dy))
+
+        elif event == cv2.EVENT_LBUTTONUP:
+            self.dragging = False
+
+    def _reset_view(self) -> None:
+        """Reset zoom and pan to default."""
+        self.zoom = 1.0
+        self.pan_x = 0.0
+        self.pan_y = 0.0
+
+    def _next_image(self) -> None:
+        """Go to next image."""
+        self.current_idx = (self.current_idx + 1) % len(self.image_paths)
+        self._reset_view()
+
+    def _prev_image(self) -> None:
+        """Go to previous image."""
+        self.current_idx = (self.current_idx - 1) % len(self.image_paths)
+        self._reset_view()
+
+    def run(self) -> None:
+        """Run the interactive viewer."""
+        cv2.namedWindow(self.window_name, cv2.WINDOW_AUTOSIZE)
+        cv2.setMouseCallback(self.window_name, self._mouse_callback)
+
+        while True:
+            # Load image if needed
+            if self.overlay_img is None and not self._load_current_image():
+                console.print(
+                    f"[yellow]Warning: Could not load "
+                    f"{self.image_paths[self.current_idx]}[/yellow]"
+                )
+                self._next_image()
+                continue
+
+            # Display image
+            display = self._get_display_image()
+            cv2.imshow(self.window_name, display)
+
+            # Wait for input (short timeout for smooth panning)
+            key = cv2.waitKey(30) & 0xFF
+
+            # Handle keyboard input
+            if key == ord("q") or key == 27:  # Q or ESC
+                break
+            elif key == ord("n") or key == 83 or key == 3:  # N or Right arrow
+                self.overlay_img = None
+                self._next_image()
+            elif key == ord("p") or key == 81 or key == 2:  # P or Left arrow
+                self.overlay_img = None
+                self._prev_image()
+            elif key == ord("r"):  # R to reset zoom
+                self._reset_view()
+            elif key == ord("t"):  # T to toggle overlay
+                self.show_overlay = not self.show_overlay
+
+        cv2.destroyAllWindows()
+
+
 def _generate_class_colors(class_names: list[str]) -> dict[str, tuple[int, int, int]]:
     """Generate consistent colors for each class name.
 
@@ -965,9 +1374,13 @@ def _draw_annotations(
             )
             # Draw label text
             cv2.putText(
-                img, label,
+                img,
+                label,
                 (x1 + 2, y1 - baseline - 2),
-                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                (255, 255, 255),
+                1,
             )
 
     return img
@@ -1026,17 +1439,25 @@ def _discover_datasets(root_path: Path, max_depth: int) -> list[Dataset]:
 
 
 def _detect_dataset(path: Path) -> Dataset | None:
-    """Try to detect a dataset at the given path."""
-    # Try YOLO first (more specific patterns)
+    """Try to detect a dataset at the given path.
+
+    Detection priority: YOLO -> COCO -> MaskDataset
+    """
+    # Try YOLO first (more specific patterns - requires data.yaml)
     dataset = YOLODataset.detect(path)
     if dataset:
         return dataset
 
-    # Try COCO
+    # Try COCO (requires annotations/*.json)
     dataset = COCODataset.detect(path)
     if dataset:
         return dataset
 
+    # Try MaskDataset (directory structure based)
+    dataset = MaskDataset.detect(path)
+    if dataset:
+        return dataset
+
     return None
 
 

argus/core/__init__.py

@@ -2,6 +2,7 @@
 
 from argus.core.base import Dataset
 from argus.core.coco import COCODataset
+from argus.core.mask import ConfigurationError, MaskDataset
 from argus.core.split import split_coco_dataset, split_yolo_dataset
 from argus.core.yolo import YOLODataset
 
@@ -9,6 +10,8 @@ __all__ = [
     "Dataset",
     "YOLODataset",
     "COCODataset",
+    "MaskDataset",
+    "ConfigurationError",
     "split_coco_dataset",
     "split_yolo_dataset",
 ]

argus/core/base.py

@@ -11,6 +11,7 @@ class DatasetFormat(str, Enum):
 
     YOLO = "yolo"
     COCO = "coco"
+    MASK = "mask"
 
 
 class TaskType(str, Enum):

argus/core/coco.py

@@ -453,12 +453,14 @@ class COCODataset(Dataset):
                         for i in range(0, len(coords), 2):
                             polygon.append((float(coords[i]), float(coords[i + 1])))
 
-                    annotations.append({
-                        "class_name": class_name,
-                        "class_id": cat_id,
-                        "bbox": bbox_tuple,
-                        "polygon": polygon,
-                    })
+                    annotations.append(
+                        {
+                            "class_name": class_name,
+                            "class_id": cat_id,
+                            "bbox": bbox_tuple,
+                            "polygon": polygon,
+                        }
+                    )
 
             except (json.JSONDecodeError, OSError):
                 continue

argus/core/mask.py

@@ -0,0 +1,648 @@
+"""Mask dataset detection and handling for semantic segmentation."""
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import cv2
+import numpy as np
+import yaml
+
+from argus.core.base import Dataset, DatasetFormat, TaskType
+
+logger = logging.getLogger(__name__)
+
+
+class ConfigurationError(Exception):
+    """Error raised when dataset configuration is invalid."""
+
+    pass
+
+
+# Directory patterns for mask dataset detection (checked in order)
+DIRECTORY_PATTERNS = [
+    ("images", "masks"),
+    ("img", "gt"),
+    ("leftImg8bit", "gtFine"),
+]
+
+# Standard image extensions
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}
+
+
+@dataclass
+class MaskDataset(Dataset):
+    """Dataset format for folder-based semantic segmentation masks.
+
+    Supports directory structures like:
+        - images/ + masks/
+        - img/ + gt/
+        - leftImg8bit/ + gtFine/ (Cityscapes-style)
+
+    Each pattern expects parallel split subdirectories:
+        dataset/
+        ├── images/
+        │   ├── train/
+        │   └── val/
+        ├── masks/
+        │   ├── train/
+        │   └── val/
+        └── classes.yaml  # Optional for grayscale, required for RGB
+    """
+
+    images_dir: str = ""
+    masks_dir: str = ""
+    class_config: dict | None = None
+    ignore_index: int | None = 255
+    format: DatasetFormat = field(default=DatasetFormat.MASK, init=False)
+    task: TaskType = field(default=TaskType.SEGMENTATION, init=False)
+
+    # Internal: maps class_id -> class_name
+    _class_mapping: dict[int, str] = field(default_factory=dict, repr=False)
+    # Internal: maps RGB tuple -> (class_id, class_name) for palette masks
+    _palette_mapping: dict[tuple[int, int, int], tuple[int, str]] = field(
+        default_factory=dict, repr=False
+    )
+    # Internal: flag for RGB palette mode
+    _is_rgb_palette: bool = field(default=False, repr=False)
+
+    @classmethod
+    def detect(cls, path: Path) -> "MaskDataset | None":
+        """Detect if the given path contains a mask dataset.
+
+        Args:
+            path: Directory path to check for dataset.
+
+        Returns:
+            MaskDataset instance if detected, None otherwise.
+        """
+        path = Path(path)
+
+        if not path.is_dir():
+            return None
+
+        # Try each directory pattern
+        for images_name, masks_name in DIRECTORY_PATTERNS:
+            images_root = path / images_name
+            masks_root = path / masks_name
+
+            if not (images_root.is_dir() and masks_root.is_dir()):
+                continue
+
+            # Check for split subdirectories with matching structure
+            splits = cls._detect_splits(images_root, masks_root)
+
+            # If no splits found, check if images and masks are directly in root
+            if not splits:
+                has_images = any(
+                    f.suffix.lower() in IMAGE_EXTENSIONS for f in images_root.iterdir()
+                )
+                has_masks = any(
+                    f.suffix.lower() == ".png" for f in masks_root.iterdir()
+                )
+
+                if has_images and has_masks:
+                    # Valid unsplit structure
+                    splits = []
+                else:
+                    continue
+
+            # Load class configuration if available
+            class_config = cls._load_class_config(path)
+
+            # Determine if masks are grayscale or RGB palette
+            is_rgb, palette_mapping = cls._detect_mask_type(path, masks_root, splits)
+
+            # RGB masks require configuration
+            if is_rgb and not class_config:
+                raise ConfigurationError(
+                    f"RGB palette masks detected in {path} but no classes.yaml found. "
+                    "RGB masks require a classes.yaml config file with palette mapping."
+                )
+
+            # Build class mapping
+            class_mapping, ignore_idx = cls._build_class_mapping(
+                path, masks_root, splits, class_config, is_rgb, palette_mapping
+            )
+
+            # Extract class names from mapping
+            class_names = [
+                class_mapping[i]
+                for i in sorted(class_mapping.keys())
+                if i != ignore_idx
+            ]
+
+            dataset = cls(
+                path=path,
+                num_classes=len(class_names),
+                class_names=class_names,
+                splits=splits,
+                images_dir=images_name,
+                masks_dir=masks_name,
+                class_config=class_config,
+                ignore_index=ignore_idx,
+            )
+            dataset._class_mapping = class_mapping
+            dataset._palette_mapping = palette_mapping if is_rgb else {}
+            dataset._is_rgb_palette = is_rgb
+
+            return dataset
+
+        return None
+
+    @classmethod
+    def _detect_splits(cls, images_root: Path, masks_root: Path) -> list[str]:
+        """Detect available splits from directory structure.
+
+        Args:
+            images_root: Root directory containing images.
+            masks_root: Root directory containing masks.
+
+        Returns:
+            List of split names found in both images and masks directories.
+        """
+        splits = []
+
+        for split_name in ["train", "val", "test"]:
+            images_split = images_root / split_name
+            masks_split = masks_root / split_name
+
+            if images_split.is_dir() and masks_split.is_dir():
+                # Verify there are actual files
+                has_images = any(
+                    f.suffix.lower() in IMAGE_EXTENSIONS for f in images_split.iterdir()
+                )
+                has_masks = any(
+                    f.suffix.lower() == ".png" for f in masks_split.iterdir()
+                )
+
+                if has_images and has_masks:
+                    splits.append(split_name)
+
+        return splits
+
+    @classmethod
+    def _load_class_config(cls, path: Path) -> dict | None:
+        """Load classes.yaml configuration if present.
+
+        Args:
+            path: Dataset root path.
+
+        Returns:
+            Parsed config dict or None if not found.
+        """
+        config_path = path / "classes.yaml"
+        if not config_path.exists():
+            config_path = path / "classes.yml"
+
+        if not config_path.exists():
+            return None
+
+        try:
+            with open(config_path, encoding="utf-8") as f:
+                return yaml.safe_load(f)
+        except (yaml.YAMLError, OSError) as e:
+            raise ConfigurationError(f"Failed to parse {config_path}: {e}") from e
+
+    @classmethod
+    def _detect_mask_type(
+        cls,
+        path: Path,
+        masks_root: Path,
+        splits: list[str],
+    ) -> tuple[bool, dict[tuple[int, int, int], tuple[int, str]]]:
+        """Detect if masks are grayscale or RGB palette format.
+
+        Args:
+            path: Dataset root path.
+            masks_root: Masks directory root.
+            splits: List of split names.
+
+        Returns:
+            Tuple of (is_rgb_palette, palette_mapping).
+        """
+        # Find sample mask files
+        sample_masks: list[Path] = []
+
+        if splits:
+            for split in splits[:1]:  # Just check first split
+                split_dir = masks_root / split
+                sample_masks.extend(list(split_dir.glob("*.png"))[:5])
+        else:
+            sample_masks.extend(list(masks_root.glob("*.png"))[:5])
+
+        if not sample_masks:
+            return False, {}
+
+        # Check first mask
+        mask = cv2.imread(str(sample_masks[0]), cv2.IMREAD_UNCHANGED)
+        if mask is None:
+            return False, {}
+
+        # Check if grayscale (single channel) or RGB (3 channels)
+        if len(mask.shape) == 2 or mask.shape[2] == 1:
+            return False, {}
+        elif mask.shape[2] >= 3:
+            return True, {}
+
+        return False, {}
+
+    @classmethod
+    def _build_class_mapping(
+        cls,
+        path: Path,
+        masks_root: Path,
+        splits: list[str],
+        class_config: dict | None,
+        is_rgb: bool,
+        palette_mapping: dict,
+    ) -> tuple[dict[int, str], int]:
+        """Build class ID to name mapping.
+
+        Args:
+            path: Dataset root path.
+            masks_root: Masks directory root.
+            splits: List of split names.
+            class_config: Parsed classes.yaml or None.
+            is_rgb: Whether masks use RGB palette encoding.
+            palette_mapping: RGB to class mapping (if RGB).
+
+        Returns:
+            Tuple of (class_id_to_name dict, ignore_index or None).
+        """
+        ignore_index: int | None = 255
+
+        if class_config:
+            # Get ignore index from config (can be null/None to disable)
+            ignore_index = class_config.get("ignore_index", 255)
+
+            # Try to get names from config
+            if "names" in class_config:
+                names = class_config["names"]
+                if isinstance(names, dict):
+                    return {int(k): v for k, v in names.items()}, ignore_index
+                elif isinstance(names, list):
+                    return {i: name for i, name in enumerate(names)}, ignore_index
+
+            # If RGB palette, build from palette config
+            if is_rgb and "palette" in class_config:
+                mapping = {}
+                for entry in class_config["palette"]:
+                    class_id = entry["id"]
+                    class_name = entry["name"]
+                    mapping[class_id] = class_name
+                return mapping, ignore_index
+
+        # Auto-detect classes from grayscale masks
+        if not is_rgb:
+            return cls._auto_detect_classes(masks_root, splits, ignore_index)
+
+        # RGB without config - error should have been raised earlier
+        return {}, ignore_index
+
+    @classmethod
+    def _auto_detect_classes(
+        cls,
+        masks_root: Path,
+        splits: list[str],
+        ignore_index: int | None,
+    ) -> tuple[dict[int, str], int | None]:
+        """Auto-detect class IDs from grayscale mask values.
+
+        Args:
+            masks_root: Masks directory root.
+            splits: List of split names.
+            ignore_index: Index to treat as ignored/void, or None.
+
+        Returns:
+            Tuple of (class_id_to_name dict, ignore_index or None).
+        """
+        unique_values: set[int] = set()
+
+        # Sample masks to find unique values
+        sample_masks: list[Path] = []
+
+        if splits:
+            for split in splits:
+                split_dir = masks_root / split
+                sample_masks.extend(list(split_dir.glob("*.png"))[:20])
+        else:
+            sample_masks.extend(list(masks_root.glob("*.png"))[:20])
+
+        for mask_path in sample_masks[:50]:  # Limit sampling
+            mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+            if mask is not None:
+                unique_values.update(np.unique(mask).tolist())
+
+        # Build mapping with auto-generated names
+        mapping = {}
+        for val in sorted(unique_values):
+            if val == ignore_index:
+                continue
+            mapping[val] = f"class_{val}"
+
+        return mapping, ignore_index
+
+    def get_image_paths(self, split: str | None = None) -> list[Path]:
+        """Get all image file paths for a split or the entire dataset.
+
+        Only returns images that have corresponding masks.
+
+        Args:
+            split: Specific split to get images from. If None, returns all images.
+
+        Returns:
+            List of image file paths sorted alphabetically.
+        """
+        images_root = self.path / self.images_dir
+        image_paths: list[Path] = []
+
+        splits_to_search = (
+            [split] if split else (self.splits if self.splits else [None])
+        )
+
+        for s in splits_to_search:
+            image_dir = images_root / s if s else images_root
+
+            if not image_dir.is_dir():
+                continue
+
+            for img_file in image_dir.iterdir():
+                if img_file.suffix.lower() not in IMAGE_EXTENSIONS:
+                    continue
+
+                # Only include if mask exists
+                mask_path = self.get_mask_path(img_file)
+                if mask_path and mask_path.exists():
+                    image_paths.append(img_file)
+                else:
+                    logger.warning(f"No mask found for image: {img_file}")
+
+        return sorted(image_paths, key=lambda p: p.name)
+
+    def get_mask_path(self, image_path: Path) -> Path | None:
+        """Return corresponding mask path for an image.
+
+        Tries multiple naming conventions in order:
+        1. Exact stem match: image.jpg -> image.png
+        2. With _mask suffix: image.jpg -> image_mask.png
+        3. With _gt suffix: image.jpg -> image_gt.png
+
+        Args:
+            image_path: Path to the image file.
+
+        Returns:
+            Path to corresponding mask, or None if not found.
+        """
+        stem = image_path.stem
+
+        # Determine the split from image path
+        image_parts = image_path.parts
+        images_dir_idx = None
+        for i, part in enumerate(image_parts):
+            if part == self.images_dir:
+                images_dir_idx = i
+                break
+
+        if images_dir_idx is None:
+            # Fallback: look in masks root
+            masks_dir = self.path / self.masks_dir
+            return self._find_mask_with_patterns(masks_dir, stem)
+
+        # Build mask directory path with same structure
+        mask_parts = list(image_parts[:-1])  # Exclude filename
+        mask_parts[images_dir_idx] = self.masks_dir
+        masks_dir = Path(*mask_parts)
+
+        return self._find_mask_with_patterns(masks_dir, stem)
+
+    def _find_mask_with_patterns(self, masks_dir: Path, stem: str) -> Path | None:
+        """Find mask file trying multiple naming patterns.
+
+        Args:
+            masks_dir: Directory containing masks.
+            stem: Image filename stem (without extension).
+
+        Returns:
+            Path to mask if found, None otherwise.
+        """
+        # Try different naming patterns in order of preference
+        patterns = [
+            f"{stem}.png",  # Exact match
+            f"{stem}_mask.png",  # Common _mask suffix
+            f"{stem}_gt.png",  # Ground truth suffix
+            f"{stem}_label.png",  # Label suffix
+        ]
+
+        for pattern in patterns:
+            mask_path = masks_dir / pattern
+            if mask_path.exists():
+                return mask_path
+
+        return None
+
+    def get_class_mapping(self) -> dict[int, str]:
+        """Return class ID to name mapping.
+
+        Returns:
+            Dictionary mapping class IDs to class names.
+        """
+        return self._class_mapping.copy()
+
+    def get_instance_counts(self) -> dict[str, dict[str, int]]:
+        """Get pixel counts per class, per split.
+
+        For mask datasets, this returns the total number of pixels
+        for each class across all masks in each split.
+
+        Returns:
+            Dictionary mapping split name to dict of class name to pixel count.
+        """
+        counts: dict[str, dict[str, int]] = {}
+
+        splits_to_process = self.splits if self.splits else ["unsplit"]
+        masks_root = self.path / self.masks_dir
+
+        for split in splits_to_process:
+            split_counts: dict[str, int] = {name: 0 for name in self.class_names}
+
+            mask_dir = masks_root if split == "unsplit" else masks_root / split
+
+            if not mask_dir.is_dir():
+                continue
+
+            for mask_path in mask_dir.glob("*.png"):
+                mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+                if mask is None:
+                    logger.warning(f"Could not read mask: {mask_path}")
+                    continue
+
+                # Count pixels for each class
+                unique, pixel_counts = np.unique(mask, return_counts=True)
+
+                for class_id, count in zip(unique, pixel_counts, strict=True):
+                    if class_id == self.ignore_index:
+                        continue
+                    class_name = self._class_mapping.get(class_id)
+                    if class_name and class_name in split_counts:
+                        split_counts[class_name] += int(count)
+
+            counts[split] = split_counts
+
+        return counts
+
+    def get_image_counts(self) -> dict[str, dict[str, int]]:
+        """Get image counts per split.
+
+        Returns:
+            Dictionary mapping split name to dict with "total" and "background" counts.
+            For mask datasets, "background" is always 0.
+        """
+        counts: dict[str, dict[str, int]] = {}
+
+        splits_to_process = self.splits if self.splits else ["unsplit"]
+
+        for split in splits_to_process:
+            image_paths = self.get_image_paths(split if split != "unsplit" else None)
+            counts[split] = {"total": len(image_paths), "background": 0}
+
+        return counts
+
+    def get_image_class_presence(self, split: str | None = None) -> dict[int, int]:
+        """Return count of images containing each class.
+
+        Args:
+            split: Specific split to analyze. If None, analyzes all splits.
+
+        Returns:
+            Dictionary mapping class ID to count of images containing that class.
+        """
+        presence: dict[int, int] = {class_id: 0 for class_id in self._class_mapping}
+
+        splits_to_process = (
+            [split] if split else (self.splits if self.splits else [None])
+        )
+        masks_root = self.path / self.masks_dir
+
+        for s in splits_to_process:
+            mask_dir = masks_root / s if s else masks_root
+
+            if not mask_dir.is_dir():
+                continue
+
+            for mask_path in mask_dir.glob("*.png"):
+                mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+                if mask is None:
+                    continue
+
+                # Find which classes are present in this mask
+                unique_values = np.unique(mask)
+                for class_id in unique_values:
+                    if class_id != self.ignore_index and class_id in presence:
+                        presence[class_id] += 1
+
+        return presence
+
+    def get_pixel_counts(self, split: str | None = None) -> dict[int, int]:
+        """Return total pixel count per class ID.
+
+        Args:
+            split: Specific split to analyze. If None, analyzes all splits.
+
+        Returns:
+            Dictionary mapping class ID to total pixel count.
+        """
+        pixel_counts: dict[int | None, int] = {
+            class_id: 0 for class_id in self._class_mapping
+        }
+        # Track ignored pixels if ignore_index is set
+        if self.ignore_index is not None:
+            pixel_counts[self.ignore_index] = 0
+
+        splits_to_process = (
+            [split] if split else (self.splits if self.splits else [None])
+        )
+        masks_root = self.path / self.masks_dir
+
+        for s in splits_to_process:
+            mask_dir = masks_root / s if s else masks_root
+
+            if not mask_dir.is_dir():
+                continue
+
+            for mask_path in mask_dir.glob("*.png"):
+                mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+                if mask is None:
+                    continue
+
+                unique, counts = np.unique(mask, return_counts=True)
+                for class_id, count in zip(unique, counts, strict=True):
+                    class_id_int = int(class_id)
+                    if class_id_int in pixel_counts:
+                        pixel_counts[class_id_int] += int(count)
+                    elif (
+                        self.ignore_index is not None
+                        and class_id_int == self.ignore_index
+                    ):
+                        pixel_counts[self.ignore_index] += int(count)
+
+        return pixel_counts
+
+    def get_annotations_for_image(self, image_path: Path) -> list[dict]:
+        """Get annotations for a specific image.
+
+        For mask datasets, this returns an empty list since annotations
+        are stored as pixel masks rather than discrete objects.
+        Use get_mask_path() to get the mask file directly.
+
+        Args:
+            image_path: Path to the image file.
+
+        Returns:
+            Empty list (masks don't have discrete annotations).
+        """
+        # Mask datasets don't have discrete annotations like detection/segmentation
+        # The mask itself IS the annotation
+        return []
+
+    def load_mask(self, image_path: Path) -> np.ndarray | None:
+        """Load the mask for a given image.
+
+        Args:
+            image_path: Path to the image file.
+
+        Returns:
+            Mask as numpy array (grayscale), or None if not found.
+        """
+        mask_path = self.get_mask_path(image_path)
+        if mask_path is None or not mask_path.exists():
+            return None
+
+        mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+        return mask
+
+    def validate_dimensions(
+        self, image_path: Path
+    ) -> tuple[bool, tuple[int, int] | None, tuple[int, int] | None]:
+        """Check if image and mask dimensions match.
+
+        Args:
+            image_path: Path to the image file.
+
+        Returns:
+            Tuple of (dimensions_match, image_shape, mask_shape).
+        """
+        mask_path = self.get_mask_path(image_path)
+        if mask_path is None or not mask_path.exists():
+            return False, None, None
+
+        img = cv2.imread(str(image_path))
+        mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
+
+        if img is None or mask is None:
+            return False, None, None
+
+        img_shape = (img.shape[0], img.shape[1])
+        mask_shape = (mask.shape[0], mask.shape[1])
+
+        return img_shape == mask_shape, img_shape, mask_shape

argus/core/yolo.py

@@ -89,6 +89,11 @@ class YOLODataset(Dataset):
                 if "names" not in config:
                     continue
 
+                # Skip if this looks like a mask dataset config
+                # (has ignore_index or palette keys which are mask-specific)
+                if "ignore_index" in config or "palette" in config:
+                    continue
+
                 names = config["names"]
 
                 # Extract class names
@@ -692,12 +697,14 @@ class YOLODataset(Dataset):
                             x = x_center - width / 2
                             y = y_center - height / 2
 
-                            annotations.append({
-                                "class_name": class_name,
-                                "class_id": class_id,
-                                "bbox": (x, y, width, height),
-                                "polygon": None,
-                            })
+                            annotations.append(
+                                {
+                                    "class_name": class_name,
+                                    "class_id": class_id,
+                                    "bbox": (x, y, width, height),
+                                    "polygon": None,
+                                }
+                            )
                         else:
                             # Segmentation: class x1 y1 x2 y2 ... xn yn
                             coords = [float(p) for p in parts[1:]]
@@ -715,12 +722,14 @@ class YOLODataset(Dataset):
                             width = max(xs) - x
                             height = max(ys) - y
 
-                            annotations.append({
-                                "class_name": class_name,
-                                "class_id": class_id,
-                                "bbox": (x, y, width, height),
-                                "polygon": polygon,
-                            })
+                            annotations.append(
+                                {
+                                    "class_name": class_name,
+                                    "class_id": class_id,
+                                    "bbox": (x, y, width, height),
+                                    "polygon": polygon,
+                                }
+                            )
 
                     except (ValueError, IndexError):
                         continue

{argus_cv-1.2.0.dist-info → argus_cv-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: argus-cv
-Version: 1.2.0
+Version: 1.3.0
 Summary: CLI tool for working with vision AI datasets
 Requires-Python: >=3.10
 Requires-Dist: numpy>=1.24.0

argus_cv-1.3.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+argus/__init__.py,sha256=277ASQvH6ZWVWUzOCVB7vAxn3OYciow4nVkjG16-kio,64
+argus/__main__.py,sha256=63ezHx8eL_lCMoZrCbKhmpao0fmdvYVw1chbknGg-oI,104
+argus/cli.py,sha256=th1Rgn1Sm9juWoavopEXLBT8XEh5lKzOMX-pccwvDgA,47904
+argus/commands/__init__.py,sha256=i2oor9hpVpF-_1qZWCGDLwwi1pZGJfZnUKJZ_NMBG18,30
+argus/core/__init__.py,sha256=II2wYJpGUUGGKOFZ5BCpMIBTfv0WP-F15U_xbpWGjk8,453
+argus/core/base.py,sha256=WBrB7XWz125YZ1UQfHQwsYAuIFY_XGEhG_0ybgPhn6s,3696
+argus/core/coco.py,sha256=atVurZV2T7cszydyD9GfDTWHGYDd-JNK5RD7nse5avc,15823
+argus/core/mask.py,sha256=m7Ztf4lAZx5ITpk3F3mETcvCC6hGydlxK0-2nCjeTfU,21835
+argus/core/split.py,sha256=kEWtbdg6bH-WiNFf83HkqZD90EL4gsavw6JiefuAETs,10776
+argus/core/yolo.py,sha256=tTc9jJzXcwa8LQ_s8nv-D_i2b9k_x-LT1O0eWr4sZ2k,28616
+argus_cv-1.3.0.dist-info/METADATA,sha256=1CCYLSnGHaAdS5jjwUuTJWRHu5OM_oFfvi1L_v3SkQw,1070
+argus_cv-1.3.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+argus_cv-1.3.0.dist-info/entry_points.txt,sha256=dvJFH7BkrOxJnifSjPhwq1YCafPaqdngWyBuFYE73yY,43
+argus_cv-1.3.0.dist-info/RECORD,,

argus_cv-1.2.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-argus/__init__.py,sha256=e8XoTWOOBcvGAX8HIUZFycf6Yi3IMFZfXSj4oiK2Cwg,64
-argus/__main__.py,sha256=63ezHx8eL_lCMoZrCbKhmpao0fmdvYVw1chbknGg-oI,104
-argus/cli.py,sha256=SHAN07n1ffqWil2qhLkc0xIdU7iXhjTPW8nFMDUhnNQ,34103
-argus/commands/__init__.py,sha256=i2oor9hpVpF-_1qZWCGDLwwi1pZGJfZnUKJZ_NMBG18,30
-argus/core/__init__.py,sha256=Plv_tk0Wq9OlGLDPOSQWxrd5cTwNK9kEZANTim3s23A,348
-argus/core/base.py,sha256=Vd_2xR6L3lhu9vHoyLeFTc0Dg59py_D9kaye1tta5Co,3678
-argus/core/coco.py,sha256=bJvOhBzwjsOU8DBijGDysnSPlprwetkPf4Z02UOmqw0,15757
-argus/core/split.py,sha256=kEWtbdg6bH-WiNFf83HkqZD90EL4gsavw6JiefuAETs,10776
-argus/core/yolo.py,sha256=W7WH7RwNYvJu8nYdNcXVNRkK5IBov-1dG23esOkCB1M,28213
-argus_cv-1.2.0.dist-info/METADATA,sha256=SP_8kVPy0Mv_9F8CM9zWTARqCOmzYFetyF7QaEgNT-Y,1070
-argus_cv-1.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-argus_cv-1.2.0.dist-info/entry_points.txt,sha256=dvJFH7BkrOxJnifSjPhwq1YCafPaqdngWyBuFYE73yY,43
-argus_cv-1.2.0.dist-info/RECORD,,