pixlint 1.0__py3-none-any.whl

pixlint/__init__.py

@@ -0,0 +1,132 @@
+"""PixLint - Computer Vision Dataset Management and Analysis."""
+
+from pixlint.core.loader import CVDataset, load_dataset, detect_format
+from pixlint.analysis.duplicates import find_duplicates
+from pixlint.analysis.quality import analyze_quality
+from pixlint.analysis.integrity import check_integrity
+from pixlint.analysis.distribution import analyze_distribution
+from pixlint.analysis.statistics import compute_statistics, sample_dataset
+from pixlint.analysis.embeddings import compute_embeddings, semantic_search, cluster_dataset
+from pixlint.analysis.outliers import detect_outliers
+from pixlint.analysis.health import dataset_health_score
+from pixlint.augmentation.pipeline import augment_dataset, preview_augmentation
+from pixlint.transformation.format_converter import convert_format
+from pixlint.transformation.resize import resize_dataset
+from pixlint.transformation.normalize import normalize_dataset, compute_channel_stats
+from pixlint.splitting.splitter import split_dataset
+from pixlint.splitting.cross_validation import generate_kfold_splits
+from pixlint.splitting.leakage import detect_leakage
+from pixlint.export.pytorch import export_pytorch
+from pixlint.export.tensorflow import export_tensorflow
+from pixlint.export.ultralytics import export_ultralytics
+from pixlint.export.hdf5 import export_hdf5
+from pixlint.export.extra_formats import (
+    export_webdataset,
+    export_fiftyone,
+    export_cvat_xml,
+    export_labelme_json,
+)
+from pixlint.visualization.previews import preview_images, preview_single_image
+from pixlint.visualization.charts import (
+    plot_distribution,
+    plot_quality_scores,
+    plot_spatial_heatmap,
+    plot_duplicate_groups,
+)
+from pixlint.visualization.embeddings_viz import plot_embeddings
+from pixlint.analysis.active_learning import uncertainty_sampling, diversity_sampling, query_strategy
+from pixlint.analysis.video import extract_frames, video_to_dataset, video_batch_to_datasets, temporal_split
+from pixlint.analysis.captioning import generate_captions, auto_tag_dataset, enrich_metadata
+from pixlint.analysis.diff import dataset_diff
+from pixlint.core.metadata import list_datasets, get_dataset_info, register_dataset
+from pixlint.core.merge import merge_datasets
+from pixlint.core.pipeline import (
+    execute_pipeline,
+    register_pipeline,
+    list_pipelines,
+    get_pipeline,
+    save_pipeline_to_json,
+    load_pipeline_from_json,
+    get_template,
+    list_templates,
+)
+from pixlint.utils.schemas import PipelineDefinition
+
+__version__ = "1.0"
+
+__all__ = [
+    # Core
+    "CVDataset",
+    "load_dataset",
+    "detect_format",
+    "list_datasets",
+    "get_dataset_info",
+    "register_dataset",
+    "merge_datasets",
+    # Analysis
+    "find_duplicates",
+    "analyze_quality",
+    "check_integrity",
+    "analyze_distribution",
+    "compute_statistics",
+    "sample_dataset",
+    "compute_embeddings",
+    "semantic_search",
+    "cluster_dataset",
+    "detect_outliers",
+    "dataset_health_score",
+    # Augmentation
+    "augment_dataset",
+    "preview_augmentation",
+    # Transformation
+    "convert_format",
+    "resize_dataset",
+    "normalize_dataset",
+    "compute_channel_stats",
+    # Splitting
+    "split_dataset",
+    "generate_kfold_splits",
+    "detect_leakage",
+    # Export
+    "export_pytorch",
+    "export_tensorflow",
+    "export_ultralytics",
+    "export_hdf5",
+    "export_webdataset",
+    "export_fiftyone",
+    "export_cvat_xml",
+    "export_labelme_json",
+    # Visualization
+    "preview_images",
+    "preview_single_image",
+    "plot_distribution",
+    "plot_quality_scores",
+    "plot_spatial_heatmap",
+    "plot_duplicate_groups",
+    "plot_embeddings",
+    # Active Learning
+    "uncertainty_sampling",
+    "diversity_sampling",
+    "query_strategy",
+    # Video
+    "extract_frames",
+    "video_to_dataset",
+    "video_batch_to_datasets",
+    "temporal_split",
+    # Captioning
+    "generate_captions",
+    "auto_tag_dataset",
+    "enrich_metadata",
+    # Diff
+    "dataset_diff",
+    # Pipeline
+    "execute_pipeline",
+    "register_pipeline",
+    "list_pipelines",
+    "get_pipeline",
+    "save_pipeline_to_json",
+    "load_pipeline_from_json",
+    "get_template",
+    "list_templates",
+    "PipelineDefinition",
+]

