Metric Engine V2 (-> feat/no-disk)

pyproject.toml

@@ -85,6 +85,7 @@ uvicorn = { extras = ["standard"], version = "^0.21.1" }
 nodejs-bin = "^16.15.1a4"
 pyjwt = "^2.7.0"
 torch = "^2.0.0"
+kornia = "^0.6.12"
 clip-ea = "^1.0"
 torchvision = "^0.15.2"
 ftfy = "^6.1.1"
@@ -142,7 +143,7 @@ profile = "black"
 plugins = "pandera.mypy"
 ignore_missing_imports = true
 enable_recursive_aliases = true
-exclude = ['examples/', "./src/encord_active/db/migrations/.*"]
+exclude = ['examples/', "./src/encord_active/db/migrations/.*", "local_tests/.*", "venv/bin/.*"]
 
 [tool.pylint.'MASTER']
 extension-pkg-allow-list = "cv2"

src/encord_active/analysis/TODO.md

@@ -0,0 +1,3 @@
+## Performance improvements
+
+- [ ] Batch images when computing embeddings. Otherwise, we won't be able to utilize the GPU

src/encord_active/analysis/base.py

@@ -0,0 +1,166 @@
+from abc import ABC, ABCMeta, abstractmethod
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from encord_active.analysis.types import (
+    AnnotationMetadata,
+    BoundingBoxBatchTensor,
+    FeatureHashBatchTensor,
+    ImageBatchTensor,
+    ImageIndexBatchTensor,
+    ImageTensor,
+    MaskBatchTensor,
+    MetricBatchDependencies,
+    MetricBatchResult,
+    MetricDependencies,
+    MetricResult,
+)
+from encord_active.db.metrics import MetricType
+
+
+@dataclass(frozen=True)  # FIXME: deprecate for new version
+class BaseFrameInput:
+    image: ImageTensor
+    image_deps: MetricDependencies
+    # key is object_hash | classification_hash
+    annotations: Dict[str, AnnotationMetadata]
+    annotations_deps: Dict[str, MetricDependencies]
+    # hash collision, if set the object or classification hash uniqueness constraint is
+    # not held for this frame
+    hash_collision: bool
+
+
+@dataclass
+class BaseFrameOutput:
+    image: Optional[MetricResult]
+    annotations: Dict[str, MetricResult]
+
+
+@dataclass
+class BaseFrameAnnotationBatchInput:
+    """
+    Dictionary to B x dependencies
+    """
+
+    objects_masks: MaskBatchTensor
+    """
+    O x Masks (O = SUM{i.annotation_count * B}
+    """
+    objects_bounding_boxes: BoundingBoxBatchTensor
+    """"
+    O x Bounding Boxes
+    """
+    objects_deps: MetricBatchDependencies
+    """
+    dependency map to Ox<value> tensors
+    """
+    objects_image_indices: ImageIndexBatchTensor
+    """
+    O x int32 index to the associated image.
+    """
+    objects_feature_hash: FeatureHashBatchTensor
+    """
+    O x int64 converted feature hash values
+    """
+    classifications_deps: MetricBatchDependencies
+    """
+    dependency map to Cx<value> tensors
+    """
+    classifications_image_indices: ImageIndexBatchTensor
+    """
+    C x int32 index to the associated image.
+    """
+    classifications_feature_hash: FeatureHashBatchTensor
+    """
+    C x int64 converted feature hash values
+    """
+
+
+@dataclass
+class BaseFrameBatchInput:
+    images: ImageBatchTensor
+    """
+    B x image tensors
+    """
+    images_deps: MetricBatchDependencies
+    """
+    Dictionary to B x dependencies
+    """
+    annotations: Optional[BaseFrameAnnotationBatchInput]
+    """
+    All annotation batches associated with this image batch.
+    Values are split into classifications & objects (objects have masks + aabb)
+    """
+
+
+@dataclass
+class BaseFrameBatchOutput:
+    images: Optional[MetricBatchResult]
+    """
+    B x image results
+    """
+    objects: Optional[MetricBatchResult]
+    """
+    O x annotation results
+    """
+    classifications: Optional[MetricBatchResult]
+    """
+    C x classification results
+    """
+
+
+class BaseEvaluation(ABC):
+    """
+    Base evaluations are the most abstract definitions of computations that need to happen.
+    They can be both:
+
+    - Quality Metrics that needs to be stored (The `BaseAnalysis` subclass)
+    - Computations the need to happen for Quality Metrics to be able to share computations.
+      For example, Embeddings, image transforms, etc.
+
+    """
+
+    def __init__(self, ident: str) -> None:
+        self.ident = ident
+
+    @abstractmethod
+    def raw_calculate(
+        self,
+        prev_frame: Optional[BaseFrameInput],
+        frame: BaseFrameInput,
+        next_frame: Optional[BaseFrameInput],
+    ) -> BaseFrameOutput:
+        """
+        Base implementation of the raw_calculate method, this api should
+        be considered unstable.
+        """
+        ...
+
+    @abstractmethod
+    def raw_calculate_batch(
+        self,
+        prev_frame: Optional[BaseFrameBatchInput],
+        frame: BaseFrameBatchInput,
+        next_frame: Optional[BaseFrameBatchInput],
+    ) -> BaseFrameBatchOutput:
+        """
+        Base implementation of batched metric calculation.
+        """
+
+
+class BaseAnalysis(BaseEvaluation, metaclass=ABCMeta):
+    """
+    The `BaseAnalysis` is all the metrics that needs to be stored in the
+    database.
+    """
+
+    def __init__(self, ident: str, long_name: str, desc: str) -> None:
+        super().__init__(ident)
+        self.long_name = long_name
+        self.desc = desc
+
+
+class BaseMetric(BaseAnalysis, metaclass=ABCMeta):
+    def __init__(self, ident: str, long_name: str, desc: str, metric_type: MetricType) -> None:
+        super().__init__(ident, long_name, desc)
+        self.metric_type = metric_type

