diff --git a/.gitignore b/.gitignore index 2308930..dd432a0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,7 @@ build/* _C.* outputs/* coreml_models/* +coreml/video_tracking/models/ +coreml/video_tracking/results/ +coreml/video_tracking/validation.json checkpoints/*.pt diff --git a/README.md b/README.md index 3ae3c12..1bcafb6 100644 --- a/README.md +++ b/README.md @@ -145,9 +145,14 @@ python ./coreml/export_to_coreml.py \ This creates three optimized CoreML models: - **Image Encoder**: Processes input images to feature embeddings (~9.6MB) -- **Prompt Encoder**: Handles user prompts (points, boxes, masks) (~2MB) +- **Prompt Encoder**: Handles user prompts (points, boxes, masks) (~2MB) - **Mask Decoder**: Generates segmentation masks from features (~8MB) +For temporal video tracking, EdgeTAM can also export four Core ML models that +preserve the video predictor's memory pipeline. See the +[Core ML video tracking guide](./coreml/video_tracking/README.md) for export, +inference, and validation instructions. + ## Performance ### Promptable Video Segmentation (PVS) diff --git a/coreml/README.md b/coreml/README.md index 0da4349..1240722 100644 --- a/coreml/README.md +++ b/coreml/README.md @@ -1,74 +1,44 @@ -# EdgeTAM CoreML Export +# EdgeTAM Core ML export -Export EdgeTAM to CoreML format for iOS/macOS deployment. - -## Quick Export +EdgeTAM provides separate Core ML pipelines for prompted image segmentation +and temporal video tracking. Install the optional dependencies before using +either exporter: ```bash -python coreml/export_to_coreml.py \ - --sam2_cfg sam2/configs/edgetam.yaml \ - --sam2_checkpoint checkpoints/edgetam.pt -``` - -This creates three CoreML models in `./coreml_models/`: -- `edgetam_image_encoder.mlpackage` (9.6MB) -- `edgetam_prompt_encoder.mlpackage` (2.0MB) -- `edgetam_mask_decoder.mlpackage` (9.8MB) - -## Usage Example - -```python -import coremltools as ct -from PIL import Image - -# Load models -image_encoder = ct.models.MLModel("coreml_models/edgetam_image_encoder.mlpackage") -prompt_encoder = ct.models.MLModel("coreml_models/edgetam_prompt_encoder.mlpackage") -mask_decoder = ct.models.MLModel("coreml_models/edgetam_mask_decoder.mlpackage") - -# Segment with point prompt -image = Image.open("image.jpg").resize((1024, 1024)) -encoder_out = image_encoder.predict({"image": image}) - -# Add your point and generate mask -# See inference_example.py for complete video tracking example +pip install -e ".[coreml]" ``` -## Video Tracking Example +## Image segmentation -The included `inference_example.py` demonstrates real-time video tracking: +Export the image encoder, prompt encoder, and mask decoder: ```bash -# Demo with default coffee video -python coreml/inference_example.py - -# Use your own video -python coreml/inference_example.py --video path/to/your/video.mp4 - -# Run different examples -python coreml/inference_example.py --example segment # Single image -python coreml/inference_example.py --example track # Real-time tracking -python coreml/inference_example.py --example demo # Video demo (default) +python coreml/export_to_coreml.py \ + --sam2_cfg sam2/configs/edgetam.yaml \ + --sam2_checkpoint checkpoints/edgetam.pt \ + --output_dir coreml_models ``` -## Performance Benchmark +See `inference_example.py` for image prompting and `benchmark_coreml.py` for a +small synthetic benchmark. -Run benchmark with: `python coreml/benchmark_coreml.py` +## Temporal video tracking -Note: This is a limited test on synthetic data. Real-world performance may vary. +The video export preserves EdgeTAM's temporal memory pipeline instead of +running image segmentation independently on every frame. It produces four +stateless Core ML packages and maintains the fixed-shape memory bank in the +client. -### Results - -| Metric | PyTorch | CoreML | Difference | -|--------|---------|--------|------------| -| Speed | 40.1ms | 39.2ms | -0.9ms | -| Quality | IoU 0.9897 | IoU 0.9893 | -0.0004 | -| Size | 54MB | 21.4MB | -32.6MB | - -## Requirements +```bash +python coreml/video_tracking/export_models.py \ + --config sam2/configs/edgetam.yaml \ + --checkpoint checkpoints/edgetam.pt \ + --output-dir coreml_models/video_tracking +``` -- PyTorch -- coremltools -- EdgeTAM checkpoint +See [video_tracking/README.md](video_tracking/README.md) for the model +architecture, Python predictor, validation command, tests, and integration +constraints. -The CoreML export maintains identical segmentation quality while being faster and 60% smaller for mobile deployment. \ No newline at end of file +Generated `.mlpackage` directories belong in `coreml_models/`, which is +excluded from version control. diff --git a/coreml/video_tracking/README.md b/coreml/video_tracking/README.md new file mode 100644 index 0000000..0315886 --- /dev/null +++ b/coreml/video_tracking/README.md @@ -0,0 +1,128 @@ +# Core ML video tracking + +This directory adds temporal video tracking to EdgeTAM's Core ML export. A +point or box prompt initializes a track on the first frame. Later frames are +processed without repeating the prompt, using the spatial memories and object +pointers produced by earlier frames. + +The exported Core ML models are stateless. The client owns the small, +fixed-shape memory bank, so the same models can be used from Python, Swift, or +another Core ML host. + +## Model pipeline + +The exporter creates four iOS 18 ML Program packages: + +| Model | Responsibility | +| --- | --- | +| `EdgeTAMVideoImageEncoder` | Produces raw, initial, and high-resolution features for each frame. | +| `EdgeTAMVideoInitializer` | Applies the first-frame point or box prompt and returns the seed mask and object pointer. | +| `EdgeTAMVideoMemoryEncoder` | Encodes the prompted mask with EdgeTAM's 2D Spatial Perceiver. | +| `EdgeTAMVideoPropagator` | Conditions the current frame on the explicit memory bank and returns the next mask, pointer, and memory. | + +The runtime keeps one conditioning memory, six recent memories, and sixteen +object pointers per tracked object. Validity tensors mask unused slots while +the bank fills. + +## Requirements + +Install EdgeTAM with its Core ML dependencies: + +```bash +pip install -e ".[coreml]" +``` + +The export targets iOS 18 and requires an EdgeTAM checkpoint. Generated model +packages are build artifacts and are not stored in the repository. + +## Export + +Run the exporter from the repository root: + +```bash +python coreml/video_tracking/export_models.py \ + --config sam2/configs/edgetam.yaml \ + --checkpoint checkpoints/edgetam.pt \ + --output-dir coreml_models/video_tracking +``` + +`--device` selects the PyTorch device used while tracing. It defaults to +`cpu`; `mps` is also useful on Apple silicon. + +## Python inference + +`CoreMLVideoPredictor` owns the explicit memory bank for one object: + +```bash +PYTHONPATH=coreml/video_tracking python +``` + +```python +from pathlib import Path + +from PIL import Image + +from edgetam_coreml_video.predictor import CoreMLVideoPredictor + +predictor = CoreMLVideoPredictor.from_directory( + Path("coreml_models/video_tracking") +) + +first_frame = Image.open("frames/00000.jpg") +result = predictor.start_track( + first_frame, + points=[[210, 350]], + labels=[1], +) + +next_frame = Image.open("frames/00001.jpg") +result = predictor.track_frame(next_frame) +binary_mask = result.mask +``` + +Prompt coordinates use the original frame's pixel coordinate system. Point +labels follow EdgeTAM conventions: `1` for a foreground point, `0` for a +background point, and `2`/`3` for the two corners of a box. One to four prompt +tokens are supported. Call `reset()` before starting a different object. + +## Numerical validation + +The validator runs the official PyTorch video predictor and the Core ML +pipeline on the same ordered JPEG frames. Filenames must have numeric stems, +such as `00000.jpg` and `00001.jpg`. + +```bash +PYTHONPATH=coreml/video_tracking \ +python coreml/video_tracking/validate_video.py \ + --frames-dir notebooks/videos/bedroom \ + --models-dir coreml_models/video_tracking \ + --checkpoint checkpoints/edgetam.pt \ + --device mps \ + --point 210 350 1 \ + --max-frames 8 \ + --json coreml/video_tracking/validation.json +``` + +For every frame, the command reports binary-mask IoU, logit cosine similarity, +mean absolute error, and maximum absolute error. + +## Tests + +```bash +PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ +PYTHONPATH=coreml/video_tracking \ +python -m pytest coreml/video_tracking/tests -q +``` + +The tests cover the fixed model contracts, prompt scaling, explicit memory-bank +updates, masked attention, sequential prediction, and validation metrics. + +## Current scope + +- Single-object, forward-only tracking. +- One to four prompt tokens on the first frame. +- Fixed 1024-by-1024 model input. +- Fixed temporal memory capacity: one conditioning frame, six recent frames, + and sixteen pointers. +- No prompt correction after initialization, reverse propagation, + quantization, or bundled Swift wrapper. diff --git a/coreml/video_tracking/edgetam_coreml_video/__init__.py b/coreml/video_tracking/edgetam_coreml_video/__init__.py new file mode 100644 index 0000000..3bbac3d --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Core ML video-tracking components for EdgeTAM.""" + +from .metrics import TensorError, binary_mask_iou, cosine_similarity, tensor_error + +__all__ = [ + "TensorError", + "binary_mask_iou", + "cosine_similarity", + "tensor_error", +] diff --git a/coreml/video_tracking/edgetam_coreml_video/explicit_memory_bank.py b/coreml/video_tracking/edgetam_coreml_video/explicit_memory_bank.py new file mode 100644 index 0000000..7a7b7a1 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/explicit_memory_bank.py @@ -0,0 +1,204 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Explicit fixed-shape NumPy memory state for stateless EdgeTAM.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +MEMORY_SHAPE = (1, 512, 64) +TEMPORAL_POSITION_SHAPE = (7, 64) +POINTER_SHAPE = (1, 256) +RECENT_SLOTS = 6 +POINTER_SLOTS = 16 +SPATIAL_TOKENS_PER_SLOT = 512 +POINTER_TOKENS_PER_SLOT = 4 +TOTAL_SPATIAL_TOKENS = 7 * SPATIAL_TOKENS_PER_SLOT +TOTAL_ATTENTION_TOKENS = TOTAL_SPATIAL_TOKENS + ( + POINTER_SLOTS * POINTER_TOKENS_PER_SLOT +) +ROTARY_TOKENS_PER_SLOT = 256 +TOTAL_ROTARY_TOKENS = 7 * ROTARY_TOKENS_PER_SLOT + + +@dataclass(frozen=True) +class MemoryBankSnapshot: + """Small validation view that does not duplicate bank tensors.""" + + recent_count: int + pointer_count: int + is_initialized: bool + + +class ExplicitMemoryBank: + """Own one object's conditioning, recent-memory, and pointer history.""" + + def __init__(self) -> None: + self.conditioning_memory = np.zeros(MEMORY_SHAPE, dtype=np.float16) + self.conditioning_position = np.zeros(MEMORY_SHAPE, dtype=np.float16) + self.recent_memory = np.zeros( + (1, RECENT_SLOTS, 512, 64), + dtype=np.float16, + ) + self.recent_positions = np.zeros_like(self.recent_memory) + self.pointer_history = np.zeros( + (1, POINTER_SLOTS, 256), + dtype=np.float16, + ) + self.temporal_positions = np.zeros( + TEMPORAL_POSITION_SHAPE, + dtype=np.float16, + ) + self.recent_count = 0 + self.pointer_count = 0 + self.is_initialized = False + + @staticmethod + def _validated( + value: np.ndarray, + shape: tuple[int, ...], + name: str, + ) -> np.ndarray: + array = np.asarray(value) + if array.shape != shape: + raise ValueError(f"{name} must have shape {shape}") + if array.dtype != np.float16: + raise ValueError(f"{name} must use float16") + return array + + def seed( + self, + memory: np.ndarray, + memory_positions: np.ndarray, + temporal_positions: np.ndarray, + pointer: np.ndarray, + ) -> None: + """Replace all state with one prompted conditioning frame.""" + + memory = self._validated(memory, MEMORY_SHAPE, "memory") + memory_positions = self._validated( + memory_positions, + MEMORY_SHAPE, + "memory_positions", + ) + temporal_positions = self._validated( + temporal_positions, + TEMPORAL_POSITION_SHAPE, + "temporal_positions", + ) + pointer = self._validated(pointer, POINTER_SHAPE, "pointer") + + self.conditioning_memory[...] = memory + self.temporal_positions[...] = temporal_positions + self.conditioning_position[...] = ( + memory_positions + temporal_positions[6].reshape(1, 1, 64) + ) + self.recent_memory.fill(0) + self.recent_positions.fill(0) + self.pointer_history.fill(0) + self.pointer_history[:, 0] = pointer + self.recent_count = 0 + self.pointer_count = 1 + self.is_initialized = True + + def commit( + self, + memory: np.ndarray, + memory_positions: np.ndarray, + pointer: np.ndarray, + ) -> None: + """Append one successful propagation result.""" + + if not self.is_initialized: + raise RuntimeError("seed must be called before commit") + memory = self._validated(memory, MEMORY_SHAPE, "memory") + memory_positions = self._validated( + memory_positions, + MEMORY_SHAPE, + "memory_positions", + ) + pointer = self._validated(pointer, POINTER_SHAPE, "pointer") + + if self.recent_count < RECENT_SLOTS: + index = self.recent_count + else: + self.recent_memory[:, :-1] = self.recent_memory[:, 1:].copy() + self.recent_positions[:, :-1] = self.recent_positions[:, 1:].copy() + index = RECENT_SLOTS - 1 + self.recent_memory[:, index] = memory + self.recent_positions[:, index] = memory_positions + self.recent_count = min(self.recent_count + 1, RECENT_SLOTS) + + self.pointer_history[:, 2:] = self.pointer_history[:, 1:-1].copy() + self.pointer_history[:, 1] = pointer + self.pointer_count = min(self.pointer_count + 1, POINTER_SLOTS) + + def model_inputs(self) -> dict[str, np.ndarray]: + """Assemble fixed Float16 tensors for one propagation call.""" + + if not self.is_initialized: + raise RuntimeError("seed must be called before model_inputs") + + spatial_bank = np.zeros((1, 7, 512, 64), dtype=np.float16) + spatial_positions = np.zeros_like(spatial_bank) + spatial_bank[:, 0] = self.conditioning_memory + spatial_positions[:, 0] = self.conditioning_position + if self.recent_count: + spatial_bank[:, 1 : self.recent_count + 1] = self.recent_memory[ + :, : self.recent_count + ] + for slot in range(self.recent_count): + temporal_index = self.recent_count - 1 - slot + spatial_positions[:, slot + 1] = ( + self.recent_positions[:, slot] + + self.temporal_positions[temporal_index].reshape(1, 1, 64) + ) + + attention_bias = np.full( + (1, 1, 1, TOTAL_ATTENTION_TOKENS), + -10000, + dtype=np.float16, + ) + valid_spatial_tokens = ( + 1 + self.recent_count + ) * SPATIAL_TOKENS_PER_SLOT + attention_bias[..., :valid_spatial_tokens] = 0 + valid_pointer_tokens = self.pointer_count * POINTER_TOKENS_PER_SLOT + attention_bias[ + ..., + TOTAL_SPATIAL_TOKENS : TOTAL_SPATIAL_TOKENS + + valid_pointer_tokens, + ] = 0 + + rotary_weight = np.zeros( + (1, TOTAL_ROTARY_TOKENS), + dtype=np.float16, + ) + valid_rotary_tokens = min( + (1 + self.recent_count) * ROTARY_TOKENS_PER_SLOT, + TOTAL_ROTARY_TOKENS, + ) + rotary_weight[:, :valid_rotary_tokens] = 1 + + return { + "spatial_bank": spatial_bank, + "spatial_positions": spatial_positions, + "pointer_bank": self.pointer_history.copy(), + "attention_bias": attention_bias, + "rotary_weight": rotary_weight, + } + + def snapshot(self) -> MemoryBankSnapshot: + """Return scalar state for diagnostics and numerical validation.""" + + return MemoryBankSnapshot( + recent_count=self.recent_count, + pointer_count=self.pointer_count, + is_initialized=self.is_initialized, + ) diff --git a/coreml/video_tracking/edgetam_coreml_video/image_encoder.py b/coreml/video_tracking/edgetam_coreml_video/image_encoder.py new file mode 100644 index 0000000..fc02a62 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/image_encoder.py @@ -0,0 +1,106 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Video-specific EdgeTAM image features for initialization and propagation.""" + +from pathlib import Path +from typing import Any + +import torch +from torch import nn + + +class VideoImageEncoder(nn.Module): + """Expose both raw and no-memory-conditioned EdgeTAM image features.""" + + def __init__(self, model: Any) -> None: + super().__init__() + self.model = model + + def forward( + self, + image: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + backbone_fpn = self.model.forward_image(image)["backbone_fpn"] + raw_vision_features = backbone_fpn[2] + initial_vision_features = raw_vision_features.flatten(2).permute(2, 0, 1) + initial_vision_features = ( + initial_vision_features + self.model.no_mem_embed + ) + initial_vision_features = initial_vision_features.permute(1, 2, 0) + initial_vision_features = initial_vision_features.reshape_as( + raw_vision_features + ) + return ( + raw_vision_features, + initial_vision_features, + backbone_fpn[0], + backbone_fpn[1], + ) + + +class CoreMLVideoImageEncoder(VideoImageEncoder): + """Add the notebook's RGB normalization inside the exported model.""" + + def __init__(self, model: Any) -> None: + super().__init__(model) + self.register_buffer( + "pixel_mean", + torch.tensor([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1), + ) + self.register_buffer( + "pixel_std", + torch.tensor([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1), + ) + + def forward( + self, + image: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + normalized_image = (image - self.pixel_mean) / self.pixel_std + return super().forward(normalized_image) + + +def export_video_image_encoder(model: Any, output_path: Path) -> Path: + """Export the video image encoder as an iOS 18 Core ML package.""" + import coremltools as ct + + wrapper = CoreMLVideoImageEncoder(model).eval() + example_image = torch.randn(1, 3, 1024, 1024) + + with torch.inference_mode(): + traced_model = torch.jit.trace( + wrapper, + example_image, + check_trace=False, + ) + + coreml_model = ct.convert( + traced_model, + inputs=[ + ct.ImageType( + name="image", + shape=(1, 3, 1024, 1024), + scale=1 / 255.0, + bias=[0, 0, 0], + color_layout=ct.colorlayout.RGB, + ) + ], + outputs=[ + ct.TensorType(name="raw_vision_features"), + ct.TensorType(name="initial_vision_features"), + ct.TensorType(name="high_res_feature_0"), + ct.TensorType(name="high_res_feature_1"), + ], + minimum_deployment_target=ct.target.iOS18, + compute_units=ct.ComputeUnit.ALL, + convert_to="mlprogram", + ) + coreml_model.author = "EdgeTAM Contributors" + coreml_model.short_description = "EdgeTAM video image encoder" + coreml_model.version = "1.0" + coreml_model.save(str(output_path)) + return output_path diff --git a/coreml/video_tracking/edgetam_coreml_video/initializer.py b/coreml/video_tracking/edgetam_coreml_video/initializer.py new file mode 100644 index 0000000..0235624 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/initializer.py @@ -0,0 +1,226 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""First-frame prompt handling for an EdgeTAM video track.""" + +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from sam2.modeling.sam2_base import NO_OBJ_SCORE + + +class VideoInitializer(nn.Module): + """Run the prompted SAM heads and expose everything needed to seed memory.""" + + def __init__(self, model: Any) -> None: + super().__init__() + self.model = model + + def forward( + self, + initial_vision_features: torch.Tensor, + high_res_feature_0: torch.Tensor, + high_res_feature_1: torch.Tensor, + point_coords: torch.Tensor, + point_labels: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + prompt_encoder = self.model.sam_prompt_encoder + point_coords = torch.cat( + [ + point_coords + 0.5, + torch.zeros( + point_coords.shape[0], + 1, + 2, + dtype=point_coords.dtype, + device=point_coords.device, + ), + ], + dim=1, + ) + point_labels = torch.cat( + [ + point_labels.to(torch.int32), + -torch.ones( + point_labels.shape[0], + 1, + dtype=torch.int32, + device=point_labels.device, + ), + ], + dim=1, + ) + + sparse_embeddings = prompt_encoder.pe_layer.forward_with_coords( + point_coords, + prompt_encoder.input_image_size, + ) + padding_weight = (point_labels == -1).unsqueeze(-1).to( + sparse_embeddings.dtype + ) + not_a_point = prompt_encoder.not_a_point_embed.weight.reshape( + 1, 1, -1 + ) + sparse_embeddings = ( + padding_weight * not_a_point + + (1.0 - padding_weight) * sparse_embeddings + ) + for label, embedding in enumerate(prompt_encoder.point_embeddings): + sparse_embeddings = sparse_embeddings + ( + (point_labels == label).unsqueeze(-1).to(sparse_embeddings.dtype) + * embedding.weight.reshape(1, 1, -1) + ) + dense_embeddings = prompt_encoder.no_mask_embed.weight.reshape( + 1, -1, 1, 1 + ).expand( + point_coords.shape[0], + -1, + prompt_encoder.image_embedding_size[0], + prompt_encoder.image_embedding_size[1], + ) + + low_res_multimasks, iou_predictions, output_tokens, score = ( + self.model.sam_mask_decoder( + image_embeddings=initial_vision_features, + image_pe=prompt_encoder.get_dense_pe(), + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=True, + repeat_image=False, + high_res_features=[high_res_feature_0, high_res_feature_1], + ) + ) + if self.model.pred_obj_scores: + object_appearing = score > 0 + appearing_weight = object_appearing[:, None, None].to( + low_res_multimasks.dtype + ) + low_res_multimasks = ( + appearing_weight * low_res_multimasks + + (1.0 - appearing_weight) * NO_OBJ_SCORE + ) + + low_res_multimasks = low_res_multimasks.float() + high_res_multimasks = F.interpolate( + low_res_multimasks, + size=(self.model.image_size, self.model.image_size), + mode="bilinear", + align_corners=False, + ) + best_indices = torch.argmax(iou_predictions, dim=-1) + low_res_mask = torch.gather( + low_res_multimasks, + 1, + best_indices.reshape(-1, 1, 1, 1).expand( + -1, + 1, + low_res_multimasks.shape[-2], + low_res_multimasks.shape[-1], + ), + ) + high_res_mask = torch.gather( + high_res_multimasks, + 1, + best_indices.reshape(-1, 1, 1, 1).expand( + -1, + 1, + high_res_multimasks.shape[-2], + high_res_multimasks.shape[-1], + ), + ) + output_token = torch.gather( + output_tokens, + 1, + best_indices.reshape(-1, 1, 1).expand( + -1, + 1, + output_tokens.shape[-1], + ), + )[:, 0] + pointer = self.model.obj_ptr_proj(output_token) + if self.model.pred_obj_scores: + if self.model.soft_no_obj_ptr: + appearing_weight = score.sigmoid() + else: + appearing_weight = object_appearing.float() + if self.model.fixed_no_obj_ptr: + pointer = appearing_weight * pointer + pointer = pointer + (1 - appearing_weight) * self.model.no_obj_ptr + + best_iou = iou_predictions.max(dim=-1).values + return low_res_mask, high_res_mask, best_iou, pointer, score + + +def export_video_initializer(model: Any, output_path: Path) -> Path: + """Export first-frame prompt initialization as an iOS 18 Core ML package.""" + import coremltools as ct + + wrapper = VideoInitializer(model).eval() + example_inputs = ( + torch.randn(1, 256, 64, 64), + torch.randn(1, 32, 256, 256), + torch.randn(1, 64, 128, 128), + torch.tensor([[[512.0, 512.0]]]), + torch.tensor([[1]], dtype=torch.int32), + ) + + with torch.inference_mode(): + traced_model = torch.jit.trace( + wrapper, + example_inputs, + check_trace=False, + ) + + point_count = ct.RangeDim(lower_bound=1, upper_bound=4, default=1) + coreml_model = ct.convert( + traced_model, + inputs=[ + ct.TensorType( + name="initial_vision_features", + shape=(1, 256, 64, 64), + ), + ct.TensorType( + name="high_res_feature_0", + shape=(1, 32, 256, 256), + ), + ct.TensorType( + name="high_res_feature_1", + shape=(1, 64, 128, 128), + ), + ct.TensorType(name="point_coords", shape=(1, point_count, 2)), + ct.TensorType( + name="point_labels", + shape=(1, point_count), + dtype=np.int32, + ), + ], + outputs=[ + ct.TensorType(name="low_res_mask"), + ct.TensorType(name="high_res_mask"), + ct.TensorType(name="best_iou"), + ct.TensorType(name="object_pointer"), + ct.TensorType(name="object_score"), + ], + minimum_deployment_target=ct.target.iOS18, + compute_units=ct.ComputeUnit.ALL, + convert_to="mlprogram", + ) + coreml_model.author = "EdgeTAM Contributors" + coreml_model.short_description = "EdgeTAM video track initializer" + coreml_model.version = "1.0" + coreml_model.save(str(output_path)) + return output_path diff --git a/coreml/video_tracking/edgetam_coreml_video/load_model.py b/coreml/video_tracking/edgetam_coreml_video/load_model.py new file mode 100644 index 0000000..b9e7568 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/load_model.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Build the PyTorch EdgeTAM reference used by export and parity checks.""" + +from pathlib import Path +from typing import Any + +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from hydra.utils import instantiate +from omegaconf import OmegaConf + +from sam2.build_sam import _load_checkpoint + + +def load_reference_model( + config: Path, + checkpoint: Path, + device: str = "cpu", +) -> Any: + """Load an eval-mode EdgeTAM model from an explicit config and checkpoint.""" + + config = Path(config).expanduser().resolve() + checkpoint = Path(checkpoint).expanduser().resolve() + if not config.is_file(): + raise FileNotFoundError(f"EdgeTAM config not found: {config}") + if not checkpoint.is_file(): + raise FileNotFoundError(f"EdgeTAM checkpoint not found: {checkpoint}") + + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(config.parent), version_base=None): + cfg = compose(config_name=config.stem) + OmegaConf.resolve(cfg) + model = instantiate(cfg.model, _recursive_=True) + + _load_checkpoint(model, str(checkpoint)) + return model.to(device).eval() diff --git a/coreml/video_tracking/edgetam_coreml_video/masked_attention.py b/coreml/video_tracking/edgetam_coreml_video/masked_attention.py new file mode 100644 index 0000000..d301673 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/masked_attention.py @@ -0,0 +1,333 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Fixed-shape EdgeTAM memory attention with explicit validity masking.""" + +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +SPATIAL_SLOTS = 7 +SPATIAL_TOKENS_PER_SLOT = 512 +ROTARY_TOKENS_PER_SLOT = 256 +NON_ROTARY_TOKENS_PER_SLOT = SPATIAL_TOKENS_PER_SLOT - ROTARY_TOKENS_PER_SLOT +POINTER_SLOTS = 16 +POINTER_CHANNELS = 256 +MEMORY_CHANNELS = 64 +POINTER_TOKENS_PER_SLOT = POINTER_CHANNELS // MEMORY_CHANNELS +TOTAL_SPATIAL_TOKENS = SPATIAL_SLOTS * SPATIAL_TOKENS_PER_SLOT +TOTAL_POINTER_TOKENS = POINTER_SLOTS * POINTER_TOKENS_PER_SLOT +TOTAL_MEMORY_TOKENS = TOTAL_SPATIAL_TOKENS + TOTAL_POINTER_TOKENS +MAXIMUM_ROTARY_TOKENS = SPATIAL_SLOTS * ROTARY_TOKENS_PER_SLOT + + +def build_attention_controls( + spatial_valid: torch.Tensor, + pointer_valid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand compact validity into fixed numeric model inputs.""" + + spatial_tokens = spatial_valid.unsqueeze(-1).expand( + -1, + -1, + SPATIAL_TOKENS_PER_SLOT, + ).flatten(1, 2) + pointer_tokens = pointer_valid.unsqueeze(-1).expand( + -1, + -1, + POINTER_TOKENS_PER_SLOT, + ).flatten(1, 2) + key_valid = torch.cat([spatial_tokens, pointer_tokens], dim=1) + attention_bias = (1.0 - key_valid).reshape( + key_valid.shape[0], + 1, + 1, + TOTAL_MEMORY_TOKENS, + ) * -10000.0 + + valid_spatial_count = spatial_valid.sum(dim=-1, keepdim=True) + token_indices = torch.arange( + MAXIMUM_ROTARY_TOKENS, + device=spatial_valid.device, + ).reshape(1, -1) + rotary_weight = ( + token_indices < valid_spatial_count * ROTARY_TOKENS_PER_SLOT + ).to(spatial_valid.dtype) + return attention_bias, rotary_weight + + +def _apply_rotary_real( + tensor: torch.Tensor, + frequency_cos: torch.Tensor, + frequency_sin: torch.Tensor, + repeat_frequencies: int, +) -> torch.Tensor: + """Apply rotary encoding without Core ML-unsupported complex tensors.""" + if repeat_frequencies > 1: + frequency_cos = ( + frequency_cos.unsqueeze(0) + .expand(repeat_frequencies, -1, -1) + .flatten(0, 1) + ) + frequency_sin = ( + frequency_sin.unsqueeze(0) + .expand(repeat_frequencies, -1, -1) + .flatten(0, 1) + ) + + pairs = tensor.float().reshape(*tensor.shape[:-1], -1, 2) + real = pairs[..., 0] + imaginary = pairs[..., 1] + frequency_cos = frequency_cos.reshape(1, 1, tensor.shape[-2], -1) + frequency_sin = frequency_sin.reshape(1, 1, tensor.shape[-2], -1) + rotated_real = real * frequency_cos - imaginary * frequency_sin + rotated_imaginary = real * frequency_sin + imaginary * frequency_cos + rotated = torch.stack([rotated_real, rotated_imaginary], dim=-1) + return rotated.flatten(-2).type_as(tensor) + + +def _apply_rotary_spatial_slots( + tensor: torch.Tensor, + frequency_cos: torch.Tensor, + frequency_sin: torch.Tensor, + rotary_weight: torch.Tensor, +) -> torch.Tensor: + """Rotate the 2D-token half of every fixed spatial-memory slot.""" + + batch, heads, _, channels = tensor.shape + spatial = tensor[:, :, :TOTAL_SPATIAL_TOKENS].reshape( + batch, + heads, + SPATIAL_SLOTS, + SPATIAL_TOKENS_PER_SLOT, + channels, + ) + non_rotary = spatial[:, :, :, :NON_ROTARY_TOKENS_PER_SLOT] + rotary = spatial[:, :, :, NON_ROTARY_TOKENS_PER_SLOT:].reshape( + batch, + heads, + MAXIMUM_ROTARY_TOKENS, + channels, + ) + rotated = _apply_rotary_real( + rotary, + frequency_cos, + frequency_sin, + SPATIAL_SLOTS, + ) + weight = rotary_weight.reshape( + rotary_weight.shape[0], + 1, + MAXIMUM_ROTARY_TOKENS, + 1, + ).to(tensor.dtype) + rotary = weight * rotated + (1.0 - weight) * rotary + rotary = rotary.reshape( + batch, + heads, + SPATIAL_SLOTS, + ROTARY_TOKENS_PER_SLOT, + channels, + ) + spatial = torch.cat([non_rotary, rotary], dim=3).flatten(2, 3) + return torch.cat([spatial, tensor[:, :, TOTAL_SPATIAL_TOKENS:]], dim=2) + + +class MaskedMemoryAttention(nn.Module): + """Run EdgeTAM memory attention over fixed state while ignoring empty slots.""" + + def __init__(self, model: Any) -> None: + super().__init__() + self.memory_attention = model.memory_attention + for index, layer in enumerate(self.memory_attention.layers): + self_frequencies = layer.self_attn.compute_cis( + end_x=64, + end_y=64, + ) + cross_query_frequencies = layer.cross_attn_image.freqs_cis_q + cross_key_frequencies = layer.cross_attn_image.freqs_cis_k + for name, frequencies in ( + ("self", self_frequencies), + ("cross_query", cross_query_frequencies), + ("cross_key", cross_key_frequencies), + ): + self.register_buffer( + f"{name}_frequency_cos_{index}", + frequencies.real.float(), + persistent=False, + ) + self.register_buffer( + f"{name}_frequency_sin_{index}", + frequencies.imag.float(), + persistent=False, + ) + + def _self_attention( + self, + attention: nn.Module, + query: torch.Tensor, + value: torch.Tensor, + layer_index: int, + ) -> torch.Tensor: + key = attention.k_proj(query) + query = attention.q_proj(query) + value = attention.v_proj(value) + + query = attention._separate_heads(query, attention.num_heads) + key = attention._separate_heads(key, attention.num_heads) + value = attention._separate_heads(value, attention.num_heads) + + frequency_cos = getattr(self, f"self_frequency_cos_{layer_index}") + frequency_sin = getattr(self, f"self_frequency_sin_{layer_index}") + query = _apply_rotary_real(query, frequency_cos, frequency_sin, 1) + key = _apply_rotary_real(key, frequency_cos, frequency_sin, 1) + attended = F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + ) + attended = attention._recombine_heads(attended) + return attention.out_proj(attended) + + def _masked_cross_attention( + self, + attention: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_bias: torch.Tensor, + rotary_weight: torch.Tensor, + layer_index: int, + ) -> torch.Tensor: + query = attention.q_proj(query) + key = attention.k_proj(key) + value = attention.v_proj(value) + + query = attention._separate_heads(query, attention.num_heads) + key = attention._separate_heads(key, attention.num_heads) + value = attention._separate_heads(value, attention.num_heads) + + query_frequency_cos = getattr( + self, + f"cross_query_frequency_cos_{layer_index}", + ) + query_frequency_sin = getattr( + self, + f"cross_query_frequency_sin_{layer_index}", + ) + key_frequency_cos = getattr( + self, + f"cross_key_frequency_cos_{layer_index}", + ) + key_frequency_sin = getattr( + self, + f"cross_key_frequency_sin_{layer_index}", + ) + query = _apply_rotary_real( + query, + query_frequency_cos, + query_frequency_sin, + 1, + ) + + key = _apply_rotary_spatial_slots( + key, + key_frequency_cos, + key_frequency_sin, + rotary_weight, + ) + + attended = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_bias.to(query.dtype), + dropout_p=0.0, + ) + attended = attention._recombine_heads(attended) + return attention.out_proj(attended) + + def forward( + self, + current_features: torch.Tensor, + current_positions: torch.Tensor, + spatial_bank: torch.Tensor, + spatial_positions: torch.Tensor, + pointer_bank: torch.Tensor, + attention_bias: torch.Tensor, + rotary_weight: torch.Tensor, + ) -> torch.Tensor: + batch, channels, height, width = current_features.shape + current = current_features.flatten(2).permute(0, 2, 1) + current_position = current_positions.flatten(2).permute(0, 2, 1) + + spatial_memory = spatial_bank.flatten(1, 2) + spatial_position = spatial_positions.flatten(1, 2) + pointer_memory = pointer_bank.reshape( + batch, + POINTER_SLOTS, + POINTER_TOKENS_PER_SLOT, + MEMORY_CHANNELS, + ).flatten(1, 2) + pointer_position = torch.zeros_like(pointer_memory) + memory = torch.cat([spatial_memory, pointer_memory], dim=1) + memory_position = torch.cat( + [spatial_position, pointer_position], + dim=1, + ) + + output = current + if self.memory_attention.pos_enc_at_input: + output = output + 0.1 * current_position + + for layer_index, layer in enumerate(self.memory_attention.layers): + normalized = layer.norm1(output) + self_query = ( + normalized + current_position + if layer.pos_enc_at_attn + else normalized + ) + self_attended = self._self_attention( + layer.self_attn, + self_query, + normalized, + layer_index, + ) + output = output + layer.dropout1(self_attended) + + normalized = layer.norm2(output) + cross_query = ( + normalized + current_position + if layer.pos_enc_at_cross_attn_queries + else normalized + ) + cross_key = ( + memory + memory_position + if layer.pos_enc_at_cross_attn_keys + else memory + ) + cross_attended = self._masked_cross_attention( + layer.cross_attn_image, + cross_query, + cross_key, + memory, + attention_bias, + rotary_weight, + layer_index, + ) + output = output + layer.dropout2(cross_attended) + + normalized = layer.norm3(output) + feed_forward = layer.linear2( + layer.dropout(layer.activation(layer.linear1(normalized))) + ) + output = output + layer.dropout3(feed_forward) + + output = self.memory_attention.norm(output) + return output.permute(0, 2, 1).reshape(batch, channels, height, width) diff --git a/coreml/video_tracking/edgetam_coreml_video/memory_encoder.py b/coreml/video_tracking/edgetam_coreml_video/memory_encoder.py new file mode 100644 index 0000000..73f5733 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/memory_encoder.py @@ -0,0 +1,102 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""EdgeTAM mask-memory encoding including the 2D Spatial Perceiver.""" + +from pathlib import Path +from typing import Any + +import torch +from torch import nn + + +class VideoMemoryEncoder(nn.Module): + """Encode one predicted mask into EdgeTAM's compact spatial memory.""" + + def __init__( + self, + model: Any, + is_mask_from_points: bool = False, + ) -> None: + super().__init__() + self.model = model + self.is_mask_from_points = is_mask_from_points + self.register_buffer( + "temporal_positions", + model.maskmem_tpos_enc[:, 0, 0].detach().clone(), + persistent=False, + ) + + def forward( + self, + raw_vision_features: torch.Tensor, + high_res_mask: torch.Tensor, + object_score: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + current_vision_features = [ + raw_vision_features.flatten(2).permute(2, 0, 1) + ] + memory_features, memory_positions = self.model._encode_new_memory( + current_vision_feats=current_vision_features, + feat_sizes=[raw_vision_features.shape[-2:]], + pred_masks_high_res=high_res_mask, + object_score_logits=object_score, + is_mask_from_pts=self.is_mask_from_points, + ) + return ( + memory_features, + memory_positions[0], + self.temporal_positions + torch.zeros_like(self.temporal_positions), + ) + + +def export_video_memory_encoder(model: Any, output_path: Path) -> Path: + """Export the Spatial Perceiver memory encoder for iOS 18.""" + import coremltools as ct + + original_binarize = model.binarize_mask_from_pts_for_mem_enc + model.binarize_mask_from_pts_for_mem_enc = True + wrapper = VideoMemoryEncoder(model, is_mask_from_points=True).eval() + example_inputs = ( + torch.randn(1, 256, 64, 64), + torch.randn(1, 1, 1024, 1024), + torch.tensor([[2.0]]), + ) + + try: + with torch.inference_mode(): + traced_model = torch.jit.trace( + wrapper, + example_inputs, + check_trace=False, + ) + finally: + model.binarize_mask_from_pts_for_mem_enc = original_binarize + + coreml_model = ct.convert( + traced_model, + inputs=[ + ct.TensorType( + name="raw_vision_features", + shape=(1, 256, 64, 64), + ), + ct.TensorType(name="high_res_mask", shape=(1, 1, 1024, 1024)), + ct.TensorType(name="object_score", shape=(1, 1)), + ], + outputs=[ + ct.TensorType(name="memory_features"), + ct.TensorType(name="memory_positions"), + ct.TensorType(name="temporal_positions"), + ], + minimum_deployment_target=ct.target.iOS18, + compute_units=ct.ComputeUnit.ALL, + convert_to="mlprogram", + ) + coreml_model.author = "EdgeTAM Contributors" + coreml_model.short_description = "EdgeTAM Spatial Perceiver memory encoder" + coreml_model.version = "1.0" + coreml_model.save(str(output_path)) + return output_path diff --git a/coreml/video_tracking/edgetam_coreml_video/metrics.py b/coreml/video_tracking/edgetam_coreml_video/metrics.py new file mode 100644 index 0000000..d2458e9 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/metrics.py @@ -0,0 +1,53 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Numerical parity metrics shared by component and video validation tests.""" + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class TensorError: + """Absolute-error summary for two tensors.""" + + max_abs: float + mean_abs: float + + +def tensor_error(reference: np.ndarray, actual: np.ndarray) -> TensorError: + """Return maximum and mean absolute error after float32 conversion.""" + + delta = np.abs( + np.asarray(reference, dtype=np.float32) + - np.asarray(actual, dtype=np.float32) + ) + return TensorError( + max_abs=float(delta.max(initial=0.0)), + mean_abs=float(delta.mean()) if delta.size else 0.0, + ) + + +def binary_mask_iou(reference: np.ndarray, actual: np.ndarray) -> float: + """Return intersection-over-union for two binary masks.""" + + reference_mask = np.asarray(reference, dtype=bool) + actual_mask = np.asarray(actual, dtype=bool) + intersection = np.logical_and(reference_mask, actual_mask).sum() + union = np.logical_or(reference_mask, actual_mask).sum() + return 1.0 if union == 0 else float(intersection / union) + + +def cosine_similarity(reference: np.ndarray, actual: np.ndarray) -> float: + """Return cosine similarity for two flattened vectors.""" + + reference_vector = np.asarray(reference, dtype=np.float32).reshape(-1) + actual_vector = np.asarray(actual, dtype=np.float32).reshape(-1) + denominator = np.linalg.norm(reference_vector) * np.linalg.norm(actual_vector) + if denominator == 0: + return 1.0 if np.array_equal(reference_vector, actual_vector) else 0.0 + return float(np.dot(reference_vector, actual_vector) / denominator) diff --git a/coreml/video_tracking/edgetam_coreml_video/predictor.py b/coreml/video_tracking/edgetam_coreml_video/predictor.py new file mode 100644 index 0000000..b20d71c --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/predictor.py @@ -0,0 +1,216 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Python runtime for the stateless EdgeTAM Core ML video pipeline.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +from PIL import Image + +from .explicit_memory_bank import ExplicitMemoryBank, MemoryBankSnapshot + +MODEL_SIZE = 1024 +MAX_POINTS = 4 + + +@dataclass(frozen=True) +class TrackingResult: + """One object's mask logits and confidence for a video frame.""" + + mask_logits: np.ndarray + mask: np.ndarray + iou: float + object_score: float + + +def prepare_points( + points: Sequence[Sequence[float]] | np.ndarray, + labels: Sequence[int] | np.ndarray, + original_size: tuple[int, int], +) -> tuple[np.ndarray, np.ndarray]: + """Scale one to four real prompt tokens to EdgeTAM's 1024 input.""" + points_array = np.asarray(points, dtype=np.float32) + labels_array = np.asarray(labels, dtype=np.int32).reshape(-1) + if points_array.ndim != 2 or points_array.shape[1] != 2: + raise ValueError("points must have shape (N, 2)") + if len(points_array) != len(labels_array): + raise ValueError("points and labels must have the same length") + if not 1 <= len(points_array) <= MAX_POINTS: + raise ValueError(f"one to {MAX_POINTS} points are supported") + + width, height = original_size + if width <= 0 or height <= 0: + raise ValueError("original_size must contain positive width and height") + + scaled_points = points_array.copy() + scaled_points[:, 0] *= MODEL_SIZE / width + scaled_points[:, 1] *= MODEL_SIZE / height + return ( + scaled_points.reshape(1, -1, 2).astype(np.float16), + labels_array.reshape(1, -1), + ) + + +class CoreMLVideoPredictor: + """Track one object with stateless models and an explicit memory bank.""" + + def __init__( + self, + image_encoder: Any, + initializer: Any, + memory_encoder: Any, + propagator: Any, + ) -> None: + self.image_encoder = image_encoder + self.initializer = initializer + self.memory_encoder = memory_encoder + self.propagator = propagator + self._bank = ExplicitMemoryBank() + self._started = False + + @classmethod + def from_directory( + cls, + model_directory: Path, + compute_units: Any = None, + ) -> CoreMLVideoPredictor: + """Load the four packages produced by ``export_models.py``.""" + import coremltools as ct + + if compute_units is None: + compute_units = ct.ComputeUnit.ALL + names = { + "image_encoder": "EdgeTAMVideoImageEncoder.mlpackage", + "initializer": "EdgeTAMVideoInitializer.mlpackage", + "memory_encoder": "EdgeTAMVideoMemoryEncoder.mlpackage", + "propagator": "EdgeTAMVideoPropagator.mlpackage", + } + models = { + name: ct.models.MLModel( + str(model_directory / package_name), + compute_units=compute_units, + ) + for name, package_name in names.items() + } + return cls(**models) + + def reset(self) -> None: + """Discard the current object track and its explicit memory bank.""" + self._bank = ExplicitMemoryBank() + self._started = False + + def debug_bank_snapshot(self) -> MemoryBankSnapshot: + """Expose scalar bank counts for parity diagnostics.""" + return self._bank.snapshot() + + @staticmethod + def _prepare_frame( + frame: Image.Image | np.ndarray, + ) -> tuple[Image.Image, tuple[int, int]]: + if isinstance(frame, np.ndarray): + frame = Image.fromarray(frame) + if not isinstance(frame, Image.Image): + raise TypeError("frame must be a PIL image or an RGB numpy array") + frame = frame.convert("RGB") + original_size = frame.size + resized = frame.resize((MODEL_SIZE, MODEL_SIZE)) + return resized, original_size + + def _encode_frame( + self, + frame: Image.Image | np.ndarray, + ) -> tuple[dict[str, np.ndarray], tuple[int, int]]: + resized, original_size = self._prepare_frame(frame) + return self.image_encoder.predict({"image": resized}), original_size + + @staticmethod + def _result( + outputs: dict[str, np.ndarray], + original_size: tuple[int, int], + ) -> TrackingResult: + low_res_mask = np.asarray(outputs["low_res_mask"]) + mask_logits_256 = low_res_mask.reshape(256, 256).astype(np.float32) + resized_logits = Image.fromarray(mask_logits_256, mode="F").resize( + original_size, + resample=Image.Resampling.BILINEAR, + ) + mask_logits = np.asarray(resized_logits, dtype=np.float32) + return TrackingResult( + mask_logits=mask_logits, + mask=mask_logits > 0.0, + iou=float(np.asarray(outputs["best_iou"]).reshape(-1)[0]), + object_score=float( + np.asarray(outputs["object_score"]).reshape(-1)[0] + ), + ) + + def start_track( + self, + frame: Image.Image | np.ndarray, + points: Sequence[Sequence[float]] | np.ndarray, + labels: Sequence[int] | np.ndarray, + ) -> TrackingResult: + """Prompt the first frame and seed this predictor's memory bank.""" + features, original_size = self._encode_frame(frame) + point_coords, point_labels = prepare_points( + points, + labels, + original_size, + ) + seed = self.initializer.predict( + { + "initial_vision_features": features["initial_vision_features"], + "high_res_feature_0": features["high_res_feature_0"], + "high_res_feature_1": features["high_res_feature_1"], + "point_coords": point_coords, + "point_labels": point_labels, + } + ) + memory = self.memory_encoder.predict( + { + "raw_vision_features": features["raw_vision_features"], + "high_res_mask": seed["high_res_mask"], + "object_score": seed["object_score"], + } + ) + self._bank.seed( + np.asarray(memory["memory_features"], dtype=np.float16), + np.asarray(memory["memory_positions"], dtype=np.float16), + np.asarray(memory["temporal_positions"], dtype=np.float16), + np.asarray(seed["object_pointer"], dtype=np.float16), + ) + self._started = True + return self._result(seed, original_size) + + def track_frame( + self, + frame: Image.Image | np.ndarray, + ) -> TrackingResult: + """Propagate the current object mask onto the next video frame.""" + if not self._started: + raise RuntimeError("start_track must be called before track_frame") + features, original_size = self._encode_frame(frame) + propagator_inputs = { + "raw_vision_features": features["raw_vision_features"], + "high_res_feature_0": features["high_res_feature_0"], + "high_res_feature_1": features["high_res_feature_1"], + **self._bank.model_inputs(), + } + outputs = self.propagator.predict(propagator_inputs) + result = self._result(outputs, original_size) + next_memory = np.asarray(outputs["memory_features"], dtype=np.float16) + next_positions = np.asarray( + outputs["memory_positions"], + dtype=np.float16, + ) + next_pointer = np.asarray(outputs["object_pointer"], dtype=np.float16) + self._bank.commit(next_memory, next_positions, next_pointer) + return result diff --git a/coreml/video_tracking/edgetam_coreml_video/propagator.py b/coreml/video_tracking/edgetam_coreml_video/propagator.py new file mode 100644 index 0000000..25d9180 --- /dev/null +++ b/coreml/video_tracking/edgetam_coreml_video/propagator.py @@ -0,0 +1,186 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""One prompt-free EdgeTAM video propagation step.""" + +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn + +from .initializer import VideoInitializer +from .masked_attention import MaskedMemoryAttention +from .memory_encoder import VideoMemoryEncoder + + +class VideoPropagator(nn.Module): + """Condition one frame on memory, decode its mask, and encode new memory.""" + + def __init__(self, model: Any) -> None: + super().__init__() + self.memory_attention = MaskedMemoryAttention(model) + self.prompt_free_head = VideoInitializer(model) + self.memory_encoder = VideoMemoryEncoder(model) + current_positions = model.image_encoder.neck.position_encoding( + torch.zeros(1, model.hidden_dim, 64, 64) + ) + self.register_buffer("current_positions", current_positions) + + def forward( + self, + raw_vision_features: torch.Tensor, + high_res_feature_0: torch.Tensor, + high_res_feature_1: torch.Tensor, + spatial_bank: torch.Tensor, + spatial_positions: torch.Tensor, + pointer_bank: torch.Tensor, + attention_bias: torch.Tensor, + rotary_weight: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + fused_features = self.memory_attention( + raw_vision_features, + self.current_positions, + spatial_bank, + spatial_positions, + pointer_bank, + attention_bias, + rotary_weight, + ) + batch = raw_vision_features.shape[0] + empty_point_coords = torch.zeros( + batch, + 1, + 2, + dtype=raw_vision_features.dtype, + device=raw_vision_features.device, + ) + empty_point_labels = -torch.ones( + batch, + 1, + dtype=torch.int32, + device=raw_vision_features.device, + ) + low_res_mask, high_res_mask, best_iou, pointer, score = ( + self.prompt_free_head( + fused_features, + high_res_feature_0, + high_res_feature_1, + empty_point_coords, + empty_point_labels, + ) + ) + memory, memory_position, _ = self.memory_encoder( + raw_vision_features, + high_res_mask, + score, + ) + return ( + low_res_mask, + high_res_mask, + best_iou, + pointer, + score, + memory, + memory_position, + ) + + +def export_video_propagator(model: Any, output_path: Path) -> Path: + """Export fixed-shape prompt-free EdgeTAM propagation for iOS 18.""" + import coremltools as ct + + propagator = VideoPropagator(model).eval() + example_inputs = ( + torch.randn(1, 256, 64, 64), + torch.randn(1, 32, 256, 256), + torch.randn(1, 64, 128, 128), + torch.randn(1, 7, 512, 64), + torch.randn(1, 7, 512, 64), + torch.randn(1, 16, 256), + torch.zeros(1, 1, 1, 3648), + torch.ones(1, 1792), + ) + + with torch.inference_mode(): + traced_model = torch.jit.trace( + propagator, + example_inputs, + check_trace=False, + ) + + coreml_model = ct.convert( + traced_model, + inputs=[ + ct.TensorType( + name="raw_vision_features", + shape=(1, 256, 64, 64), + dtype=np.float16, + ), + ct.TensorType( + name="high_res_feature_0", + shape=(1, 32, 256, 256), + dtype=np.float16, + ), + ct.TensorType( + name="high_res_feature_1", + shape=(1, 64, 128, 128), + dtype=np.float16, + ), + ct.TensorType( + name="spatial_bank", + shape=(1, 7, 512, 64), + dtype=np.float16, + ), + ct.TensorType( + name="spatial_positions", + shape=(1, 7, 512, 64), + dtype=np.float16, + ), + ct.TensorType( + name="pointer_bank", + shape=(1, 16, 256), + dtype=np.float16, + ), + ct.TensorType( + name="attention_bias", + shape=(1, 1, 1, 3648), + dtype=np.float16, + ), + ct.TensorType( + name="rotary_weight", + shape=(1, 1792), + dtype=np.float16, + ), + ], + outputs=[ + ct.TensorType(name="low_res_mask"), + ct.TensorType(name="high_res_mask"), + ct.TensorType(name="best_iou"), + ct.TensorType(name="object_pointer"), + ct.TensorType(name="object_score"), + ct.TensorType(name="memory_features"), + ct.TensorType(name="memory_positions"), + ], + minimum_deployment_target=ct.target.iOS18, + compute_units=ct.ComputeUnit.ALL, + compute_precision=ct.precision.FLOAT16, + convert_to="mlprogram", + ) + coreml_model.author = "EdgeTAM Contributors" + coreml_model.short_description = "Stateless EdgeTAM video propagation" + coreml_model.version = "1.0" + coreml_model.save(str(output_path)) + return output_path diff --git a/coreml/video_tracking/export_models.py b/coreml/video_tracking/export_models.py new file mode 100644 index 0000000..631c656 --- /dev/null +++ b/coreml/video_tracking/export_models.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Export the EdgeTAM video tracking pipeline as iOS 18 Core ML packages.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from edgetam_coreml_video.image_encoder import export_video_image_encoder +from edgetam_coreml_video.initializer import export_video_initializer +from edgetam_coreml_video.load_model import load_reference_model +from edgetam_coreml_video.memory_encoder import export_video_memory_encoder +from edgetam_coreml_video.propagator import export_video_propagator + + +def model_output_paths(output_directory: Path) -> dict[str, Path]: + """Return the stable package names used by Python and iOS clients.""" + return { + "image_encoder": output_directory / "EdgeTAMVideoImageEncoder.mlpackage", + "initializer": output_directory / "EdgeTAMVideoInitializer.mlpackage", + "memory_encoder": output_directory / "EdgeTAMVideoMemoryEncoder.mlpackage", + "propagator": output_directory / "EdgeTAMVideoPropagator.mlpackage", + } + + +def export_all_models( + config: Path, + checkpoint: Path, + output_directory: Path, + device: str = "cpu", +) -> dict[str, Path]: + """Load EdgeTAM once and export every video tracking package.""" + output_directory.mkdir(parents=True, exist_ok=True) + paths = model_output_paths(output_directory) + model = load_reference_model(config, checkpoint, device=device) + + exporters = ( + ("image_encoder", export_video_image_encoder), + ("initializer", export_video_initializer), + ("memory_encoder", export_video_memory_encoder), + ("propagator", export_video_propagator), + ) + for name, exporter in exporters: + print(f"Exporting {name} -> {paths[name]}") + exporter(model, paths[name]) + return paths + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + type=Path, + default=REPO_ROOT / "sam2/configs/edgetam.yaml", + ) + parser.add_argument( + "--checkpoint", + type=Path, + default=REPO_ROOT / "checkpoints/edgetam.pt", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path(__file__).resolve().parent / "models", + ) + parser.add_argument("--device", default="cpu") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + exported = export_all_models( + config=args.config, + checkpoint=args.checkpoint, + output_directory=args.output_dir, + device=args.device, + ) + print("Exported Core ML video packages:") + for path in exported.values(): + print(path) + + +if __name__ == "__main__": + main() diff --git a/coreml/video_tracking/tests/conftest.py b/coreml/video_tracking/tests/conftest.py new file mode 100644 index 0000000..15707c6 --- /dev/null +++ b/coreml/video_tracking/tests/conftest.py @@ -0,0 +1,21 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import pytest + +from edgetam_coreml_video.load_model import load_reference_model + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +CONFIG_PATH = REPOSITORY_ROOT / "sam2" / "configs" / "edgetam.yaml" +CHECKPOINT_PATH = REPOSITORY_ROOT / "checkpoints" / "edgetam.pt" + + +@pytest.fixture(scope="session") +def reference_model(): + return load_reference_model(CONFIG_PATH, CHECKPOINT_PATH) diff --git a/coreml/video_tracking/tests/test_explicit_memory_bank.py b/coreml/video_tracking/tests/test_explicit_memory_bank.py new file mode 100644 index 0000000..62b1c96 --- /dev/null +++ b/coreml/video_tracking/tests/test_explicit_memory_bank.py @@ -0,0 +1,98 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import pytest + +from edgetam_coreml_video.explicit_memory_bank import ExplicitMemoryBank + + +def _memory(value: float) -> np.ndarray: + return np.full((1, 512, 64), value, dtype=np.float16) + + +def _position(value: float) -> np.ndarray: + return np.full((1, 512, 64), value, dtype=np.float16) + + +def _pointer(value: float) -> np.ndarray: + return np.full((1, 256), value, dtype=np.float16) + + +def _temporal_positions() -> np.ndarray: + return np.broadcast_to( + np.arange(7, dtype=np.float16).reshape(7, 1), + (7, 64), + ).copy() + + +def test_explicit_memory_bank_rolls_recent_memory_and_pointer_history(): + bank = ExplicitMemoryBank() + bank.seed( + _memory(10), + _position(100), + _temporal_positions(), + _pointer(10), + ) + for value in range(1, 8): + bank.commit( + _memory(float(value)), + _position(float(value * 10)), + _pointer(float(value)), + ) + + inputs = bank.model_inputs() + snapshot = bank.snapshot() + + assert inputs["spatial_bank"][:, :, 0, 0].tolist() == [ + [10.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0] + ] + assert inputs["spatial_positions"][:, :, 0, 0].tolist() == [ + [106.0, 25.0, 34.0, 43.0, 52.0, 61.0, 70.0] + ] + assert inputs["pointer_bank"][:, :8, 0].tolist() == [ + [10.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0] + ] + assert np.all(inputs["attention_bias"][..., :3616] == 0) + assert np.all(inputs["attention_bias"][..., 3616:] == -10000) + assert np.all(inputs["rotary_weight"] == 1) + assert snapshot.recent_count == 6 + assert snapshot.pointer_count == 8 + assert snapshot.is_initialized is True + + +def test_explicit_memory_bank_warmup_masks_unused_fixed_slots(): + bank = ExplicitMemoryBank() + bank.seed( + _memory(10), + _position(100), + _temporal_positions(), + _pointer(10), + ) + bank.commit(_memory(1), _position(10), _pointer(1)) + + inputs = bank.model_inputs() + + assert np.all(inputs["attention_bias"][..., :1024] == 0) + assert np.all(inputs["attention_bias"][..., 1024:3584] == -10000) + assert np.all(inputs["attention_bias"][..., 3584:3592] == 0) + assert np.all(inputs["attention_bias"][..., 3592:] == -10000) + assert np.all(inputs["rotary_weight"][:, :512] == 1) + assert np.all(inputs["rotary_weight"][:, 512:] == 0) + assert inputs["spatial_positions"][0, 0, 0, 0] == 106 + assert inputs["spatial_positions"][0, 1, 0, 0] == 10 + + +def test_explicit_memory_bank_rejects_wrong_seed_shapes(): + bank = ExplicitMemoryBank() + + with pytest.raises(ValueError, match="memory must have shape"): + bank.seed( + np.zeros((1, 511, 64), dtype=np.float16), + _position(0), + _temporal_positions(), + _pointer(0), + ) diff --git a/coreml/video_tracking/tests/test_export_models.py b/coreml/video_tracking/tests/test_export_models.py new file mode 100644 index 0000000..762ac47 --- /dev/null +++ b/coreml/video_tracking/tests/test_export_models.py @@ -0,0 +1,18 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +from export_models import model_output_paths + + +def test_model_output_paths_are_separate_coreml_packages(tmp_path: Path): + assert model_output_paths(tmp_path) == { + "image_encoder": tmp_path / "EdgeTAMVideoImageEncoder.mlpackage", + "initializer": tmp_path / "EdgeTAMVideoInitializer.mlpackage", + "memory_encoder": tmp_path / "EdgeTAMVideoMemoryEncoder.mlpackage", + "propagator": tmp_path / "EdgeTAMVideoPropagator.mlpackage", + } diff --git a/coreml/video_tracking/tests/test_load_model.py b/coreml/video_tracking/tests/test_load_model.py new file mode 100644 index 0000000..e4aa0a8 --- /dev/null +++ b/coreml/video_tracking/tests/test_load_model.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import pytest + +from edgetam_coreml_video.load_model import load_reference_model + + +def test_load_reference_model_rejects_missing_config(tmp_path: Path): + checkpoint = tmp_path / "edgetam.pt" + checkpoint.touch() + + with pytest.raises(FileNotFoundError, match="config"): + load_reference_model(tmp_path / "missing.yaml", checkpoint) + + +def test_load_reference_model_rejects_missing_checkpoint(tmp_path: Path): + config = tmp_path / "edgetam.yaml" + config.touch() + + with pytest.raises(FileNotFoundError, match="checkpoint"): + load_reference_model(config, tmp_path / "missing.pt") diff --git a/coreml/video_tracking/tests/test_masked_memory_attention.py b/coreml/video_tracking/tests/test_masked_memory_attention.py new file mode 100644 index 0000000..dec4b2d --- /dev/null +++ b/coreml/video_tracking/tests/test_masked_memory_attention.py @@ -0,0 +1,149 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import inspect + +import torch + +from edgetam_coreml_video.masked_attention import ( + MaskedMemoryAttention, + _apply_rotary_spatial_slots, + build_attention_controls, +) +from sam2.modeling.position_encoding import apply_rotary_enc_v2 + + +def test_attention_controls_expand_slot_validity_without_model_graph_logic(): + spatial_valid = torch.tensor( + [[1, 1, 0, 0, 0, 0, 0]], + dtype=torch.float16, + ) + pointer_valid = torch.tensor( + [[1, 1] + [0] * 14], + dtype=torch.float16, + ) + + attention_bias, rotary_weight = build_attention_controls( + spatial_valid, + pointer_valid, + ) + + assert attention_bias.shape == (1, 1, 1, 3648) + assert rotary_weight.shape == (1, 1792) + assert attention_bias.dtype == torch.float16 + assert rotary_weight.dtype == torch.float16 + assert torch.all(attention_bias[..., :1024] == 0) + assert torch.all(attention_bias[..., 1024:3584] == -10000) + assert torch.all(attention_bias[..., 3584:3592] == 0) + assert torch.all(attention_bias[..., 3592:] == -10000) + assert torch.all(rotary_weight[:, :512] == 1) + assert torch.all(rotary_weight[:, 512:] == 0) + + +def test_converted_attention_path_does_not_construct_validity_from_counts(): + source = inspect.getsource(MaskedMemoryAttention.forward) + source += inspect.getsource(MaskedMemoryAttention._masked_cross_attention) + + assert "spatial_valid" not in source + assert "pointer_valid" not in source + assert "torch.arange" not in source + assert ".sum(" not in source + + +def test_rotary_encoding_targets_last_half_of_each_spatial_slot(reference_model): + torch.manual_seed(31) + attention = reference_model.memory_attention.layers[0].cross_attn_image + key = torch.randn(1, 1, 3648, 256) + spatial_valid = torch.tensor( + [[1, 1, 0, 0, 0, 0, 0]], + dtype=torch.float32, + ) + pointer_valid = torch.tensor( + [[1, 1] + [0] * 14], + dtype=torch.float32, + ) + _, rotary_weight = build_attention_controls(spatial_valid, pointer_valid) + + actual = _apply_rotary_spatial_slots( + key, + attention.freqs_cis_k.real.float(), + attention.freqs_cis_k.imag.float(), + rotary_weight, + ) + expected = key.clone() + expected[:, :, :1024] = apply_rotary_enc_v2( + key[:, :, :1024], + attention.freqs_cis_k, + repeat_freqs=2, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def _pointer_tokens(pointer_bank: torch.Tensor, count: int) -> torch.Tensor: + pointers = pointer_bank[:, :count] + batch, pointer_count, channels = pointers.shape + return pointers.reshape(batch, pointer_count, channels // 64, 64).flatten(1, 2) + + +def test_masked_memory_attention_matches_compact_reference_memory(reference_model): + torch.manual_seed(3) + spatial_count = 3 + pointer_count = 2 + current = torch.randn(1, 256, 64, 64) + current_position = reference_model.image_encoder.neck.position_encoding(current) + spatial_bank = torch.randn(1, 7, 512, 64) + spatial_positions = torch.randn(1, 7, 512, 64) + spatial_valid = torch.tensor([[1, 1, 1, 0, 0, 0, 0]], dtype=torch.float32) + pointer_bank = torch.randn(1, 16, 256) + pointer_valid = torch.tensor( + [[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], + dtype=torch.float32, + ) + attention_bias, rotary_weight = build_attention_controls( + spatial_valid, + pointer_valid, + ) + wrapper = MaskedMemoryAttention(reference_model).eval() + + compact_spatial = spatial_bank[:, :spatial_count].flatten(1, 2) + compact_positions = spatial_positions[:, :spatial_count].flatten(1, 2) + compact_pointers = _pointer_tokens(pointer_bank, pointer_count) + compact_pointer_positions = torch.zeros_like(compact_pointers) + compact_memory = torch.cat([compact_spatial, compact_pointers], dim=1) + compact_memory_positions = torch.cat( + [compact_positions, compact_pointer_positions], dim=1 + ) + + with torch.inference_mode(): + actual = wrapper( + current, + current_position, + spatial_bank, + spatial_positions, + pointer_bank, + attention_bias, + rotary_weight, + ) + expected = reference_model.memory_attention( + curr=current.flatten(2).permute(2, 0, 1), + curr_pos=current_position.flatten(2).permute(2, 0, 1), + memory=compact_memory.permute(1, 0, 2), + memory_pos=compact_memory_positions.permute(1, 0, 2), + num_obj_ptr_tokens=pointer_count * 4, + num_spatial_mem=spatial_count, + ) + expected = expected.permute(1, 2, 0).reshape_as(current) + + absolute_error = (actual - expected).abs() + cosine = torch.nn.functional.cosine_similarity( + actual.flatten(), + expected.flatten(), + dim=0, + ) + assert absolute_error.mean().item() < 0.02 + assert absolute_error.max().item() < 0.5 + assert cosine.item() > 0.999 diff --git a/coreml/video_tracking/tests/test_metrics.py b/coreml/video_tracking/tests/test_metrics.py new file mode 100644 index 0000000..84e74c9 --- /dev/null +++ b/coreml/video_tracking/tests/test_metrics.py @@ -0,0 +1,41 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np + +from edgetam_coreml_video.metrics import ( + binary_mask_iou, + cosine_similarity, + tensor_error, +) + + +def test_binary_mask_iou_is_one_for_identical_masks(): + mask = np.array([[True, False], [True, True]]) + + assert binary_mask_iou(mask, mask) == 1.0 + + +def test_binary_mask_iou_is_one_for_two_empty_masks(): + mask = np.zeros((2, 2), dtype=bool) + + assert binary_mask_iou(mask, mask) == 1.0 + + +def test_tensor_error_reports_maximum_and_mean_absolute_error(): + error = tensor_error( + np.array([0.0, 2.0], dtype=np.float32), + np.array([1.0, 2.0], dtype=np.float32), + ) + + assert error.max_abs == 1.0 + assert error.mean_abs == 0.5 + + +def test_cosine_similarity_is_one_for_equal_vectors(): + vector = np.array([1.0, 2.0, 3.0], dtype=np.float32) + + assert np.isclose(cosine_similarity(vector, vector), 1.0) diff --git a/coreml/video_tracking/tests/test_predictor.py b/coreml/video_tracking/tests/test_predictor.py new file mode 100644 index 0000000..272d0e2 --- /dev/null +++ b/coreml/video_tracking/tests/test_predictor.py @@ -0,0 +1,146 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import numpy as np +from PIL import Image + +from edgetam_coreml_video.predictor import ( + CoreMLVideoPredictor, + prepare_points, +) + + +class FakeModel: + def __init__(self, output): + self.output = output + self.calls = [] + + def predict(self, inputs, state=None): + self.calls.append((inputs, state)) + if callable(self.output): + return self.output(inputs) + return self.output + + +class FakePropagator(FakeModel): + def __init__(self): + super().__init__(self._output) + self.step = 0 + + def _output(self, inputs): + assert "mode" not in inputs + assert not any(name.startswith("seed_") for name in inputs) + self.step += 1 + propagated = float(self.step + 1) + return { + "low_res_mask": np.full((1, 1, 256, 256), propagated), + "high_res_mask": np.full((1, 1, 1024, 1024), propagated + 10), + "best_iou": np.array([0.75]), + "object_pointer": np.full((1, 256), propagated, dtype=np.float16), + "object_score": np.array([[1.0]]), + "memory_features": np.full( + (1, 512, 64), + propagated, + dtype=np.float16, + ), + "memory_positions": np.full( + (1, 512, 64), + propagated * 10, + dtype=np.float16, + ), + } + + +def _image_features(): + return { + "raw_vision_features": np.zeros((1, 256, 64, 64)), + "initial_vision_features": np.zeros((1, 256, 64, 64)), + "high_res_feature_0": np.zeros((1, 32, 256, 256)), + "high_res_feature_1": np.zeros((1, 64, 128, 128)), + } + + +def test_prepare_points_scales_to_1024_without_transformer_visible_padding(): + coords, labels = prepare_points( + points=np.array([[160.0, 120.0]]), + labels=np.array([1]), + original_size=(640, 480), + ) + + np.testing.assert_allclose(coords[0, 0], [256.0, 256.0]) + assert coords.shape == (1, 1, 2) + assert labels.tolist() == [[1]] + assert labels.dtype == np.int32 + + +def test_prepare_frame_matches_video_predictor_pillow_resize(): + pixels = np.array( + [ + [[255, 0, 0], [0, 255, 0]], + [[0, 0, 255], [255, 255, 255]], + ], + dtype=np.uint8, + ) + frame = Image.fromarray(pixels, mode="RGB") + + actual, original_size = CoreMLVideoPredictor._prepare_frame(frame) + expected = frame.resize((1024, 1024)) + + assert original_size == (2, 2) + np.testing.assert_array_equal(np.asarray(actual), np.asarray(expected)) + + +def test_predictor_seeds_and_propagates_with_explicit_stateless_bank( + tmp_path: Path, +): + image_encoder = FakeModel(_image_features()) + initializer = FakeModel( + { + "low_res_mask": np.ones((1, 1, 256, 256)), + "high_res_mask": np.ones((1, 1, 1024, 1024)), + "best_iou": np.array([0.8]), + "object_pointer": np.full((1, 256), 10, dtype=np.float16), + "object_score": np.array([[1.5]]), + } + ) + memory_encoder = FakeModel( + { + "memory_features": np.zeros((1, 512, 64)), + "memory_positions": np.zeros((1, 512, 64)), + "temporal_positions": np.zeros((7, 64), dtype=np.float16), + } + ) + propagator = FakePropagator() + predictor = CoreMLVideoPredictor( + image_encoder=image_encoder, + initializer=initializer, + memory_encoder=memory_encoder, + propagator=propagator, + ) + frame = Image.new("RGB", (640, 480)) + + seed = predictor.start_track(frame, [[160.0, 120.0]], [1]) + first_propagated = predictor.track_frame(frame) + second_propagated = predictor.track_frame(frame) + snapshot = predictor.debug_bank_snapshot() + + assert seed.mask_logits.shape == (480, 640) + assert first_propagated.mask_logits.shape == (480, 640) + assert second_propagated.mask_logits.shape == (480, 640) + assert np.all(seed.mask_logits == 1.0) + assert np.all(first_propagated.mask_logits == 2.0) + assert np.all(second_propagated.mask_logits == 3.0) + assert np.all(seed.mask) + assert np.all(first_propagated.mask) + assert np.all(second_propagated.mask) + assert len(propagator.calls) == 2 + assert all(state is None for _, state in propagator.calls) + assert "attention_bias" in propagator.calls[0][0] + assert "rotary_weight" in propagator.calls[0][0] + assert snapshot.recent_count == 2 + assert snapshot.pointer_count == 3 diff --git a/coreml/video_tracking/tests/test_validate_video.py b/coreml/video_tracking/tests/test_validate_video.py new file mode 100644 index 0000000..b42168e --- /dev/null +++ b/coreml/video_tracking/tests/test_validate_video.py @@ -0,0 +1,187 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import json +from argparse import Namespace +from types import SimpleNamespace + +import numpy as np +from PIL import Image + +import validate_video as validate_video_module +from validate_video import ( + _compare_coreml_frames, + compare_mask_logits, + parse_args, + run_validation, +) + + +def test_compare_mask_logits_reports_tensor_error_and_binary_iou(): + reference = np.array([[1.0, 1.0], [-1.0, -1.0]]) + candidate = np.array([[1.0, -1.0], [-1.0, -1.0]]) + + comparison = compare_mask_logits(reference, candidate) + + assert comparison.max_abs == 2.0 + assert comparison.mean_abs == 0.5 + assert comparison.mask_iou == 0.5 + + +class _Result: + def __init__(self, mask_logits, iou=0.75, object_score=1.25): + self.mask_logits = mask_logits + self.iou = iou + self.object_score = object_score + + +class _Predictor: + def __init__(self, logits): + self.logits = iter(logits) + self.start_calls = 0 + self.track_calls = 0 + self.recent_count = 0 + self.pointer_count = 0 + + def start_track(self, frame, points, labels): + self.start_calls += 1 + self.recent_count = 0 + self.pointer_count = 1 + return _Result(next(self.logits)) + + def track_frame(self, frame): + self.track_calls += 1 + self.recent_count = min(self.recent_count + 1, 6) + self.pointer_count = min(self.pointer_count + 1, 16) + return _Result(next(self.logits)) + + def debug_bank_snapshot(self): + return SimpleNamespace( + recent_count=self.recent_count, + pointer_count=self.pointer_count, + ) + + +def test_compare_coreml_frames_returns_one_row_per_frame(tmp_path): + frame_paths = [] + for frame_index in range(2): + frame_path = tmp_path / f"{frame_index:05d}.jpg" + Image.new("RGB", (8, 6), (frame_index * 20, 0, 0)).save(frame_path) + frame_paths.append(frame_path) + reference = { + 0: np.array([[1.0, -1.0]], dtype=np.float32), + 1: np.array([[1.0, 1.0]], dtype=np.float32), + } + predictor = _Predictor( + [ + np.array([[1.0, -1.0]], dtype=np.float32), + np.array([[1.0, -1.0]], dtype=np.float32), + ] + ) + + rows = _compare_coreml_frames( + frame_paths, + reference, + predictor, + np.array([[2.0, 3.0]], dtype=np.float32), + np.array([1], dtype=np.int32), + ) + + assert predictor.start_calls == 1 + assert predictor.track_calls == 1 + assert [row["frame"] for row in rows] == [0, 1] + assert rows[0]["mask_iou"] == 1.0 + assert rows[1]["mask_iou"] == 0.5 + assert rows[0]["predicted_iou"] == 0.75 + assert rows[0]["object_score"] == 1.25 + assert rows[0]["recent_count"] == 0 + assert rows[0]["pointer_count"] == 1 + assert rows[1]["recent_count"] == 1 + assert rows[1]["pointer_count"] == 2 + + +def test_compare_coreml_frames_reports_saturated_bank_counts(tmp_path): + frame_paths = [] + logits = [] + reference = {} + for frame_index in range(20): + frame_path = tmp_path / f"{frame_index:05d}.jpg" + Image.new("RGB", (2, 2)).save(frame_path) + frame_paths.append(frame_path) + frame_logits = np.ones((2, 2), dtype=np.float32) + logits.append(frame_logits) + reference[frame_index] = frame_logits + + rows = _compare_coreml_frames( + frame_paths, + reference, + _Predictor(logits), + np.array([[1.0, 1.0]], dtype=np.float32), + np.array([1], dtype=np.int32), + ) + + assert [row["recent_count"] for row in rows] == [ + 0, 1, 2, 3, 4, 5, 6, 6, 6, 6, + 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, + ] + assert [row["pointer_count"] for row in rows] == [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, + 11, 12, 13, 14, 15, 16, 16, 16, 16, 16, + ] + + +def test_parse_args_accepts_json_output(monkeypatch, tmp_path): + json_path = tmp_path / "validation.json" + monkeypatch.setattr( + "sys.argv", + [ + "validate_video.py", + "--frames-dir", + "frames", + "--json", + str(json_path), + ], + ) + + args = parse_args() + + assert args.json == json_path + + +def test_run_validation_writes_requested_json(monkeypatch, tmp_path): + rows = [ + { + "frame": 0, + "mask_iou": 0.99, + "cosine": 0.999, + "mean_abs": 0.1, + "max_abs": 1.0, + } + ] + captured = {} + + def fake_validate_video(**kwargs): + captured.update(kwargs) + return rows + + monkeypatch.setattr(validate_video_module, "validate_video", fake_validate_video) + json_path = tmp_path / "validation.json" + args = Namespace( + frames_dir=tmp_path / "frames", + models_dir=tmp_path / "models", + checkpoint=tmp_path / "edgetam.pt", + config_name="edgetam.yaml", + device="mps", + max_frames=8, + point=[[210.0, 350.0, 1.0]], + json=json_path, + ) + + actual_rows = run_validation(args) + + assert actual_rows == rows + assert captured["max_frames"] == 8 + assert json.loads(json_path.read_text()) == rows diff --git a/coreml/video_tracking/tests/test_video_image_encoder.py b/coreml/video_tracking/tests/test_video_image_encoder.py new file mode 100644 index 0000000..3fc0679 --- /dev/null +++ b/coreml/video_tracking/tests/test_video_image_encoder.py @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import coremltools as ct +import torch + +from edgetam_coreml_video.image_encoder import ( + CoreMLVideoImageEncoder, + VideoImageEncoder, + export_video_image_encoder, +) + + +def test_video_image_encoder_returns_raw_and_initial_features(reference_model): + torch.manual_seed(0) + image = torch.randn(1, 3, 256, 256) + wrapper = VideoImageEncoder(reference_model).eval() + + with torch.inference_mode(): + raw, initial, high_res_0, high_res_1 = wrapper(image) + expected_fpn = reference_model.forward_image(image)["backbone_fpn"] + + torch.testing.assert_close(raw, expected_fpn[2]) + expected_initial = raw.flatten(2).permute(2, 0, 1) + expected_initial = expected_initial + reference_model.no_mem_embed + expected_initial = expected_initial.permute(1, 2, 0).reshape_as(raw) + torch.testing.assert_close(initial, expected_initial) + torch.testing.assert_close(high_res_0, expected_fpn[0]) + torch.testing.assert_close(high_res_1, expected_fpn[1]) + + +def test_video_image_encoder_keeps_raw_features_unconditioned(reference_model): + image = torch.zeros(1, 3, 256, 256) + wrapper = VideoImageEncoder(reference_model).eval() + + with torch.inference_mode(): + raw, initial, _, _ = wrapper(image) + + assert not torch.equal(raw, initial) + + +def test_coreml_video_image_encoder_applies_notebook_normalization( + reference_model, +): + torch.manual_seed(5) + image_0_to_1 = torch.rand(1, 3, 256, 256) + mean = torch.tensor([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1) + std = torch.tensor([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1) + coreml_wrapper = CoreMLVideoImageEncoder(reference_model).eval() + normalized_wrapper = VideoImageEncoder(reference_model).eval() + + with torch.inference_mode(): + actual = coreml_wrapper(image_0_to_1) + expected = normalized_wrapper((image_0_to_1 - mean) / std) + + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close(actual_tensor, expected_tensor) + + +def test_export_video_image_encoder_declares_video_feature_outputs( + reference_model, + tmp_path: Path, +): + output_path = tmp_path / "edgetam_video_image_encoder.mlpackage" + + export_video_image_encoder(reference_model, output_path) + spec = ct.models.MLModel(str(output_path), skip_model_load=True).get_spec() + + assert [feature.name for feature in spec.description.input] == ["image"] + assert [feature.name for feature in spec.description.output] == [ + "raw_vision_features", + "initial_vision_features", + "high_res_feature_0", + "high_res_feature_1", + ] diff --git a/coreml/video_tracking/tests/test_video_initializer.py b/coreml/video_tracking/tests/test_video_initializer.py new file mode 100644 index 0000000..2837959 --- /dev/null +++ b/coreml/video_tracking/tests/test_video_initializer.py @@ -0,0 +1,87 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import coremltools as ct +import torch + +from edgetam_coreml_video.image_encoder import VideoImageEncoder +from edgetam_coreml_video.initializer import ( + VideoInitializer, + export_video_initializer, +) + + +def test_video_initializer_matches_prompted_sam_heads(reference_model): + torch.manual_seed(1) + image = torch.randn(1, 3, 1024, 1024) + point_coords = torch.tensor([[[210.0, 350.0]]]) + point_labels = torch.tensor([[1]], dtype=torch.int32) + image_encoder = VideoImageEncoder(reference_model).eval() + initializer = VideoInitializer(reference_model).eval() + + with torch.inference_mode(): + raw, initial, high_res_0, high_res_1 = image_encoder(image) + actual = initializer( + initial, + high_res_0, + high_res_1, + point_coords, + point_labels, + ) + expected = reference_model._forward_sam_heads( + backbone_features=initial, + point_inputs={ + "point_coords": point_coords, + "point_labels": point_labels, + }, + high_res_features=[high_res_0, high_res_1], + multimask_output=True, + ) + + low_res_mask, high_res_mask, iou, object_pointer, object_score = actual + torch.testing.assert_close(low_res_mask, expected[3]) + torch.testing.assert_close(high_res_mask, expected[4]) + torch.testing.assert_close(iou, expected[2].max(dim=-1).values) + torch.testing.assert_close(object_pointer, expected[5]) + torch.testing.assert_close(object_score, expected[6]) + assert raw.shape == (1, 256, 64, 64) + + +def test_export_video_initializer_declares_tracking_seed_outputs( + reference_model, + tmp_path: Path, +): + output_path = tmp_path / "edgetam_video_initializer.mlpackage" + + export_video_initializer(reference_model, output_path) + spec = ct.models.MLModel(str(output_path), skip_model_load=True).get_spec() + + assert [feature.name for feature in spec.description.input] == [ + "initial_vision_features", + "high_res_feature_0", + "high_res_feature_1", + "point_coords", + "point_labels", + ] + assert [feature.name for feature in spec.description.output] == [ + "low_res_mask", + "high_res_mask", + "best_iou", + "object_pointer", + "object_score", + ] + coords_range = spec.description.input[3].type.multiArrayType.shapeRange + labels_range = spec.description.input[4].type.multiArrayType.shapeRange + assert [ + (item.lowerBound, item.upperBound) + for item in coords_range.sizeRanges + ] == [(1, 1), (1, 4), (2, 2)] + assert [ + (item.lowerBound, item.upperBound) + for item in labels_range.sizeRanges + ] == [(1, 1), (1, 4)] diff --git a/coreml/video_tracking/tests/test_video_memory_encoder.py b/coreml/video_tracking/tests/test_video_memory_encoder.py new file mode 100644 index 0000000..bb3a6c4 --- /dev/null +++ b/coreml/video_tracking/tests/test_video_memory_encoder.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import coremltools as ct +import torch + +from edgetam_coreml_video.memory_encoder import ( + VideoMemoryEncoder, + export_video_memory_encoder, +) + + +def test_video_memory_encoder_matches_edgetam_memory_path(reference_model): + torch.manual_seed(2) + raw_features = torch.randn(1, 256, 64, 64) + high_res_mask = torch.randn(1, 1, 1024, 1024) + object_score = torch.tensor([[2.0]]) + wrapper = VideoMemoryEncoder(reference_model).eval() + + with torch.inference_mode(): + actual_features, actual_positions, actual_temporal_positions = wrapper( + raw_features, + high_res_mask, + object_score, + ) + vision_features = [raw_features.flatten(2).permute(2, 0, 1)] + expected_features, expected_positions = reference_model._encode_new_memory( + current_vision_feats=vision_features, + feat_sizes=[(64, 64)], + pred_masks_high_res=high_res_mask, + object_score_logits=object_score, + is_mask_from_pts=False, + ) + + torch.testing.assert_close(actual_features, expected_features) + torch.testing.assert_close(actual_positions, expected_positions[0]) + torch.testing.assert_close( + actual_temporal_positions, + reference_model.maskmem_tpos_enc[:, 0, 0], + ) + assert actual_features.shape == (1, 512, 64) + assert actual_positions.shape == (1, 512, 64) + assert actual_temporal_positions.shape == (7, 64) + + +def test_seed_memory_encoder_binarizes_prompted_mask_like_video_predictor( + reference_model, +): + raw_features = torch.randn(1, 256, 64, 64) + high_res_mask = torch.linspace(-1, 1, 1024 * 1024).reshape( + 1, 1, 1024, 1024 + ) + object_score = torch.tensor([[2.0]]) + original_binarize = reference_model.binarize_mask_from_pts_for_mem_enc + reference_model.binarize_mask_from_pts_for_mem_enc = True + wrapper = VideoMemoryEncoder( + reference_model, + is_mask_from_points=True, + ).eval() + + try: + with torch.inference_mode(): + actual = wrapper(raw_features, high_res_mask, object_score) + expected_features, expected_positions = ( + reference_model._encode_new_memory( + current_vision_feats=[ + raw_features.flatten(2).permute(2, 0, 1) + ], + feat_sizes=[(64, 64)], + pred_masks_high_res=high_res_mask, + object_score_logits=object_score, + is_mask_from_pts=True, + ) + ) + finally: + reference_model.binarize_mask_from_pts_for_mem_enc = original_binarize + + torch.testing.assert_close(actual[0], expected_features) + torch.testing.assert_close(actual[1], expected_positions[0]) + torch.testing.assert_close( + actual[2], + reference_model.maskmem_tpos_enc[:, 0, 0], + ) + + +def test_export_video_memory_encoder_declares_spatial_memory_outputs( + reference_model, + tmp_path: Path, +): + output_path = tmp_path / "edgetam_video_memory_encoder.mlpackage" + + export_video_memory_encoder(reference_model, output_path) + spec = ct.models.MLModel(str(output_path), skip_model_load=True).get_spec() + + assert [feature.name for feature in spec.description.input] == [ + "raw_vision_features", + "high_res_mask", + "object_score", + ] + assert [feature.name for feature in spec.description.output] == [ + "memory_features", + "memory_positions", + "temporal_positions", + ] diff --git a/coreml/video_tracking/tests/test_video_propagator.py b/coreml/video_tracking/tests/test_video_propagator.py new file mode 100644 index 0000000..72351e8 --- /dev/null +++ b/coreml/video_tracking/tests/test_video_propagator.py @@ -0,0 +1,140 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import coremltools as ct +import torch + +from edgetam_coreml_video.masked_attention import build_attention_controls +from edgetam_coreml_video.propagator import ( + VideoPropagator, + export_video_propagator, +) + + +def _operation_types(spec) -> list[str]: + def walk_block(block): + for operation in block.operations: + yield operation.type + for nested_block in operation.blocks: + yield from walk_block(nested_block) + + operation_types = [] + for function in spec.mlProgram.functions.values(): + for block in function.block_specializations.values(): + operation_types.extend(walk_block(block)) + return operation_types + + +def test_export_video_propagator_declares_fixed_stateless_contract( + reference_model, + tmp_path: Path, +): + output_path = tmp_path / "EdgeTAMVideoPropagator.mlpackage" + + export_video_propagator(reference_model, output_path) + spec = ct.models.MLModel(str(output_path), skip_model_load=True).get_spec() + + assert [feature.name for feature in spec.description.input] == [ + "raw_vision_features", + "high_res_feature_0", + "high_res_feature_1", + "spatial_bank", + "spatial_positions", + "pointer_bank", + "attention_bias", + "rotary_weight", + ] + assert { + feature.name: list(feature.type.multiArrayType.shape) + for feature in spec.description.input + } == { + "raw_vision_features": [1, 256, 64, 64], + "high_res_feature_0": [1, 32, 256, 256], + "high_res_feature_1": [1, 64, 128, 128], + "spatial_bank": [1, 7, 512, 64], + "spatial_positions": [1, 7, 512, 64], + "pointer_bank": [1, 16, 256], + "attention_bias": [1, 1, 1, 3648], + "rotary_weight": [1, 1792], + } + assert [feature.name for feature in spec.description.output] == [ + "low_res_mask", + "high_res_mask", + "best_iou", + "object_pointer", + "object_score", + "memory_features", + "memory_positions", + ] + assert not spec.description.state + assert "select" not in _operation_types(spec) + + +def test_video_propagator_runs_prompt_free_heads_and_encodes_new_memory( + reference_model, +): + torch.manual_seed(4) + raw = torch.randn(1, 256, 64, 64) + high_res_0 = torch.randn(1, 32, 256, 256) + high_res_1 = torch.randn(1, 64, 128, 128) + spatial_bank = torch.randn(1, 7, 512, 64) + spatial_positions = torch.randn(1, 7, 512, 64) + spatial_valid = torch.tensor([[1, 1, 0, 0, 0, 0, 0]], dtype=torch.float32) + pointer_bank = torch.randn(1, 16, 256) + pointer_valid = torch.tensor( + [[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], + dtype=torch.float32, + ) + attention_bias, rotary_weight = build_attention_controls( + spatial_valid, + pointer_valid, + ) + propagator = VideoPropagator(reference_model).eval() + + with torch.inference_mode(): + actual = propagator( + raw, + high_res_0, + high_res_1, + spatial_bank, + spatial_positions, + pointer_bank, + attention_bias, + rotary_weight, + ) + fused = propagator.memory_attention( + raw, + propagator.current_positions, + spatial_bank, + spatial_positions, + pointer_bank, + attention_bias, + rotary_weight, + ) + expected_heads = reference_model._forward_sam_heads( + backbone_features=fused, + point_inputs=None, + high_res_features=[high_res_0, high_res_1], + multimask_output=True, + ) + expected_memory = reference_model._encode_new_memory( + current_vision_feats=[raw.flatten(2).permute(2, 0, 1)], + feat_sizes=[(64, 64)], + pred_masks_high_res=expected_heads[4], + object_score_logits=expected_heads[6], + is_mask_from_pts=False, + ) + + low, high, iou, pointer, score, memory, memory_position = actual + torch.testing.assert_close(low, expected_heads[3]) + torch.testing.assert_close(high, expected_heads[4]) + torch.testing.assert_close(iou, expected_heads[2].max(dim=-1).values) + torch.testing.assert_close(pointer, expected_heads[5]) + torch.testing.assert_close(score, expected_heads[6]) + torch.testing.assert_close(memory, expected_memory[0]) + torch.testing.assert_close(memory_position, expected_memory[1][0]) diff --git a/coreml/video_tracking/validate_video.py b/coreml/video_tracking/validate_video.py new file mode 100644 index 0000000..2dfebb9 --- /dev/null +++ b/coreml/video_tracking/validate_video.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Compare Core ML video tracking with EdgeTAM's PyTorch predictor.""" + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np +from PIL import Image + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from edgetam_coreml_video.metrics import ( + binary_mask_iou, + cosine_similarity, + tensor_error, +) +from edgetam_coreml_video.predictor import CoreMLVideoPredictor +from sam2.build_sam import build_sam2_video_predictor + + +@dataclass(frozen=True) +class MaskComparison: + max_abs: float + mean_abs: float + cosine: float + mask_iou: float + + +def compare_mask_logits( + reference: np.ndarray, + candidate: np.ndarray, +) -> MaskComparison: + """Compare continuous logits and their zero-threshold binary masks.""" + reference = np.asarray(reference, dtype=np.float32) + candidate = np.asarray(candidate, dtype=np.float32) + if reference.shape != candidate.shape: + raise ValueError( + f"mask shape mismatch: {reference.shape} != {candidate.shape}" + ) + error = tensor_error(reference, candidate) + return MaskComparison( + max_abs=error.max_abs, + mean_abs=error.mean_abs, + cosine=cosine_similarity(reference, candidate), + mask_iou=binary_mask_iou(reference > 0, candidate > 0), + ) + + +def _frame_paths(frame_directory: Path, max_frames: int) -> list[Path]: + paths = [ + path + for path in frame_directory.iterdir() + if path.suffix.lower() in {".jpg", ".jpeg"} + ] + try: + paths.sort(key=lambda path: int(path.stem)) + except ValueError as error: + raise ValueError("JPEG frame filenames must use numeric stems") from error + if not paths: + raise ValueError(f"no JPEG frames found in {frame_directory}") + if max_frames <= 0: + raise ValueError("max_frames must be positive") + return paths[:max_frames] + + +def _pytorch_logits( + frame_directory: Path, + checkpoint: Path, + config_name: str, + device: str, + points: np.ndarray, + labels: np.ndarray, + frame_count: int, +) -> dict[int, np.ndarray]: + predictor = build_sam2_video_predictor( + config_name, + str(checkpoint), + device=device, + apply_postprocessing=False, + hydra_overrides_extra=[ + "++model.binarize_mask_from_pts_for_mem_enc=true", + ], + ) + inference_state = predictor.init_state( + video_path=str(frame_directory), + offload_video_to_cpu=device != "cpu", + ) + predictor.reset_state(inference_state) + predictor.add_new_points_or_box( + inference_state=inference_state, + frame_idx=0, + obj_id=1, + points=points, + labels=labels, + ) + logits = {} + for frame_index, _, mask_logits in predictor.propagate_in_video( + inference_state, + max_frame_num_to_track=frame_count - 1, + ): + if frame_index >= frame_count: + break + logits[frame_index] = mask_logits[0, 0].detach().float().cpu().numpy() + return logits + + +def _compare_coreml_frames( + frame_paths: list[Path], + reference: dict[int, np.ndarray], + predictor: CoreMLVideoPredictor, + points: np.ndarray, + labels: np.ndarray, +) -> list[dict[str, float | int]]: + """Run Core ML sequentially and return one comparison row per frame.""" + rows = [] + for frame_index, frame_path in enumerate(frame_paths): + with Image.open(frame_path) as frame: + result = ( + predictor.start_track(frame, points, labels) + if frame_index == 0 + else predictor.track_frame(frame) + ) + comparison = compare_mask_logits( + reference[frame_index], + result.mask_logits, + ) + snapshot = predictor.debug_bank_snapshot() + rows.append( + { + "frame": frame_index, + **asdict(comparison), + "object_score": result.object_score, + "predicted_iou": result.iou, + "recent_count": snapshot.recent_count, + "pointer_count": snapshot.pointer_count, + } + ) + return rows + + +def validate_video( + frame_directory: Path, + model_directory: Path, + checkpoint: Path, + config_name: str, + device: str, + points: np.ndarray, + labels: np.ndarray, + max_frames: int, +) -> list[dict[str, float | int]]: + """Run both predictors and return per-frame numerical comparisons.""" + frames = _frame_paths(frame_directory, max_frames) + reference = _pytorch_logits( + frame_directory, + checkpoint, + config_name, + device, + points, + labels, + len(frames), + ) + predictor = CoreMLVideoPredictor.from_directory(model_directory) + return _compare_coreml_frames( + frames, + reference, + predictor, + points, + labels, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--frames-dir", type=Path, required=True) + parser.add_argument( + "--models-dir", + type=Path, + default=Path(__file__).resolve().parent / "models", + ) + parser.add_argument( + "--checkpoint", + type=Path, + default=REPO_ROOT / "checkpoints/edgetam.pt", + ) + parser.add_argument("--config-name", default="edgetam.yaml") + parser.add_argument("--device", default="cpu") + parser.add_argument("--max-frames", type=int, default=8) + parser.add_argument( + "--point", + type=float, + nargs=3, + action="append", + metavar=("X", "Y", "LABEL"), + help="Repeat for up to four prompts; defaults to 210 350 1.", + ) + parser.add_argument( + "--json", + type=Path, + help="Optionally write all per-frame comparisons as JSON.", + ) + return parser.parse_args() + + +def run_validation( + args: argparse.Namespace, +) -> list[dict[str, float | int]]: + """Run validation from parsed arguments and write optional JSON.""" + prompt_rows = args.point or [[210.0, 350.0, 1.0]] + points = np.asarray([row[:2] for row in prompt_rows], dtype=np.float32) + labels = np.asarray([int(row[2]) for row in prompt_rows], dtype=np.int32) + rows = validate_video( + frame_directory=args.frames_dir, + model_directory=args.models_dir, + checkpoint=args.checkpoint, + config_name=args.config_name, + device=args.device, + points=points, + labels=labels, + max_frames=args.max_frames, + ) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(rows, indent=2) + "\n") + return rows + + +def main() -> None: + rows = run_validation(parse_args()) + + print("frame mask_iou cosine mean_abs max_abs") + for row in rows: + print( + f"{row['frame']:5d} {row['mask_iou']:.6f} " + f"{row['cosine']:.6f} {row['mean_abs']:.6f} " + f"{row['max_abs']:.6f}" + ) + + +if __name__ == "__main__": + main()