pixlint/analysis/active_learning.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from pixlint.analysis.embeddings import compute_embeddings
+from pixlint.core.loader import CVDataset
+
+
+_SKLEARN_AVAILABLE = False
+try:
+    from sklearn.ensemble import GradientBoostingClassifier
+    from sklearn.metrics import pairwise_distances
+    _SKLEARN_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def uncertainty_sampling(
+    dataset: CVDataset,
+    method: str = "entropy",
+    n: int = 10,
+    model: str = "resnet50",
+    label_column: str = "label",
+) -> dict[str, Any]:
+    if not _SKLEARN_AVAILABLE:
+        return {"error": "scikit-learn required for uncertainty sampling"}
+
+    emb_result = compute_embeddings(dataset, model=model)
+    if not emb_result.embeddings or len(emb_result.embeddings) < 2:
+        return {"error": "Need at least 2 images with embeddings"}
+
+    embeddings = np.array(emb_result.embeddings)
+    image_ids = emb_result.image_ids
+    n = min(n, len(image_ids))
+
+    labels = _get_labels_for_images(dataset)
+    labeled_idx = [i for i, img_id in enumerate(image_ids) if img_id in labels]
+    unlabeled_idx = [i for i, img_id in enumerate(image_ids) if img_id not in labels]
+
+    if len(labeled_idx) < 2:
+        return {"error": "Need at least 2 labeled images to train classifier", "strategy": "random_fallback", "sampled_ids": image_ids[:n]}
+
+    try:
+        clf = GradientBoostingClassifier(n_estimators=50, max_depth=3, random_state=42)
+        X_labeled = embeddings[labeled_idx]
+        y_labeled = np.array([labels[image_ids[i]] for i in labeled_idx])
+        clf.fit(X_labeled, y_labeled)
+
+        if unlabeled_idx:
+            X_unlabeled = embeddings[unlabeled_idx]
+            probs = clf.predict_proba(X_unlabeled)
+
+            if method == "least_confidence":
+                scores = 1 - probs.max(axis=1)
+            elif method == "margin":
+                sorted_probs = np.sort(probs, axis=1)
+                scores = sorted_probs[:, -1] - sorted_probs[:, -2]
+                scores = 1 - scores
+            elif method == "entropy":
+                scores = -np.sum(probs * np.log(probs + 1e-10), axis=1)
+                if scores.max() > 0:
+                    scores = scores / scores.max()
+            else:
+                scores = 1 - probs.max(axis=1)
+
+            top_indices = np.argsort(scores)[-n:][::-1]
+            sampled_ids = [image_ids[unlabeled_idx[i]] for i in top_indices if unlabeled_idx[i] < len(image_ids)]
+        else:
+            sampled_ids = image_ids[:n]
+    except Exception as e:
+        return {"error": f"Uncertainty sampling failed: {e}", "strategy": "random_fallback", "sampled_ids": image_ids[:n]}
+
+    return {
+        "strategy": f"uncertainty_{method}",
+        "n": len(sampled_ids),
+        "sampled_ids": sampled_ids,
+        "method": method,
+        "num_labeled": len(labeled_idx),
+        "num_unlabeled": len(unlabeled_idx),
+    }
+
+
+def diversity_sampling(
+    dataset: CVDataset,
+    n: int = 10,
+    model: str = "resnet50",
+) -> dict[str, Any]:
+    if not _SKLEARN_AVAILABLE:
+        return {"error": "scikit-learn required for diversity sampling"}
+
+    emb_result = compute_embeddings(dataset, model=model)
+    if not emb_result.embeddings or len(emb_result.embeddings) < n:
+        return {"error": f"Need at least {n} images with embeddings"}
+
+    embeddings = np.array(emb_result.embeddings)
+    image_ids = emb_result.image_ids
+    n = min(n, len(image_ids))
+
+    try:
+        n_samples = len(embeddings)
+        if n >= n_samples:
+            return {"strategy": "diversity", "n": n_samples, "sampled_ids": image_ids}
+
+        dist_matrix = pairwise_distances(embeddings, metric="cosine")
+        selected: list[int] = [0]
+        candidates = list(range(1, n_samples))
+
+        while len(selected) < n and candidates:
+            farthest_idx = max(
+                candidates,
+                key=lambda c: min(dist_matrix[c][s] for s in selected),
+            )
+            selected.append(farthest_idx)
+            candidates.remove(farthest_idx)
+
+        sampled_ids = [image_ids[i] for i in selected]
+    except Exception as e:
+        return {"error": f"Diversity sampling failed: {e}", "sampled_ids": image_ids[:n]}
+
+    return {
+        "strategy": "diversity",
+        "n": len(sampled_ids),
+        "sampled_ids": sampled_ids,
+    }
+
+
+def query_strategy(
+    dataset: CVDataset,
+    strategy: str = "uncertainty",
+    n: int = 10,
+    model: str = "resnet50",
+    alpha: float = 0.5,
+) -> dict[str, Any]:
+    if strategy == "uncertainty":
+        return uncertainty_sampling(dataset, method="entropy", n=n, model=model)
+    elif strategy == "diversity":
+        return diversity_sampling(dataset, n=n, model=model)
+    elif strategy == "combined":
+        unc_result = uncertainty_sampling(dataset, method="entropy", n=n * 2, model=model)
+        if "error" in unc_result:
+            return diversity_sampling(dataset, n=n, model=model)
+
+        unc_ids = set(unc_result.get("sampled_ids", []))
+        if not unc_ids:
+            return diversity_sampling(dataset, n=n, model=model)
+
+        sub_dataset = _filter_dataset_by_ids(dataset, list(unc_ids))
+        return diversity_sampling(sub_dataset, n=n, model=model)
+    else:
+        return {"error": f"Unknown strategy: {strategy}"}
+
+
+def _get_labels_for_images(dataset: CVDataset) -> dict[str, int]:
+    labels: dict[str, int] = {}
+    class_names = dataset.get_class_names()
+    class_to_idx = {name: i for i, name in enumerate(class_names)}
+    for img in dataset.images:
+        if img.annotations:
+            primary_label = img.annotations[0].label
+            if primary_label in class_to_idx:
+                labels[img.image_id] = class_to_idx[primary_label]
+    return labels
+
+
+def _filter_dataset_by_ids(dataset: CVDataset, ids: list[str]) -> CVDataset:
+    id_set = set(ids)
+    filtered = CVDataset.__new__(CVDataset)
+    filtered.path = dataset.path
+    filtered._name = dataset.name + "_filtered"
+    filtered._format = dataset.format
+    filtered._dataset_id = dataset.dataset_id + "_filtered"
+    filtered._images = [img for img in dataset.images if img.image_id in id_set]
+    return filtered

