Add Core ML video tracking export

.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

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)

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.
+Generated `.mlpackage` directories belong in `coreml_models/`, which is
+excluded from version control.

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.

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",
+]