#!/usr/bin/env python3 """ Depth Estimation Privacy Skill — Monocular depth maps via Depth Anything v2. Backend selection: macOS → CoreML (.mlpackage via coremltools) — runs on Neural Engine Other → PyTorch (depth_anything_v2 pip package + HF weights) — runs on CUDA/MPS/CPU Implements the TransformSkillBase interface to provide real-time depth map overlays on camera feeds. When used as a privacy skill, the depth-only mode anonymizes the scene while preserving spatial layout and activity recognition. Usage: python transform.py --model depth-anything-v2-small --device auto python transform.py --config config.json """ import sys import os import platform import argparse from pathlib import Path # Import the base class from the same directory _script_dir = Path(__file__).resolve().parent sys.path.insert(0, str(_script_dir)) from transform_base import TransformSkillBase, _log # noqa: E402 COLORMAP_MAP = { "inferno": 1, # cv2.COLORMAP_INFERNO "viridis": 16, # cv2.COLORMAP_VIRIDIS "plasma": 13, # cv2.COLORMAP_PLASMA "magma": 12, # cv2.COLORMAP_MAGMA "jet": 2, # cv2.COLORMAP_JET "turbo": 18, # cv2.COLORMAP_TURBO "hot": 11, # cv2.COLORMAP_HOT "cool": 8, # cv2.COLORMAP_COOL } # CoreML model registry — mirrors apple/coreml-depth-anything-v2-small HF repo COREML_VARIANTS = { "DepthAnythingV2SmallF16": { "precision": "float16", "size_mb": 49.8, "description": "Float16 — optimized for Neural Engine", }, "DepthAnythingV2SmallF16INT8": { "precision": "float16_int8", "size_mb": 25.0, "description": "Float16 + INT8 quantization — smallest", }, "DepthAnythingV2SmallF32": { "precision": "float32", "size_mb": 99.2, "description": "Float32 — highest precision", }, } # Default CoreML variant (best balance of speed + quality on Neural Engine) DEFAULT_COREML_VARIANT = "DepthAnythingV2SmallF16" # HuggingFace repo for CoreML models COREML_HF_REPO = "apple/coreml-depth-anything-v2-small" # CoreML input size — MUST match model exactly (multiples of 14 for ViT) COREML_INPUT_SIZE = (518, 392) # width, height # Where Aegis DepthVisionStudio stores downloaded models MODELS_DIR = Path.home() / ".aegis-ai" / "models" / "feature-extraction" # PyTorch model configs (fallback on non-macOS) PYTORCH_CONFIGS = { "depth-anything-v2-small": { "encoder": "vits", "features": 64, "out_channels": [48, 96, 192, 384], "repo": "depth-anything/Depth-Anything-V2-Small", "filename": "depth_anything_v2_vits.pth", }, "depth-anything-v2-base": { "encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768], "repo": "depth-anything/Depth-Anything-V2-Base", "filename": "depth_anything_v2_vitb.pth", }, "depth-anything-v2-large": { "encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024], "repo": "depth-anything/Depth-Anything-V2-Large", "filename": "depth_anything_v2_vitl.pth", }, } class DepthEstimationSkill(TransformSkillBase): """ Depth estimation using Depth Anything v2. Produces colorized depth maps that can be blended with the original frame (overlay mode), shown side-by-side, or displayed as depth-only anonymized view. """ def __init__(self): super().__init__() self._tag = "DepthEstimation" self.model = None self.backend = None # "coreml" or "pytorch" self.colormap_id = 16 # cv2.COLORMAP_VIRIDIS self.opacity = 0.5 self.blend_mode = "depth_only" # Default for privacy: depth_only anonymizes self._coreml_input_size = COREML_INPUT_SIZE def parse_extra_args(self, parser: argparse.ArgumentParser): parser.add_argument("--model", type=str, default="depth-anything-v2-small", choices=["depth-anything-v2-small", "depth-anything-v2-base", "depth-anything-v2-large"]) parser.add_argument("--variant", type=str, default=DEFAULT_COREML_VARIANT, help="CoreML variant ID (macOS only)") parser.add_argument("--colormap", type=str, default="viridis", choices=list(COLORMAP_MAP.keys())) parser.add_argument("--blend-mode", type=str, default="depth_only", choices=["overlay", "side_by_side", "depth_only"]) parser.add_argument("--opacity", type=float, default=0.5) def load_model(self, config: dict) -> dict: model_name = config.get("model", "depth-anything-v2-small") self.colormap_id = COLORMAP_MAP.get(config.get("colormap", "viridis"), 16) self.opacity = config.get("opacity", 0.5) self.blend_mode = config.get("blend_mode", "depth_only") # Try CoreML first on macOS if platform.system() == "Darwin": try: info = self._load_coreml(config) return info except Exception as e: _log(f"CoreML load failed ({e}), falling back to PyTorch", self._tag) # Fallback: PyTorch return self._load_pytorch(model_name, config) # ── CoreML backend (macOS) ──────────────────────────────────────── def _load_coreml(self, config: dict) -> dict: """Load CoreML .mlpackage model — runs on Apple Neural Engine.""" import coremltools as ct variant_id = config.get("variant", DEFAULT_COREML_VARIANT) model_path = MODELS_DIR / f"{variant_id}.mlpackage" # Auto-download from HuggingFace if not present if not model_path.exists(): _log(f"CoreML model not found at {model_path}, downloading from HF...", self._tag) self._download_coreml_model(variant_id) if not model_path.exists(): raise FileNotFoundError(f"CoreML model not found: {model_path}") _log(f"Loading CoreML model: {variant_id} (Neural Engine)", self._tag) self.model = ct.models.MLModel(str(model_path), compute_units=ct.ComputeUnit.ALL) self.backend = "coreml" _log(f"CoreML model loaded: {variant_id}", self._tag) return { "model": f"coreml-{variant_id}", "device": "neural_engine", "blend_mode": self.blend_mode, "colormap": config.get("colormap", "viridis"), "backend": "coreml", } def _download_coreml_model(self, variant_id: str): """Download CoreML .mlpackage from HuggingFace using huggingface_hub.""" try: from huggingface_hub import snapshot_download MODELS_DIR.mkdir(parents=True, exist_ok=True) mlpackage_name = f"{variant_id}.mlpackage" _log(f"Downloading {mlpackage_name} from {COREML_HF_REPO}...", self._tag) # Download only the specific variant's .mlpackage directory snapshot_download( COREML_HF_REPO, local_dir=str(MODELS_DIR), allow_patterns=[f"{mlpackage_name}/**"], ) model_path = MODELS_DIR / mlpackage_name if model_path.exists(): _log(f"Downloaded CoreML model: {model_path}", self._tag) else: _log(f"Download completed but model not found at {model_path}", self._tag) except Exception as e: _log(f"CoreML model download failed: {e}", self._tag) raise # ── PyTorch backend (fallback) ──────────────────────────────────── def _load_pytorch(self, model_name: str, config: dict) -> dict: """Load PyTorch model — fallback for non-macOS or when CoreML is unavailable.""" import torch from depth_anything_v2.dpt import DepthAnythingV2 from huggingface_hub import hf_hub_download _log(f"Loading {model_name} on {self.device} (PyTorch)", self._tag) cfg = PYTORCH_CONFIGS.get(model_name) if not cfg: raise ValueError(f"Unknown model: {model_name}. Choose from: {list(PYTORCH_CONFIGS.keys())}") # Download weights from HuggingFace Hub (cached after first download) _log(f"Downloading weights from HF: {cfg['repo']}", self._tag) weights_path = hf_hub_download(cfg["repo"], cfg["filename"]) # Build model from pip package self.model = DepthAnythingV2( encoder=cfg["encoder"], features=cfg["features"], out_channels=cfg["out_channels"], ) self.model.load_state_dict(torch.load(weights_path, map_location=self.device, weights_only=True)) self.model.to(self.device) self.model.eval() self.backend = "pytorch" _log(f"PyTorch model loaded: {model_name} on {self.device}", self._tag) return { "model": model_name, "device": self.device, "blend_mode": self.blend_mode, "colormap": config.get("colormap", "viridis"), "backend": "pytorch", } # ── Frame transform ─────────────────────────────────────────────── def transform_frame(self, image, metadata: dict): import cv2 import numpy as np if self.backend == "coreml": depth_colored = self._infer_coreml(image) else: depth_colored = self._infer_pytorch(image) if self.blend_mode == "overlay": output = cv2.addWeighted(image, 1 - self.opacity, depth_colored, self.opacity, 0) elif self.blend_mode == "side_by_side": output = np.hstack([image, depth_colored]) else: # depth_only — full anonymization output = depth_colored return output def _infer_coreml(self, image): """Run CoreML inference and return colorized depth map (BGR, original size).""" import cv2 import numpy as np from PIL import Image original_h, original_w = image.shape[:2] input_w, input_h = self._coreml_input_size # BGR → RGB → resize to model input → PIL rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) resized = cv2.resize(rgb, (input_w, input_h), interpolation=cv2.INTER_LINEAR) pil_image = Image.fromarray(resized, mode="RGB") # Inference prediction = self.model.predict({"image": pil_image}) # Extract depth map (first output key) output_key = list(prediction.keys())[0] depth_map = prediction[output_key] # Convert PIL Image to numpy if needed if isinstance(depth_map, Image.Image): depth_map = np.array(depth_map) depth_map = np.array(depth_map) if depth_map.ndim > 2: depth_map = np.squeeze(depth_map) # Normalize → uint8 → colormap → resize back depth_norm = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min() + 1e-8) depth_uint8 = (depth_norm * 255).astype(np.uint8) depth_colored = cv2.applyColorMap(depth_uint8, self.colormap_id) depth_colored = cv2.resize(depth_colored, (original_w, original_h)) return depth_colored def _infer_pytorch(self, image): """Run PyTorch inference and return colorized depth map (BGR, original size).""" import torch import cv2 import numpy as np rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) with torch.no_grad(): depth = self.model.infer_image(rgb) d_min, d_max = depth.min(), depth.max() depth_norm = ((depth - d_min) / (d_max - d_min + 1e-8) * 255).astype(np.uint8) depth_colored = cv2.applyColorMap(depth_norm, self.colormap_id) return depth_colored # ── Config updates ──────────────────────────────────────────────── def on_config_update(self, config: dict): """Handle live config updates from Aegis.""" if "colormap" in config: self.colormap_id = COLORMAP_MAP.get(config["colormap"], self.colormap_id) _log(f"Colormap updated: {config['colormap']}", self._tag) if "opacity" in config: self.opacity = float(config["opacity"]) _log(f"Opacity updated: {self.opacity}", self._tag) if "blend_mode" in config: self.blend_mode = config["blend_mode"] _log(f"Blend mode updated: {self.blend_mode}", self._tag) def get_output_mode(self) -> str: """Use base64 for privacy transforms — avoids temp file cleanup issues.""" return "base64" if __name__ == "__main__": DepthEstimationSkill().run()