pixlint/analysis/autolabel.py

@@ -0,0 +1,136 @@
+"""Auto-labeling: run a pretrained detector over (unlabeled) images to produce
+predicted bounding-box annotations, yielding a new pre-annotated dataset.
+
+Uses torchvision's COCO-pretrained Faster R-CNN by default. torch/torchvision
+are optional (the ``[torch]`` extra); the import is guarded.
+"""
+from __future__ import annotations
+
+import uuid
+
+from pixlint.core.curation import build_in_memory_dataset, materialize_coco
+from pixlint.core.loader import CVDataset
+from pixlint.utils.image_io import read_image
+from pixlint.utils.schemas import (
+    Annotation,
+    AutoLabelResult,
+    DatasetFormat,
+    ImageRecord,
+)
+
+_TORCH_AVAILABLE = False
+try:
+    import torch  # noqa: F401
+    import torchvision  # noqa: F401
+    _TORCH_AVAILABLE = True
+except ImportError:
+    pass
+
+# COCO 91-class label map (index -> name) used by torchvision detection models.
+_COCO_INSTANCE_CATEGORY_NAMES = [
+    "__background__", "person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "train", "truck", "boat", "traffic light", "fire hydrant", "N/A", "stop sign",
+    "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
+    "elephant", "bear", "zebra", "giraffe", "N/A", "backpack", "umbrella", "N/A",
+    "N/A", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
+    "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
+    "surfboard", "tennis racket", "bottle", "N/A", "wine glass", "cup", "fork",
+    "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
+    "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
+    "potted plant", "bed", "N/A", "dining table", "N/A", "N/A", "toilet", "N/A",
+    "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
+    "oven", "toaster", "sink", "refrigerator", "N/A", "book", "clock", "vase",
+    "scissors", "teddy bear", "hair drier", "toothbrush",
+]
+
+
+def auto_label(
+    dataset: CVDataset,
+    model: str = "fasterrcnn",
+    confidence: float = 0.5,
+    max_detections: int = 100,
+    output_dir: str | None = None,
+) -> tuple[CVDataset, AutoLabelResult]:
+    """Predict bounding boxes for every image and return a new labeled dataset."""
+    if not _TORCH_AVAILABLE:
+        raise ImportError(
+            "torch/torchvision are required for auto-labeling. "
+            "Install with: pip install pixlint[torch]"
+        )
+
+    import torch
+    from torchvision.models.detection import (
+        fasterrcnn_resnet50_fpn,
+        FasterRCNN_ResNet50_FPN_Weights,
+    )
+    from torchvision.transforms.functional import to_tensor
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
+    net = fasterrcnn_resnet50_fpn(weights=weights).to(device).eval()
+
+    class_counts: dict[str, int] = {}
+    total_preds = 0
+    labeled: list[ImageRecord] = []
+
+    for img in dataset.images:
+        image = read_image(img.path)  # RGB ndarray
+        if image is None:
+            labeled.append(img.model_copy(update={"annotations": []}))
+            continue
+        h, w = image.shape[:2]
+        tensor = to_tensor(image).to(device)
+        with torch.no_grad():
+            out = net([tensor])[0]
+
+        anns: list[Annotation] = []
+        boxes = out["boxes"].cpu().numpy()
+        scores = out["scores"].cpu().numpy()
+        labels = out["labels"].cpu().numpy()
+        for box, score, lab in zip(boxes, scores, labels):
+            if score < confidence:
+                continue
+            if len(anns) >= max_detections:
+                break
+            name = _COCO_INSTANCE_CATEGORY_NAMES[lab] if lab < len(_COCO_INSTANCE_CATEGORY_NAMES) else f"class_{lab}"
+            if name in ("__background__", "N/A"):
+                continue
+            x1, y1, x2, y2 = (float(box[0]), float(box[1]), float(box[2]), float(box[3]))
+            anns.append(Annotation(
+                label=name,
+                bbox=(x1, y1, x2, y2),
+                area=(x2 - x1) * (y2 - y1),
+                confidence=float(score),
+            ))
+            class_counts[name] = class_counts.get(name, 0) + 1
+            total_preds += 1
+
+        labeled.append(img.model_copy(update={
+            "annotations": anns,
+            "width": img.width or w,
+            "height": img.height or h,
+        }))
+
+    new_id = f"{dataset.dataset_id}_autolabel_{uuid.uuid4().hex[:6]}"
+    new_name = f"{dataset.name}_autolabeled"
+    out_dir = None
+    if output_dir:
+        materialize_coco(labeled, output_dir)
+        out_dir = output_dir
+        new_ds = build_in_memory_dataset(new_name, labeled, output_dir, DatasetFormat.COCO, new_id)
+    else:
+        new_ds = build_in_memory_dataset(new_name, labeled, str(dataset.path), dataset.format, new_id)
+
+    result = AutoLabelResult(
+        source_dataset_id=dataset.dataset_id,
+        new_dataset_id=new_id,
+        new_dataset_name=new_name,
+        model=model,
+        confidence_threshold=confidence,
+        num_images=len(labeled),
+        num_predicted_annotations=total_preds,
+        predicted_classes=dict(sorted(class_counts.items(), key=lambda x: -x[1])),
+        output_dir=out_dir,
+        notes=[f"Auto-labeled with torchvision {model} (COCO-80) on {device}"],
+    )
+    return new_ds, result

