PyTorch DataLoader Integration

VigilCV's VisionSentinel integrates directly into PyTorch Dataset.__getitem__ to validate images on the fly during training.

Quality-Guarded Dataset

from pathlib import Path
from collections.abc import Callable
from torch.utils.data import Dataset
from PIL import Image
import torch

from vigilcv import VisionSentinel
from vigilcv.exceptions import VigilCVError

sentinel = VisionSentinel(blur_threshold=80.0, min_entropy=2.5)


class QualityGuardedDataset(Dataset):
    """
    A PyTorch Dataset that validates images with VigilCV before loading.
    Corrupted or degraded images are replaced with a fallback tensor
    instead of crashing the DataLoader batch collation.
    """

    def __init__(
        self,
        image_dir: str | Path,
        transform: Callable | None = None,
    ) -> None:
        self.image_paths = sorted(Path(image_dir).glob("*.jpg"))
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        path = self.image_paths[idx]

        try:
            # Step 1: Pre-flight quality gate
            sentinel.guard(path)

            # Step 2: Safe to load and transform
            with Image.open(path) as img:
                img_rgb = img.convert("RGB")
                tensor = self.transform(img_rgb) if self.transform else torch.zeros(3, 224, 224)
                return tensor, 1  # label=1 for valid

        except VigilCVError:
            # Return fallback tensor — preserves batch collation
            return torch.zeros(3, 224, 224), -1  # label=-1 for invalid

Usage with DataLoader

from torch.utils.data import DataLoader
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

dataset = QualityGuardedDataset("dataset/train/", transform=transform)
loader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

for images, labels in loader:
    valid_mask = labels != -1
    if valid_mask.sum() == 0:
        continue  # Skip fully-degraded batches
    
    images = images[valid_mask]
    labels = labels[valid_mask]
    
    outputs = model(images)
    loss = criterion(outputs, labels)
    loss.backward()

Why Not Filter in __init__?

Pre-filtering image_paths at startup scans the entire dataset once. This is fine for small datasets but causes a long cold start for 100k+ image datasets. VigilCV's per-item lazy validation ensures:

1. No cold-start penalty

2. Handles files that become corrupt during training (e.g., NFS issues)

3. Degraded images return fallback tensors — no crashes

Performance Impact

The overhead is negligible relative to GPU batch transfer time (~40ms).