src/encord_active/analysis/config.py

@@ -0,0 +1,171 @@
+import os
+from typing import Union
+
+import torch
+
+from .base import BaseEvaluation
+from .conversions.grayscale import RGBToGray
+from .conversions.hsv import RGBToHSV
+from .embedding import NearestImageEmbeddingQuery
+from .embeddings.clip import ClipImgEmbedding
+from .embeddings.hu_moment import HuMomentEmbeddings
+from .metric import DerivedMetric, RandomSamplingQuery
+from .metrics.active_learning.uniqueness import ImageUniqueness
+from .metrics.annotation.count import ObjectCountMetric
+from .metrics.annotation.density import ObjectDensityMetric
+from .metrics.annotation.distance_to_border import DistanceToBorderMetric
+from .metrics.annotation.maximum_label_iou import MaximumLabelIOUMetric
+from .metrics.annotation.nearest_neighbor_agreement import NearestNeighborAgreement
+from .metrics.image.area import AreaMetric, AreaRelativeMetric
+from .metrics.image.aspect_ratio import AspectRatioMetric
+from .metrics.image.brightness import BrightnessMetric
+from .metrics.image.colors import HSVColorMetric
+from .metrics.image.contrast import ContrastMetric
+from .metrics.image.height import HeightMetric
+from .metrics.image.random_value import RandomMetric
+from .metrics.image.sharpness import SharpnessMetric
+from .metrics.image.width import WidthMetric
+
+"""
+Operations to support tracking for:
+1. Add new input (image/video/...):
+    - calculate new data
+    - invalidate nearby embeddings
+2. Delete input
+    - cascade delete data
+    - invalidate nearby embeddings
+3. Add object
+    - see new input
+    - calculate new data
+    - invalidate nearby embeddings
+4. Update object
+    - internally implement as delete + add (as 1 transaction in db state)
+5. Delete object
+    - see new output
+    - cascade delete data
+    - invalidate nearby embeddings
+"""
+
+
+class AnalysisConfig:
+    """
+    Config for the complete analysis state.
+    This should be generated once and re-used for all metric analysis
+    Based upon the results.
+    """
+
+    def __init__(
+        self,
+        analysis: list[BaseEvaluation],
+        derived_embeddings: list[Union[NearestImageEmbeddingQuery, RandomSamplingQuery]],
+        derived_metrics: list[DerivedMetric],
+        device: torch.device,
+    ) -> None:
+        """
+        Args:
+            analysis:
+                List of pure analysis that only depend on the input state, or
+                earlier analysis results in the list. This can generate
+                temporary tensors, embeddings and metrics.
+            derive_embeddings:
+                List of embeddings based properties that are derived from input embeddings,
+                all operations:
+                    - Nearest Embedding
+                    - Nearest N Embeddings (where N is hardcoded constant).
+            derive_metrics:
+                List of metrics that only take exported results from earlier analysis and
+                derived_embedding results.
+        """
+        self.analysis = analysis
+        self.derived_embeddings = derived_embeddings
+        self.derived_metrics = derived_metrics
+        self.device = device
+
+    def validate(self) -> None:
+        """
+        Asserts that the given analysis config is valid for execution.
+        """
+        ...
+
+    def analysis_versions(self) -> dict[str, int]:
+        """
+        Returns:
+            The version for each exported metric or embedding,
+            this value should be incremented whenever the code
+            is changed in a way that would result in different results
+            being returned if the function were to be re-executed.
+            Exported features missing
+        """
+        ...
+
+
+def default_torch_device() -> torch.device:
+    """Default torch device to use for the analysis config"""
+    torch_device_str = "cpu"
+    # FIXME: currently CPU only due to GPU being too slow due to lack of batch support.
+    #  keep this until batch functions are properly implemented.
+    if torch.cuda.is_available() and os.environ.get("ACTIVE_ALLOW_TORCH_CUDA_BACKEND", "0") == "1":
+        torch_device_str = "cuda"
+    elif torch.backends.mps.is_available() and os.environ.get("ACTIVE_ALLOW_TORCH_MPS_BACKEND", "0") == "1":
+        torch_device_str = "mps"
+    return torch.device(torch_device_str)
+
+
+def create_analysis(device: torch.device) -> AnalysisConfig:
+    """List of all analysis passes to be executed by encord active"""
+    analysis: list[BaseEvaluation] = [
+        # Generate embeddings
+        ClipImgEmbedding(device, "embedding_clip", "ViT-B/32"),
+        HuMomentEmbeddings(),
+        # FIXME: HuMomentEmbeddings("embedding_hu_moments"),
+        # Image transformations
+        RGBToHSV(),
+        RGBToGray(),
+        # Unified Image or Annotation Feature Metrics
+        AspectRatioMetric(),  # metric_aspect_ratio
+        AreaMetric(),  # metric_area
+        BrightnessMetric(),  # metric_brightness
+        ContrastMetric(),  # metric_contrast
+        SharpnessMetric(),  # metric_sharpness
+        HSVColorMetric("red", hue_query=0.0),  # metric_red
+        HSVColorMetric("green", hue_query=1 / 3.0),  # metric_green
+        HSVColorMetric("blue", hue_query=2 / 3.0),  # metric_blue
+        HeightMetric(),  # metric_height
+        WidthMetric(),  # metric_width
+        RandomMetric(),  # metric_random
+        # Annotation only metrics
+        AreaRelativeMetric(),  # metric_area_relative
+        MaximumLabelIOUMetric(),  # metric_max_iou
+        DistanceToBorderMetric(),  # metric_label_border_closeness
+        # Annotation based Frame Metrics
+        ObjectCountMetric(),  # metric_object_count
+        ObjectDensityMetric(),  # metric_object_density ("Frame object density")
+        # ?????
+        # FIXME: metric_image_difficulty ("Image Difficulty")  FIXME: KMeans!!??
+        # FIXME: metric_label_inconsistent_classification_and_track ("Inconsistent Object Classification and Track IDs")
+        # FIXME: metric_label_shape_outlier ("Shape outlier detection")
+        # FIXME: metric_seq_occlusion_detection (ALSO - MISSING COLUMN NAME, WILL NEED ALEMBIC MIGRATION)
+        # Temporal metrics
+        # FIXME: DISABLED TemporalShapeChange(),  # metric_label_poly_similarity ("Polygon Shape Similarity")
+        # FIXME: DISABLED TemporalMissingObjectsAndWrongTracks(),  # metric_label_missing_or_broken_tracks ("Missing Objects and Broken Tracks")
+    ]
+    derived_embeddings: list[Union[NearestImageEmbeddingQuery, RandomSamplingQuery]] = [
+        # Derive properties from embedding
+        NearestImageEmbeddingQuery("derived_clip_nearest", "embedding_clip"),
+    ]
+    derived_metrics: list[DerivedMetric] = [
+        # Metrics depending on derived embedding / metric properties ONLY
+        NearestNeighborAgreement(),  # metric_annotation_quality
+        ImageUniqueness(),  # metric_image_uniqueness ("Image Singularity")
+    ]
+    return AnalysisConfig(
+        analysis=analysis,
+        # FIXME?: derived_temporal (derived metrics that have access to analysis values in N frame range)
+        derived_embeddings=derived_embeddings,
+        derived_metrics=derived_metrics,
+        device=device,
+    )
+
+
+def get_default_analysis():
+    analysis = create_analysis(torch.device("cpu"))