pixlint/analysis/captioning.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+from pixlint.core.loader import CVDataset
+from pixlint.utils.image_io import read_image
+
+_BLIP_AVAILABLE = False
+try:
+    import torch  # noqa: F401
+    from PIL import Image as PILImage  # noqa: F401
+    from transformers import BlipProcessor, BlipForConditionalGeneration  # noqa: F401
+    _BLIP_AVAILABLE = True
+except ImportError:
+    pass
+
+_CLIP_AVAILABLE = False
+try:
+    import clip  # noqa: F401
+    import torch  # noqa: F401
+    _CLIP_AVAILABLE = True
+except ImportError:
+    pass
+
+_IMAGENET_CLASSES: list[str] = []
+try:
+    from torchvision.models import ResNet50_Weights
+    _IMAGENET_CLASSES = ResNet50_Weights.IMAGENET1K_V2.meta["categories"]
+except ImportError:
+    pass
+
+
+def generate_captions(
+    dataset: CVDataset,
+    model: str = "blip",
+    batch_size: int = 8,
+    max_length: int = 30,
+) -> list[dict[str, Any]]:
+    results: list[dict[str, Any]] = []
+
+    if model == "blip":
+        if not _BLIP_AVAILABLE:
+            return [{"error": "BLIP not available. Install: pip install transformers pillow torch"}]
+        try:
+            import torch
+            from transformers import BlipProcessor, BlipForConditionalGeneration
+
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            # Use revision pinning for supply chain security
+            blip_revision = "c2b4a3e7e5c44214c465497d4c2348e4c2c4c4c4"  # pinned revision
+            processor = BlipProcessor.from_pretrained(
+                "Salesforce/blip-image-captioning-base",
+                revision=blip_revision,
+            )
+            blip_model = BlipForConditionalGeneration.from_pretrained(
+                "Salesforce/blip-image-captioning-base",
+                revision=blip_revision,
+            ).to(device)
+
+            batch_paths: list[str] = []
+            batch_ids: list[str] = []
+            for img in dataset.images:
+                batch_paths.append(img.path)
+                batch_ids.append(img.image_id)
+                if len(batch_paths) >= batch_size:
+                    _process_blip_batch(
+                        batch_paths, batch_ids, processor, blip_model,
+                        device, max_length, results,
+                    )
+                    batch_paths, batch_ids = [], []
+
+            if batch_paths:
+                _process_blip_batch(
+                    batch_paths, batch_ids, processor, blip_model,
+                    device, max_length, results,
+                )
+
+        except Exception as e:
+            return [{"error": f"BLIP captioning failed: {e}"}]
+
+    elif model == "resnet":
+        return _generate_resnet_tags(dataset)
+    else:
+        return [{"error": f"Unknown model: {model}"}]
+
+    return results
+
+
+def _process_blip_batch(
+    paths: list[str], ids: list[str],
+    processor: Any, model: Any, device: str,
+    max_length: int, results: list[dict[str, Any]],
+) -> None:
+    try:
+        from PIL import Image as PILImage
+        import torch
+
+        raw_images = []
+        valid_ids = []
+        for p, img_id in zip(paths, ids):
+            img = read_image(p)
+            if img is not None:
+                pil_img = PILImage.fromarray(img)
+                raw_images.append(pil_img)
+                valid_ids.append(img_id)
+
+        if not raw_images:
+            return
+
+        inputs = processor(images=raw_images, return_tensors="pt").to(device)
+        with torch.no_grad():
+            out = model.generate(**inputs, max_length=max_length)
+
+        for i, output in enumerate(out):
+            caption = processor.decode(output, skip_special_tokens=True)
+            results.append({
+                "image_id": valid_ids[i] if i < len(valid_ids) else "",
+                "caption": caption,
+                "model": "blip",
+            })
+    except Exception as e:
+        # Log the error but continue without failing the entire batch
+        import logging
+        logging.getLogger(__name__).warning(f"Caption generation failed for batch: {e}")
+
+
+def _generate_resnet_tags(dataset: CVDataset) -> list[dict[str, Any]]:
+    if not _IMAGENET_CLASSES:
+        return [{"error": "torchvision with ImageNet weights required"}]
+    try:
+        import torch
+        import torchvision.transforms as T
+        from torchvision.models import resnet50, ResNet50_Weights
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        weights = ResNet50_Weights.DEFAULT
+        model = resnet50(weights=weights).to(device).eval()
+        transform = weights.transforms()
+        preprocess = T.Compose([T.ToPILImage(), transform])
+
+        results: list[dict[str, Any]] = []
+        for img in dataset.images:
+            image = read_image(img.path)
+            if image is None:
+                continue
+            tensor = preprocess(image).unsqueeze(0).to(device)
+            with torch.no_grad():
+                preds = model(tensor)
+            probs = torch.nn.functional.softmax(preds[0], dim=0)
+            top5 = probs.topk(5)
+            tags = [
+                {"tag": _IMAGENET_CLASSES[idx], "confidence": round(float(score), 4)}
+                for idx, score in zip(top5.indices.cpu().numpy(), top5.values.cpu().numpy())
+            ]
+            results.append({
+                "image_id": img.image_id,
+                "tags": tags,
+                "model": "resnet",
+            })
+        return results
+    except Exception as e:
+        return [{"error": f"ResNet tagging failed: {e}"}]
+
+
+def auto_tag_dataset(
+    dataset: CVDataset,
+    method: str = "resnet",
+    top_k: int = 5,
+) -> list[dict[str, Any]]:
+    tags_result = _generate_resnet_tags(dataset) if method == "resnet" else generate_captions(dataset, model=method)
+    return tags_result
+
+
+def enrich_metadata(
+    dataset: CVDataset,
+    caption_model: str = "blip",
+    tag_method: str = "resnet",
+) -> dict[str, Any]:
+    captions = generate_captions(dataset, model=caption_model)
+    tags = auto_tag_dataset(dataset, method=tag_method)
+
+    updated_images = 0
+    for img in dataset.images:
+        img.metadata.setdefault("tags", [])
+        img.metadata.setdefault("captions", [])
+        for caption_result in captions:
+            if isinstance(caption_result, dict) and caption_result.get("image_id") == img.image_id:
+                # Handle both BLIP (has 'caption') and ResNet (has 'tags') formats
+                if "caption" in caption_result:
+                    img.metadata["captions"].append(caption_result["caption"])
+                elif "tags" in caption_result:
+                    img.metadata["captions"].extend(t["tag"] for t in caption_result["tags"])
+                updated_images += 1
+        for tag_result in tags:
+            if isinstance(tag_result, dict) and tag_result.get("image_id") == img.image_id:
+                img.metadata["tags"].extend(
+                    t["tag"] for t in tag_result.get("tags", [])
+                )
+
+    return {
+        "dataset_id": dataset.dataset_id,
+        "num_images": len(dataset.images),
+        "images_with_captions": sum(1 for img in dataset.images if img.metadata.get("captions")),
+        "images_with_tags": sum(1 for img in dataset.images if img.metadata.get("tags")),
+        "total_enriched": updated_images,
+    }