diff --git a/.github/workflows/pets-dino-release-publish.yml b/.github/workflows/pets-dino-release-publish.yml new file mode 100644 index 000000000..8f2827e23 --- /dev/null +++ b/.github/workflows/pets-dino-release-publish.yml @@ -0,0 +1,98 @@ +name: Pets DINO Release publish + +on: + workflow_dispatch: + inputs: + artifact_run_id: + description: Successful Pets DINO candidate workflow run ID + required: true + type: string + +permissions: + contents: read + +env: + CANDIDATE_WORKFLOW_ID: "340207631" + CANDIDATE_WORKFLOW_PATH: ".github/workflows/pets-dino-source-contract.yml" + +jobs: + publish-release: + runs-on: ubuntu-latest + environment: pets-model-release + permissions: + actions: read + contents: write + steps: + - name: Require an explicit artifact run + run: | + if [ -z "${{ inputs.artifact_run_id }}" ]; then + echo "artifact_run_id is required" >&2 + exit 1 + fi + + - name: Resolve candidate run provenance + id: candidate + env: + GH_TOKEN: ${{ github.token }} + run: | + gh api "repos/${GITHUB_REPOSITORY}/actions/runs/${{ inputs.artifact_run_id }}" > candidate-run.json + head_sha=$(python -c 'import json; print(json.load(open("candidate-run.json"))["head_sha"])') + echo "head_sha=${head_sha}" >> "${GITHUB_OUTPUT}" + + - name: Check out the exact builder commit + uses: actions/checkout@v4 + with: + ref: ${{ steps.candidate.outputs.head_sha }} + path: builder + + - name: Check out the release manifest commit + uses: actions/checkout@v4 + with: + ref: ${{ github.sha }} + path: release + + - uses: actions/download-artifact@v4 + with: + name: pet-models-v1-candidate + path: model-dist + run-id: ${{ inputs.artifact_run_id }} + github-token: ${{ github.token }} + + - name: Verify artifact, manifest, and builder provenance + run: | + builder_commit=$(git -C builder rev-parse HEAD) + python release/tools/model_release_provenance.py \ + --run-json candidate-run.json \ + --build-manifest model-dist/dinov2_vits14.build.json \ + --builder-commit "${builder_commit}" \ + --expected-repository "${GITHUB_REPOSITORY}" \ + --expected-workflow-id "${CANDIDATE_WORKFLOW_ID}" \ + --expected-workflow-path "${CANDIDATE_WORKFLOW_PATH}" + python release/tools/validate_dinov2_torchscript.py \ + model-dist/dinov2_vits14.pt \ + model-dist/dinov2_vits14.pt.metadata.json \ + --manifest release/src/iPhoto/pets/model_manifest.json \ + --metadata-only + + - name: Refuse overwrite of an immutable v1 + env: + GH_TOKEN: ${{ github.token }} + run: | + if gh release view pet-models-v1 --repo "${GITHUB_REPOSITORY}" >/dev/null 2>&1; then + echo "pet-models-v1 already exists; refusing to overwrite it" >&2 + exit 1 + fi + + - name: Create v1 when no Release exists + env: + GH_TOKEN: ${{ github.token }} + run: >- + gh release create pet-models-v1 + model-dist/dinov2_vits14.pt + model-dist/dinov2_vits14.pt.metadata.json + model-dist/dinov2_vits14.build.json + model-dist/SHA256SUMS + --repo "${GITHUB_REPOSITORY}" + --target "${{ steps.candidate.outputs.head_sha }}" + --title "iPhotron Pets models v1" + --notes "Pinned DINOv2 ViT-S/14 TorchScript artifact validated on Ubuntu, macOS, and Windows." diff --git a/.github/workflows/pets-dino-source-contract.yml b/.github/workflows/pets-dino-source-contract.yml new file mode 100644 index 000000000..8d80193ff --- /dev/null +++ b/.github/workflows/pets-dino-source-contract.yml @@ -0,0 +1,120 @@ +name: Pets DINO Release contract + +on: + workflow_dispatch: + pull_request: + branches: [main, edit-base] + +permissions: + contents: read + +env: + PYTHON_VERSION: "3.12" + TORCH_VERSION: "2.12.1" + TORCHVISION_VERSION: "0.27.1" + DINO_REVISION: "7764ea0f912e53c92e82eb78a2a1631e92725fc8" + DINO_WEIGHTS_URL: "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pth" + +jobs: + build-release-candidate: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.head.sha || github.sha }} + - uses: actions/checkout@v4 + with: + repository: facebookresearch/dinov2 + ref: ${{ env.DINO_REVISION }} + path: .model-build/dinov2 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + cache: pip + cache-dependency-path: pyproject.toml + - name: Install pinned CPU build runtime + run: >- + python -m pip install torch==${TORCH_VERSION} + torchvision==${TORCHVISION_VERSION} + --index-url https://download.pytorch.org/whl/cpu + - name: Download pinned official checkpoint + run: >- + curl --fail --location --retry 3 "${DINO_WEIGHTS_URL}" + --output .model-build/dinov2_vits14_pretrain.pth + - name: Build one verified TorchScript candidate + env: + XFORMERS_DISABLED: "1" + IPHOTO_MODEL_BUILD_COMMIT: ${{ github.event.pull_request.head.sha || github.sha }} + run: >- + python tools/convert_dinov2_torchscript.py + model-dist/dinov2_vits14.pt + --source-dir .model-build/dinov2 + --checkpoint .model-build/dinov2_vits14_pretrain.pth + --runtime-metadata model-dist/dinov2_vits14.pt.metadata.json + --build-manifest model-dist/dinov2_vits14.build.json + - name: Record artifact checksum + run: sha256sum model-dist/dinov2_vits14.pt > model-dist/SHA256SUMS + - uses: actions/upload-artifact@v4 + with: + name: pet-models-v1-candidate + path: model-dist/* + if-no-files-found: error + retention-days: 14 + + validate-release-candidate: + needs: build-release-candidate + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + - name: Install pinned CPU runtime on Linux and Windows + if: runner.os != 'macOS' + run: | + python -m pip install numpy + python -m pip install torch==${{ env.TORCH_VERSION }} --index-url https://download.pytorch.org/whl/cpu + - name: Install pinned runtime on macOS + if: runner.os == 'macOS' + run: python -m pip install numpy torch==${{ env.TORCH_VERSION }} + - uses: actions/download-artifact@v4 + with: + name: pet-models-v1-candidate + path: model-dist + - name: Validate identical release bytes and inference + shell: bash + run: >- + python tools/validate_dinov2_torchscript.py + model-dist/dinov2_vits14.pt + model-dist/dinov2_vits14.pt.metadata.json + | tee "model-dist/validation-${RUNNER_OS}.json" + - uses: actions/upload-artifact@v4 + with: + name: pet-models-v1-validation-${{ runner.os }} + path: model-dist/validation-*.json + if-no-files-found: error + + published-release-contract: + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ env.PYTHON_VERSION }} + - name: Install pinned published-artifact runtime + run: | + python -m pip install certifi numpy pillow pytest + python -m pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu + - name: Validate the immutable public Release asset + env: + PYTHONPATH: src + IPHOTO_RUN_PETS_DINO_RELEASE_CONTRACT: "1" + run: python -m pytest -q -s tests/contracts/test_pets_dino_source_contract.py diff --git a/docs/development.md b/docs/development.md index 34727c09e..fb4b9586d 100644 --- a/docs/development.md +++ b/docs/development.md @@ -235,12 +235,14 @@ YOLOX release asset and can be overridden with: export IPHOTO_PET_DETECTOR_MODEL_URL="https://example.invalid/yolox_nano.onnx" ``` -Production never executes Torch Hub. It loads a packaged DINOv2 TorchScript model, -or downloads the fixed HTTPS artifact declared by SHA-256 and exact byte size in -`src/iPhoto/pets/model_manifest.json`. Release engineering may regenerate the -artifact from the pinned source revision with -`tools/convert_dinov2_torchscript.py`; that tool also checks eager/TorchScript -numeric equivalence before publishing. +When DINOv2 is missing, iPhotron downloads the fixed `pet-models-v1` TorchScript +Release declared by HTTPS URL, SHA-256, exact byte size, cache schema, and +producer version in `src/iPhoto/pets/model_manifest.json`. It validates CPU +loading and output shape before publishing the model and metadata atomically. +Production never downloads or executes DINOv2 source and never traces a model on +the user's machine. `tools/convert_dinov2_torchscript.py` is release-only: it +loads a local checkout of the pinned source commit and verified official +checkpoint, then checks eager/TorchScript numeric equivalence. For offline or packaged validation, disable first-use downloads with: diff --git a/docs/misc/BUILD_DEB.md b/docs/misc/BUILD_DEB.md index 28b094744..19580e24c 100644 --- a/docs/misc/BUILD_DEB.md +++ b/docs/misc/BUILD_DEB.md @@ -15,7 +15,7 @@ runtime depends on the helper binary plus the shared libraries under Builds that ship People/Pets recognition must also preserve the selected AI runtimes from the standalone bundle. People needs `insightface` and -`onnxruntime`; Pets needs `onnxruntime`, `torch`, `torchvision`, `usearch`, and +`onnxruntime`; Pets needs `onnxruntime`, `torch==2.12.1`, `usearch`, and `certifi`. Offline builds also retain the shared `extension/models` cache. These are added at the Nuitka stage described in [`BUILD_EXE.md`](BUILD_EXE.md); the `.deb` stage must not strip them from @@ -126,7 +126,7 @@ Description: Folder-native local photo album manager ```bash find "$APP_ROOT" -path '*insightface*' -o -path '*onnxruntime*' find "$APP_ROOT/extension/models" -name 'det_500m.onnx' -o -name 'w600k_mbf.onnx' - find "$APP_ROOT" -path '*torch*' -o -path '*torchvision*' -o -path '*usearch*' + find "$APP_ROOT" -path '*torch*' -o -path '*usearch*' find "$APP_ROOT/extension/models/pets" -name 'yolox_nano_coco.onnx' -o -name 'dinov2_vits14.pt' ``` @@ -194,4 +194,4 @@ sudo apt remove iPhotron | Native maps fail with GLX/XCB startup errors | The runtime was installed correctly, but the desktop session lacks XWayland/XCB GL integration | Install/enable XWayland and rerun, or set `IPHOTO_PREFER_OSMAND_NATIVE_WIDGET=0` to force the helper-backed Python OBF path | | People scan is unavailable in the installed app | The standalone build was produced without the optional face runtime | Rebuild the standalone app with `insightface`, `onnxruntime`, and `src/extension/models` included before staging the `.deb` | | People scan starts but never creates clusters | The model cache or an InsightFace submodel/dependency is missing from `/opt/iPhotron/` | Verify `extension/models`, exclude unused `albumentations`/`pydantic` packages at the Nuitka stage, and keep InsightFace limited to detection and recognition | -| Pets scan is unavailable in the installed app | The standalone build omitted `pets-ai` packages or `extension/models/pets` | Rebuild the standalone app with `onnxruntime`, `torch`, `torchvision`, `usearch`, `certifi`, and both Pets model files before staging the `.deb` | +| Pets scan is unavailable in the installed app | The standalone build omitted `pets-ai` packages or `extension/models/pets` | Rebuild the standalone app with `onnxruntime`, `torch==2.12.1`, `usearch`, `certifi`, and both Pets model files before staging the `.deb` | diff --git a/docs/misc/BUILD_EXE.md b/docs/misc/BUILD_EXE.md index 43a0e7c96..49a48d10f 100644 --- a/docs/misc/BUILD_EXE.md +++ b/docs/misc/BUILD_EXE.md @@ -176,8 +176,9 @@ not enable it. Build environments that promise Pets support must install: python -m pip install -e ".[pets-ai]" ``` -The standalone bundle must retain `onnxruntime`, `torch`, `torchvision`, -`usearch`, and `certifi`. An offline-ready build must also include: +The standalone bundle must retain `onnxruntime`, `torch==2.12.1`, `usearch`, and +`certifi`. `torchvision` is a model-build dependency, not a production runtime +dependency. An offline-ready build must also include: ```text extension/models/pets/ @@ -192,14 +193,13 @@ that enable Pets should also include the optional runtime explicitly: ```bash --include-package=onnxruntime --include-package=torch ---include-package=torchvision --include-package=usearch --include-package=certifi --include-data-dir=src/extension/models=extension/models ``` The current platform build scripts explicitly include the People runtime but do -not yet add `torch`, `torchvision`, or `usearch` flags. Therefore a stock script +not yet add `torch` or `usearch` flags. Therefore a stock script build must not be advertised as Pets-enabled merely because the model directory was copied; add the flags above (or update the script) and perform the Pets smoke test before release. diff --git a/docs/misc/PETS_RECOGNITION_RUNTIME.md b/docs/misc/PETS_RECOGNITION_RUNTIME.md index 8fc24f750..b7aef8843 100644 --- a/docs/misc/PETS_RECOGNITION_RUNTIME.md +++ b/docs/misc/PETS_RECOGNITION_RUNTIME.md @@ -38,29 +38,44 @@ Install the optional runtime with: pip install -e ".[pets-ai]" ``` -The extra provides `onnxruntime`, `torch`, `torchvision`, `usearch`, and -`certifi`. Bundled models are read-only fallbacks; downloads are written to the +The extra provides `onnxruntime`, `torch==2.12.1`, `usearch`, and `certifi`. +`torchvision==0.27.1` is confined to the controlled model build workflow. +Bundled models are read-only fallbacks; downloads are written to the platform user cache: ```text src/extension/models/pets/ ├── detector/yolox_nano_coco.onnx -└── embedding/dinov2_vits14/dinov2_vits14.pt +└── embedding/dinov2_vits14/ + ├── dinov2_vits14.pt + └── dinov2_vits14.pt.metadata.json ``` -`IPHOTO_PET_MODEL_DIR` overrides that root. Missing models may be populated on -first use unless `IPHOTO_PET_MODEL_AUTO_DOWNLOAD=0`. The detector URL defaults -to the upstream YOLOX release and can be overridden with -`IPHOTO_PET_DETECTOR_MODEL_URL`. Production does not execute Torch Hub. DINOv2 -must be supplied as the hash- and size-verified TorchScript artifact declared -in `iPhoto/pets/model_manifest.json`; Torch Hub is restricted to the release -conversion tool. `IPHOTO_PET_SCAN_DISABLED=1` disables the worker without -disabling the rest of the application. +Lookup uses the bundled extension first and then the platform user cache. A +writable development extension directory is preferred for installation; signed +macOS app bundles install directly into the user cache without a writability +probe. `IPHOTO_PET_MODEL_DIR` is authoritative: when set, both lookup and +installation use only that root. + +Missing models may be populated lazily on the first non-empty scan unless +`IPHOTO_PET_MODEL_AUTO_DOWNLOAD=0`. The detector URL defaults to the upstream +YOLOX release and can be overridden with `IPHOTO_PET_DETECTOR_MODEL_URL`. +DINOv2 acquisition downloads only the fixed `pet-models-v1` TorchScript Release, +validates its SHA-256, exact size, producer/cache schema, CPU load, and output +shape, then publishes metadata first and the model as the visibility point. +Production never invokes Torch Hub, downloads source, imports xFormers, or traces +TorchScript. Legacy derived caches are replaced under the acquisition lock; +bundled artifacts are never deleted. Only local `EACCES`, `EPERM`, or `EROFS` +storage failures fall back to the user cache; network, TLS, disk-full, I/O, hash, +or compatibility failures do not masquerade as storage fallback. +`IPHOTO_PET_SCAN_DISABLED=1` disables the worker without disabling the rest of +the application. Packaged/offline builds that promise Pets support must include the Python AI -runtime and the two model files under `extension/models/pets`. A build that -omits them must preserve graceful degradation: core browsing, People, Maps, -editing, and library state remain usable. +runtime, the YOLOX detector, the DINOv2 TorchScript artifact, and its +`dinov2_vits14.pt.metadata.json` sidecar under `extension/models/pets`. A build +that omits them must preserve graceful degradation: core browsing, People, +Maps, editing, and library state remain usable. ## Detection And Clustering Contract @@ -125,16 +140,14 @@ starts before the previous drain finishes. | `done` | Detection completed, including valid images with no pets. | | `skipped` | Video, non-primary Live Photo component, or another ineligible asset. | -Interactive scans start Face and Pet workers alongside metadata scanning and -enqueue rows only after their asset batches commit. When a saved library needs -a startup metadata scan, startup first warms the gallery, runs that scan, then -starts both AI workers with closed input so they drain persisted -`pending`/`retry` rows. This avoids model initialization and competing AI work -on the first-frame path. If the metadata scan scope is already complete, startup -still starts the Pet backfill worker whenever persisted `pending` or `retry` rows -need draining. With no metadata scan and no queued AI work, startup does not -launch scan workers; an explicit rescan is only needed to reset or rediscover -otherwise completed/failed assets. +Interactive rescans start Face and Pet workers alongside metadata scanning and +enqueue rows only after their asset batches commit. Desktop startup keeps the +first-frame and metadata-scan paths AI-free. After the startup scan succeeds—or +immediately when its scope was already complete—a 1500 ms interaction-idle gate +starts both closed-input workers at `LowestPriority`. Click, wheel, key, drag, +touch, and gesture input restart the gate. Missing People/Pets models may then +download automatically; switching libraries, shutdown, cancellation, or a new +generation cancels pending activation. The Pet worker uses small batches and queue top-up from the asset repository. Missing dependencies/models are runtime-availability failures: pending rows are diff --git a/docs/requirements/startup-chain-optimization/ENGINEERING_CLOSURE_REPORT.md b/docs/requirements/startup-chain-optimization/ENGINEERING_CLOSURE_REPORT.md index 56b6cd63e..6928fb582 100644 --- a/docs/requirements/startup-chain-optimization/ENGINEERING_CLOSURE_REPORT.md +++ b/docs/requirements/startup-chain-optimization/ENGINEERING_CLOSURE_REPORT.md @@ -19,8 +19,9 @@ 因此之前记录的 Pets 自动化完成结论已经撤回。对 Pets 模型契约、大图库 复杂度、跨库一致性、升级回填、打包模型目录和身份语义的复审已完成本机自动化 核心修复。完整问题、裁决和验收要求见 `PETS_REVIEW_REMEDIATION_LEDGER.md`。项目固定 -`pet-models-v1` TorchScript Release 尚未发布；远端 CI 已在 -`e0001ee646e95fadc33659ebe277eb067e79a084` 上 9/9 成功。只读 packaged 安装、 +[`pet-models-v1`](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/releases/tag/pet-models-v1) +TorchScript Release 已发布，并由 [run 32807111211](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/actions/runs/32807111211) +在 Ubuntu、macOS、Windows 验证同一资产。只读 packaged 安装、 真实升级库、网络失败和跨平台 50k 报告仍为 `manual_validation_pending`，因此不得恢复 `engineering_complete`。 @@ -45,7 +46,7 @@ Windows 实机复测表明前次 WIC fresh-decode 修复会被已持久化的错 - `DesktopCoordinatorRuntime` 是唯一桌面组合根，Recognition、Location/Info、Edit 与地图能力均延迟到首次使用；People dashboard 快照只在 People feature 首次创建时预热。 - settings/shell 同步初始化和 Windows/Linux pre-show Detail 异常也进入唯一 terminal 协议；pre-show Detail 使用可重试降级窗口。 - 模块预载使用 generation-aware owner 和完成信号，不再依赖持续轮询 timer；退出时等待预载线程收口。 -- People/Pets 的模型扫描不再在 startup completed 后自动启动；首次进入识别功能才构造服务与 worker，消除快速关窗的 QThread 竞争。 +- People/Pets 在主 metadata scan 成功后通过 1500 ms 交互空闲门控自动启动；worker 使用最低优先级，切库、取消和 shutdown generation 阻止迟到启动。 ### Probe、数据库与慢存储 diff --git a/docs/requirements/startup-chain-optimization/MANUAL_VALIDATION_MATRIX.md b/docs/requirements/startup-chain-optimization/MANUAL_VALIDATION_MATRIX.md index f597bd5a0..e2f08de2f 100644 --- a/docs/requirements/startup-chain-optimization/MANUAL_VALIDATION_MATRIX.md +++ b/docs/requirements/startup-chain-optimization/MANUAL_VALIDATION_MATRIX.md @@ -70,7 +70,7 @@ Metal 与 OpenGL 分开采集： | --- | --- | --- | --- | | 只读安装目录与首次模型落盘 | macOS Apple Silicon、macOS Intel、Windows `Program Files`、Linux AppImage | bundled root 保持只读；查找顺序为 override、用户 cache、bundled；缺失模型只写用户 cache；首次识别成功 | `pending_manual_validation` | | 模型获取与自愈 | 在线下载、离线 bundled fallback、损坏 cache、代理失败、证书失败 | 在线文件 hash/size/shape 正确；离线可用 bundled；损坏 cache 可重取；代理/证书失败给出可操作提示且不污染 cache | `pending_manual_validation` | -| 固定 DINO Release 产物 | 使用固定 source commit 的受控构建环境和本仓库不可变 `pet-models-v1` 标签 | 生产 Torch Hub 路径已删除；仍须发布 `dinov2_vits14.pt`、记录 artifact/build manifest SHA、填写 Release HTTPS URL，并重新通过 hash/shape/packaging 门禁 | `pending_manual_validation` | +| 固定 DINO Release 产物 | 使用固定 source commit 的受控构建环境和本仓库不可变 `pet-models-v1` 标签 | 生产 Torch Hub 路径已删除；Release 包含 `.pt`、runtime metadata、build manifest 与 SHA 清单；run head/build manifest/tag target 均为 `ceb24fa0…`，公开 URL/SHA/size 与三平台 load/shape 已复核 | `automated_release_contract_pass`（[Release](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/releases/tag/pet-models-v1)，[run 32807111211](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/actions/runs/32807111211)） | | 真实旧图库升级 | 含 name、cover、hidden、rejection、pet merge、跨类型 merge 的脱敏副本 | interactive 不被 backfill 阻塞；后台清空 pending/retry；durable state 不丢；失败资产明确显示 stale/source generation | `pending_manual_validation` | | 真实照片识别 | 多宠照片、小目标 tile、People overlap、大狗+远处小猫、重叠 cat/dog | 类别、bbox、去重和 People 优先级符合契约；旧 source annotation 与 canonical identity 显示一致 | `pending_manual_validation` | | Windows Live Photo 静态图方向 | Windows packaged，优先复测 `IMG_3684.HEIC` 的脱敏副本，并覆盖 iPhone Orientation 5/6/7/8 的 HEIC+MOV/JPEG+MOV | 静态图与动态视频视觉方向一致；静态图无二次 EXIF 旋转；JPEG 等 WIC 未预转正格式仍正确应用 EXIF；首次展示、Live Motion 返回静帧和连续切换均通过 | `user_verified_original_sample_pass / formal_artifact_evidence_pending` | diff --git a/docs/requirements/startup-chain-optimization/NEXT_DEVELOPMENT_HANDOFF.md b/docs/requirements/startup-chain-optimization/NEXT_DEVELOPMENT_HANDOFF.md index cb61b62a1..77093dcdc 100644 --- a/docs/requirements/startup-chain-optimization/NEXT_DEVELOPMENT_HANDOFF.md +++ b/docs/requirements/startup-chain-optimization/NEXT_DEVELOPMENT_HANDOFF.md @@ -20,7 +20,7 @@ 本轮复审发现并修复了启动 terminal 边界、后台 import 资源所有权、 Recognition 启动期回退和 packaged A/B 同构校验四类问题。用户确认的产品策略为： -1. People/Pets dashboard 只在首次进入 People 时异步预热。 +1. 主 metadata scan 成功后等待 1500 ms 无交互，再以最低优先级自动启动 People/Pets；用户交互、切库和 shutdown 会延后或取消启动。 2. Windows/Linux pre-show Detail 创建失败时显示基础窗口和非模态恢复面板，允许新 generation 重试。 ## 已完成的修复 @@ -57,9 +57,9 @@ PR workflow 已完成自动化修复。当前全量 `2756 passed, 14 skipped`， 远端 CI 已在 `e0001ee646e95fadc33659ebe277eb067e79a084` 上 9/9 成功，证据为 [run 30263442145](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/actions/runs/30263442145)。 -仍有一个不可伪造为完成的交付边界：仓库不可变 `pet-models-v1` TorchScript Release -尚未发布，因此生产 Torch Hub 转换路径仍保留。该项及全部真实平台项目继续记录 -在人工矩阵中。 +仓库固定 [`pet-models-v1`](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/releases/tag/pet-models-v1) +TorchScript Release 已发布；生产仅下载 manifest 固定的 URL/SHA/size，不再执行 +Torch Hub、DINO source 或用户侧 trace。其余真实平台项目继续记录在人工矩阵中。 定向验收覆盖 terminal 唯一性、取消后的迟到 import、Recognition 首用预热、构建清单和 A/B 拒绝路径。完整平台性能数字仍按 `STARTUP_BENCHMARK_RUNBOOK.md` 采集。 diff --git a/docs/requirements/startup-chain-optimization/PETS_REVIEW_REMEDIATION_LEDGER.md b/docs/requirements/startup-chain-optimization/PETS_REVIEW_REMEDIATION_LEDGER.md index 92c11a65b..8fb90c814 100644 --- a/docs/requirements/startup-chain-optimization/PETS_REVIEW_REMEDIATION_LEDGER.md +++ b/docs/requirements/startup-chain-optimization/PETS_REVIEW_REMEDIATION_LEDGER.md @@ -37,8 +37,11 @@ | R2-09 | 每批加载全部 profiles 并重建 ANN；现有规模测试只预装 50k 后测 +2。 | session 级、按 contract/species 分区的增量 USearch index；只读取和更新受影响候选；初始、重建和 batch 16 更新均使用批量 `index.add`。 | `pets-scale-contract` 覆盖空库 batch 16 增长到 1k/10k/50k 及 50k+2，`1 passed in 65.08s`。 | `automated_pass` | | R2-10 | journal recovery 仅在新扫描批次前触发，其他 mutation 可越过旧 applying 操作。 | 初始化和每个 public mutation 前按创建顺序恢复；失败时拒绝新 mutation。 | `test_public_mutation_cannot_overtake_unrecovered_journal_owner` 通过。 | `automated_pass` | | R2-11 | `_publish_staged_thumbnails` 的逐文件 replace 没有内部补偿或预登记正式目标。 | publish 前 journal 记录清单；部分失败反向清理，进程恢复也能识别 orphan。 | `test_thumbnail_publish_compensates_when_later_replace_fails` 及 scan recovery 回归通过。 | `automated_pass` | -| R2-12 | DINO 仅固定 Torch Hub source revision，首次权重内容没有发布前 SHA。 | 生产运行时改用项目固定 Release TorchScript，manifest 固定 URL/SHA/size/shape；Torch Hub 仅保留开发转换工具。 | 已固定并在首次加载前验证官方 checkpoint SHA/size，生成 cache 复核 metadata/hash；仓库 `pet-models-v1` 不可变 TorchScript Release 尚未发布，生产 Torch Hub 路径尚不能删除。 | `manual_pending` | +| R2-12 | DINO 仅固定 Torch Hub source revision，首次权重内容没有发布前 SHA。 | 生产运行时改用项目固定 Release TorchScript，manifest 固定 URL/SHA/size/shape；Torch Hub 仅保留发布转换工具。 | [`pet-models-v1`](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/releases/tag/pet-models-v1) 已发布；候选 [run 32807111211](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/actions/runs/32807111211) 使用固定 source/checkpoint/Torch 并在 Ubuntu、macOS、Windows 验证同一资产。生产源码无 Torch Hub/xFormers 分支。 | `automated_pass` | | R2-13 | `.github/workflows/test.yml` 的 PR base 仅允许 `main`，当前 stacked PR 没有 head status。 | 同时允许 `main` 与 `codex/startup-chain-optimization`，增加手工触发，并为独立 Pets job 配置 Linux Qt headless runtime。 | `e0001ee646e95fadc33659ebe277eb067e79a084` 的 9 个 GitHub Actions job 全部成功：[run 30263442145](https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/actions/runs/30263442145)。 | `automated_pass` | +| R2-14 | 损坏的 authoritative YOLOX override 因 SHA `RuntimeError` 无法自愈。 | detector 在通用跨进程 acquisition lock 内重验；允许下载时只在原 root 删除并原子重取，关闭下载时保留文件；override 永不 fallback。 | override repair、download-disabled、permission/non-permission 与双进程单次下载合同通过。 | `automated_pass` | +| R2-15 | v1 artifact run、build manifest 和 Release tag provenance 不一致。 | builder 显式 checkout head；发布入口校验 run API/build manifest/builder checkout；v1 一次性受保护修正后永久拒绝覆盖。 | run `32807111211`、build manifest、tag 和 Release target 均固定 `ceb24fa0…`；公开 SHA `7f8e204c…ab13`。 | `automated_pass` | +| R2-16 | 自动 recognition 属于启动策略变化，缺少 +1.5s/+5s 资源和 lifecycle 证据。 | 增加调度 profile、CPU/RSS/I/O 采样、同 commit policy 开关和安全逐样本图库恢复。 | [PR #917 benchmark](PR917_RECOGNITION_IDLE_BENCHMARK.md)：30-sample resource、10-pair alternating A/B、30 quick-close、5 missing-model 均通过；其它 packaged 平台保留人工项。 | `automated_source_evidence / packaged_manual_pending` | 二次审查的工程关闭上限为 `automated_remediation_complete / manual_validation_pending`。远端 CI、真实安装权限、 diff --git a/docs/requirements/startup-chain-optimization/PR917_RECOGNITION_IDLE_BENCHMARK.md b/docs/requirements/startup-chain-optimization/PR917_RECOGNITION_IDLE_BENCHMARK.md new file mode 100644 index 000000000..1be8d5ab3 --- /dev/null +++ b/docs/requirements/startup-chain-optimization/PR917_RECOGNITION_IDLE_BENCHMARK.md @@ -0,0 +1,59 @@ +# PR #917 Recognition Idle-Start Benchmark + +## Context + +- Runtime: source checkout `e91920f4`, macOS arm64, Cocoa/Metal, hot cache. +- Library: disposable `/tmp` copy restored before every sample from the same + five-image template; the operator library was never modified. +- Policy arms: default `auto_after_metadata_idle` and the same commit with + `IPHOTO_STARTUP_RECOGNITION_AUTO_START=0`. +- Sampling: 100 ms child-process CPU/RSS sampling. macOS psutil does not expose + per-process I/O counters, so read/write bytes are explicitly `null`. + +Two 30-sample auto batches and two 30-sample feature-scoped batches were +collected. Because whole-batch ordering showed thermal/cache bias, an additional +10-pair alternating A/B batch is the causal timing comparison below. + +## Alternating A/B timing + +| Metric | Feature-scoped P50/P95 | Auto P50/P95 | Auto delta P50/P95 | +|---|---:|---:|---:| +| show → interactive | 19.945 / 21.592 ms | 19.825 / 21.815 ms | -0.60% / +1.03% | +| process → first gallery | 798.130 / 1064.118 ms | 797.375 / 1057.369 ms | -0.09% / -0.63% | +| process → first usable thumbnail | 3259.572 / 3614.575 ms | 3257.639 / 3514.421 ms | -0.06% / -2.77% | +| max post-interactive GUI job | 139.608 / 186.723 ms | 140.911 / 183.887 ms | +0.93% / -1.52% | + +Recognition activation occurred after the first usable thumbnail. Automatic +startup introduced no measurable first-frame/gallery regression in the +alternating comparison. The absolute post-interactive job duration is a +pre-existing coordinator construction cost present in both arms; post- +recognition GUI job stall was `0 ms` in all 30 auto samples. + +## Auto-start resource envelope + +The second 30-sample auto batch recorded: + +| Snapshot | CPU P50/P95 | RSS P50/P95 | +|---|---:|---:| +| interactive | 358.603 / 379.743 ms | 184.4 / 186.7 MiB | +| recognition activation | 1851.815 / 1986.214 ms | 536.3 / 543.4 MiB | +| activation +1.5 s | 2593.641 / 2671.004 ms | 656.1 / 672.4 MiB | +| activation +5 s | 8510.109 / 8786.367 ms | 1395.0 / 1469.3 MiB | + +These values describe intentional background AI work, not the first-frame path. +The worker threads run at `LowestPriority`; input before activation resets the +1500 ms gate. + +## Failure and lifecycle scenarios + +- Quick-close: 30/30 valid with `--allow-degraded`; no sample emitted + `recognition.worker.started` and no late-QThread/fatal diagnostics appeared. +- Missing Pets models with downloads disabled: 5/5 valid; each run reported the + missing YOLOX model, retained pending work, remained interactive, and had + `0 ms` post-recognition GUI job stall. +- 50k backlog/index behavior remains covered by the cross-platform + `pets-production-shape-contract`; startup generation, cancellation, and + worker admission are covered by the three-platform startup contracts. + +This is controlled source-runtime evidence on Apple Silicon. Packaged Windows, +Linux AppImage, and macOS Intel resource figures remain manual validation items. diff --git a/docs/requirements/startup-chain-optimization/STARTUP_BENCHMARK_RUNBOOK.md b/docs/requirements/startup-chain-optimization/STARTUP_BENCHMARK_RUNBOOK.md index feaaee500..c24af48b9 100644 --- a/docs/requirements/startup-chain-optimization/STARTUP_BENCHMARK_RUNBOOK.md +++ b/docs/requirements/startup-chain-optimization/STARTUP_BENCHMARK_RUNBOOK.md @@ -62,6 +62,52 @@ and add both `--cache-eviction-method METHOD` and `--confirm-controlled-cold-cache`. The tool records but intentionally does not elevate privileges or purge caches itself. +### Recognition idle-start evidence + +Use the same candidate commit for both arms. The default is the automatic +policy; add `--set-env IPHOTO_STARTUP_RECOGNITION_AUTO_START=0` for the +feature-scoped baseline. Explicit People-page activation remains enabled. Keep +the application alive long enough to observe five seconds after recognition +activation and enable child-process resource sampling: + +```bash +.venv/bin/python tools/startup_benchmark.py collect \ + --revision CURRENT_SHA \ + --scenario recognition-auto-models-present \ + --runtime source \ + --qt-backend cocoa \ + --graphics-backend metal \ + --cache-state hot \ + --samples 30 \ + --auto-exit-delay-ms 10000 \ + --timeout-seconds 30 \ + --sample-resources \ + --resource-sample-interval-ms 100 \ + --library /absolute/path/to/dedicated-recognition-library \ + --confirm-dedicated-library \ + --output-dir benchmark-output/recognition/candidate/models-present \ + -- .venv/bin/python -m iPhoto.gui.main +``` + +For repeated recognition runs, preserve one read-only prepared template and use +`--library-template TEMPLATE --library OUTPUT_DIR/active-library +--confirm-template-restore`. The collector refuses any other restore target, +recreates only that disposable path before each sample, and never mutates the +template or operator photo library. + +Repeat with scenarios `recognition-auto-missing-models` (empty isolated model +root plus `--set-env IPHOTO_PET_MODEL_AUTO_DOWNLOAD=0`), +`recognition-auto-50k-pending` (prepared 50k +status backlog), and `recognition-quick-close` (`--auto-exit-delay-ms 250`). +Quick-close collection also uses `--allow-degraded`, because it intentionally +exits before a usable thumbnail is required. +The JSON/Markdown summary records CPU time, RSS, read bytes, and write bytes at +interactive, recognition activation, activation +1.5 s, and activation +5 s. +First-gallery/thumbnail P50 and P95 must not regress, post-interactive GUI jobs +are reported against the same-commit baseline, post-recognition GUI jobs remain +below 100 ms, and quick-close must have no recognition worker start or late +QThread diagnostics. + Windows packaged runs use the same CLI under PowerShell, with the Nuitka `.exe` after `--`. Keep Defender enabled. Collect separate `local-ssd-indexed`, `offline-removable`, and `delayed-smb` scenarios and verify after the batch that diff --git a/pyproject.toml b/pyproject.toml index 0c7c68edf..fe8507610 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,12 +38,14 @@ test = [ "pytest-mock", "pytest-qt", "pytest-timeout", + "psutil>=5.9", ] dev = [ "pytest", "pytest-mock", "pytest-qt", "pytest-timeout", + "psutil>=5.9", "ruff", "black", "mypy", @@ -57,8 +59,7 @@ ai-demo = [ pets-ai = [ "certifi>=2024", "onnxruntime>=1.18,<2", - "torch", - "torchvision", + "torch==2.12.1", "usearch>=2.26,<3", ] diff --git a/src/iPhoto/bootstrap/runtime_context.py b/src/iPhoto/bootstrap/runtime_context.py index 6c2659ccd..cf77df468 100644 --- a/src/iPhoto/bootstrap/runtime_context.py +++ b/src/iPhoto/bootstrap/runtime_context.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING @@ -21,6 +22,7 @@ from .library_session import LibrarySession _logger = logging.getLogger(__name__) +STARTUP_RECOGNITION_AUTO_START_ENV = "IPHOTO_STARTUP_RECOGNITION_AUTO_START" def _create_settings_manager() -> "SettingsManager": @@ -277,6 +279,12 @@ def open_initial_collection(self) -> Path | None: def schedule_idle_startup_jobs(self) -> None: self.start_deferred_startup_scan() + raw_policy = str(os.environ.get(STARTUP_RECOGNITION_AUTO_START_ENV, "1")).strip().lower() + if raw_policy in {"0", "false", "no", "off"}: + return + requester = getattr(self.library, "request_startup_recognition_after_idle", None) + if callable(requester): + requester() def start_deferred_startup_scan(self) -> None: """Start a scan that was intentionally delayed until after first gallery load.""" diff --git a/src/iPhoto/gui/main.py b/src/iPhoto/gui/main.py index fd7c0ab83..7574b15a8 100644 --- a/src/iPhoto/gui/main.py +++ b/src/iPhoto/gui/main.py @@ -274,6 +274,115 @@ def _belongs_to_window(self, watched: QObject) -> bool: return False +class _RecognitionIdleActivityFilter(QObject): + """Reset the startup recognition idle gate without consuming user input.""" + + def __init__( + self, + window: QObject, + app: QApplication, + callback: Callable[[], None], + ) -> None: + try: + super().__init__(window) + except TypeError: + super().__init__() + self._window = window + self._app = app + self._callback = callback + self._installed = False + + def install(self) -> None: + if self._installed: + return + install_filter = getattr(self._app, "installEventFilter", None) + if not callable(install_filter): + return + install_filter(self) + self._installed = True + + def release(self) -> None: + if not self._installed: + return + remove_filter = getattr(self._app, "removeEventFilter", None) + if callable(remove_filter): + try: + remove_filter(self) + except RuntimeError: + pass + self._installed = False + + def eventFilter(self, watched: QObject, event: QEvent) -> bool: # noqa: N802 + if not self._installed or event.type() not in _STARTUP_INPUT_EVENT_TYPES: + return False + if not self._belongs_to_window(watched): + return False + if event.type() == QEvent.Type.MouseMove: + buttons = getattr(event, "buttons", None) + if not callable(buttons) or buttons() == Qt.MouseButton.NoButton: + return False + self._callback() + return False + + def _belongs_to_window(self, watched: QObject) -> bool: + if self._object_is_owned_by_window(watched): + return True + for getter_name in ("activePopupWidget", "activeModalWidget"): + getter = getattr(self._app, getter_name, None) + active = getter() if callable(getter) else None + if active is None or not self._object_contains(active, watched): + continue + if self._object_is_owned_by_window(active): + return True + return False + + def _object_is_owned_by_window(self, value: object) -> bool: + main_handle_getter = getattr(self._window, "windowHandle", None) + main_handle = main_handle_getter() if callable(main_handle_getter) else None + pending = [value] + seen: set[int] = set() + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + if current is self._window or current is main_handle: + return True + is_ancestor = getattr(self._window, "isAncestorOf", None) + if callable(is_ancestor): + try: + if is_ancestor(current): + return True + except (RuntimeError, TypeError): + pass + for getter_name in ("parent", "parentWidget", "transientParent"): + getter = getattr(current, getter_name, None) + if callable(getter): + try: + pending.append(getter()) + except RuntimeError: + pass + window_handle_getter = getattr(current, "windowHandle", None) + if callable(window_handle_getter): + try: + pending.append(window_handle_getter()) + except RuntimeError: + pass + return False + + @staticmethod + def _object_contains(container: object, watched: object) -> bool: + if container is watched: + return True + is_ancestor = getattr(container, "isAncestorOf", None) + if not callable(is_ancestor): + return False + try: + return bool(is_ancestor(watched)) + except (RuntimeError, TypeError): + return False + + def _bootstrap_macos_external_tool_path() -> None: """Expose common Homebrew/MacPorts tool paths to GUI-launched app bundles.""" @@ -687,6 +796,13 @@ def _handle_startup_job_failure(failure) -> None: set_startup_orchestrator(startup) startup_input_guard = _StartupInputGuard(window, app) startup_input_guard.install() + recognition_activity_filter = _RecognitionIdleActivityFilter( + window, + app, + lambda: context.library.notify_user_activity(), + ) + recognition_activity_filter.install() + setattr(window, "_recognition_idle_activity_filter", recognition_activity_filter) from iPhoto.bootstrap.library_probe import LibraryProbeController diff --git a/src/iPhoto/library/runtime_controller.py b/src/iPhoto/library/runtime_controller.py index ef2d460d5..c324ec88c 100644 --- a/src/iPhoto/library/runtime_controller.py +++ b/src/iPhoto/library/runtime_controller.py @@ -26,6 +26,7 @@ Signal, ) +from ..bootstrap.startup_profile import mark from ..errors import LibraryUnavailableError from ..utils.logging import get_logger @@ -44,6 +45,7 @@ from .watch_service import LibraryWatchResult, LibraryWatchService LOGGER = get_logger() +_STARTUP_RECOGNITION_IDLE_MS = 1500 if TYPE_CHECKING: # pragma: no cover from ..application.ports import ( @@ -142,6 +144,13 @@ def __init__(self, parent: QObject | None = None) -> None: self._recognition_services_root: Path | None = None self._recognition_scans_root: Path | None = None self._recognition_generation = 0 + self._startup_recognition_request: tuple[Path, int] | None = None + self._startup_recognition_timer = QTimer(self) + self._startup_recognition_timer.setSingleShot(True) + self._startup_recognition_timer.setInterval(_STARTUP_RECOGNITION_IDLE_MS) + self._startup_recognition_timer.timeout.connect( + self._activate_startup_recognition_after_idle + ) self._delivered_recognition_event_ids: set[str] = set() self._retiring_recognition_workers: set[QThread] = set() self._library_session: "LibrarySession | None" = None @@ -331,6 +340,7 @@ def shutdown(self) -> None: """Stop background workers and watchers during application shutdown.""" had_scanner_worker = self._current_scanner_worker is not None + self._cancel_startup_recognition_request() self.stop_scanning(wait=True) self._recognition_services_root = None self._recognition_scans_root = None @@ -657,6 +667,12 @@ def activate_recognition_scans(self) -> None: """Start model workers after a recognition viewport is usable.""" root = self._root + scanner = self._current_scanner_worker + if scanner is not None and getattr( + scanner, "_defer_ai_workers_until_scan_finished", False + ): + self.request_startup_recognition_after_idle() + return if ( root is None or root != self._recognition_services_root @@ -670,6 +686,85 @@ def activate_recognition_scans(self) -> None: if self._start_ai_scan_workers(root, startup=True): self._recognition_scans_root = root + def request_startup_recognition_after_idle(self) -> None: + """Start People/Pets after startup metadata is complete and input is idle.""" + + root = self._root + if root is None: + return + generation = int(self._recognition_generation) + self._startup_recognition_request = (Path(root), generation) + scanner = self._current_scanner_worker + waits_for_scan = scanner is not None and getattr( + scanner, "_defer_ai_workers_until_scan_finished", False + ) + mark( + "recognition.startup.requested", + generation=generation, + waits_for_scan=waits_for_scan, + ) + if waits_for_scan: + return + self._arm_startup_recognition_idle_timer(Path(root), generation) + + def notify_user_activity(self) -> None: + """Postpone a pending startup recognition scan without consuming input.""" + + request = self._startup_recognition_request + if request is None or not self._startup_recognition_timer.isActive(): + return + self._startup_recognition_timer.start(_STARTUP_RECOGNITION_IDLE_MS) + mark( + "recognition.startup.idle_reset", + generation=request[1], + delay_ms=_STARTUP_RECOGNITION_IDLE_MS, + ) + + def _arm_startup_recognition_idle_timer(self, root: Path, generation: int) -> None: + request = self._startup_recognition_request + if ( + request != (Path(root), int(generation)) + or self._root != Path(root) + or int(self._recognition_generation) != int(generation) + ): + return + self._startup_recognition_timer.start(_STARTUP_RECOGNITION_IDLE_MS) + mark( + "recognition.startup.idle_armed", + generation=generation, + delay_ms=_STARTUP_RECOGNITION_IDLE_MS, + ) + + def _cancel_startup_recognition_request(self, *, reason: str = "cancelled") -> None: + request = self._startup_recognition_request + self._startup_recognition_timer.stop() + self._startup_recognition_request = None + if request is not None: + mark( + "recognition.startup.cancelled", + generation=request[1], + reason=reason, + ) + + def _activate_startup_recognition_after_idle(self) -> None: + request = self._startup_recognition_request + if request is None: + return + root, generation = request + self._startup_recognition_request = None + if ( + self._root != root + or int(self._recognition_generation) != generation + or self._current_scanner_worker is not None + ): + return + session = self._library_session + if session is None or Path(session.library_root) != root: + return + mark("recognition.startup.activated", generation=generation) + self.bind_recognition_services(session.people, session.pets) + self.activate_recognition_scans() + def activate_recognition_services( self, people_service: PeopleService | None, diff --git a/src/iPhoto/library/scan_coordinator.py b/src/iPhoto/library/scan_coordinator.py index 9aba4ffa8..8b8516b10 100644 --- a/src/iPhoto/library/scan_coordinator.py +++ b/src/iPhoto/library/scan_coordinator.py @@ -19,6 +19,18 @@ LOGGER = get_logger() +def _prepare_face_runtime_imports() -> None: + """Serialize cv2 initialization before handing work to a Face QThread.""" + + try: + from ..people.pipeline import prepare_face_runtime_imports + + prepare_face_runtime_imports() + except ImportError: + # The worker preserves the existing graceful optional-runtime error. + LOGGER.debug("Face runtime import preflight is unavailable", exc_info=True) + + class _PairingWorker(QRunnable): """Run live-photo pairing off the main thread after a scan completes.""" @@ -77,6 +89,12 @@ def start_scanning( All scanned assets are written to the global database at the library root. """ + if not startup: + cancel_startup_recognition = getattr( + self, "_cancel_startup_recognition_request", None + ) + if callable(cancel_startup_recognition): + cancel_startup_recognition() # Scanner and face-recognition workers bring Pillow/NumPy and optional # AI runtimes into the process. Import them only when a scan actually # starts, never while constructing the first window frame. @@ -158,6 +176,7 @@ def _start_ai_scan_workers(self, library_root: Path, *, startup: bool = False) - started: list[object] = [] if self._current_face_scanner is None: + _prepare_face_runtime_imports() from .workers.face_scan_worker import FaceScanWorker face_worker = FaceScanWorker( @@ -210,7 +229,17 @@ def _start_ai_scan_workers(self, library_root: Path, *, startup: bool = False) - try: if startup: worker.finish_input() - worker.start() + worker.start(QThread.Priority.LowestPriority) + else: + worker.start() + mark( + "recognition.worker.started", + generation=generation, + worker=( + "face" if worker is self._current_face_scanner else "pet" + ), + startup=bool(startup), + ) except Exception: # noqa: BLE001 all_started = False if self._current_face_scanner is worker: @@ -263,6 +292,11 @@ def _start_pet_backfill_worker(self, library_root: Path) -> None: def stop_scanning(self, *, wait: bool = False, timeout_ms: int = 2000) -> None: """Cancel the currently running scan, if any.""" + cancel_startup_recognition = getattr( + self, "_cancel_startup_recognition_request", None + ) + if callable(cancel_startup_recognition): + cancel_startup_recognition() _locker = QMutexLocker(self._scan_buffer_lock) scanner_worker = self._current_scanner_worker face_scanner = self._current_face_scanner @@ -284,10 +318,20 @@ def stop_scanning(self, *, wait: bool = False, timeout_ms: int = 2000) -> None: if deferred_queue is not None: deferred_queue.clear() if self._current_face_scanner is not None: + mark( + "recognition.worker.cancelled", + generation=self._recognition_generation, + worker="face", + ) self._current_face_scanner.cancel() self._retiring_recognition_workers.add(self._current_face_scanner) self._current_face_scanner = None if self._current_pet_scanner is not None: + mark( + "recognition.worker.cancelled", + generation=self._recognition_generation, + worker="pet", + ) self._current_pet_scanner.cancel() self._retiring_recognition_workers.add(self._current_pet_scanner) self._current_pet_scanner = None @@ -602,10 +646,14 @@ def _on_scan_finished( pet_scanner.finish_input() if worker.cancelled: + if defer_ai_workers: + self._cancel_startup_recognition_request() self.scanFinished.emit(root, False) self._start_next_deferred_scan() return if worker.failed: + if defer_ai_workers: + self._cancel_startup_recognition_request() self.scanFinished.emit(root, False) self._start_next_deferred_scan() return @@ -623,21 +671,46 @@ def _on_scan_finished( ) except Exception as exc: LOGGER.warning("Failed to persist scan finalization for %s: %s", root, exc) + if defer_ai_workers: + self._cancel_startup_recognition_request() self.scanFinished.emit(root, False) self._start_next_deferred_scan() return - # A startup metadata scan only prepares the Gallery index. People and - # Pets are feature-scoped services and their model workers are started - # by ``activate_recognition_services`` on first use. Starting them here - # used to add model pressure immediately after startup completion and - # made a quick window close race QThread destruction. + # Startup metadata enumeration stays AI-free. Once it commits, arm the + # interaction-idle gate for both People and Pets. + startup_recognition_generation = ( + int(getattr(self, "_recognition_generation", 0)) + if defer_ai_workers + else None + ) + if startup_recognition_generation is not None: + mark( + "recognition.startup.scan_ready", + generation=startup_recognition_generation, + ) # Emit immediately so the UI (status bar, map refresh) can react without # waiting for the potentially slow live-photo pairing step. self.scanFinished.emit(root, True) self._start_next_deferred_scan() + # scanFinished listeners may synchronously close/rebind the runtime. In + # that case stop_scanning() advances the recognition generation; do not + # create a new QThread during teardown. A queued metadata scan also gets + # priority and will schedule its own recognition work when appropriate. + if ( + startup_recognition_generation is not None + and int(getattr(self, "_recognition_generation", 0)) + == startup_recognition_generation + and self._current_scanner_worker is None + ): + library_root = getattr(self, "_root", None) + if library_root is not None: + arm_idle = getattr(self, "_arm_startup_recognition_idle_timer", None) + if callable(arm_idle): + arm_idle(Path(library_root), startup_recognition_generation) + # Persist live-photo pairings in the background to avoid blocking the # main thread while downstream listeners start refreshing. self._scan_thread_pool.start(_PairingWorker(root, scan_service)) @@ -651,6 +724,12 @@ def _on_scan_error(self, worker: ScannerWorker, root: Path, message: str) -> Non pet_scanner = self._current_pet_scanner self._live_scan_root = None del locker + if getattr(worker, "_defer_ai_workers_until_scan_finished", False): + cancel_startup_recognition = getattr( + self, "_cancel_startup_recognition_request", None + ) + if callable(cancel_startup_recognition): + cancel_startup_recognition() if face_scanner is not None: face_scanner.finish_input() if pet_scanner is not None: diff --git a/src/iPhoto/people/pipeline.py b/src/iPhoto/people/pipeline.py index f1da6a923..730948fcb 100644 --- a/src/iPhoto/people/pipeline.py +++ b/src/iPhoto/people/pipeline.py @@ -7,6 +7,7 @@ import logging import os import sys +import threading import typing import uuid from collections import Counter, defaultdict, deque @@ -48,6 +49,24 @@ DEFAULT_FACE_TINY_AREA_RATIO = 0.0005 DEFAULT_FACE_SMALL_AREA_RATIO = 0.005 DEFAULT_FACE_RELATIVE_AREA_RATIO = 0.15 +_FACE_RUNTIME_IMPORT_LOCK = threading.Lock() +_FACE_RUNTIME_IMPORT_READY = False + + +def prepare_face_runtime_imports() -> None: + """Complete cv2's process-global import before a Face QThread starts.""" + + global _FACE_RUNTIME_IMPORT_READY + if _FACE_RUNTIME_IMPORT_READY: + return + with _FACE_RUNTIME_IMPORT_LOCK: + if _FACE_RUNTIME_IMPORT_READY: + return + # Importing cv2 concurrently with PySide/Shiboken signature mapping can + # terminate the process on macOS (KeyError: cv2.Error). Model creation + # and downloads remain lazy; only the native module import is serialized. + __import__("cv2") + _FACE_RUNTIME_IMPORT_READY = True @dataclass(frozen=True) diff --git a/src/iPhoto/pets/_pipeline_impl.py b/src/iPhoto/pets/_pipeline_impl.py new file mode 100644 index 000000000..3f9745702 --- /dev/null +++ b/src/iPhoto/pets/_pipeline_impl.py @@ -0,0 +1,2501 @@ +"""Pet detection, embedding, clustering, and identity helpers.""" + +from __future__ import annotations + +import errno +import hashlib +import json +import logging +import math +import os +import ssl +import sys +import tempfile +import time +import uuid +from collections import Counter, defaultdict +from collections.abc import Callable, Sequence +from contextlib import contextmanager, suppress +from dataclasses import dataclass, replace +from enum import StrEnum +from pathlib import Path +from urllib import request +from urllib.parse import urlparse + +import numpy as np +from PIL import Image + +from iPhoto.utils.pathutils import LibraryAssetPathError, resolve_library_asset_path + +from .errors import ( + PetInferenceError, + PetModelUnavailableError, + PetPipelineInvariantError, + PetRuntimeUnavailableError, +) +from .image_utils import ( + PetImageLoadError, + crop_pet_region, + image_to_chw_float, + load_image_rgb, + save_pet_thumbnail, +) +from .records import PetDetectionRecord, PetProfile, PetRecord +from .repository_utils import ( + compute_cluster_center, + cosine_distance, + cosine_distance_matrix, + key_detection_sort_key, + normalize_vector, + profile_state_for_sample_count, + utc_now_iso, +) +from .state_repository import PetStateRepository + + +class _ModelStoragePermissionError(OSError): + pass + + +_MODEL_STORAGE_ERRNOS = { + errno.EACCES, + errno.EPERM, + errno.EROFS, +} + + +def _raise_if_model_storage_error(exc: OSError, path: Path) -> None: + if exc.errno in _MODEL_STORAGE_ERRNOS: + raise _ModelStoragePermissionError( + exc.errno, + f"model storage is not writable: {path}", + ) from exc + raise exc + + +def _load_pet_model_manifest() -> dict: + manifest_path = Path(__file__).with_name("model_manifest.json") + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + detector = manifest["detector"] + embedder = manifest["embedder"] + except (OSError, json.JSONDecodeError, KeyError, TypeError) as exc: + raise PetPipelineInvariantError(f"Invalid Pets model manifest: {manifest_path}") from exc + if int(manifest.get("schema_version") or 0) != 2: + raise PetPipelineInvariantError(f"Unsupported Pets model manifest: {manifest_path}") + if urlparse(str(detector.get("url") or "")).scheme.lower() != "https": + raise PetPipelineInvariantError("Pets detector manifest URL must use HTTPS.") + if detector.get("input") != { + "layout": "NCHW", + "channel_order": "BGR", + "dtype": "float32", + "range": [0, 255], + "shape": [1, 3, 416, 416], + }: + raise PetPipelineInvariantError("Pets detector manifest input contract is invalid.") + if embedder.get("input_shape") != [1, 3, 224, 224]: + raise PetPipelineInvariantError("Pets embedder manifest input contract is invalid.") + if embedder.get("source_repository") != "facebookresearch/dinov2": + raise PetPipelineInvariantError("Pets embedder source repository is invalid.") + if embedder.get("source_revision") != ( + "7764ea0f912e53c92e82eb78a2a1631e92725fc8" + ): + raise PetPipelineInvariantError("Pets embedder source revision is invalid.") + if embedder.get("source_tree_sha1") != "2a27257b79b0633b027a21014bc9360e3c1b3f43": + raise PetPipelineInvariantError("Pets embedder source tree is invalid.") + weights_url = str(embedder.get("weights_url") or "").strip() + if urlparse(weights_url).scheme.lower() != "https": + raise PetPipelineInvariantError("Pets embedder checkpoint URL must use HTTPS.") + weights_sha256 = str(embedder.get("weights_sha256") or "") + if len(weights_sha256) != 64 or any( + char not in "0123456789abcdef" for char in weights_sha256.lower() + ): + raise PetPipelineInvariantError("Pets embedder checkpoint SHA-256 is invalid.") + if int(embedder.get("weights_size") or 0) <= 0: + raise PetPipelineInvariantError("Pets embedder checkpoint size is invalid.") + if embedder.get("artifact_kind") != "release_torchscript": + raise PetPipelineInvariantError("Pets embedder must be a release TorchScript artifact.") + if int(embedder.get("cache_schema_version") or 0) != 2: + raise PetPipelineInvariantError("Pets embedder cache schema is invalid.") + if embedder.get("release_tag") != "pet-models-v1": + raise PetPipelineInvariantError("Pets embedder release tag is invalid.") + if embedder.get("producer_torch_version") != "2.12.1": + raise PetPipelineInvariantError("Pets embedder Torch producer version is invalid.") + if embedder.get("producer_torchvision_version") != "0.27.1": + raise PetPipelineInvariantError("Pets embedder Torchvision producer version is invalid.") + if embedder.get("producer_python_version") != "3.12": + raise PetPipelineInvariantError("Pets embedder Python producer version is invalid.") + torchscript_url = str(embedder.get("torchscript_url") or "").strip() + expected_torchscript_url = ( + "https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/" + "releases/download/pet-models-v1/dinov2_vits14.pt" + ) + if ( + urlparse(torchscript_url).scheme.lower() != "https" + or torchscript_url != expected_torchscript_url + ): + raise PetPipelineInvariantError("Pets embedder TorchScript URL must use HTTPS.") + if len(str(embedder.get("torchscript_sha256") or "")) != 64: + raise PetPipelineInvariantError("Pets embedder TorchScript SHA-256 is invalid.") + if int(embedder.get("torchscript_size") or 0) <= 0: + raise PetPipelineInvariantError("Pets embedder TorchScript size is invalid.") + return manifest + + +PET_MODEL_MANIFEST = _load_pet_model_manifest() +_DETECTOR_MANIFEST = PET_MODEL_MANIFEST["detector"] +_EMBEDDER_MANIFEST = PET_MODEL_MANIFEST["embedder"] +SUPPORTED_DEFAULT_SPECIES = frozenset({"cat", "dog"}) +PET_DETECTOR_PIPELINE_VERSION = "yolox-letterbox-tiles-people-priority-v6" +PET_CLUSTERING_PIPELINE_VERSION = "species-bounded-single-link-v3" +PET_EMBEDDING_PIPELINE_VERSION = "dinov2-vits14-imagenet-normalized-v1" +PET_KEY_VERSION = "v2" +PET_DETECTOR_KEY_VERSION = "yolox-nano-coco-0.1.1rc0-raw-bgr-v1" +DEFAULT_PET_DISTANCE_THRESHOLD = 0.42 +PET_CLUSTER_DIAMETER_MULTIPLIER = 1.5 +PET_PET_IOU_THRESHOLD = 0.50 +PET_PET_SMALLER_BOX_COVERAGE_THRESHOLD = 0.90 +PET_PET_NORMALIZED_CENTER_DISTANCE_THRESHOLD = 0.40 +PET_PET_CROSS_SPECIES_MUTUAL_COVERAGE_THRESHOLD = 0.90 +PET_PEOPLE_IOU_THRESHOLD = 0.50 +PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD = 0.90 +PET_PEOPLE_LARGER_PET_RATIO = 1.50 +PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD = 0.60 +PET_CANDIDATE_QUALITY_VERSION = "tiny-low-confidence-v1" +DEFAULT_PET_TINY_AREA_RATIO = 0.001 +DEFAULT_PET_TINY_MAX_CONFIDENCE = 0.45 +DEFAULT_PET_DETECTOR_MODEL_URL = str(_DETECTOR_MANIFEST["url"]) +PET_MODEL_AUTO_DOWNLOAD_ENV = "IPHOTO_PET_MODEL_AUTO_DOWNLOAD" +IPHOTO_PET_MODEL_DIR_ENV = "IPHOTO_PET_MODEL_DIR" +PET_DETECTOR_MODEL_URL_ENV = "IPHOTO_PET_DETECTOR_MODEL_URL" +PET_DETECTOR_MODEL_SHA256_ENV = "IPHOTO_PET_DETECTOR_MODEL_SHA256" +DEFAULT_PET_DETECTOR_MODEL_SHA256 = str(_DETECTOR_MANIFEST["sha256"]) +DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES = int(_DETECTOR_MANIFEST["max_bytes"]) +# The development conversion tool uses this immutable source revision. The +# production runtime builds a local cache from the fixed, hash-verified Meta +# checkpoint when no valid bundled artifact is available. +_DINO_SOURCE_REVISION = str(_EMBEDDER_MANIFEST["source_revision"]) +_DINO_WEIGHTS_URL = str(_EMBEDDER_MANIFEST["weights_url"]) +_DINO_WEIGHTS_SHA256 = str(_EMBEDDER_MANIFEST["weights_sha256"]).lower() +_DINO_WEIGHTS_SIZE = int(_EMBEDDER_MANIFEST["weights_size"]) +_DINO_TORCHSCRIPT_URL = str(_EMBEDDER_MANIFEST["torchscript_url"]) +_DINO_TORCHSCRIPT_SHA256 = str(_EMBEDDER_MANIFEST["torchscript_sha256"]).lower() +_DINO_TORCHSCRIPT_SIZE = int(_EMBEDDER_MANIFEST["torchscript_size"]) +_DINO_CACHE_SCHEMA_VERSION = int(_EMBEDDER_MANIFEST["cache_schema_version"]) +_DINO_TORCH_VERSION = str(_EMBEDDER_MANIFEST["producer_torch_version"]) +_DOWNLOAD_TIMEOUT_SECONDS = 60 +_DOWNLOAD_CHUNK_SIZE = 1024 * 256 +_MODEL_LOCK_RETRY_ERRNOS = {errno.EACCES, errno.EAGAIN, errno.EDEADLK} +_YOLOX_STRIDES = (8, 16, 32) +_YOLOX_RAW_COORD_LIMIT = 32.0 +_LOGGER = logging.getLogger(__name__) +COCO_ANIMAL_LABELS = { + 15: "cat", + 16: "dog", + 17: "horse", + 18: "sheep", + 19: "cow", + 20: "elephant", + 21: "bear", + 22: "zebra", + 23: "giraffe", +} + + +@dataclass(frozen=True) +class DetectedAssetPets: + asset_id: str + asset_rel: str + detections: list[PetDetectionRecord] + error: str | None = None + + +class PetIdentityResolutionSource(StrEnum): + KEY = "key" + REDIRECT_KEY = "redirect_key" + PROFILE = "profile" + REDIRECT_PROFILE = "redirect_profile" + NEW = "new" + + +@dataclass(frozen=True) +class PetIdentityResolution: + raw_pet_id: str + canonical_pet_id: str + source: PetIdentityResolutionSource + + @property + def is_redirect_alias(self) -> bool: + return self.source in { + PetIdentityResolutionSource.REDIRECT_KEY, + PetIdentityResolutionSource.REDIRECT_PROFILE, + } + + +@dataclass(frozen=True) +class _DetectedPetBox: + bbox: tuple[int, int, int, int] + confidence: float + species_label: str + quality_score: float = 0.0 + + +@dataclass(frozen=True) +class _YoloxPreprocessResult: + tensor: np.ndarray + resize_ratio: float + pad_left: int = 0 + pad_top: int = 0 + + +@dataclass(frozen=True) +class PetScanMetrics: + candidate_boxes: int = 0 + unsupported_species: int = 0 + too_small: int = 0 + pet_quality_rejected: int = 0 + people_overlaps: int = 0 + accepted_detections: int = 0 + pet_candidate_identities: int = 0 + pet_promotions: int = 0 + same_asset_cannot_link_hits: int = 0 + same_asset_manual_conflicts: int = 0 + + +@dataclass(frozen=True) +class _PetPeopleOverlapDecision: + suppressed: bool + reason: str = "" + pet_to_face_area_ratio: float = 0.0 + pet_image_coverage: float = 0.0 + + +class PetClusterPipeline: + def __init__( + self, + *, + model_root: Path, + detector_model_name: str = "yolox_nano_coco.onnx", + embedding_model_name: str = "dinov2_vits14", + allow_model_download: bool | None = None, + distance_threshold: float = DEFAULT_PET_DISTANCE_THRESHOLD, + min_pet_size: int = 48, + supported_species: frozenset[str] = SUPPORTED_DEFAULT_SPECIES, + detector_score_threshold: float = 0.30, + enable_tiled_detection: bool = True, + tile_scan_min_confidence: float | None = None, + tiny_area_ratio: float = DEFAULT_PET_TINY_AREA_RATIO, + tiny_max_confidence: float = DEFAULT_PET_TINY_MAX_CONFIDENCE, + ) -> None: + self._model_root = Path(model_root) + self._detector_model_name = detector_model_name + self._embedding_model_name = embedding_model_name + self._allow_model_download = ( + pet_model_auto_download_enabled() + if allow_model_download is None + else bool(allow_model_download) + ) + self._distance_threshold = float(distance_threshold) + self._min_pet_size = int(min_pet_size) + self._supported_species = frozenset(supported_species) + self._detector_score_threshold = float(detector_score_threshold) + self._enable_tiled_detection = bool(enable_tiled_detection) + self._tile_scan_min_confidence = ( + self._detector_score_threshold + if tile_scan_min_confidence is None + else float(tile_scan_min_confidence) + ) + self._tiny_area_ratio = float(tiny_area_ratio) + self._tiny_max_confidence = float(tiny_max_confidence) + self._detector: _YoloxOnnxPetDetector | None = None + self._embedder: _DinoV2Embedder | None = None + self._last_scan_metrics = PetScanMetrics() + + @property + def distance_threshold(self) -> float: + return self._distance_threshold + + @property + def detector_pipeline_version(self) -> str: + return PET_DETECTOR_PIPELINE_VERSION + + @property + def candidate_quality_version(self) -> str: + return PET_CANDIDATE_QUALITY_VERSION + + @property + def last_scan_metrics(self) -> PetScanMetrics: + return self._last_scan_metrics + + def detect_pets_for_rows( + self, + rows: list[dict], + *, + library_root: Path, + thumbnail_dir: Path, + published_thumbnail_dir: Path | None = None, + is_cancelled: Callable[[], bool] | None = None, + people_boxes_by_asset_id: dict[str, Sequence[tuple[int, int, int, int]]] | None = None, + ) -> list[DetectedAssetPets]: + if not rows: + return [] + embedder = self._ensure_embedder() + detector = self._ensure_detector() + cancellation_requested = is_cancelled or (lambda: False) + stored_thumbnail_dir = Path(published_thumbnail_dir or thumbnail_dir) + results: list[DetectedAssetPets] = [] + candidate_boxes = 0 + unsupported_species = 0 + too_small = 0 + pet_quality_rejected = 0 + people_overlaps = 0 + accepted_detections = 0 + excluded_people_boxes = people_boxes_by_asset_id or {} + for row in rows: + if cancellation_requested(): + break + asset_id = str(row.get("id") or "") + asset_rel = Path(str(row.get("rel") or "")).as_posix() + try: + image_path = resolve_library_asset_path(library_root, asset_rel) + image = load_image_rgb(image_path) + boxes = detector.detect(image) + except LibraryAssetPathError as exc: + results.append( + DetectedAssetPets( + asset_id=asset_id, + asset_rel=asset_rel, + detections=[], + error=str(exc), + ) + ) + continue + except PetImageLoadError as exc: + if cancellation_requested(): + break + results.append( + DetectedAssetPets( + asset_id=asset_id, + asset_rel=asset_rel, + detections=[], + error=str(exc).strip() or exc.__class__.__name__, + ) + ) + continue + except PetPipelineInvariantError: + raise + except Exception as exc: # noqa: BLE001 + if cancellation_requested(): + break + results.append( + DetectedAssetPets( + asset_id=asset_id, + asset_rel=asset_rel, + detections=[], + error=str(exc).strip() or exc.__class__.__name__, + ) + ) + continue + + image_width, image_height = image.size + detections: list[PetDetectionRecord] = [] + created_thumbnail_paths: list[Path] = [] + supported_boxes: list[_DetectedPetBox] = [] + quality_candidates: list[ + tuple[tuple[int, int, int, int], float, str, float] + ] = [] + candidate_boxes += len(boxes) + for detected in boxes: + if detected.species_label not in self._supported_species: + unsupported_species += 1 + continue + bbox = _normalize_bbox( + detected.bbox, + image_width=image_width, + image_height=image_height, + ) + if bbox[2] < self._min_pet_size or bbox[3] < self._min_pet_size: + too_small += 1 + continue + area_ratio = (bbox[2] * bbox[3]) / max(1, image_width * image_height) + if ( + area_ratio < self._tiny_area_ratio + and float(detected.confidence) < self._tiny_max_confidence + ): + pet_quality_rejected += 1 + continue + quality_candidates.append( + (bbox, float(detected.confidence), detected.species_label, area_ratio) + ) + + largest_pet_area_ratio = max( + (area_ratio for _bbox, _confidence, _species, area_ratio in quality_candidates), + default=0.0, + ) + for bbox, confidence, species_label, area_ratio in quality_candidates: + supported_boxes.append( + _DetectedPetBox( + bbox=bbox, + confidence=confidence, + species_label=species_label, + quality_score=_pet_candidate_quality_score( + confidence=confidence, + relative_area_ratio=area_ratio + / max(largest_pet_area_ratio, np.finfo(np.float32).eps), + ), + ) + ) + + deduped_boxes = _dedupe_supported_species_boxes(supported_boxes) + people_boxes = excluded_people_boxes.get(asset_id, ()) + accepted_boxes: list[_DetectedPetBox] = [] + for detected in deduped_boxes: + decision = _pet_people_overlap_decision( + detected.bbox, + people_boxes, + image_dimensions=(image_width, image_height), + ) + if decision.suppressed: + people_overlaps += 1 + _LOGGER.info( + "Suppressed pet candidate for asset %s: reason=%s " + "pet_to_face_area_ratio=%.3f pet_image_coverage=%.3f " + "iou_threshold=%.2f smaller_box_coverage_threshold=%.2f " + "larger_pet_ratio=%.2f mural_image_coverage_threshold=%.2f", + asset_id, + decision.reason, + decision.pet_to_face_area_ratio, + decision.pet_image_coverage, + PET_PEOPLE_IOU_THRESHOLD, + PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, + PET_PEOPLE_LARGER_PET_RATIO, + PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD, + ) + continue + accepted_boxes.append(detected) + + for detected in accepted_boxes: + bbox = detected.bbox + detection_id = uuid.uuid4().hex + thumbnail_path = thumbnail_dir / f"{detection_id}.png" + try: + crop = crop_pet_region(image, bbox, padding_ratio=0.08) + embedding = embedder.embed(crop) + save_pet_thumbnail(image, bbox, thumbnail_path, padding_ratio=0.08) + created_thumbnail_paths.append(thumbnail_path) + except PetPipelineInvariantError: + for created_path in created_thumbnail_paths: + created_path.unlink(missing_ok=True) + raise + except Exception as exc: # noqa: BLE001 + for created_path in created_thumbnail_paths: + try: + created_path.unlink(missing_ok=True) + except OSError: + _LOGGER.warning( + "Failed to roll back pet thumbnail %s", + created_path, + exc_info=True, + ) + results.append( + DetectedAssetPets( + asset_id=asset_id, + asset_rel=asset_rel, + detections=[], + error=str(exc).strip() or exc.__class__.__name__, + ) + ) + detections = [] + break + detections.append( + PetDetectionRecord( + detection_id=detection_id, + pet_key=build_pet_key( + asset_id=asset_id, + bbox=bbox, + image_width=image_width, + image_height=image_height, + species_label=detected.species_label, + ), + asset_id=asset_id, + asset_rel=asset_rel, + box_x=bbox[0], + box_y=bbox[1], + box_w=bbox[2], + box_h=bbox[3], + confidence=float(detected.confidence), + embedding=embedding, + embedding_dim=int(embedding.shape[0]), + embedding_model=self._embedding_model_name, + detector_model=self._detector_model_name, + thumbnail_path=(stored_thumbnail_dir / thumbnail_path.name) + .relative_to(stored_thumbnail_dir.parent) + .as_posix(), + pet_id=None, + detected_at=utc_now_iso(), + image_width=image_width, + image_height=image_height, + species_label=detected.species_label, + quality_score=detected.quality_score, + pet_key_version=PET_KEY_VERSION, + embedding_pipeline_version=PET_EMBEDDING_PIPELINE_VERSION, + ) + ) + has_asset_error = any( + result.asset_id == asset_id and result.error for result in results + ) + if detections or not has_asset_error: + accepted_detections += len(detections) + results.append( + DetectedAssetPets( + asset_id=asset_id, + asset_rel=asset_rel, + detections=detections, + ) + ) + self._last_scan_metrics = PetScanMetrics( + candidate_boxes=candidate_boxes, + unsupported_species=unsupported_species, + too_small=too_small, + pet_quality_rejected=pet_quality_rejected, + people_overlaps=people_overlaps, + accepted_detections=accepted_detections, + ) + return results + + def _ensure_detector(self) -> _YoloxOnnxPetDetector: + if self._detector is None: + model_path = self._resolve_model_path(Path("detector") / self._detector_model_name) + try: + self._detector = _YoloxOnnxPetDetector( + model_path, + score_threshold=self._detector_score_threshold, + allow_model_download=self._allow_model_download, + enable_tiled_detection=self._enable_tiled_detection, + tile_scan_min_confidence=self._tile_scan_min_confidence, + tile_species=self._supported_species, + ) + except ( + PetRuntimeUnavailableError, + PetModelUnavailableError, + PetPipelineInvariantError, + ): + raise + except RuntimeError as exc: + raise PetModelUnavailableError(str(exc)) from exc + return self._detector + + def _ensure_embedder(self) -> _DinoV2Embedder: + if self._embedder is None: + model_dir = self._resolve_model_path( + Path("embedding") / self._embedding_model_name, + directory=True, + ) + try: + self._embedder = _DinoV2Embedder( + model_dir, + model_name=self._embedding_model_name, + allow_model_download=self._allow_model_download, + ) + except ( + PetRuntimeUnavailableError, + PetModelUnavailableError, + PetPipelineInvariantError, + ): + raise + except RuntimeError as exc: + raise PetModelUnavailableError(str(exc)) from exc + return self._embedder + + def _resolve_model_path(self, relative_path: Path, *, directory: bool = False) -> Path: + override = pet_model_override_dir() + if override is not None and self._model_root != override: + raise PetModelUnavailableError( + "Pet scanning unavailable: model root does not match " + f"{IPHOTO_PET_MODEL_DIR_ENV}." + ) + if self._model_root == default_pet_model_dir(): + return resolve_pet_model_path(relative_path, directory=directory) + return self._model_root / relative_path + + +def build_pet_key( + *, + asset_id: str, + bbox: tuple[int, int, int, int], + image_width: int, + image_height: int, + species_label: str | None = None, + detector_key_version: str = PET_DETECTOR_KEY_VERSION, + quantization: int = 12, +) -> str: + x, y, width, height = bbox + center_x = x + width / 2.0 + center_y = y + height / 2.0 + quantized = ( + _quantize_value(center_x, quantization), + _quantize_value(center_y, quantization), + _quantize_value(width, quantization), + _quantize_value(height, quantization), + ) + species = _normalize_species_label(species_label) or "unknown" + payload = ( + f"{PET_KEY_VERSION}|{detector_key_version}|{asset_id}|" + f"{image_width}x{image_height}|{species}|" + f"{quantized[0]}|{quantized[1]}|{quantized[2]}|{quantized[3]}" + ) + digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() + return f"{PET_KEY_VERSION}:{digest}" + + +def cluster_pet_records( + detections: list[PetDetectionRecord], + *, + distance_threshold: float = 0.42, +) -> tuple[list[PetDetectionRecord], list[PetRecord]]: + if not detections: + return [], [] + + updated_detections = list(detections) + pets: list[PetRecord] = [] + labels = _cluster_pet_detection_labels( + detections, + distance_threshold=distance_threshold, + ) + grouped_indices: dict[str, list[int]] = defaultdict(list) + for index, label in enumerate(labels.tolist()): + grouped_indices[f"cluster-{label}"].append(index) + + for grouped in grouped_indices.values(): + members = [detections[index] for index in grouped] + key_detection = max(members, key=key_detection_sort_key) + pet_id = uuid.uuid4().hex + center_embedding = compute_cluster_center( + np.stack([member.embedding for member in members], axis=0) + ) + timestamp = utc_now_iso() + evidence_asset_count = len({member.asset_id for member in members if member.asset_id}) + pets.append( + PetRecord( + pet_id=pet_id, + name=None, + key_detection_id=key_detection.detection_id, + detection_count=len(members), + center_embedding=center_embedding, + embedding_dim=int(center_embedding.shape[0]), + created_at=timestamp, + updated_at=timestamp, + sample_count=len(members), + profile_state=profile_state_for_sample_count(evidence_asset_count), + species_label=_dominant_species_label(members), + embedding_pipeline_version=members[0].embedding_pipeline_version, + generation_id=members[0].generation_id, + boundary_embeddings=_boundary_embeddings(members, center_embedding), + evidence_asset_count=evidence_asset_count, + ) + ) + for index in grouped: + updated_detections[index] = replace(updated_detections[index], pet_id=pet_id) + + pets.sort(key=lambda pet: (-pet.detection_count, pet.created_at, pet.pet_id)) + return updated_detections, pets + + +def build_pet_records_from_detections( + detections: Sequence[PetDetectionRecord], + *, + names_by_pet_id: dict[str, str | None] | None = None, + created_at_by_pet_id: dict[str, str] | None = None, + allow_mixed_identity_members: bool = False, +) -> list[PetRecord]: + grouped: dict[str, list[PetDetectionRecord]] = defaultdict(list) + for detection in detections: + if detection.pet_id: + grouped[str(detection.pet_id)].append(detection) + + names = dict(names_by_pet_id or {}) + created = dict(created_at_by_pet_id or {}) + updated_at = utc_now_iso() + pets: list[PetRecord] = [] + for pet_id, members in grouped.items(): + species_labels = { + label + for label in (_normalize_species_label(member.species_label) for member in members) + if label is not None + } + if len(species_labels) > 1: + if not allow_mixed_identity_members: + raise ValueError( + f"Pet {pet_id} mixes incompatible species labels: " + f"{sorted(species_labels)}" + ) + _LOGGER.info( + "Preserving mixed-species Pet identity %s: species=%s", + pet_id, + sorted(species_labels), + ) + contract_groups: dict[tuple[str, int, int], list[PetDetectionRecord]] = defaultdict(list) + for member in members: + contract_groups[ + ( + str(member.embedding_pipeline_version or ""), + int(member.embedding_dim), + int(member.generation_id), + ) + ].append(member) + if len(contract_groups) != 1: + if not allow_mixed_identity_members: + raise ValueError( + f"Pet {pet_id} mixes incompatible embedding contracts: " + f"{sorted(contract_groups)}" + ) + _LOGGER.info( + "Preserving mixed-contract Pet identity %s: contracts=%s", + pet_id, + sorted(contract_groups), + ) + _profile_contract, profile_members = max( + contract_groups.items(), + key=lambda item: ( + len(item[1]), + item[0][2], + item[0][0], + item[0][1], + ), + ) + key_detection = max(members, key=key_detection_sort_key) + center_embedding = compute_cluster_center( + np.stack([member.embedding for member in profile_members], axis=0) + ) + sample_count = len(members) + evidence_asset_count = len({member.asset_id for member in members if member.asset_id}) + pets.append( + PetRecord( + pet_id=pet_id, + name=names.get(pet_id), + key_detection_id=key_detection.detection_id, + detection_count=sample_count, + center_embedding=center_embedding, + embedding_dim=int(center_embedding.shape[0]), + created_at=created.get( + pet_id, + min((member.detected_at for member in members), default=updated_at), + ), + updated_at=updated_at, + sample_count=sample_count, + profile_state=profile_state_for_sample_count(evidence_asset_count), + species_label=_dominant_species_label(members), + embedding_pipeline_version=profile_members[0].embedding_pipeline_version, + generation_id=profile_members[0].generation_id, + boundary_embeddings=_boundary_embeddings(profile_members, center_embedding), + evidence_asset_count=evidence_asset_count, + ) + ) + pets.sort(key=lambda pet: (-pet.detection_count, pet.created_at, pet.pet_id)) + return pets + + +def _boundary_embeddings( + members: Sequence[PetDetectionRecord], + center_embedding: np.ndarray, +) -> tuple[np.ndarray, ...]: + ranked = sorted( + members, + key=lambda member: ( + -cosine_distance(member.embedding, center_embedding), + member.detection_id, + ), + ) + return tuple(normalize_vector(member.embedding) for member in ranked[:8]) + + +def canonicalize_pet_identities( + detections: list[PetDetectionRecord], + pets: list[PetRecord], + state_repository: PetStateRepository, + *, + distance_threshold: float, +) -> tuple[list[PetDetectionRecord], list[PetRecord]]: + if not detections or not pets: + return detections, pets + + profiles = {profile.pet_id: profile for profile in state_repository.get_identity_profiles()} + redirects = state_repository.get_merge_redirect_map() + pet_key_map = state_repository.get_pet_key_map(detection.pet_key for detection in detections) + detections_by_pet_id: dict[str, list[PetDetectionRecord]] = defaultdict(list) + for detection in detections: + if detection.pet_id is not None: + detections_by_pet_id[detection.pet_id].append(detection) + + canonical_members: dict[str, list[PetDetectionRecord]] = defaultdict(list) + canonical_names: dict[str, str | None] = {} + canonical_created_at: dict[str, str] = {} + direct_anchors: set[str] = set() + + for pet in pets: + members = detections_by_pet_id.get(pet.pet_id, []) + resolution = resolve_canonical_pet_id( + pet, + members, + profiles=profiles, + pet_key_map=pet_key_map, + redirects=redirects, + distance_threshold=distance_threshold, + ) + canonical_id = resolution.canonical_pet_id + is_incompatible = bool( + canonical_members.get(canonical_id) + and not _detection_groups_compatible( + canonical_members[canonical_id], + members, + distance_threshold=distance_threshold, + ) + ) + same_asset_conflict = bool( + canonical_members.get(canonical_id) + and _detection_groups_share_asset(canonical_members[canonical_id], members) + ) + if is_incompatible and ( + same_asset_conflict + or (not resolution.is_redirect_alias and canonical_id in direct_anchors) + ): + canonical_id = uuid.uuid4().hex + resolution = PetIdentityResolution( + raw_pet_id=canonical_id, + canonical_pet_id=canonical_id, + source=PetIdentityResolutionSource.NEW, + ) + if not resolution.is_redirect_alias: + direct_anchors.add(canonical_id) + profile = profiles.get(canonical_id) + canonical_members[canonical_id].extend(members) + canonical_names.setdefault(canonical_id, profile.name if profile is not None else None) + canonical_created_at.setdefault( + canonical_id, + profile.created_at if profile is not None else pet.created_at, + ) + + updated = list(detections) + index_by_detection_id = { + detection.detection_id: index for index, detection in enumerate(detections) + } + for canonical_id, members in canonical_members.items(): + for member in members: + updated[index_by_detection_id[member.detection_id]] = replace( + member, + pet_id=canonical_id, + ) + canonical_pets = build_pet_records_from_detections( + updated, + names_by_pet_id=canonical_names, + created_at_by_pet_id=canonical_created_at, + ) + return updated, canonical_pets + + +def resolve_canonical_pet_id( + pet: PetRecord, + members: list[PetDetectionRecord], + *, + profiles: dict[str, PetProfile], + pet_key_map: dict[str, str], + redirects: dict[str, str], + distance_threshold: float, +) -> PetIdentityResolution: + vote_counter = Counter( + pet_key_map[member.pet_key] for member in members if member.pet_key in pet_key_map + ) + if vote_counter: + raw_pet_id = max( + vote_counter.items(), + key=lambda item: ( + item[1], + profiles[item[0]].updated_at if item[0] in profiles else "", + item[0], + ), + )[0] + canonical_pet_id = redirects.get(raw_pet_id, raw_pet_id) + return PetIdentityResolution( + raw_pet_id=raw_pet_id, + canonical_pet_id=canonical_pet_id, + source=( + PetIdentityResolutionSource.REDIRECT_KEY + if canonical_pet_id != raw_pet_id + else PetIdentityResolutionSource.KEY + ), + ) + + best_profile_id: str | None = None + best_distance = float("inf") + pet_species = _normalize_species_label(pet.species_label) + for profile in profiles.values(): + is_redirect_alias = profile.pet_id in redirects + if not is_redirect_alias and str(profile.profile_state or "unstable") != "stable": + continue + profile_species = _normalize_species_label(profile.species_label) + if pet_species != profile_species: + continue + if profile.embedding_dim <= 0 or profile.center_embedding.size == 0: + continue + if profile.center_embedding.shape != pet.center_embedding.shape: + continue + distance = cosine_distance(pet.center_embedding, profile.center_embedding) + if distance < best_distance: + best_distance = distance + best_profile_id = profile.pet_id + + if best_profile_id is not None and best_distance <= distance_threshold: + canonical_pet_id = redirects.get(best_profile_id, best_profile_id) + return PetIdentityResolution( + raw_pet_id=best_profile_id, + canonical_pet_id=canonical_pet_id, + source=( + PetIdentityResolutionSource.REDIRECT_PROFILE + if canonical_pet_id != best_profile_id + else PetIdentityResolutionSource.PROFILE + ), + ) + new_pet_id = uuid.uuid4().hex + return PetIdentityResolution( + raw_pet_id=new_pet_id, + canonical_pet_id=new_pet_id, + source=PetIdentityResolutionSource.NEW, + ) + + +def _cluster_pet_detection_labels( + detections: Sequence[PetDetectionRecord], + *, + distance_threshold: float, +) -> np.ndarray: + if not detections: + return np.empty((0,), dtype=np.int32) + embeddings = np.stack([detection.embedding for detection in detections], axis=0).astype( + np.float32 + ) + return _cluster_embeddings_bounded_single_link( + embeddings, + compatibility_keys=[ + ( + _normalize_species_label(detection.species_label), + str(detection.embedding_pipeline_version or ""), + int(detection.embedding_dim), + int(detection.generation_id), + ) + for detection in detections + ], + member_keys=[detection.detection_id for detection in detections], + cannot_link_keys=[detection.asset_id for detection in detections], + distance_threshold=distance_threshold, + ) + + +def _cluster_embeddings_bounded_single_link( + embeddings: np.ndarray, + *, + compatibility_keys: Sequence[object], + member_keys: Sequence[str] | None = None, + cannot_link_keys: Sequence[str] | None = None, + distance_threshold: float, +) -> np.ndarray: + count = int(embeddings.shape[0]) + if count == 0: + return np.empty((0,), dtype=np.int32) + return _cluster_distance_matrix_bounded_single_link( + cosine_distance_matrix(embeddings), + compatibility_keys=compatibility_keys, + member_keys=member_keys, + cannot_link_keys=cannot_link_keys, + link_threshold=distance_threshold, + diameter_threshold=distance_threshold * PET_CLUSTER_DIAMETER_MULTIPLIER, + ) + + +def _cluster_distance_matrix_bounded_single_link( + distance_matrix: np.ndarray, + *, + compatibility_keys: Sequence[object], + link_threshold: float, + diameter_threshold: float, + member_keys: Sequence[str] | None = None, + cannot_link_keys: Sequence[str] | None = None, +) -> np.ndarray: + """Cluster nearest-neighbour links without allowing unbounded similarity chains.""" + + count = int(distance_matrix.shape[0]) + if count == 0: + return np.empty((0,), dtype=np.int32) + if distance_matrix.shape != (count, count): + raise ValueError("Pet distance matrix must be square.") + if len(compatibility_keys) != count: + raise ValueError("Pet compatibility key count must match the distance matrix.") + stable_keys = tuple(member_keys or (str(index) for index in range(count))) + if len(stable_keys) != count: + raise ValueError("Pet member key count must match the distance matrix.") + resolved_cannot_link_keys = tuple(cannot_link_keys or ("" for _ in range(count))) + if len(resolved_cannot_link_keys) != count: + raise ValueError("Pet cannot-link key count must match the distance matrix.") + + clusters: list[list[int]] = [[index] for index in range(count)] + diameters: list[float] = [0.0] * count + cannot_link_sets: list[set[str]] = [ + ({key} if key else set()) for key in resolved_cannot_link_keys + ] + + while True: + best_pair: tuple[int, int] | None = None + best_key: tuple[float, float, tuple[str, ...], tuple[str, ...]] | None = None + best_diameter = 0.0 + for left_index in range(len(clusters)): + for right_index in range(left_index + 1, len(clusters)): + left = clusters[left_index] + right = clusters[right_index] + if not _cluster_keys_compatible(left, right, compatibility_keys): + continue + if cannot_link_sets[left_index] & cannot_link_sets[right_index]: + _LOGGER.debug( + "Pet clustering constraint hit: same_asset_cannot_link_hits=1" + ) + continue + cross_distances = [ + float(distance_matrix[left_member, right_member]) + for left_member in left + for right_member in right + ] + connection_distance = min(cross_distances) + if connection_distance > link_threshold: + continue + merged_diameter = max( + diameters[left_index], + diameters[right_index], + max(cross_distances), + ) + if merged_diameter > diameter_threshold: + continue + left_keys = tuple(sorted(stable_keys[index] for index in left)) + right_keys = tuple(sorted(stable_keys[index] for index in right)) + ordered_cluster_keys = tuple(sorted((left_keys, right_keys))) + tie_key = ( + connection_distance, + merged_diameter, + ordered_cluster_keys[0], + ordered_cluster_keys[1], + ) + if best_key is None or tie_key < best_key: + best_key = tie_key + best_pair = (left_index, right_index) + best_diameter = merged_diameter + if best_pair is None: + break + left_index, right_index = best_pair + merged = sorted(clusters[left_index] + clusters[right_index]) + clusters[left_index] = merged + diameters[left_index] = best_diameter + cannot_link_sets[left_index].update(cannot_link_sets[right_index]) + del clusters[right_index] + del diameters[right_index] + del cannot_link_sets[right_index] + ordering = sorted( + range(len(clusters)), + key=lambda cluster_index: tuple( + sorted(stable_keys[index] for index in clusters[cluster_index]) + ), + ) + clusters = [clusters[index] for index in ordering] + diameters = [diameters[index] for index in ordering] + cannot_link_sets = [cannot_link_sets[index] for index in ordering] + + labels = np.empty((count,), dtype=np.int32) + for cluster_id, members in enumerate(clusters): + for member in members: + labels[member] = cluster_id + return labels + + +def _cluster_keys_compatible( + left: Sequence[int], + right: Sequence[int], + compatibility_keys: Sequence[object], +) -> bool: + keys = {compatibility_keys[index] for index in [*left, *right]} + return len(keys) == 1 + + +def _detection_species_compatible( + left: Sequence[PetDetectionRecord], + right: Sequence[PetDetectionRecord], +) -> bool: + labels = [ + *(_normalize_species_label(detection.species_label) for detection in left), + *(_normalize_species_label(detection.species_label) for detection in right), + ] + return len(set(labels)) <= 1 + + +def _detection_contracts_compatible( + left: Sequence[PetDetectionRecord], + right: Sequence[PetDetectionRecord], +) -> bool: + contracts = { + ( + str(detection.embedding_pipeline_version or ""), + int(detection.embedding_dim), + int(detection.generation_id), + ) + for detection in [*left, *right] + } + return len(contracts) <= 1 + + +def _detection_groups_compatible( + left: Sequence[PetDetectionRecord], + right: Sequence[PetDetectionRecord], + *, + distance_threshold: float, +) -> bool: + if not _detection_species_compatible(left, right) or not _detection_contracts_compatible( + left, right + ): + return False + if _detection_groups_share_asset(left, right): + return False + if not left or not right: + return True + cross_distances = [ + cosine_distance(left_detection.embedding, right_detection.embedding) + for left_detection in left + for right_detection in right + ] + if min(cross_distances) > distance_threshold: + return False + members = [*left, *right] + distance_matrix = cosine_distance_matrix( + np.stack([member.embedding for member in members], axis=0) + ) + return float(distance_matrix.max()) <= (distance_threshold * PET_CLUSTER_DIAMETER_MULTIPLIER) + + +def _detection_groups_share_asset( + left: Sequence[PetDetectionRecord], + right: Sequence[PetDetectionRecord], +) -> bool: + left_assets = {detection.asset_id for detection in left if detection.asset_id} + right_assets = {detection.asset_id for detection in right if detection.asset_id} + return bool(left_assets & right_assets) + + +def _dominant_species_label(detections: Sequence[PetDetectionRecord]) -> str | None: + counter = Counter( + label + for label in (_normalize_species_label(detection.species_label) for detection in detections) + if label is not None + ) + if not counter: + return None + return max(counter.items(), key=lambda item: (item[1], item[0]))[0] + + +def _normalize_species_label(value: object) -> str | None: + if value is None: + return None + label = str(value).strip().lower() + return label or None + + +def _pet_candidate_quality_score( + *, + confidence: float, + relative_area_ratio: float, +) -> float: + """Rank retained detections without turning relative size into a hard gate.""" + + normalized_area = min(1.0, math.sqrt(max(0.0, float(relative_area_ratio)))) + return float(0.75 * max(0.0, min(1.0, confidence)) + 0.25 * normalized_area) + + +def _normalize_bbox( + raw_bbox, + *, + image_width: int, + image_height: int, +) -> tuple[int, int, int, int]: + box = np.asarray(raw_bbox, dtype=np.float32).flatten().tolist() + x, y, width, height = [round(value) for value in box[:4]] + x = max(0, min(x, image_width - 1)) + y = max(0, min(y, image_height - 1)) + width = max(1, min(width, image_width - x)) + height = max(1, min(height, image_height - y)) + return x, y, width, height + + +def _quantize_value(value: float, step: int) -> int: + step = max(1, int(step)) + return int(round(float(value) / step) * step) + + +class _YoloxOnnxPetDetector: + def __init__( + self, + model_path: Path, + *, + score_threshold: float = 0.30, + allow_model_download: bool = True, + enable_tiled_detection: bool = True, + tile_scan_min_confidence: float | None = None, + tile_species: frozenset[str] = SUPPORTED_DEFAULT_SPECIES, + execution_providers: Sequence[str] | None = None, + ) -> None: + self._model_path = Path(model_path) + self._score_threshold = float(score_threshold) + self._enable_tiled_detection = bool(enable_tiled_detection) + self._tile_scan_min_confidence = ( + self._score_threshold + if tile_scan_min_confidence is None + else float(tile_scan_min_confidence) + ) + self._tile_species = frozenset(tile_species) + try: + import onnxruntime as ort + except ImportError as exc: + raise PetRuntimeUnavailableError( + "Pet scanning unavailable: missing onnxruntime. Install the optional " + 'Pets AI runtime with: pip install -e ".[pets-ai]"' + ) from exc + self._model_path = ensure_pet_detector_model( + self._model_path, + allow_model_download=allow_model_download, + ) + try: + providers = list(execution_providers or _resolve_execution_providers(ort)) + self._session = ort.InferenceSession(str(self._model_path), providers=providers) + self._input_name = self._session.get_inputs()[0].name + shape = self._session.get_inputs()[0].shape + self._input_size = _input_size_from_shape(shape) + _validate_yolox_session_contract(self._session) + except Exception as exc: # noqa: BLE001 - provider failures vary by backend + if isinstance(exc, PetPipelineInvariantError): + raise + raise PetModelUnavailableError( + "Pet scanning unavailable: failed to initialize YOLOX detector model at " + f"{self._model_path} ({_error_reason(exc)}). Check the model cache, " + "disable unsupported execution providers, or reinstall the Pets AI runtime." + ) from exc + + def detect(self, image) -> list[_DetectedPetBox]: + boxes = self._detect_single_image(image) + if not self._enable_tiled_detection: + return _dedupe_supported_species_boxes(boxes) + + image_width, image_height = image.size + for crop_box in _select_uncovered_tile_regions( + image_width, + image_height, + boxes, + max_regions=4, + ): + left, top, *_ = crop_box + crop = image.crop(crop_box) + boxes.extend(self._detect_single_image(crop, offset=(left, top))) + return _dedupe_supported_species_boxes(boxes) + + def _detect_single_image( + self, + image, + *, + offset: tuple[int, int] = (0, 0), + ) -> list[_DetectedPetBox]: + image_width, image_height = image.size + input_width, input_height = self._input_size + preprocessed = _preprocess_yolox( + image, + input_width=input_width, + input_height=input_height, + ) + try: + outputs = self._session.run(None, {self._input_name: preprocessed.tensor}) + except Exception as exc: # noqa: BLE001 - provider failures vary by backend + raise PetInferenceError(f"Pet detector inference failed: {_error_reason(exc)}") from exc + predictions = _flatten_predictions(outputs) + boxes: list[_DetectedPetBox] = [] + for x0, y0, x1, y1, confidence, class_id in _decode_yolox_predictions( + predictions, + input_size=self._input_size, + ): + species = COCO_ANIMAL_LABELS.get(int(class_id)) + if species is None or confidence < self._score_threshold: + continue + bbox = _map_yolox_box_to_source( + (x0, y0, x1, y1), + preprocessed=preprocessed, + image_width=image_width, + image_height=image_height, + offset=offset, + ) + boxes.append( + _DetectedPetBox( + bbox=bbox, + confidence=float(confidence), + species_label=species, + ) + ) + return boxes + + def _has_tile_species_box(self, boxes: list[_DetectedPetBox]) -> bool: + return any( + box.species_label in self._tile_species + and box.confidence >= self._tile_scan_min_confidence + for box in boxes + ) + + +class _DinoV2Embedder: + def __init__( + self, + model_dir: Path, + *, + model_name: str, + allow_model_download: bool = True, + ) -> None: + self._model_dir = Path(model_dir) + self._model_name = model_name + try: + import torch + except ImportError as exc: + raise PetRuntimeUnavailableError( + "Pet scanning unavailable: missing torch for DINOv2 pet embeddings. " + 'Install the optional Pets AI runtime with: pip install -e ".[pets-ai]"' + ) from exc + self._torch = torch + runtime_version = str(getattr(torch, "__version__", "")).split("+", 1)[0] + if runtime_version != _DINO_TORCH_VERSION: + raise PetRuntimeUnavailableError( + "Pet scanning unavailable: DINOv2 requires " + f"torch=={_DINO_TORCH_VERSION}, but found {runtime_version or 'unknown'}. " + 'Reinstall the optional runtime with: pip install -e ".[pets-ai]"' + ) + self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model_path = self._model_dir / f"{model_name}.pt" + if model_path.is_file(): + try: + _validate_dinov2_cache_metadata(model_path, model_name=model_name) + self._model = torch.jit.load(str(model_path), map_location=self._device) + except (OSError, RuntimeError): + if not allow_model_download: + raise + self._model = self._download_dinov2_model(model_path) + elif allow_model_download: + self._model = self._download_dinov2_model(model_path) + else: + raise PetModelUnavailableError( + "Pet scanning unavailable: missing DINOv2 TorchScript model at " + f"{model_path}. Set IPHOTO_PET_MODEL_DIR or enable pet model downloads." + ) + self._model.eval() + + def embed(self, image) -> np.ndarray: + tensor = image_to_chw_float(image, (224, 224)) + torch = self._torch + try: + with torch.no_grad(): + input_tensor = torch.from_numpy(tensor).to(self._device) + output = self._model(input_tensor) + if isinstance(output, (list, tuple)): + output = output[0] + vector = output.detach().cpu().numpy().reshape(-1) + except Exception as exc: # noqa: BLE001 - backend failures vary by device/runtime + raise PetInferenceError( + f"Pet embedding inference failed: {_error_reason(exc)}" + ) from exc + expected_dimension = int(_EMBEDDER_MANIFEST["output_shape"][-1]) + if vector.size != expected_dimension: + raise PetPipelineInvariantError( + "Pet scanning unavailable: DINOv2 output contract mismatch " + f"({vector.size} != {expected_dimension})." + ) + return normalize_vector(vector.astype(np.float32)) + + def _download_dinov2_model(self, model_path: Path): + try: + return self._build_dinov2_cache(model_path) + except _ModelStoragePermissionError as exc: + fallback = _model_storage_fallback_path(model_path) + if fallback is None: + raise PetModelUnavailableError( + "Pet scanning unavailable: model storage is not writable at " + f"{model_path.parent}." + ) from exc + _LOGGER.warning( + "Falling back to the user Pets model cache after %s was not writable", + model_path.parent, + exc_info=exc, + ) + try: + return self._build_dinov2_cache(fallback) + except _ModelStoragePermissionError as fallback_exc: + raise PetModelUnavailableError( + "Pet scanning unavailable: model storage is not writable at " + f"{fallback.parent}." + ) from fallback_exc + + def _build_dinov2_cache(self, model_path: Path): + try: + loaded = _acquire_dinov2_release( + self._torch, + model_path, + model_name=self._model_name, + download_file=_download_file, + ) + loaded.eval() + loaded.to(self._device) + return loaded + except _ModelStoragePermissionError: + raise + except Exception as exc: + raise PetModelUnavailableError( + "Pet scanning unavailable: failed to acquire the verified DINOv2 " + f"Release artifact ({_error_reason(exc)})." + ) from exc + + +def _resolve_execution_providers(ort) -> list[str]: + available = set(ort.get_available_providers()) + preferred = [ + "CUDAExecutionProvider", + "CoreMLExecutionProvider", + "OpenVINOExecutionProvider", + "CPUExecutionProvider", + ] + providers = [provider for provider in preferred if provider in available] + return providers or ["CPUExecutionProvider"] + + +def _input_size_from_shape(shape: Sequence[object]) -> tuple[int, int]: + if len(shape) >= 4 and isinstance(shape[2], int) and isinstance(shape[3], int): + return int(shape[3]), int(shape[2]) + return 640, 640 + + +def _validate_yolox_session_contract(session) -> None: + inputs = session.get_inputs() + outputs = session.get_outputs() + if len(inputs) != 1 or len(inputs[0].shape) != 4: + raise RuntimeError("Pet scanning unavailable: YOLOX input contract is invalid.") + if inputs[0].shape[1] not in {3, "3"}: + raise RuntimeError("Pet scanning unavailable: YOLOX input must have three channels.") + concrete_input = tuple( + int(value) if isinstance(value, (int, np.integer)) else None for value in inputs[0].shape + ) + expected_input = tuple(int(value) for value in _DETECTOR_MANIFEST["input"]["shape"]) + if concrete_input != expected_input: + raise RuntimeError( + "Pet scanning unavailable: YOLOX input shape does not match the manifest." + ) + if not outputs or len(outputs[0].shape) < 2: + raise RuntimeError("Pet scanning unavailable: YOLOX output contract is invalid.") + last_output_dim = outputs[0].shape[-1] + expected_output = tuple(int(value) for value in _DETECTOR_MANIFEST["output_shape"]) + concrete_output = tuple( + int(value) if isinstance(value, (int, np.integer)) else None for value in outputs[0].shape + ) + if isinstance(last_output_dim, (int, np.integer)) and concrete_output != expected_output: + raise RuntimeError( + "Pet scanning unavailable: YOLOX output shape does not match the manifest." + ) + + +def _preprocess_yolox( + image, + *, + input_width: int, + input_height: int, +) -> _YoloxPreprocessResult: + image_width, image_height = image.size + resize_ratio = min( + input_width / float(max(1, image_width)), + input_height / float(max(1, image_height)), + ) + resized_width = max(1, int(image_width * resize_ratio)) + resized_height = max(1, int(image_height * resize_ratio)) + resized = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR) + canvas = Image.new("RGB", (input_width, input_height), (114, 114, 114)) + canvas.paste(resized, (0, 0)) + array = np.asarray(canvas, dtype=np.float32) + # YOLOX 0.1.1rc0 deployment weights consume OpenCV-style BGR bytes. + # This release explicitly removed legacy mean/std normalization. + array = array[:, :, ::-1] + array = np.transpose(array, (2, 0, 1))[None, :, :, :] + return _YoloxPreprocessResult( + tensor=np.ascontiguousarray(array), + resize_ratio=float(resize_ratio), + ) + + +def _map_yolox_box_to_source( + box: tuple[float, float, float, float], + *, + preprocessed: _YoloxPreprocessResult, + image_width: int, + image_height: int, + offset: tuple[int, int] = (0, 0), +) -> tuple[int, int, int, int]: + ratio = max(float(preprocessed.resize_ratio), 1e-6) + offset_x, offset_y = offset + x0, y0, x1, y1 = box + left = round((x0 - preprocessed.pad_left) / ratio) + top = round((y0 - preprocessed.pad_top) / ratio) + right = round((x1 - preprocessed.pad_left) / ratio) + bottom = round((y1 - preprocessed.pad_top) / ratio) + left = max(0, min(left, image_width - 1)) + top = max(0, min(top, image_height - 1)) + right = max(left + 1, min(right, image_width)) + bottom = max(top + 1, min(bottom, image_height)) + return ( + left + offset_x, + top + offset_y, + max(1, right - left), + max(1, bottom - top), + ) + + +def _tile_scan_regions(image_width: int, image_height: int) -> list[tuple[int, int, int, int]]: + width = max(1, int(image_width)) + height = max(1, int(image_height)) + + def region(left: float, top: float, right: float, bottom: float) -> tuple[int, int, int, int]: + x0 = max(0, min(round(width * left), width - 1)) + y0 = max(0, min(round(height * top), height - 1)) + x1 = max(x0 + 1, min(round(width * right), width)) + y1 = max(y0 + 1, min(round(height * bottom), height)) + return x0, y0, x1, y1 + + candidates = [ + region(0.0, 0.0, 0.70, 1.0), + region(0.30, 0.0, 1.0, 1.0), + region(0.0, 0.0, 1.0, 0.70), + region(0.0, 0.30, 1.0, 1.0), + region(0.0, 0.0, 0.65, 0.65), + region(0.35, 0.0, 1.0, 0.65), + region(0.15, 0.15, 0.85, 0.85), + ] + return list(dict.fromkeys(candidates)) + + +def _select_uncovered_tile_regions( + image_width: int, + image_height: int, + boxes: Sequence[_DetectedPetBox], + *, + max_regions: int, +) -> list[tuple[int, int, int, int]]: + """Choose a bounded set of tiles by area not covered by full-frame pets.""" + + supported = [box for box in boxes if box.species_label in SUPPORTED_DEFAULT_SPECIES] + ranked: list[tuple[float, tuple[int, int, int, int]]] = [] + for region in _tile_scan_regions(image_width, image_height): + x0, y0, x1, y1 = region + tile_box = (x0, y0, x1 - x0, y1 - y0) + tile_area = max(1, tile_box[2] * tile_box[3]) + covered = min( + tile_area, + sum(_bbox_intersection_area(tile_box, box.bbox) for box in supported), + ) + uncovered_ratio = 1.0 - (covered / float(tile_area)) + if uncovered_ratio >= 0.20: + ranked.append((uncovered_ratio, region)) + ranked.sort(key=lambda item: (-item[0], item[1])) + return [region for _, region in ranked[: max(0, int(max_regions))]] + + +def _flatten_predictions(outputs: Sequence[np.ndarray]) -> np.ndarray: + if not outputs: + return np.empty((0, 0), dtype=np.float32) + prediction = np.asarray(outputs[0], dtype=np.float32) + return prediction.reshape(-1, prediction.shape[-1]) + + +def _decode_yolox_predictions( + predictions: np.ndarray, + *, + input_size: tuple[int, int], +) -> list[tuple[float, float, float, float, float, int]]: + if predictions.size == 0: + return [] + decoded = np.asarray(predictions, dtype=np.float32) + if _looks_like_raw_yolox_output(decoded, input_size=input_size): + decoded = _decode_raw_yolox_output(decoded, input_size=input_size) + return [_decode_prediction(prediction) for prediction in decoded if prediction.shape[0] >= 6] + + +def _decode_prediction(prediction: np.ndarray) -> tuple[float, float, float, float, float, int]: + if prediction.shape[0] >= 85: + cx, cy, width, height = [float(value) for value in prediction[:4]] + object_score = float(prediction[4]) + class_scores = prediction[5:] + class_index = int(np.argmax(class_scores)) + confidence = object_score * float(class_scores[class_index]) + x0 = cx - width / 2.0 + y0 = cy - height / 2.0 + x1 = cx + width / 2.0 + y1 = cy + height / 2.0 + return x0, y0, x1, y1, confidence, class_index + x0, y0, x1, y1 = [float(value) for value in prediction[:4]] + confidence = float(prediction[4]) + class_id = round(float(prediction[5])) + return x0, y0, x1, y1, confidence, class_id + + +def _looks_like_raw_yolox_output( + predictions: np.ndarray, + *, + input_size: tuple[int, int], +) -> bool: + if predictions.ndim != 2 or predictions.shape[1] < 85: + return False + grids, _strides = _yolox_grids(input_size) + if predictions.shape[0] != grids.shape[0]: + return False + coord_max = float(np.nanmax(np.abs(predictions[:, :4]))) if predictions.size else 0.0 + return coord_max <= _YOLOX_RAW_COORD_LIMIT + + +def _decode_raw_yolox_output( + predictions: np.ndarray, + *, + input_size: tuple[int, int], +) -> np.ndarray: + grids, strides = _yolox_grids(input_size) + decoded = np.array(predictions, dtype=np.float32, copy=True) + decoded[:, :2] = (decoded[:, :2] + grids) * strides + decoded[:, 2:4] = np.exp(np.clip(decoded[:, 2:4], -20.0, 20.0)) * strides + return decoded + + +def _yolox_grids(input_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]: + input_width, input_height = input_size + grid_parts: list[np.ndarray] = [] + stride_parts: list[np.ndarray] = [] + for stride in _YOLOX_STRIDES: + grid_height = int(input_height) // stride + grid_width = int(input_width) // stride + yv, xv = np.meshgrid( + np.arange(grid_height, dtype=np.float32), + np.arange(grid_width, dtype=np.float32), + indexing="ij", + ) + grid = np.stack((xv, yv), axis=-1).reshape(-1, 2) + grid_parts.append(grid) + stride_parts.append(np.full((grid.shape[0], 1), stride, dtype=np.float32)) + return np.concatenate(grid_parts, axis=0), np.concatenate(stride_parts, axis=0) + + +def _dedupe_supported_species_boxes( + boxes: list[_DetectedPetBox], + *, + threshold: float = PET_PET_IOU_THRESHOLD, + smaller_box_coverage_threshold: float = PET_PET_SMALLER_BOX_COVERAGE_THRESHOLD, + normalized_center_distance_threshold: float = PET_PET_NORMALIZED_CENTER_DISTANCE_THRESHOLD, + cross_species_mutual_coverage_threshold: float = ( + PET_PET_CROSS_SPECIES_MUTUAL_COVERAGE_THRESHOLD + ), + cross_species_threshold: float = 0.90, + cross_species_score_margin: float = 0.25, +) -> list[_DetectedPetBox]: + selected: list[_DetectedPetBox] = [] + for box in sorted( + boxes, + key=lambda item: (item.quality_score, item.confidence), + reverse=True, + ): + suppress = False + for existing in selected: + overlap = _bbox_iou(existing.bbox, box.bbox) + existing_box_coverage, candidate_box_coverage = _bbox_pair_coverages( + existing.bbox, + box.bbox, + ) + smaller_box_coverage = max(existing_box_coverage, candidate_box_coverage) + normalized_center_distance = _bbox_normalized_center_distance( + existing.bbox, + box.bbox, + ) + reason = "" + if existing.species_label == box.species_label: + if overlap >= threshold: + reason = "same_species_iou" + elif ( + smaller_box_coverage >= smaller_box_coverage_threshold + and normalized_center_distance <= normalized_center_distance_threshold + ): + reason = "same_species_containment" + elif ( + existing.species_label in SUPPORTED_DEFAULT_SPECIES + and box.species_label in SUPPORTED_DEFAULT_SPECIES + and existing_box_coverage >= cross_species_mutual_coverage_threshold + and candidate_box_coverage >= cross_species_mutual_coverage_threshold + ): + reason = "cross_species_mutual_coverage" + elif ( + existing.species_label != box.species_label + and overlap >= cross_species_threshold + and existing.confidence - box.confidence >= cross_species_score_margin + ): + reason = "cross_species_iou" + + if reason: + _LOGGER.debug( + "Suppressed pet box: reason=%s species=%s candidate_confidence=%.3f " + "candidate_bbox=%s kept_species=%s kept_confidence=%.3f kept_bbox=%s " + "iou=%.3f smaller_box_coverage=%.3f kept_box_coverage=%.3f " + "candidate_box_coverage=%.3f " + "normalized_center_distance=%.3f", + reason, + box.species_label, + box.confidence, + box.bbox, + existing.species_label, + existing.confidence, + existing.bbox, + overlap, + smaller_box_coverage, + existing_box_coverage, + candidate_box_coverage, + normalized_center_distance, + ) + suppress = True + break + if suppress: + continue + selected.append(box) + return selected + + +def _bbox_iou(left: tuple[int, int, int, int], right: tuple[int, int, int, int]) -> float: + intersection = _bbox_intersection_area(left, right) + left_area = max(0, left[2]) * max(0, left[3]) + right_area = max(0, right[2]) * max(0, right[3]) + union = left_area + right_area - intersection + if union <= 0: + return 0.0 + return intersection / float(union) + + +def _bbox_pair_coverages( + left: tuple[int, int, int, int], + right: tuple[int, int, int, int], +) -> tuple[float, float]: + left_area = max(0, left[2]) * max(0, left[3]) + right_area = max(0, right[2]) * max(0, right[3]) + if left_area <= 0 or right_area <= 0: + return 0.0, 0.0 + intersection = _bbox_intersection_area(left, right) + return intersection / float(left_area), intersection / float(right_area) + + +def _bbox_normalized_center_distance( + left: tuple[int, int, int, int], + right: tuple[int, int, int, int], +) -> float: + left_area = max(0, left[2]) * max(0, left[3]) + right_area = max(0, right[2]) * max(0, right[3]) + smaller_area = min(left_area, right_area) + if smaller_area <= 0: + return float("inf") + left_center = (left[0] + left[2] / 2.0, left[1] + left[3] / 2.0) + right_center = (right[0] + right[2] / 2.0, right[1] + right[3] / 2.0) + return math.hypot( + left_center[0] - right_center[0], + left_center[1] - right_center[1], + ) / math.sqrt(smaller_area) + + +def _bbox_intersection_area( + left: tuple[int, int, int, int], + right: tuple[int, int, int, int], +) -> int: + lx, ly, lw, lh = left + rx, ry, rw, rh = right + left_x2 = lx + lw + left_y2 = ly + lh + right_x2 = rx + rw + right_y2 = ry + rh + inter_left = max(lx, rx) + inter_top = max(ly, ry) + inter_right = min(left_x2, right_x2) + inter_bottom = min(left_y2, right_y2) + inter_width = max(0, inter_right - inter_left) + inter_height = max(0, inter_bottom - inter_top) + return inter_width * inter_height + + +def _pet_box_overlaps_people_boxes( + pet_box: tuple[int, int, int, int], + people_boxes: Sequence[tuple[int, int, int, int]], + *, + image_dimensions: tuple[int, int] | None = None, + iou_threshold: float = PET_PEOPLE_IOU_THRESHOLD, + smaller_box_coverage_threshold: float = PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, +) -> bool: + """Return whether a pet detection conflicts with a People face region.""" + + return _pet_people_overlap_decision( + pet_box, + people_boxes, + image_dimensions=image_dimensions, + iou_threshold=iou_threshold, + smaller_box_coverage_threshold=smaller_box_coverage_threshold, + ).suppressed + + +def _pet_people_overlap_decision( + pet_box: tuple[int, int, int, int], + people_boxes: Sequence[tuple[int, int, int, int]], + *, + image_dimensions: tuple[int, int] | None = None, + iou_threshold: float = PET_PEOPLE_IOU_THRESHOLD, + smaller_box_coverage_threshold: float = PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, + larger_pet_ratio: float = PET_PEOPLE_LARGER_PET_RATIO, + mural_image_coverage_threshold: float = PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD, +) -> _PetPeopleOverlapDecision: + """Classify face/pet overlap without losing pets held by people. + + A pet body may legitimately contain a much smaller human face. Preserve + that candidate unless it also spans most of the image, which is the shape + produced by the known wall-mural false-positive regression. + """ + + pet_area = max(0, pet_box[2]) * max(0, pet_box[3]) + if pet_area <= 0: + return _PetPeopleOverlapDecision(False) + image_area = 0 + if image_dimensions is not None: + image_area = max(0, int(image_dimensions[0])) * max(0, int(image_dimensions[1])) + pet_image_coverage = pet_area / float(image_area) if image_area else 0.0 + for people_box in people_boxes: + people_area = max(0, people_box[2]) * max(0, people_box[3]) + if people_area <= 0: + continue + intersection = _bbox_intersection_area(pet_box, people_box) + if intersection <= 0: + continue + pet_to_face_ratio = pet_area / float(people_area) + preserve_larger_pet = ( + image_area > 0 + and pet_to_face_ratio > larger_pet_ratio + and pet_image_coverage < mural_image_coverage_threshold + ) + if preserve_larger_pet: + continue + if _bbox_iou(pet_box, people_box) >= iou_threshold: + return _PetPeopleOverlapDecision( + True, + "iou", + pet_to_face_ratio, + pet_image_coverage, + ) + smaller_area = min(pet_area, people_area) + if intersection / float(smaller_area) >= smaller_box_coverage_threshold: + return _PetPeopleOverlapDecision( + True, + "smaller_box_coverage", + pet_to_face_ratio, + pet_image_coverage, + ) + return _PetPeopleOverlapDecision(False, pet_image_coverage=pet_image_coverage) + + +def pet_model_auto_download_enabled() -> bool: + raw = str(os.environ.get(PET_MODEL_AUTO_DOWNLOAD_ENV, "")).strip().lower() + return raw not in {"0", "false", "no", "off"} + + +def ensure_pet_detector_model( + model_path: Path, + *, + allow_model_download: bool = True, + model_url: str | None = None, +) -> Path: + target = Path(model_path) + custom_url = str(model_url or os.environ.get(PET_DETECTOR_MODEL_URL_ENV) or "").strip() + expected_sha256 = ( + str(os.environ.get(PET_DETECTOR_MODEL_SHA256_ENV) or "").strip().lower() + if custom_url + else DEFAULT_PET_DETECTOR_MODEL_SHA256 + ) + if custom_url and not expected_sha256: + raise RuntimeError( + "Pet scanning unavailable: a custom detector URL requires " + f"{PET_DETECTOR_MODEL_SHA256_ENV}." + ) + url = str(custom_url or DEFAULT_PET_DETECTOR_MODEL_URL).strip() + try: + with _model_acquisition_lock(target): + if target.is_file(): + try: + _validate_downloaded_file( + target, + label="YOLOX pet detector model", + expected_sha256=expected_sha256, + max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, + ) + return target + except RuntimeError: + if not allow_model_download: + raise + try: + target.unlink() + except OSError as exc: + _raise_if_model_storage_error(exc, target) + except OSError as exc: + _raise_if_model_storage_error(exc, target) + if not allow_model_download: + raise RuntimeError( + "Pet scanning unavailable: missing YOLOX model at " + f"{target}. Set IPHOTO_PET_MODEL_DIR or enable pet model downloads." + ) + if not url: + raise RuntimeError( + "Pet scanning unavailable: missing YOLOX model at " + f"{target} and no pet detector download URL is configured." + ) + return _download_file( + url, + target, + label="YOLOX pet detector model", + expected_sha256=expected_sha256, + max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, + ) + except _ModelStoragePermissionError as exc: + fallback = _model_storage_fallback_path(target) + if fallback is None: + raise RuntimeError( + "Pet scanning unavailable: model storage is not writable at " + f"{target.parent}." + ) from exc + _LOGGER.warning( + "Falling back to %s after model storage failure", + fallback.parent, + exc_info=exc, + ) + try: + return _download_file( + url, + fallback, + label="YOLOX pet detector model", + expected_sha256=expected_sha256, + max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, + ) + except _ModelStoragePermissionError as fallback_exc: + raise RuntimeError( + "Pet scanning unavailable: model storage is not writable at " + f"{fallback.parent}." + ) from fallback_exc + + +def default_pet_model_dir() -> Path: + return user_pet_model_cache_dir() + + +def bundled_pet_model_dir() -> Path: + package_root = Path(__file__).resolve().parents[2] + return package_root / "extension" / "models" / "pets" + + +def user_pet_model_cache_dir() -> Path: + if sys.platform == "darwin": + base = Path.home() / "Library" / "Caches" + elif os.name == "nt": + base = Path(os.environ.get("LOCALAPPDATA") or Path.home() / "AppData" / "Local") + else: + base = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") + return base / "iPhoto" / "models" / "pets" + + +def pet_model_search_roots() -> tuple[Path, ...]: + override = pet_model_override_dir() + if override is not None: + return (override,) + return (bundled_pet_model_dir(), user_pet_model_cache_dir()) + + +def pet_model_override_dir() -> Path | None: + override = str(os.environ.get(IPHOTO_PET_MODEL_DIR_ENV) or "").strip() + if not override: + return None + return Path(override).expanduser() + + +def _is_packaged_macos_app_path(path: Path) -> bool: + if sys.platform != "darwin": + return False + parts = path.parts + return any( + part.lower().endswith(".app") + and index + 1 < len(parts) + and parts[index + 1] == "Contents" + for index, part in enumerate(parts) + ) + + +def _directory_is_writable(path: Path) -> bool: + try: + path.mkdir(parents=True, exist_ok=True) + probe = path / f".iphoto-write-probe-{uuid.uuid4().hex}" + with probe.open("xb"): + pass + probe.unlink(missing_ok=True) + return True + except OSError as exc: + if exc.errno in _MODEL_STORAGE_ERRNOS: + return False + raise + + +def pet_model_install_root() -> Path: + override = pet_model_override_dir() + if override is not None: + return override + + bundled = bundled_pet_model_dir() + if not _is_packaged_macos_app_path(bundled) and _directory_is_writable(bundled): + return bundled + return user_pet_model_cache_dir() + + +def _model_storage_fallback_path(path: Path) -> Path | None: + if pet_model_override_dir() is not None: + return None + target = Path(path) + bundled = bundled_pet_model_dir() + try: + relative = target.relative_to(bundled) + except ValueError: + return None + fallback = user_pet_model_cache_dir() / relative + if fallback == target: + return None + return fallback + + +def resolve_pet_model_path(relative_path: Path, *, directory: bool = False) -> Path: + relative = Path(relative_path) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("Pet model path must be relative to a configured model root.") + override = pet_model_override_dir() + if override is not None and not directory: + # Override lookup and repair are authoritative. Read-only resolution must + # never delete the candidate or consult another model root. + return override / relative + user_cache = user_pet_model_cache_dir() + bundled_invalid = False + invalid_user_candidate: Path | None = None + for root in pet_model_search_roots(): + candidate = root / relative + exists = candidate.is_dir() if directory else candidate.is_file() + if not exists: + continue + try: + if directory: + model_name = relative.name + model_path = candidate / f"{model_name}.pt" + if not model_path.is_file(): + raise RuntimeError("DINOv2 model file is missing") + _validate_dinov2_cache_metadata(model_path, model_name=model_name) + else: + _validate_downloaded_file( + candidate, + label="YOLOX pet detector model", + expected_sha256=( + str( + os.environ.get(PET_DETECTOR_MODEL_SHA256_ENV) + or DEFAULT_PET_DETECTOR_MODEL_SHA256 + ).lower() + ), + max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, + ) + return candidate + except (OSError, RuntimeError) as exc: + if override is not None and root == override: + raise RuntimeError( + f"Pet scanning unavailable: invalid model override artifact at {candidate}." + ) from exc + if root == user_cache and not directory: + invalid_user_candidate = candidate + continue + if root == user_cache: + model_path = candidate / f"{relative.name}.pt" + try: + model_path.unlink(missing_ok=True) + _dinov2_metadata_path(model_path).unlink(missing_ok=True) + except OSError: + _LOGGER.warning( + "Failed to quarantine invalid Pets model cache %s", + candidate, + exc_info=True, + ) + if root == bundled_pet_model_dir(): + bundled_invalid = True + if override is not None: + return override / relative + if invalid_user_candidate is not None: + return invalid_user_candidate + if bundled_invalid: + return user_cache / relative + return pet_model_install_root() / relative + + +def _download_file( + url: str, + destination: Path, + *, + label: str, + expected_sha256: str, + max_bytes: int, + exact_size: int | None = None, +) -> Path: + destination = Path(destination) + if urlparse(url).scheme.lower() != "https": + raise RuntimeError(f"Pet scanning unavailable: {label} URL must use HTTPS.") + try: + try: + destination.parent.mkdir(parents=True, exist_ok=True) + except OSError as exc: + _raise_if_model_storage_error(exc, destination.parent) + try: + temp_context = tempfile.TemporaryDirectory( + prefix="iphoto-pet-model-", + dir=destination.parent, + ) + except OSError as exc: + _raise_if_model_storage_error(exc, destination.parent) + with temp_context as tmp_dir: + tmp_path = Path(tmp_dir) / destination.name + try: + handle = tmp_path.open("wb") + except OSError as exc: + _raise_if_model_storage_error(exc, tmp_path) + with handle: + with request.urlopen( # noqa: S310 + url, + timeout=_DOWNLOAD_TIMEOUT_SECONDS, + context=_download_ssl_context(url), + ) as response: + total = 0 + while True: + chunk = response.read(_DOWNLOAD_CHUNK_SIZE) + if not chunk: + break + total += len(chunk) + if total > int(max_bytes): + raise RuntimeError(f"Downloaded {label} exceeds its size limit.") + try: + handle.write(chunk) + except OSError as exc: + _raise_if_model_storage_error(exc, tmp_path) + try: + _validate_downloaded_file( + tmp_path, + label=label, + expected_sha256=expected_sha256, + max_bytes=max_bytes, + exact_size=exact_size, + ) + except OSError as exc: + _raise_if_model_storage_error(exc, tmp_path) + try: + tmp_path.replace(destination) + except OSError as exc: + _raise_if_model_storage_error(exc, destination) + return destination + except _ModelStoragePermissionError: + raise + except TimeoutError as exc: + raise RuntimeError( + f"Pet scanning unavailable: downloading {label} timed out. " + "Check your network connection or install the model manually." + ) from exc + except OSError as exc: + raise RuntimeError( + f"Pet scanning unavailable: failed to download {label} from {url} " + f"({_error_reason(exc)}). Check your network connection, set " + f"{PET_DETECTOR_MODEL_URL_ENV}, or install the model manually." + ) from exc + except Exception as exc: + if isinstance(exc, RuntimeError) and str(exc).startswith("Pet scanning unavailable:"): + raise + raise RuntimeError( + f"Pet scanning unavailable: failed to download {label} from {url} " + f"({_error_reason(exc)}). Check your network connection, set " + f"{PET_DETECTOR_MODEL_URL_ENV}, or install the model manually." + ) from exc + + +def _validate_downloaded_file( + path: Path, + *, + label: str, + expected_sha256: str, + max_bytes: int, + exact_size: int | None = None, +) -> None: + size = Path(path).stat().st_size + if size <= 0: + raise RuntimeError(f"Downloaded {label} is empty.") + if size > int(max_bytes): + raise RuntimeError(f"Downloaded {label} exceeds its size limit.") + if exact_size is not None and size != int(exact_size): + raise RuntimeError(f"Downloaded {label} has the wrong file size.") + digest = _file_sha256(path) + if not expected_sha256 or digest != expected_sha256.lower(): + raise RuntimeError(f"Downloaded {label} failed SHA-256 verification.") + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for chunk in iter(lambda: handle.read(_DOWNLOAD_CHUNK_SIZE), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _dinov2_metadata_path(model_path: Path) -> Path: + return Path(model_path).with_suffix(f"{Path(model_path).suffix}.metadata.json") + + +@contextmanager +def _model_acquisition_lock(model_path: Path): + """Serialize one cache artifact across threads and processes.""" + + lock_path = Path(model_path).with_suffix(f"{Path(model_path).suffix}.acquire.lock") + try: + lock_path.parent.mkdir(parents=True, exist_ok=True) + handle = lock_path.open("a+b") + except OSError as exc: + _raise_if_model_storage_error(exc, lock_path.parent) + + try: + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"\0") + handle.flush() + handle.seek(0) + if os.name == "nt": + import msvcrt + + while True: + try: + handle.seek(0) + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + break + except OSError as exc: + if exc.errno not in _MODEL_LOCK_RETRY_ERRNOS: + raise + time.sleep(0.1) + try: + yield + finally: + handle.seek(0) + with suppress(OSError): + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + with suppress(OSError): + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + finally: + handle.close() + + +_dinov2_acquisition_lock = _model_acquisition_lock + + +def _publish_dinov2_cache_pair( + candidate: Path, + metadata_path: Path, + model_path: Path, +) -> None: + candidate = Path(candidate) + metadata_path = Path(metadata_path) + model_path = Path(model_path) + final_metadata_path = _dinov2_metadata_path(model_path) + published_metadata = False + published_model = False + try: + metadata_path.replace(final_metadata_path) + published_metadata = True + candidate.replace(model_path) + published_model = True + except OSError as exc: + if published_model: + with suppress(OSError): + model_path.unlink(missing_ok=True) + if published_metadata: + with suppress(OSError): + final_metadata_path.unlink(missing_ok=True) + _raise_if_model_storage_error(exc, model_path.parent) + + +def _dinov2_release_metadata(*, model_name: str) -> dict: + return { + "artifact_kind": "release_torchscript", + "cache_schema_version": _DINO_CACHE_SCHEMA_VERSION, + "model_name": model_name, + "release_tag": _EMBEDDER_MANIFEST["release_tag"], + "source_repository": _EMBEDDER_MANIFEST["source_repository"], + "source_revision": _DINO_SOURCE_REVISION, + "source_tree_sha1": _EMBEDDER_MANIFEST["source_tree_sha1"], + "weights_sha256": _DINO_WEIGHTS_SHA256, + "weights_size": _DINO_WEIGHTS_SIZE, + "producer_python_version": _EMBEDDER_MANIFEST["producer_python_version"], + "producer_torch_version": _DINO_TORCH_VERSION, + "producer_torchvision_version": _EMBEDDER_MANIFEST[ + "producer_torchvision_version" + ], + "torchscript_url": _DINO_TORCHSCRIPT_URL, + "torchscript_sha256": _DINO_TORCHSCRIPT_SHA256, + "torchscript_size": _DINO_TORCHSCRIPT_SIZE, + "input_shape": _EMBEDDER_MANIFEST["input_shape"], + "output_shape": _EMBEDDER_MANIFEST["output_shape"], + } + + +def _validate_dinov2_cache_metadata(model_path: Path, *, model_name: str) -> None: + metadata_path = _dinov2_metadata_path(model_path) + if not metadata_path.is_file(): + raise RuntimeError( + "Pet scanning unavailable: DINOv2 TorchScript metadata is missing for " + f"{model_path}. Remove the incomplete cache so it can be rebuilt." + ) + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"Pet scanning unavailable: invalid DINOv2 metadata at {metadata_path}." + ) from exc + expected = _dinov2_release_metadata(model_name=model_name) + if any(metadata.get(key) != value for key, value in expected.items()): + raise RuntimeError( + f"Pet scanning unavailable: DINOv2 metadata contract mismatch at {metadata_path}." + ) + size = int(metadata["torchscript_size"]) + digest = str(metadata["torchscript_sha256"]) + size_matches = size == model_path.stat().st_size + hash_matches = digest.lower() == _file_sha256(model_path) + if not size_matches or not hash_matches: + raise RuntimeError( + f"Pet scanning unavailable: DINOv2 cache integrity check failed for {model_path}." + ) + + +def _verify_dinov2_release_candidate(torch_runtime, candidate: Path): + model = torch_runtime.jit.load(str(candidate), map_location="cpu").eval() + shape = tuple(_EMBEDDER_MANIFEST["input_shape"]) + tensor_factory = getattr(torch_runtime, "zeros", None) or torch_runtime.randn + example = tensor_factory(shape, dtype=torch_runtime.float32) + with torch_runtime.no_grad(): + output = model(example) + if isinstance(output, (list, tuple)): + output = output[0] + expected_shape = tuple(_EMBEDDER_MANIFEST["output_shape"]) + actual_shape = tuple(output.shape) + if actual_shape != expected_shape: + raise RuntimeError( + f"DINOv2 output shape mismatch: {actual_shape} != {expected_shape}" + ) + return model + + +def _acquire_dinov2_release( + torch_runtime, + model_path: Path, + *, + model_name: str, + download_file, + validate_metadata=_validate_dinov2_cache_metadata, +): + """Download, validate, and atomically publish the fixed Release artifact.""" + + _install_certifi_environment() + model_path = Path(model_path) + try: + model_path.parent.mkdir(parents=True, exist_ok=True) + except OSError as exc: + _raise_if_model_storage_error(exc, model_path.parent) + + with _dinov2_acquisition_lock(model_path): + if model_path.is_file(): + try: + validate_metadata(model_path, model_name=model_name) + return torch_runtime.jit.load(str(model_path), map_location="cpu") + except (OSError, RuntimeError): + pass + + for stale_path in (model_path, _dinov2_metadata_path(model_path)): + try: + stale_path.unlink(missing_ok=True) + except OSError as exc: + _raise_if_model_storage_error(exc, stale_path) + + try: + temp_context = tempfile.TemporaryDirectory( + prefix="iphoto-dinov2-release-", + dir=model_path.parent, + ) + except OSError as exc: + _raise_if_model_storage_error(exc, model_path.parent) + + with temp_context as temp_dir: + candidate = Path(temp_dir) / model_path.name + metadata_path = _dinov2_metadata_path(candidate) + download_file( + _DINO_TORCHSCRIPT_URL, + candidate, + label="DINOv2 TorchScript Release artifact", + expected_sha256=_DINO_TORCHSCRIPT_SHA256, + max_bytes=_DINO_TORCHSCRIPT_SIZE, + exact_size=_DINO_TORCHSCRIPT_SIZE, + ) + _verify_dinov2_release_candidate(torch_runtime, candidate) + try: + metadata_path.write_text( + json.dumps( + _dinov2_release_metadata(model_name=model_name), + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + except OSError as exc: + _raise_if_model_storage_error(exc, metadata_path) + _publish_dinov2_cache_pair(candidate, metadata_path, model_path) + validate_metadata(model_path, model_name=model_name) + return torch_runtime.jit.load(str(model_path), map_location="cpu") + + +def _error_reason(exc: Exception) -> str: + return str(exc).strip() or exc.__class__.__name__ + + +def _download_ssl_context(url: str) -> ssl.SSLContext | None: + if not url.lower().startswith("https://"): + return None + try: + import certifi + except ImportError: + return ssl.create_default_context() + _install_certifi_environment() + return ssl.create_default_context(cafile=certifi.where()) + + +def _install_certifi_environment() -> None: + try: + import certifi + except ImportError: + return + cafile = certifi.where() + os.environ.setdefault("SSL_CERT_FILE", cafile) + os.environ.setdefault("REQUESTS_CA_BUNDLE", cafile) diff --git a/src/iPhoto/pets/model_manifest.json b/src/iPhoto/pets/model_manifest.json index f1d8367b4..1727b0b62 100644 --- a/src/iPhoto/pets/model_manifest.json +++ b/src/iPhoto/pets/model_manifest.json @@ -1,5 +1,5 @@ { - "schema_version": 1, + "schema_version": 2, "detector": { "filename": "detector/yolox_nano_coco.onnx", "url": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.onnx", @@ -19,9 +19,19 @@ "model_name": "dinov2_vits14", "source_repository": "facebookresearch/dinov2", "source_revision": "7764ea0f912e53c92e82eb78a2a1631e92725fc8", - "torchscript_url": null, - "torchscript_sha256": "36e01591cfd52d194845d5b23b870dcb698f79c33f6140144180568f2155ab94", - "torchscript_size": 88587286, + "source_tree_sha1": "2a27257b79b0633b027a21014bc9360e3c1b3f43", + "weights_url": "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/dinov2_vits14_pretrain.pth", + "weights_sha256": "b938bf1bc15cd2ec0feacfe3a1bb553fe8ea9ca46a7e1d8d00217f29aef60cd9", + "weights_size": 88283115, + "artifact_kind": "release_torchscript", + "release_tag": "pet-models-v1", + "cache_schema_version": 2, + "producer_python_version": "3.12", + "producer_torch_version": "2.12.1", + "producer_torchvision_version": "0.27.1", + "torchscript_url": "https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/releases/download/pet-models-v1/dinov2_vits14.pt", + "torchscript_sha256": "7f8e204c6222662ee53ef61fb02b1278781e0ab70a57fa73a63cdef9ed32ab13", + "torchscript_size": 88566102, "input_shape": [1, 3, 224, 224], "output_shape": [1, 384] } diff --git a/src/iPhoto/pets/pipeline.py b/src/iPhoto/pets/pipeline.py index bc4d46348..ada6c0f5c 100644 --- a/src/iPhoto/pets/pipeline.py +++ b/src/iPhoto/pets/pipeline.py @@ -1,2103 +1,232 @@ -"""Pet detection, embedding, clustering, and identity helpers.""" +"""Pets recognition pipeline compatibility surface. + +The stable public import path remains ``iPhoto.pets.pipeline`` while the bulk of +legacy recognition logic lives in ``_pipeline_impl``. First-use model +acquisition hardening is expressed here with normal inheritance and composition; +this module never replaces itself in ``sys.modules`` and never mutates classes or +function globals in the implementation module. +""" from __future__ import annotations -import hashlib -import json -import logging -import math -import os -import ssl -import sys -import tempfile -import uuid -from collections import Counter, defaultdict -from collections.abc import Callable, Sequence -from dataclasses import dataclass, replace -from enum import StrEnum from pathlib import Path -from urllib import request -from urllib.parse import urlparse - -import numpy as np -from PIL import Image - -from iPhoto.utils.pathutils import LibraryAssetPathError, resolve_library_asset_path - -from .errors import ( - PetInferenceError, - PetModelUnavailableError, - PetPipelineInvariantError, - PetRuntimeUnavailableError, -) -from .image_utils import ( - PetImageLoadError, - crop_pet_region, - image_to_chw_float, - load_image_rgb, - save_pet_thumbnail, -) -from .records import PetDetectionRecord, PetProfile, PetRecord -from .repository_utils import ( - compute_cluster_center, - cosine_distance, - cosine_distance_matrix, - key_detection_sort_key, - normalize_vector, - profile_state_for_sample_count, - utc_now_iso, -) -from .state_repository import PetStateRepository - -def _load_pet_model_manifest() -> dict: - manifest_path = Path(__file__).with_name("model_manifest.json") +from . import _pipeline_impl as _impl +from ._pipeline_impl import * # noqa: F403 + + +# Private implementation helpers intentionally re-exported for the established +# test/debug surface. Keeping real module attributes here also means targeted +# monkeypatches used by model-acquisition tests affect the hardened code below +# without modifying the implementation module globally. +_EMBEDDER_MANIFEST = _impl._EMBEDDER_MANIFEST +_DINO_SOURCE_REVISION = _impl._DINO_SOURCE_REVISION +_DINO_WEIGHTS_URL = _impl._DINO_WEIGHTS_URL +_DINO_WEIGHTS_SHA256 = _impl._DINO_WEIGHTS_SHA256 +_DINO_WEIGHTS_SIZE = _impl._DINO_WEIGHTS_SIZE +_DINO_TORCHSCRIPT_URL = _impl._DINO_TORCHSCRIPT_URL +_DINO_TORCHSCRIPT_SHA256 = _impl._DINO_TORCHSCRIPT_SHA256 +_DINO_TORCHSCRIPT_SIZE = _impl._DINO_TORCHSCRIPT_SIZE +_DINO_CACHE_SCHEMA_VERSION = _impl._DINO_CACHE_SCHEMA_VERSION +_DINO_TORCH_VERSION = _impl._DINO_TORCH_VERSION +_ModelStoragePermissionError = _impl._ModelStoragePermissionError +_raise_if_model_storage_error = _impl._raise_if_model_storage_error +_model_storage_fallback_path = _impl._model_storage_fallback_path +_install_certifi_environment = _impl._install_certifi_environment +_dinov2_metadata_path = _impl._dinov2_metadata_path +_validate_dinov2_cache_metadata = _impl._validate_dinov2_cache_metadata +_dinov2_release_metadata = _impl._dinov2_release_metadata +_model_acquisition_lock = _impl._model_acquisition_lock +_dinov2_acquisition_lock = _impl._dinov2_acquisition_lock +_publish_dinov2_cache_pair = _impl._publish_dinov2_cache_pair +_file_sha256 = _impl._file_sha256 +_error_reason = _impl._error_reason +_original_download_file = _impl._download_file + +def _download_file(*args, **kwargs): + """Keep detector-specific remediation out of DINOv2 download failures.""" + + label = str(kwargs.get("label") or "") try: - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - detector = manifest["detector"] - embedder = manifest["embedder"] - except (OSError, json.JSONDecodeError, KeyError, TypeError) as exc: - raise PetPipelineInvariantError(f"Invalid Pets model manifest: {manifest_path}") from exc - if int(manifest.get("schema_version") or 0) != 1: - raise PetPipelineInvariantError(f"Unsupported Pets model manifest: {manifest_path}") - if urlparse(str(detector.get("url") or "")).scheme.lower() != "https": - raise PetPipelineInvariantError("Pets detector manifest URL must use HTTPS.") - if detector.get("input") != { - "layout": "NCHW", - "channel_order": "BGR", - "dtype": "float32", - "range": [0, 255], - "shape": [1, 3, 416, 416], - }: - raise PetPipelineInvariantError("Pets detector manifest input contract is invalid.") - if embedder.get("input_shape") != [1, 3, 224, 224]: - raise PetPipelineInvariantError("Pets embedder manifest input contract is invalid.") - torchscript_url = str(embedder.get("torchscript_url") or "").strip() - if torchscript_url and urlparse(torchscript_url).scheme.lower() != "https": - raise PetPipelineInvariantError("Pets embedder TorchScript URL must use HTTPS.") - if len(str(embedder.get("torchscript_sha256") or "")) != 64: - raise PetPipelineInvariantError("Pets embedder TorchScript SHA-256 is invalid.") - if int(embedder.get("torchscript_size") or 0) <= 0: - raise PetPipelineInvariantError("Pets embedder TorchScript size is invalid.") - return manifest - - -PET_MODEL_MANIFEST = _load_pet_model_manifest() -_DETECTOR_MANIFEST = PET_MODEL_MANIFEST["detector"] -_EMBEDDER_MANIFEST = PET_MODEL_MANIFEST["embedder"] -SUPPORTED_DEFAULT_SPECIES = frozenset({"cat", "dog"}) -PET_DETECTOR_PIPELINE_VERSION = "yolox-letterbox-tiles-people-priority-v6" -PET_CLUSTERING_PIPELINE_VERSION = "species-bounded-single-link-v3" -PET_EMBEDDING_PIPELINE_VERSION = "dinov2-vits14-imagenet-normalized-v1" -PET_KEY_VERSION = "v2" -PET_DETECTOR_KEY_VERSION = "yolox-nano-coco-0.1.1rc0-raw-bgr-v1" -DEFAULT_PET_DISTANCE_THRESHOLD = 0.42 -PET_CLUSTER_DIAMETER_MULTIPLIER = 1.5 -PET_PET_IOU_THRESHOLD = 0.50 -PET_PET_SMALLER_BOX_COVERAGE_THRESHOLD = 0.90 -PET_PET_NORMALIZED_CENTER_DISTANCE_THRESHOLD = 0.40 -PET_PET_CROSS_SPECIES_MUTUAL_COVERAGE_THRESHOLD = 0.90 -PET_PEOPLE_IOU_THRESHOLD = 0.50 -PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD = 0.90 -PET_PEOPLE_LARGER_PET_RATIO = 1.50 -PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD = 0.60 -PET_CANDIDATE_QUALITY_VERSION = "tiny-low-confidence-v1" -DEFAULT_PET_TINY_AREA_RATIO = 0.001 -DEFAULT_PET_TINY_MAX_CONFIDENCE = 0.45 -DEFAULT_PET_DETECTOR_MODEL_URL = str(_DETECTOR_MANIFEST["url"]) -PET_MODEL_AUTO_DOWNLOAD_ENV = "IPHOTO_PET_MODEL_AUTO_DOWNLOAD" -PET_DETECTOR_MODEL_URL_ENV = "IPHOTO_PET_DETECTOR_MODEL_URL" -PET_DETECTOR_MODEL_SHA256_ENV = "IPHOTO_PET_DETECTOR_MODEL_SHA256" -DEFAULT_PET_DETECTOR_MODEL_SHA256 = str(_DETECTOR_MANIFEST["sha256"]) -DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES = int(_DETECTOR_MANIFEST["max_bytes"]) -# The development conversion tool uses this immutable source revision. The -# production runtime only loads the fixed, hash-verified TorchScript artifact. -_DINO_SOURCE_REVISION = str(_EMBEDDER_MANIFEST["source_revision"]) -_DOWNLOAD_TIMEOUT_SECONDS = 60 -_DOWNLOAD_CHUNK_SIZE = 1024 * 256 -_YOLOX_STRIDES = (8, 16, 32) -_YOLOX_RAW_COORD_LIMIT = 32.0 -_LOGGER = logging.getLogger(__name__) -COCO_ANIMAL_LABELS = { - 15: "cat", - 16: "dog", - 17: "horse", - 18: "sheep", - 19: "cow", - 20: "elephant", - 21: "bear", - 22: "zebra", - 23: "giraffe", -} - - -@dataclass(frozen=True) -class DetectedAssetPets: - asset_id: str - asset_rel: str - detections: list[PetDetectionRecord] - error: str | None = None - - -class PetIdentityResolutionSource(StrEnum): - KEY = "key" - REDIRECT_KEY = "redirect_key" - PROFILE = "profile" - REDIRECT_PROFILE = "redirect_profile" - NEW = "new" - - -@dataclass(frozen=True) -class PetIdentityResolution: - raw_pet_id: str - canonical_pet_id: str - source: PetIdentityResolutionSource - - @property - def is_redirect_alias(self) -> bool: - return self.source in { - PetIdentityResolutionSource.REDIRECT_KEY, - PetIdentityResolutionSource.REDIRECT_PROFILE, - } - - -@dataclass(frozen=True) -class _DetectedPetBox: - bbox: tuple[int, int, int, int] - confidence: float - species_label: str - quality_score: float = 0.0 - - -@dataclass(frozen=True) -class _YoloxPreprocessResult: - tensor: np.ndarray - resize_ratio: float - pad_left: int = 0 - pad_top: int = 0 - - -@dataclass(frozen=True) -class PetScanMetrics: - candidate_boxes: int = 0 - unsupported_species: int = 0 - too_small: int = 0 - pet_quality_rejected: int = 0 - people_overlaps: int = 0 - accepted_detections: int = 0 - pet_candidate_identities: int = 0 - pet_promotions: int = 0 - same_asset_cannot_link_hits: int = 0 - same_asset_manual_conflicts: int = 0 - - -@dataclass(frozen=True) -class _PetPeopleOverlapDecision: - suppressed: bool - reason: str = "" - pet_to_face_area_ratio: float = 0.0 - pet_image_coverage: float = 0.0 - - -class PetClusterPipeline: - def __init__( - self, - *, - model_root: Path, - detector_model_name: str = "yolox_nano_coco.onnx", - embedding_model_name: str = "dinov2_vits14", - allow_model_download: bool | None = None, - distance_threshold: float = DEFAULT_PET_DISTANCE_THRESHOLD, - min_pet_size: int = 48, - supported_species: frozenset[str] = SUPPORTED_DEFAULT_SPECIES, - detector_score_threshold: float = 0.30, - enable_tiled_detection: bool = True, - tile_scan_min_confidence: float | None = None, - tiny_area_ratio: float = DEFAULT_PET_TINY_AREA_RATIO, - tiny_max_confidence: float = DEFAULT_PET_TINY_MAX_CONFIDENCE, - ) -> None: - self._model_root = Path(model_root) - self._detector_model_name = detector_model_name - self._embedding_model_name = embedding_model_name - self._allow_model_download = ( - pet_model_auto_download_enabled() - if allow_model_download is None - else bool(allow_model_download) - ) - self._distance_threshold = float(distance_threshold) - self._min_pet_size = int(min_pet_size) - self._supported_species = frozenset(supported_species) - self._detector_score_threshold = float(detector_score_threshold) - self._enable_tiled_detection = bool(enable_tiled_detection) - self._tile_scan_min_confidence = ( - self._detector_score_threshold - if tile_scan_min_confidence is None - else float(tile_scan_min_confidence) - ) - self._tiny_area_ratio = float(tiny_area_ratio) - self._tiny_max_confidence = float(tiny_max_confidence) - self._detector: _YoloxOnnxPetDetector | None = None - self._embedder: _DinoV2Embedder | None = None - self._last_scan_metrics = PetScanMetrics() - - @property - def distance_threshold(self) -> float: - return self._distance_threshold - - @property - def detector_pipeline_version(self) -> str: - return PET_DETECTOR_PIPELINE_VERSION - - @property - def candidate_quality_version(self) -> str: - return PET_CANDIDATE_QUALITY_VERSION - - @property - def last_scan_metrics(self) -> PetScanMetrics: - return self._last_scan_metrics - - def detect_pets_for_rows( - self, - rows: list[dict], - *, - library_root: Path, - thumbnail_dir: Path, - published_thumbnail_dir: Path | None = None, - is_cancelled: Callable[[], bool] | None = None, - people_boxes_by_asset_id: dict[str, Sequence[tuple[int, int, int, int]]] | None = None, - ) -> list[DetectedAssetPets]: - if not rows: - return [] - embedder = self._ensure_embedder() - detector = self._ensure_detector() - cancellation_requested = is_cancelled or (lambda: False) - stored_thumbnail_dir = Path(published_thumbnail_dir or thumbnail_dir) - results: list[DetectedAssetPets] = [] - candidate_boxes = 0 - unsupported_species = 0 - too_small = 0 - pet_quality_rejected = 0 - people_overlaps = 0 - accepted_detections = 0 - excluded_people_boxes = people_boxes_by_asset_id or {} - for row in rows: - if cancellation_requested(): - break - asset_id = str(row.get("id") or "") - asset_rel = Path(str(row.get("rel") or "")).as_posix() - try: - image_path = resolve_library_asset_path(library_root, asset_rel) - image = load_image_rgb(image_path) - boxes = detector.detect(image) - except LibraryAssetPathError as exc: - results.append( - DetectedAssetPets( - asset_id=asset_id, - asset_rel=asset_rel, - detections=[], - error=str(exc), - ) - ) - continue - except PetImageLoadError as exc: - if cancellation_requested(): - break - results.append( - DetectedAssetPets( - asset_id=asset_id, - asset_rel=asset_rel, - detections=[], - error=str(exc).strip() or exc.__class__.__name__, - ) - ) - continue - except PetPipelineInvariantError: - raise - except Exception as exc: # noqa: BLE001 - if cancellation_requested(): - break - results.append( - DetectedAssetPets( - asset_id=asset_id, - asset_rel=asset_rel, - detections=[], - error=str(exc).strip() or exc.__class__.__name__, - ) - ) - continue - - image_width, image_height = image.size - detections: list[PetDetectionRecord] = [] - created_thumbnail_paths: list[Path] = [] - supported_boxes: list[_DetectedPetBox] = [] - quality_candidates: list[ - tuple[tuple[int, int, int, int], float, str, float] - ] = [] - candidate_boxes += len(boxes) - for detected in boxes: - if detected.species_label not in self._supported_species: - unsupported_species += 1 - continue - bbox = _normalize_bbox( - detected.bbox, - image_width=image_width, - image_height=image_height, - ) - if bbox[2] < self._min_pet_size or bbox[3] < self._min_pet_size: - too_small += 1 - continue - area_ratio = (bbox[2] * bbox[3]) / max(1, image_width * image_height) - if ( - area_ratio < self._tiny_area_ratio - and float(detected.confidence) < self._tiny_max_confidence - ): - pet_quality_rejected += 1 - continue - quality_candidates.append( - (bbox, float(detected.confidence), detected.species_label, area_ratio) - ) - - largest_pet_area_ratio = max( - (area_ratio for _bbox, _confidence, _species, area_ratio in quality_candidates), - default=0.0, - ) - for bbox, confidence, species_label, area_ratio in quality_candidates: - supported_boxes.append( - _DetectedPetBox( - bbox=bbox, - confidence=confidence, - species_label=species_label, - quality_score=_pet_candidate_quality_score( - confidence=confidence, - relative_area_ratio=area_ratio - / max(largest_pet_area_ratio, np.finfo(np.float32).eps), - ), - ) - ) - - deduped_boxes = _dedupe_supported_species_boxes(supported_boxes) - people_boxes = excluded_people_boxes.get(asset_id, ()) - accepted_boxes: list[_DetectedPetBox] = [] - for detected in deduped_boxes: - decision = _pet_people_overlap_decision( - detected.bbox, - people_boxes, - image_dimensions=(image_width, image_height), - ) - if decision.suppressed: - people_overlaps += 1 - _LOGGER.info( - "Suppressed pet candidate for asset %s: reason=%s " - "pet_to_face_area_ratio=%.3f pet_image_coverage=%.3f " - "iou_threshold=%.2f smaller_box_coverage_threshold=%.2f " - "larger_pet_ratio=%.2f mural_image_coverage_threshold=%.2f", - asset_id, - decision.reason, - decision.pet_to_face_area_ratio, - decision.pet_image_coverage, - PET_PEOPLE_IOU_THRESHOLD, - PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, - PET_PEOPLE_LARGER_PET_RATIO, - PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD, - ) - continue - accepted_boxes.append(detected) - - for detected in accepted_boxes: - bbox = detected.bbox - detection_id = uuid.uuid4().hex - thumbnail_path = thumbnail_dir / f"{detection_id}.png" - try: - crop = crop_pet_region(image, bbox, padding_ratio=0.08) - embedding = embedder.embed(crop) - save_pet_thumbnail(image, bbox, thumbnail_path, padding_ratio=0.08) - created_thumbnail_paths.append(thumbnail_path) - except PetPipelineInvariantError: - for created_path in created_thumbnail_paths: - created_path.unlink(missing_ok=True) - raise - except Exception as exc: # noqa: BLE001 - for created_path in created_thumbnail_paths: - try: - created_path.unlink(missing_ok=True) - except OSError: - _LOGGER.warning( - "Failed to roll back pet thumbnail %s", - created_path, - exc_info=True, - ) - results.append( - DetectedAssetPets( - asset_id=asset_id, - asset_rel=asset_rel, - detections=[], - error=str(exc).strip() or exc.__class__.__name__, - ) - ) - detections = [] - break - detections.append( - PetDetectionRecord( - detection_id=detection_id, - pet_key=build_pet_key( - asset_id=asset_id, - bbox=bbox, - image_width=image_width, - image_height=image_height, - species_label=detected.species_label, - ), - asset_id=asset_id, - asset_rel=asset_rel, - box_x=bbox[0], - box_y=bbox[1], - box_w=bbox[2], - box_h=bbox[3], - confidence=float(detected.confidence), - embedding=embedding, - embedding_dim=int(embedding.shape[0]), - embedding_model=self._embedding_model_name, - detector_model=self._detector_model_name, - thumbnail_path=(stored_thumbnail_dir / thumbnail_path.name) - .relative_to(stored_thumbnail_dir.parent) - .as_posix(), - pet_id=None, - detected_at=utc_now_iso(), - image_width=image_width, - image_height=image_height, - species_label=detected.species_label, - quality_score=detected.quality_score, - pet_key_version=PET_KEY_VERSION, - embedding_pipeline_version=PET_EMBEDDING_PIPELINE_VERSION, - ) - ) - has_asset_error = any( - result.asset_id == asset_id and result.error for result in results - ) - if detections or not has_asset_error: - accepted_detections += len(detections) - results.append( - DetectedAssetPets( - asset_id=asset_id, - asset_rel=asset_rel, - detections=detections, - ) - ) - self._last_scan_metrics = PetScanMetrics( - candidate_boxes=candidate_boxes, - unsupported_species=unsupported_species, - too_small=too_small, - pet_quality_rejected=pet_quality_rejected, - people_overlaps=people_overlaps, - accepted_detections=accepted_detections, - ) - return results - - def _ensure_detector(self) -> _YoloxOnnxPetDetector: - if self._detector is None: - model_path = self._resolve_model_path(Path("detector") / self._detector_model_name) - try: - self._detector = _YoloxOnnxPetDetector( - model_path, - score_threshold=self._detector_score_threshold, - allow_model_download=self._allow_model_download, - enable_tiled_detection=self._enable_tiled_detection, - tile_scan_min_confidence=self._tile_scan_min_confidence, - tile_species=self._supported_species, - ) - except ( - PetRuntimeUnavailableError, - PetModelUnavailableError, - PetPipelineInvariantError, - ): - raise - except RuntimeError as exc: - raise PetModelUnavailableError(str(exc)) from exc - return self._detector - - def _ensure_embedder(self) -> _DinoV2Embedder: - if self._embedder is None: - model_dir = self._resolve_model_path( - Path("embedding") / self._embedding_model_name, - directory=True, - ) - try: - self._embedder = _DinoV2Embedder( - model_dir, - model_name=self._embedding_model_name, - allow_model_download=self._allow_model_download, - ) - except ( - PetRuntimeUnavailableError, - PetModelUnavailableError, - PetPipelineInvariantError, - ): - raise - except RuntimeError as exc: - raise PetModelUnavailableError(str(exc)) from exc - return self._embedder - - def _resolve_model_path(self, relative_path: Path, *, directory: bool = False) -> Path: - if self._model_root == default_pet_model_dir(): - return resolve_pet_model_path(relative_path, directory=directory) - return self._model_root / relative_path - - -def build_pet_key( - *, - asset_id: str, - bbox: tuple[int, int, int, int], - image_width: int, - image_height: int, - species_label: str | None = None, - detector_key_version: str = PET_DETECTOR_KEY_VERSION, - quantization: int = 12, -) -> str: - x, y, width, height = bbox - center_x = x + width / 2.0 - center_y = y + height / 2.0 - quantized = ( - _quantize_value(center_x, quantization), - _quantize_value(center_y, quantization), - _quantize_value(width, quantization), - _quantize_value(height, quantization), - ) - species = _normalize_species_label(species_label) or "unknown" - payload = ( - f"{PET_KEY_VERSION}|{detector_key_version}|{asset_id}|" - f"{image_width}x{image_height}|{species}|" - f"{quantized[0]}|{quantized[1]}|{quantized[2]}|{quantized[3]}" - ) - digest = hashlib.sha256(payload.encode("utf-8")).hexdigest() - return f"{PET_KEY_VERSION}:{digest}" - - -def cluster_pet_records( - detections: list[PetDetectionRecord], - *, - distance_threshold: float = 0.42, -) -> tuple[list[PetDetectionRecord], list[PetRecord]]: - if not detections: - return [], [] - - updated_detections = list(detections) - pets: list[PetRecord] = [] - labels = _cluster_pet_detection_labels( - detections, - distance_threshold=distance_threshold, - ) - grouped_indices: dict[str, list[int]] = defaultdict(list) - for index, label in enumerate(labels.tolist()): - grouped_indices[f"cluster-{label}"].append(index) - - for grouped in grouped_indices.values(): - members = [detections[index] for index in grouped] - key_detection = max(members, key=key_detection_sort_key) - pet_id = uuid.uuid4().hex - center_embedding = compute_cluster_center( - np.stack([member.embedding for member in members], axis=0) - ) - timestamp = utc_now_iso() - evidence_asset_count = len({member.asset_id for member in members if member.asset_id}) - pets.append( - PetRecord( - pet_id=pet_id, - name=None, - key_detection_id=key_detection.detection_id, - detection_count=len(members), - center_embedding=center_embedding, - embedding_dim=int(center_embedding.shape[0]), - created_at=timestamp, - updated_at=timestamp, - sample_count=len(members), - profile_state=profile_state_for_sample_count(evidence_asset_count), - species_label=_dominant_species_label(members), - embedding_pipeline_version=members[0].embedding_pipeline_version, - generation_id=members[0].generation_id, - boundary_embeddings=_boundary_embeddings(members, center_embedding), - evidence_asset_count=evidence_asset_count, - ) - ) - for index in grouped: - updated_detections[index] = replace(updated_detections[index], pet_id=pet_id) - - pets.sort(key=lambda pet: (-pet.detection_count, pet.created_at, pet.pet_id)) - return updated_detections, pets - - -def build_pet_records_from_detections( - detections: Sequence[PetDetectionRecord], - *, - names_by_pet_id: dict[str, str | None] | None = None, - created_at_by_pet_id: dict[str, str] | None = None, - allow_mixed_identity_members: bool = False, -) -> list[PetRecord]: - grouped: dict[str, list[PetDetectionRecord]] = defaultdict(list) - for detection in detections: - if detection.pet_id: - grouped[str(detection.pet_id)].append(detection) - - names = dict(names_by_pet_id or {}) - created = dict(created_at_by_pet_id or {}) - updated_at = utc_now_iso() - pets: list[PetRecord] = [] - for pet_id, members in grouped.items(): - species_labels = { - label - for label in (_normalize_species_label(member.species_label) for member in members) - if label is not None - } - if len(species_labels) > 1: - if not allow_mixed_identity_members: - raise ValueError( - f"Pet {pet_id} mixes incompatible species labels: " - f"{sorted(species_labels)}" - ) - _LOGGER.info( - "Preserving mixed-species Pet identity %s: species=%s", - pet_id, - sorted(species_labels), - ) - contract_groups: dict[tuple[str, int, int], list[PetDetectionRecord]] = defaultdict(list) - for member in members: - contract_groups[ - ( - str(member.embedding_pipeline_version or ""), - int(member.embedding_dim), - int(member.generation_id), - ) - ].append(member) - if len(contract_groups) != 1: - if not allow_mixed_identity_members: - raise ValueError( - f"Pet {pet_id} mixes incompatible embedding contracts: " - f"{sorted(contract_groups)}" - ) - _LOGGER.info( - "Preserving mixed-contract Pet identity %s: contracts=%s", - pet_id, - sorted(contract_groups), - ) - _profile_contract, profile_members = max( - contract_groups.items(), - key=lambda item: ( - len(item[1]), - item[0][2], - item[0][0], - item[0][1], - ), - ) - key_detection = max(members, key=key_detection_sort_key) - center_embedding = compute_cluster_center( - np.stack([member.embedding for member in profile_members], axis=0) + return _original_download_file(*args, **kwargs) + except RuntimeError as exc: + if not label.startswith("DINOv2"): + raise + message = str(exc) + detector_hint = ( + "Check your network connection, set " + f"{_impl.PET_DETECTOR_MODEL_URL_ENV}, or install the model manually." ) - sample_count = len(members) - evidence_asset_count = len({member.asset_id for member in members if member.asset_id}) - pets.append( - PetRecord( - pet_id=pet_id, - name=names.get(pet_id), - key_detection_id=key_detection.detection_id, - detection_count=sample_count, - center_embedding=center_embedding, - embedding_dim=int(center_embedding.shape[0]), - created_at=created.get( - pet_id, - min((member.detected_at for member in members), default=updated_at), - ), - updated_at=updated_at, - sample_count=sample_count, - profile_state=profile_state_for_sample_count(evidence_asset_count), - species_label=_dominant_species_label(members), - embedding_pipeline_version=profile_members[0].embedding_pipeline_version, - generation_id=profile_members[0].generation_id, - boundary_embeddings=_boundary_embeddings(profile_members, center_embedding), - evidence_asset_count=evidence_asset_count, + if detector_hint not in message: + raise + raise RuntimeError( + message.replace( + detector_hint, + "Check your network connection or install the model manually.", ) - ) - pets.sort(key=lambda pet: (-pet.detection_count, pet.created_at, pet.pet_id)) - return pets - - -def _boundary_embeddings( - members: Sequence[PetDetectionRecord], - center_embedding: np.ndarray, -) -> tuple[np.ndarray, ...]: - ranked = sorted( - members, - key=lambda member: ( - -cosine_distance(member.embedding, center_embedding), - member.detection_id, - ), - ) - return tuple(normalize_vector(member.embedding) for member in ranked[:8]) - - -def canonicalize_pet_identities( - detections: list[PetDetectionRecord], - pets: list[PetRecord], - state_repository: PetStateRepository, - *, - distance_threshold: float, -) -> tuple[list[PetDetectionRecord], list[PetRecord]]: - if not detections or not pets: - return detections, pets - - profiles = {profile.pet_id: profile for profile in state_repository.get_identity_profiles()} - redirects = state_repository.get_merge_redirect_map() - pet_key_map = state_repository.get_pet_key_map(detection.pet_key for detection in detections) - detections_by_pet_id: dict[str, list[PetDetectionRecord]] = defaultdict(list) - for detection in detections: - if detection.pet_id is not None: - detections_by_pet_id[detection.pet_id].append(detection) + ) from exc - canonical_members: dict[str, list[PetDetectionRecord]] = defaultdict(list) - canonical_names: dict[str, str | None] = {} - canonical_created_at: dict[str, str] = {} - direct_anchors: set[str] = set() - for pet in pets: - members = detections_by_pet_id.get(pet.pet_id, []) - resolution = resolve_canonical_pet_id( - pet, - members, - profiles=profiles, - pet_key_map=pet_key_map, - redirects=redirects, - distance_threshold=distance_threshold, - ) - canonical_id = resolution.canonical_pet_id - is_incompatible = bool( - canonical_members.get(canonical_id) - and not _detection_groups_compatible( - canonical_members[canonical_id], - members, - distance_threshold=distance_threshold, - ) - ) - same_asset_conflict = bool( - canonical_members.get(canonical_id) - and _detection_groups_share_asset(canonical_members[canonical_id], members) - ) - if is_incompatible and ( - same_asset_conflict - or (not resolution.is_redirect_alias and canonical_id in direct_anchors) - ): - canonical_id = uuid.uuid4().hex - resolution = PetIdentityResolution( - raw_pet_id=canonical_id, - canonical_pet_id=canonical_id, - source=PetIdentityResolutionSource.NEW, - ) - if not resolution.is_redirect_alias: - direct_anchors.add(canonical_id) - profile = profiles.get(canonical_id) - canonical_members[canonical_id].extend(members) - canonical_names.setdefault(canonical_id, profile.name if profile is not None else None) - canonical_created_at.setdefault( - canonical_id, - profile.created_at if profile is not None else pet.created_at, - ) - - updated = list(detections) - index_by_detection_id = { - detection.detection_id: index for index, detection in enumerate(detections) - } - for canonical_id, members in canonical_members.items(): - for member in members: - updated[index_by_detection_id[member.detection_id]] = replace( - member, - pet_id=canonical_id, - ) - canonical_pets = build_pet_records_from_detections( - updated, - names_by_pet_id=canonical_names, - created_at_by_pet_id=canonical_created_at, - ) - return updated, canonical_pets +def resolve_pet_model_path(relative_path: Path, *, directory: bool = False) -> Path: + """Resolve Pets models without mutating a DINOv2 cache from the read path.""" + if not directory: + return _impl.resolve_pet_model_path(relative_path, directory=False) -def resolve_canonical_pet_id( - pet: PetRecord, - members: list[PetDetectionRecord], - *, - profiles: dict[str, PetProfile], - pet_key_map: dict[str, str], - redirects: dict[str, str], - distance_threshold: float, -) -> PetIdentityResolution: - vote_counter = Counter( - pet_key_map[member.pet_key] for member in members if member.pet_key in pet_key_map - ) - if vote_counter: - raw_pet_id = max( - vote_counter.items(), - key=lambda item: ( - item[1], - profiles[item[0]].updated_at if item[0] in profiles else "", - item[0], - ), - )[0] - canonical_pet_id = redirects.get(raw_pet_id, raw_pet_id) - return PetIdentityResolution( - raw_pet_id=raw_pet_id, - canonical_pet_id=canonical_pet_id, - source=( - PetIdentityResolutionSource.REDIRECT_KEY - if canonical_pet_id != raw_pet_id - else PetIdentityResolutionSource.KEY - ), - ) + relative = Path(relative_path) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError("Pet model path must be relative to a configured model root.") - best_profile_id: str | None = None - best_distance = float("inf") - pet_species = _normalize_species_label(pet.species_label) - for profile in profiles.values(): - is_redirect_alias = profile.pet_id in redirects - if not is_redirect_alias and str(profile.profile_state or "unstable") != "stable": - continue - profile_species = _normalize_species_label(profile.species_label) - if pet_species != profile_species: - continue - if profile.embedding_dim <= 0 or profile.center_embedding.size == 0: - continue - if profile.center_embedding.shape != pet.center_embedding.shape: + override = _impl.pet_model_override_dir() + if override is not None: + candidate = override / relative + model_path = candidate / f"{relative.name}.pt" + if model_path.is_file(): + try: + _validate_dinov2_cache_metadata(model_path, model_name=relative.name) + except (OSError, RuntimeError): + # Acquisition repairs this exact authoritative location under lock. + pass + return candidate + + user_cache = _impl.user_pet_model_cache_dir() + bundled = _impl.bundled_pet_model_dir() + bundled_invalid = False + for root in _impl.pet_model_search_roots(): + candidate = root / relative + if not candidate.is_dir(): continue - distance = cosine_distance(pet.center_embedding, profile.center_embedding) - if distance < best_distance: - best_distance = distance - best_profile_id = profile.pet_id - - if best_profile_id is not None and best_distance <= distance_threshold: - canonical_pet_id = redirects.get(best_profile_id, best_profile_id) - return PetIdentityResolution( - raw_pet_id=best_profile_id, - canonical_pet_id=canonical_pet_id, - source=( - PetIdentityResolutionSource.REDIRECT_PROFILE - if canonical_pet_id != best_profile_id - else PetIdentityResolutionSource.PROFILE - ), - ) - new_pet_id = uuid.uuid4().hex - return PetIdentityResolution( - raw_pet_id=new_pet_id, - canonical_pet_id=new_pet_id, - source=PetIdentityResolutionSource.NEW, - ) - - -def _cluster_pet_detection_labels( - detections: Sequence[PetDetectionRecord], - *, - distance_threshold: float, -) -> np.ndarray: - if not detections: - return np.empty((0,), dtype=np.int32) - embeddings = np.stack([detection.embedding for detection in detections], axis=0).astype( - np.float32 - ) - return _cluster_embeddings_bounded_single_link( - embeddings, - compatibility_keys=[ - ( - _normalize_species_label(detection.species_label), - str(detection.embedding_pipeline_version or ""), - int(detection.embedding_dim), - int(detection.generation_id), - ) - for detection in detections - ], - member_keys=[detection.detection_id for detection in detections], - cannot_link_keys=[detection.asset_id for detection in detections], - distance_threshold=distance_threshold, - ) - - -def _cluster_embeddings_bounded_single_link( - embeddings: np.ndarray, - *, - compatibility_keys: Sequence[object], - member_keys: Sequence[str] | None = None, - cannot_link_keys: Sequence[str] | None = None, - distance_threshold: float, -) -> np.ndarray: - count = int(embeddings.shape[0]) - if count == 0: - return np.empty((0,), dtype=np.int32) - return _cluster_distance_matrix_bounded_single_link( - cosine_distance_matrix(embeddings), - compatibility_keys=compatibility_keys, - member_keys=member_keys, - cannot_link_keys=cannot_link_keys, - link_threshold=distance_threshold, - diameter_threshold=distance_threshold * PET_CLUSTER_DIAMETER_MULTIPLIER, - ) - - -def _cluster_distance_matrix_bounded_single_link( - distance_matrix: np.ndarray, - *, - compatibility_keys: Sequence[object], - link_threshold: float, - diameter_threshold: float, - member_keys: Sequence[str] | None = None, - cannot_link_keys: Sequence[str] | None = None, -) -> np.ndarray: - """Cluster nearest-neighbour links without allowing unbounded similarity chains.""" - - count = int(distance_matrix.shape[0]) - if count == 0: - return np.empty((0,), dtype=np.int32) - if distance_matrix.shape != (count, count): - raise ValueError("Pet distance matrix must be square.") - if len(compatibility_keys) != count: - raise ValueError("Pet compatibility key count must match the distance matrix.") - stable_keys = tuple(member_keys or (str(index) for index in range(count))) - if len(stable_keys) != count: - raise ValueError("Pet member key count must match the distance matrix.") - resolved_cannot_link_keys = tuple(cannot_link_keys or ("" for _ in range(count))) - if len(resolved_cannot_link_keys) != count: - raise ValueError("Pet cannot-link key count must match the distance matrix.") - - clusters: list[list[int]] = [[index] for index in range(count)] - diameters: list[float] = [0.0] * count - cannot_link_sets: list[set[str]] = [ - ({key} if key else set()) for key in resolved_cannot_link_keys - ] - - while True: - best_pair: tuple[int, int] | None = None - best_key: tuple[float, float, tuple[str, ...], tuple[str, ...]] | None = None - best_diameter = 0.0 - for left_index in range(len(clusters)): - for right_index in range(left_index + 1, len(clusters)): - left = clusters[left_index] - right = clusters[right_index] - if not _cluster_keys_compatible(left, right, compatibility_keys): - continue - if cannot_link_sets[left_index] & cannot_link_sets[right_index]: - _LOGGER.debug( - "Pet clustering constraint hit: same_asset_cannot_link_hits=1" - ) - continue - cross_distances = [ - float(distance_matrix[left_member, right_member]) - for left_member in left - for right_member in right - ] - connection_distance = min(cross_distances) - if connection_distance > link_threshold: - continue - merged_diameter = max( - diameters[left_index], - diameters[right_index], - max(cross_distances), - ) - if merged_diameter > diameter_threshold: - continue - left_keys = tuple(sorted(stable_keys[index] for index in left)) - right_keys = tuple(sorted(stable_keys[index] for index in right)) - ordered_cluster_keys = tuple(sorted((left_keys, right_keys))) - tie_key = ( - connection_distance, - merged_diameter, - ordered_cluster_keys[0], - ordered_cluster_keys[1], - ) - if best_key is None or tie_key < best_key: - best_key = tie_key - best_pair = (left_index, right_index) - best_diameter = merged_diameter - if best_pair is None: - break - left_index, right_index = best_pair - merged = sorted(clusters[left_index] + clusters[right_index]) - clusters[left_index] = merged - diameters[left_index] = best_diameter - cannot_link_sets[left_index].update(cannot_link_sets[right_index]) - del clusters[right_index] - del diameters[right_index] - del cannot_link_sets[right_index] - ordering = sorted( - range(len(clusters)), - key=lambda cluster_index: tuple( - sorted(stable_keys[index] for index in clusters[cluster_index]) - ), - ) - clusters = [clusters[index] for index in ordering] - diameters = [diameters[index] for index in ordering] - cannot_link_sets = [cannot_link_sets[index] for index in ordering] - - labels = np.empty((count,), dtype=np.int32) - for cluster_id, members in enumerate(clusters): - for member in members: - labels[member] = cluster_id - return labels - - -def _cluster_keys_compatible( - left: Sequence[int], - right: Sequence[int], - compatibility_keys: Sequence[object], -) -> bool: - keys = {compatibility_keys[index] for index in [*left, *right]} - return len(keys) == 1 - - -def _detection_species_compatible( - left: Sequence[PetDetectionRecord], - right: Sequence[PetDetectionRecord], -) -> bool: - labels = [ - *(_normalize_species_label(detection.species_label) for detection in left), - *(_normalize_species_label(detection.species_label) for detection in right), - ] - return len(set(labels)) <= 1 - - -def _detection_contracts_compatible( - left: Sequence[PetDetectionRecord], - right: Sequence[PetDetectionRecord], -) -> bool: - contracts = { - ( - str(detection.embedding_pipeline_version or ""), - int(detection.embedding_dim), - int(detection.generation_id), - ) - for detection in [*left, *right] - } - return len(contracts) <= 1 - - -def _detection_groups_compatible( - left: Sequence[PetDetectionRecord], - right: Sequence[PetDetectionRecord], - *, - distance_threshold: float, -) -> bool: - if not _detection_species_compatible(left, right) or not _detection_contracts_compatible( - left, right - ): - return False - if _detection_groups_share_asset(left, right): - return False - if not left or not right: - return True - cross_distances = [ - cosine_distance(left_detection.embedding, right_detection.embedding) - for left_detection in left - for right_detection in right - ] - if min(cross_distances) > distance_threshold: - return False - members = [*left, *right] - distance_matrix = cosine_distance_matrix( - np.stack([member.embedding for member in members], axis=0) - ) - return float(distance_matrix.max()) <= (distance_threshold * PET_CLUSTER_DIAMETER_MULTIPLIER) - - -def _detection_groups_share_asset( - left: Sequence[PetDetectionRecord], - right: Sequence[PetDetectionRecord], -) -> bool: - left_assets = {detection.asset_id for detection in left if detection.asset_id} - right_assets = {detection.asset_id for detection in right if detection.asset_id} - return bool(left_assets & right_assets) - - -def _dominant_species_label(detections: Sequence[PetDetectionRecord]) -> str | None: - counter = Counter( - label - for label in (_normalize_species_label(detection.species_label) for detection in detections) - if label is not None - ) - if not counter: - return None - return max(counter.items(), key=lambda item: (item[1], item[0]))[0] - - -def _normalize_species_label(value: object) -> str | None: - if value is None: - return None - label = str(value).strip().lower() - return label or None - - -def _pet_candidate_quality_score( - *, - confidence: float, - relative_area_ratio: float, -) -> float: - """Rank retained detections without turning relative size into a hard gate.""" - - normalized_area = min(1.0, math.sqrt(max(0.0, float(relative_area_ratio)))) - return float(0.75 * max(0.0, min(1.0, confidence)) + 0.25 * normalized_area) - - -def _normalize_bbox( - raw_bbox, - *, - image_width: int, - image_height: int, -) -> tuple[int, int, int, int]: - box = np.asarray(raw_bbox, dtype=np.float32).flatten().tolist() - x, y, width, height = [round(value) for value in box[:4]] - x = max(0, min(x, image_width - 1)) - y = max(0, min(y, image_height - 1)) - width = max(1, min(width, image_width - x)) - height = max(1, min(height, image_height - y)) - return x, y, width, height - - -def _quantize_value(value: float, step: int) -> int: - step = max(1, int(step)) - return int(round(float(value) / step) * step) - - -class _YoloxOnnxPetDetector: - def __init__( - self, - model_path: Path, - *, - score_threshold: float = 0.30, - allow_model_download: bool = True, - enable_tiled_detection: bool = True, - tile_scan_min_confidence: float | None = None, - tile_species: frozenset[str] = SUPPORTED_DEFAULT_SPECIES, - execution_providers: Sequence[str] | None = None, - ) -> None: - self._model_path = Path(model_path) - self._score_threshold = float(score_threshold) - self._enable_tiled_detection = bool(enable_tiled_detection) - self._tile_scan_min_confidence = ( - self._score_threshold - if tile_scan_min_confidence is None - else float(tile_scan_min_confidence) - ) - self._tile_species = frozenset(tile_species) - try: - import onnxruntime as ort - except ImportError as exc: - raise PetRuntimeUnavailableError( - "Pet scanning unavailable: missing onnxruntime. Install the optional " - 'Pets AI runtime with: pip install -e ".[pets-ai]"' - ) from exc - ensure_pet_detector_model( - self._model_path, - allow_model_download=allow_model_download, - ) try: - providers = list(execution_providers or _resolve_execution_providers(ort)) - self._session = ort.InferenceSession(str(self._model_path), providers=providers) - self._input_name = self._session.get_inputs()[0].name - shape = self._session.get_inputs()[0].shape - self._input_size = _input_size_from_shape(shape) - _validate_yolox_session_contract(self._session) - except Exception as exc: - if isinstance(exc, PetPipelineInvariantError): - raise - raise PetModelUnavailableError( - "Pet scanning unavailable: failed to initialize YOLOX detector model at " - f"{self._model_path} ({_error_reason(exc)}). Check the model cache, " - "disable unsupported execution providers, or reinstall the Pets AI runtime." - ) from exc + model_name = relative.name + model_path = candidate / f"{model_name}.pt" + if not model_path.is_file(): + raise RuntimeError("DINOv2 model file is missing") + _validate_dinov2_cache_metadata(model_path, model_name=model_name) + return candidate + except (OSError, RuntimeError): + # DINOv2 cache cleanup is intentionally deferred to the acquisition + # owner. A resolver can race with metadata-first publication, so it + # must never unlink either side of the cache pair here. + if root == bundled: + bundled_invalid = True - def detect(self, image) -> list[_DetectedPetBox]: - boxes = self._detect_single_image(image) - if not self._enable_tiled_detection: - return _dedupe_supported_species_boxes(boxes) + if bundled_invalid: + return user_cache / relative + return _impl.pet_model_install_root() / relative - image_width, image_height = image.size - for crop_box in _select_uncovered_tile_regions( - image_width, - image_height, - boxes, - max_regions=4, - ): - left, top, *_ = crop_box - crop = image.crop(crop_box) - boxes.extend(self._detect_single_image(crop, offset=(left, top))) - return _dedupe_supported_species_boxes(boxes) - def _detect_single_image( - self, - image, - *, - offset: tuple[int, int] = (0, 0), - ) -> list[_DetectedPetBox]: - image_width, image_height = image.size - input_width, input_height = self._input_size - preprocessed = _preprocess_yolox( - image, - input_width=input_width, - input_height=input_height, - ) - try: - outputs = self._session.run(None, {self._input_name: preprocessed.tensor}) - except Exception as exc: # noqa: BLE001 - provider failures vary by backend - raise PetInferenceError(f"Pet detector inference failed: {_error_reason(exc)}") from exc - predictions = _flatten_predictions(outputs) - boxes: list[_DetectedPetBox] = [] - for x0, y0, x1, y1, confidence, class_id in _decode_yolox_predictions( - predictions, - input_size=self._input_size, - ): - species = COCO_ANIMAL_LABELS.get(int(class_id)) - if species is None or confidence < self._score_threshold: - continue - bbox = _map_yolox_box_to_source( - (x0, y0, x1, y1), - preprocessed=preprocessed, - image_width=image_width, - image_height=image_height, - offset=offset, - ) - boxes.append( - _DetectedPetBox( - bbox=bbox, - confidence=float(confidence), - species_label=species, - ) - ) - return boxes +_LegacyDinoV2Embedder = _impl._DinoV2Embedder - def _has_tile_species_box(self, boxes: list[_DetectedPetBox]) -> bool: - return any( - box.species_label in self._tile_species - and box.confidence >= self._tile_scan_min_confidence - for box in boxes - ) +class _DinoV2Embedder(_LegacyDinoV2Embedder): + """Acquire the fixed DINOv2 Release without coupling it to device activation.""" -class _DinoV2Embedder: - def __init__( - self, - model_dir: Path, - *, - model_name: str, - allow_model_download: bool = True, - ) -> None: - self._model_dir = Path(model_dir) - self._model_name = model_name + def _build_dinov2_cache(self, model_path: Path): + loaded = self._build_verified_dinov2_cpu_cache(model_path) try: - import torch - except ImportError as exc: - raise PetRuntimeUnavailableError( - "Pet scanning unavailable: missing torch for DINOv2 pet embeddings. " - 'Install the optional Pets AI runtime with: pip install -e ".[pets-ai]"' + loaded.eval() + loaded.to(self._device) + except Exception as exc: # noqa: BLE001 - backend failures vary by runtime + raise _impl.PetModelUnavailableError( + "Pet scanning unavailable: DINOv2 cache was built and verified, " + f"but the runtime device could not load it ({_error_reason(exc)})." ) from exc - self._torch = torch - self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model_path = self._model_dir / f"{model_name}.pt" - if model_path.is_file(): - _validate_dinov2_cache_metadata(model_path, model_name=model_name) - self._model = torch.jit.load(str(model_path), map_location=self._device) - elif allow_model_download: - self._model = self._download_dinov2_model(model_path) - else: - raise PetModelUnavailableError( - "Pet scanning unavailable: missing DINOv2 TorchScript model at " - f"{model_path}. Set IPHOTO_PET_MODEL_DIR or enable pet model downloads." - ) - self._model.eval() + return loaded - def embed(self, image) -> np.ndarray: - tensor = image_to_chw_float(image, (224, 224)) - torch = self._torch + def _load_verified_dinov2_cpu_cache(self, model_path: Path): + if not model_path.is_file(): + return None try: - with torch.no_grad(): - input_tensor = torch.from_numpy(tensor).to(self._device) - output = self._model(input_tensor) - if isinstance(output, (list, tuple)): - output = output[0] - vector = output.detach().cpu().numpy().reshape(-1) - except Exception as exc: # noqa: BLE001 - provider failures vary by backend - raise PetInferenceError( - f"Pet embedding inference failed: {_error_reason(exc)}" - ) from exc - expected_dimension = int(_EMBEDDER_MANIFEST["output_shape"][-1]) - if vector.size != expected_dimension: - raise PetPipelineInvariantError( - "Pet scanning unavailable: DINOv2 output contract mismatch " - f"({vector.size} != {expected_dimension})." - ) - return normalize_vector(vector.astype(np.float32)) + _validate_dinov2_cache_metadata(model_path, model_name=self._model_name) + except RuntimeError: + return None + return self._torch.jit.load(str(model_path), map_location="cpu") + + def _build_verified_dinov2_cpu_cache(self, model_path: Path): + """Download, publish, validate, and reload the Release artifact on CPU.""" - def _download_dinov2_model(self, model_path: Path): - url = str(_EMBEDDER_MANIFEST.get("torchscript_url") or "").strip() - if not url: - raise PetModelUnavailableError( - "Pet scanning unavailable: the fixed DINOv2 TorchScript release " - "artifact has not been configured. Install a package containing the " - "verified model or set IPHOTO_PET_MODEL_DIR." - ) try: - _install_certifi_environment() - _download_file( - url, - model_path, - label="DINOv2 TorchScript model", - expected_sha256=str(_EMBEDDER_MANIFEST["torchscript_sha256"]), - max_bytes=int(_EMBEDDER_MANIFEST["torchscript_size"]), - ) - _dinov2_metadata_path(model_path).write_text( - json.dumps( - { - "model_name": self._model_name, - "source_repository": _EMBEDDER_MANIFEST["source_repository"], - "source_revision": _DINO_SOURCE_REVISION, - "torchscript_sha256": _EMBEDDER_MANIFEST["torchscript_sha256"], - "torchscript_size": _EMBEDDER_MANIFEST["torchscript_size"], - "input_shape": _EMBEDDER_MANIFEST["input_shape"], - "output_shape": _EMBEDDER_MANIFEST["output_shape"], - }, - indent=2, - sort_keys=True, - ) - + "\n", - encoding="utf-8", - ) - _validate_dinov2_cache_metadata( + return _impl._acquire_dinov2_release( + self._torch, model_path, model_name=self._model_name, + download_file=_download_file, + validate_metadata=_validate_dinov2_cache_metadata, ) - model = self._torch.jit.load(str(model_path), map_location=self._device) + except _ModelStoragePermissionError: + raise except Exception as exc: - model_path.unlink(missing_ok=True) - _dinov2_metadata_path(model_path).unlink(missing_ok=True) - raise PetModelUnavailableError( - "Pet scanning unavailable: failed to download the verified DINOv2 " - f"TorchScript model ({_error_reason(exc)})." + raise _impl.PetModelUnavailableError( + "Pet scanning unavailable: failed to acquire the verified DINOv2 " + f"Release artifact ({_error_reason(exc)})." ) from exc - model.eval() - model.to(self._device) - return model - - -def _resolve_execution_providers(ort) -> list[str]: - available = set(ort.get_available_providers()) - preferred = [ - "CUDAExecutionProvider", - "CoreMLExecutionProvider", - "OpenVINOExecutionProvider", - "CPUExecutionProvider", - ] - providers = [provider for provider in preferred if provider in available] - return providers or ["CPUExecutionProvider"] - - -def _input_size_from_shape(shape: Sequence[object]) -> tuple[int, int]: - if len(shape) >= 4 and isinstance(shape[2], int) and isinstance(shape[3], int): - return int(shape[3]), int(shape[2]) - return 640, 640 - - -def _validate_yolox_session_contract(session) -> None: - inputs = session.get_inputs() - outputs = session.get_outputs() - if len(inputs) != 1 or len(inputs[0].shape) != 4: - raise RuntimeError("Pet scanning unavailable: YOLOX input contract is invalid.") - if inputs[0].shape[1] not in {3, "3"}: - raise RuntimeError("Pet scanning unavailable: YOLOX input must have three channels.") - concrete_input = tuple( - int(value) if isinstance(value, (int, np.integer)) else None for value in inputs[0].shape - ) - expected_input = tuple(int(value) for value in _DETECTOR_MANIFEST["input"]["shape"]) - if concrete_input != expected_input: - raise RuntimeError( - "Pet scanning unavailable: YOLOX input shape does not match the manifest." - ) - if not outputs or len(outputs[0].shape) < 2: - raise RuntimeError("Pet scanning unavailable: YOLOX output contract is invalid.") - last_output_dim = outputs[0].shape[-1] - expected_output = tuple(int(value) for value in _DETECTOR_MANIFEST["output_shape"]) - concrete_output = tuple( - int(value) if isinstance(value, (int, np.integer)) else None for value in outputs[0].shape - ) - if isinstance(last_output_dim, (int, np.integer)) and concrete_output != expected_output: - raise RuntimeError( - "Pet scanning unavailable: YOLOX output shape does not match the manifest." - ) - -def _preprocess_yolox( - image, - *, - input_width: int, - input_height: int, -) -> _YoloxPreprocessResult: - image_width, image_height = image.size - resize_ratio = min( - input_width / float(max(1, image_width)), - input_height / float(max(1, image_height)), - ) - resized_width = max(1, int(image_width * resize_ratio)) - resized_height = max(1, int(image_height * resize_ratio)) - resized = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR) - canvas = Image.new("RGB", (input_width, input_height), (114, 114, 114)) - canvas.paste(resized, (0, 0)) - array = np.asarray(canvas, dtype=np.float32) - # YOLOX 0.1.1rc0 deployment weights consume OpenCV-style BGR bytes. - # This release explicitly removed legacy mean/std normalization. - array = array[:, :, ::-1] - array = np.transpose(array, (2, 0, 1))[None, :, :, :] - return _YoloxPreprocessResult( - tensor=np.ascontiguousarray(array), - resize_ratio=float(resize_ratio), - ) +class _LazyDinoV2Embedder: + """Delay torch import/model acquisition until YOLOX accepted a pet crop.""" + def __init__(self, factory) -> None: + self._factory = factory + self._delegate = None -def _map_yolox_box_to_source( - box: tuple[float, float, float, float], - *, - preprocessed: _YoloxPreprocessResult, - image_width: int, - image_height: int, - offset: tuple[int, int] = (0, 0), -) -> tuple[int, int, int, int]: - ratio = max(float(preprocessed.resize_ratio), 1e-6) - offset_x, offset_y = offset - x0, y0, x1, y1 = box - left = round((x0 - preprocessed.pad_left) / ratio) - top = round((y0 - preprocessed.pad_top) / ratio) - right = round((x1 - preprocessed.pad_left) / ratio) - bottom = round((y1 - preprocessed.pad_top) / ratio) - left = max(0, min(left, image_width - 1)) - top = max(0, min(top, image_height - 1)) - right = max(left + 1, min(right, image_width)) - bottom = max(top + 1, min(bottom, image_height)) - return ( - left + offset_x, - top + offset_y, - max(1, right - left), - max(1, bottom - top), - ) + def _ensure_delegate(self): + if self._delegate is None: + self._delegate = self._factory() + return self._delegate + def embed(self, image): + return self._ensure_delegate().embed(image) -def _tile_scan_regions(image_width: int, image_height: int) -> list[tuple[int, int, int, int]]: - width = max(1, int(image_width)) - height = max(1, int(image_height)) - def region(left: float, top: float, right: float, bottom: float) -> tuple[int, int, int, int]: - x0 = max(0, min(round(width * left), width - 1)) - y0 = max(0, min(round(height * top), height - 1)) - x1 = max(x0 + 1, min(round(width * right), width)) - y1 = max(y0 + 1, min(round(height * bottom), height)) - return x0, y0, x1, y1 +class PetClusterPipeline(_impl.PetClusterPipeline): + """Pipeline with DINOv2 construction deferred until the first accepted crop.""" - candidates = [ - region(0.0, 0.0, 0.70, 1.0), - region(0.30, 0.0, 1.0, 1.0), - region(0.0, 0.0, 1.0, 0.70), - region(0.0, 0.30, 1.0, 1.0), - region(0.0, 0.0, 0.65, 0.65), - region(0.35, 0.0, 1.0, 0.65), - region(0.15, 0.15, 0.85, 0.85), - ] - return list(dict.fromkeys(candidates)) - - -def _select_uncovered_tile_regions( - image_width: int, - image_height: int, - boxes: Sequence[_DetectedPetBox], - *, - max_regions: int, -) -> list[tuple[int, int, int, int]]: - """Choose a bounded set of tiles by area not covered by full-frame pets.""" - - supported = [box for box in boxes if box.species_label in SUPPORTED_DEFAULT_SPECIES] - ranked: list[tuple[float, tuple[int, int, int, int]]] = [] - for region in _tile_scan_regions(image_width, image_height): - x0, y0, x1, y1 = region - tile_box = (x0, y0, x1 - x0, y1 - y0) - tile_area = max(1, tile_box[2] * tile_box[3]) - covered = min( - tile_area, - sum(_bbox_intersection_area(tile_box, box.bbox) for box in supported), - ) - uncovered_ratio = 1.0 - (covered / float(tile_area)) - if uncovered_ratio >= 0.20: - ranked.append((uncovered_ratio, region)) - ranked.sort(key=lambda item: (-item[0], item[1])) - return [region for _, region in ranked[: max(0, int(max_regions))]] - - -def _flatten_predictions(outputs: Sequence[np.ndarray]) -> np.ndarray: - if not outputs: - return np.empty((0, 0), dtype=np.float32) - prediction = np.asarray(outputs[0], dtype=np.float32) - return prediction.reshape(-1, prediction.shape[-1]) - - -def _decode_yolox_predictions( - predictions: np.ndarray, - *, - input_size: tuple[int, int], -) -> list[tuple[float, float, float, float, float, int]]: - if predictions.size == 0: - return [] - decoded = np.asarray(predictions, dtype=np.float32) - if _looks_like_raw_yolox_output(decoded, input_size=input_size): - decoded = _decode_raw_yolox_output(decoded, input_size=input_size) - return [_decode_prediction(prediction) for prediction in decoded if prediction.shape[0] >= 6] - - -def _decode_prediction(prediction: np.ndarray) -> tuple[float, float, float, float, float, int]: - if prediction.shape[0] >= 85: - cx, cy, width, height = [float(value) for value in prediction[:4]] - object_score = float(prediction[4]) - class_scores = prediction[5:] - class_index = int(np.argmax(class_scores)) - confidence = object_score * float(class_scores[class_index]) - x0 = cx - width / 2.0 - y0 = cy - height / 2.0 - x1 = cx + width / 2.0 - y1 = cy + height / 2.0 - return x0, y0, x1, y1, confidence, class_index - x0, y0, x1, y1 = [float(value) for value in prediction[:4]] - confidence = float(prediction[4]) - class_id = round(float(prediction[5])) - return x0, y0, x1, y1, confidence, class_id - - -def _looks_like_raw_yolox_output( - predictions: np.ndarray, - *, - input_size: tuple[int, int], -) -> bool: - if predictions.ndim != 2 or predictions.shape[1] < 85: - return False - grids, _strides = _yolox_grids(input_size) - if predictions.shape[0] != grids.shape[0]: - return False - coord_max = float(np.nanmax(np.abs(predictions[:, :4]))) if predictions.size else 0.0 - return coord_max <= _YOLOX_RAW_COORD_LIMIT - - -def _decode_raw_yolox_output( - predictions: np.ndarray, - *, - input_size: tuple[int, int], -) -> np.ndarray: - grids, strides = _yolox_grids(input_size) - decoded = np.array(predictions, dtype=np.float32, copy=True) - decoded[:, :2] = (decoded[:, :2] + grids) * strides - decoded[:, 2:4] = np.exp(np.clip(decoded[:, 2:4], -20.0, 20.0)) * strides - return decoded - - -def _yolox_grids(input_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]: - input_width, input_height = input_size - grid_parts: list[np.ndarray] = [] - stride_parts: list[np.ndarray] = [] - for stride in _YOLOX_STRIDES: - grid_height = int(input_height) // stride - grid_width = int(input_width) // stride - yv, xv = np.meshgrid( - np.arange(grid_height, dtype=np.float32), - np.arange(grid_width, dtype=np.float32), - indexing="ij", - ) - grid = np.stack((xv, yv), axis=-1).reshape(-1, 2) - grid_parts.append(grid) - stride_parts.append(np.full((grid.shape[0], 1), stride, dtype=np.float32)) - return np.concatenate(grid_parts, axis=0), np.concatenate(stride_parts, axis=0) - - -def _dedupe_supported_species_boxes( - boxes: list[_DetectedPetBox], - *, - threshold: float = PET_PET_IOU_THRESHOLD, - smaller_box_coverage_threshold: float = PET_PET_SMALLER_BOX_COVERAGE_THRESHOLD, - normalized_center_distance_threshold: float = PET_PET_NORMALIZED_CENTER_DISTANCE_THRESHOLD, - cross_species_mutual_coverage_threshold: float = ( - PET_PET_CROSS_SPECIES_MUTUAL_COVERAGE_THRESHOLD - ), - cross_species_threshold: float = 0.90, - cross_species_score_margin: float = 0.25, -) -> list[_DetectedPetBox]: - selected: list[_DetectedPetBox] = [] - for box in sorted( - boxes, - key=lambda item: (item.quality_score, item.confidence), - reverse=True, - ): - suppress = False - for existing in selected: - overlap = _bbox_iou(existing.bbox, box.bbox) - existing_box_coverage, candidate_box_coverage = _bbox_pair_coverages( - existing.bbox, - box.bbox, - ) - smaller_box_coverage = max(existing_box_coverage, candidate_box_coverage) - normalized_center_distance = _bbox_normalized_center_distance( - existing.bbox, - box.bbox, + def _resolve_model_path(self, relative_path: Path, *, directory: bool = False) -> Path: + override = _impl.pet_model_override_dir() + if override is not None and self._model_root != override: + raise _impl.PetModelUnavailableError( + "Pet scanning unavailable: model root does not match " + f"{_impl.IPHOTO_PET_MODEL_DIR_ENV}." ) - reason = "" - if existing.species_label == box.species_label: - if overlap >= threshold: - reason = "same_species_iou" - elif ( - smaller_box_coverage >= smaller_box_coverage_threshold - and normalized_center_distance <= normalized_center_distance_threshold - ): - reason = "same_species_containment" - elif ( - existing.species_label in SUPPORTED_DEFAULT_SPECIES - and box.species_label in SUPPORTED_DEFAULT_SPECIES - and existing_box_coverage >= cross_species_mutual_coverage_threshold - and candidate_box_coverage >= cross_species_mutual_coverage_threshold - ): - reason = "cross_species_mutual_coverage" - elif ( - existing.species_label != box.species_label - and overlap >= cross_species_threshold - and existing.confidence - box.confidence >= cross_species_score_margin - ): - reason = "cross_species_iou" - - if reason: - _LOGGER.debug( - "Suppressed pet box: reason=%s species=%s candidate_confidence=%.3f " - "candidate_bbox=%s kept_species=%s kept_confidence=%.3f kept_bbox=%s " - "iou=%.3f smaller_box_coverage=%.3f kept_box_coverage=%.3f " - "candidate_box_coverage=%.3f " - "normalized_center_distance=%.3f", - reason, - box.species_label, - box.confidence, - box.bbox, - existing.species_label, - existing.confidence, - existing.bbox, - overlap, - smaller_box_coverage, - existing_box_coverage, - candidate_box_coverage, - normalized_center_distance, - ) - suppress = True - break - if suppress: - continue - selected.append(box) - return selected - - -def _bbox_iou(left: tuple[int, int, int, int], right: tuple[int, int, int, int]) -> float: - intersection = _bbox_intersection_area(left, right) - left_area = max(0, left[2]) * max(0, left[3]) - right_area = max(0, right[2]) * max(0, right[3]) - union = left_area + right_area - intersection - if union <= 0: - return 0.0 - return intersection / float(union) - - -def _bbox_pair_coverages( - left: tuple[int, int, int, int], - right: tuple[int, int, int, int], -) -> tuple[float, float]: - left_area = max(0, left[2]) * max(0, left[3]) - right_area = max(0, right[2]) * max(0, right[3]) - if left_area <= 0 or right_area <= 0: - return 0.0, 0.0 - intersection = _bbox_intersection_area(left, right) - return intersection / float(left_area), intersection / float(right_area) - - -def _bbox_normalized_center_distance( - left: tuple[int, int, int, int], - right: tuple[int, int, int, int], -) -> float: - left_area = max(0, left[2]) * max(0, left[3]) - right_area = max(0, right[2]) * max(0, right[3]) - smaller_area = min(left_area, right_area) - if smaller_area <= 0: - return float("inf") - left_center = (left[0] + left[2] / 2.0, left[1] + left[3] / 2.0) - right_center = (right[0] + right[2] / 2.0, right[1] + right[3] / 2.0) - return math.hypot( - left_center[0] - right_center[0], - left_center[1] - right_center[1], - ) / math.sqrt(smaller_area) - - -def _bbox_intersection_area( - left: tuple[int, int, int, int], - right: tuple[int, int, int, int], -) -> int: - lx, ly, lw, lh = left - rx, ry, rw, rh = right - left_x2 = lx + lw - left_y2 = ly + lh - right_x2 = rx + rw - right_y2 = ry + rh - inter_left = max(lx, rx) - inter_top = max(ly, ry) - inter_right = min(left_x2, right_x2) - inter_bottom = min(left_y2, right_y2) - inter_width = max(0, inter_right - inter_left) - inter_height = max(0, inter_bottom - inter_top) - return inter_width * inter_height - - -def _pet_box_overlaps_people_boxes( - pet_box: tuple[int, int, int, int], - people_boxes: Sequence[tuple[int, int, int, int]], - *, - image_dimensions: tuple[int, int] | None = None, - iou_threshold: float = PET_PEOPLE_IOU_THRESHOLD, - smaller_box_coverage_threshold: float = PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, -) -> bool: - """Return whether a pet detection conflicts with a People face region.""" - - return _pet_people_overlap_decision( - pet_box, - people_boxes, - image_dimensions=image_dimensions, - iou_threshold=iou_threshold, - smaller_box_coverage_threshold=smaller_box_coverage_threshold, - ).suppressed - - -def _pet_people_overlap_decision( - pet_box: tuple[int, int, int, int], - people_boxes: Sequence[tuple[int, int, int, int]], - *, - image_dimensions: tuple[int, int] | None = None, - iou_threshold: float = PET_PEOPLE_IOU_THRESHOLD, - smaller_box_coverage_threshold: float = PET_PEOPLE_SMALLER_BOX_COVERAGE_THRESHOLD, - larger_pet_ratio: float = PET_PEOPLE_LARGER_PET_RATIO, - mural_image_coverage_threshold: float = PET_PEOPLE_MURAL_IMAGE_COVERAGE_THRESHOLD, -) -> _PetPeopleOverlapDecision: - """Classify face/pet overlap without losing pets held by people. - - A pet body may legitimately contain a much smaller human face. Preserve - that candidate unless it also spans most of the image, which is the shape - produced by the known wall-mural false-positive regression. - """ + if self._model_root == _impl.default_pet_model_dir(): + return resolve_pet_model_path(relative_path, directory=directory) + return self._model_root / relative_path - pet_area = max(0, pet_box[2]) * max(0, pet_box[3]) - if pet_area <= 0: - return _PetPeopleOverlapDecision(False) - image_area = 0 - if image_dimensions is not None: - image_area = max(0, int(image_dimensions[0])) * max(0, int(image_dimensions[1])) - pet_image_coverage = pet_area / float(image_area) if image_area else 0.0 - for people_box in people_boxes: - people_area = max(0, people_box[2]) * max(0, people_box[3]) - if people_area <= 0: - continue - intersection = _bbox_intersection_area(pet_box, people_box) - if intersection <= 0: - continue - pet_to_face_ratio = pet_area / float(people_area) - preserve_larger_pet = ( - image_area > 0 - and pet_to_face_ratio > larger_pet_ratio - and pet_image_coverage < mural_image_coverage_threshold - ) - if preserve_larger_pet: - continue - if _bbox_iou(pet_box, people_box) >= iou_threshold: - return _PetPeopleOverlapDecision( - True, - "iou", - pet_to_face_ratio, - pet_image_coverage, - ) - smaller_area = min(pet_area, people_area) - if intersection / float(smaller_area) >= smaller_box_coverage_threshold: - return _PetPeopleOverlapDecision( - True, - "smaller_box_coverage", - pet_to_face_ratio, - pet_image_coverage, + def _ensure_embedder(self): + if self._embedder is None: + model_dir = self._resolve_model_path( + Path("embedding") / self._embedding_model_name, + directory=True, ) - return _PetPeopleOverlapDecision(False, pet_image_coverage=pet_image_coverage) - - -def pet_model_auto_download_enabled() -> bool: - raw = str(os.environ.get(PET_MODEL_AUTO_DOWNLOAD_ENV, "")).strip().lower() - return raw not in {"0", "false", "no", "off"} - -def ensure_pet_detector_model( - model_path: Path, - *, - allow_model_download: bool = True, - model_url: str | None = None, -) -> Path: - target = Path(model_path) - custom_url = str(model_url or os.environ.get(PET_DETECTOR_MODEL_URL_ENV) or "").strip() - expected_sha256 = ( - str(os.environ.get(PET_DETECTOR_MODEL_SHA256_ENV) or "").strip().lower() - if custom_url - else DEFAULT_PET_DETECTOR_MODEL_SHA256 - ) - if custom_url and not expected_sha256: - raise RuntimeError( - "Pet scanning unavailable: a custom detector URL requires " - f"{PET_DETECTOR_MODEL_SHA256_ENV}." - ) - if target.is_file(): - _validate_downloaded_file( - target, - label="YOLOX pet detector model", - expected_sha256=expected_sha256, - max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, - ) - return target - if not allow_model_download: - raise RuntimeError( - "Pet scanning unavailable: missing YOLOX model at " - f"{target}. Set IPHOTO_PET_MODEL_DIR or enable pet model downloads." - ) - - url = str(custom_url or DEFAULT_PET_DETECTOR_MODEL_URL).strip() - if not url: - raise RuntimeError( - "Pet scanning unavailable: missing YOLOX model at " - f"{target} and no pet detector download URL is configured." - ) - _download_file( - url, - target, - label="YOLOX pet detector model", - expected_sha256=expected_sha256, - max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, - ) - return target - - -def default_pet_model_dir() -> Path: - return user_pet_model_cache_dir() - - -def bundled_pet_model_dir() -> Path: - package_root = Path(__file__).resolve().parents[2] - return package_root / "extension" / "models" / "pets" - - -def user_pet_model_cache_dir() -> Path: - if sys.platform == "darwin": - base = Path.home() / "Library" / "Caches" - elif os.name == "nt": - base = Path(os.environ.get("LOCALAPPDATA") or Path.home() / "AppData" / "Local") - else: - base = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") - return base / "iPhoto" / "models" / "pets" - - -def pet_model_search_roots() -> tuple[Path, ...]: - roots: list[Path] = [] - override = str(os.environ.get("IPHOTO_PET_MODEL_DIR") or "").strip() - if override: - roots.append(Path(override).expanduser()) - roots.extend((user_pet_model_cache_dir(), bundled_pet_model_dir())) - return tuple(dict.fromkeys(roots)) - - -def resolve_pet_model_path(relative_path: Path, *, directory: bool = False) -> Path: - relative = Path(relative_path) - if relative.is_absolute() or ".." in relative.parts: - raise ValueError("Pet model path must be relative to a configured model root.") - override = str(os.environ.get("IPHOTO_PET_MODEL_DIR") or "").strip() - user_cache = user_pet_model_cache_dir() - for root in pet_model_search_roots(): - candidate = root / relative - exists = candidate.is_dir() if directory else candidate.is_file() - if not exists: - continue - try: - if directory: - model_name = relative.name - model_path = candidate / f"{model_name}.pt" - if not model_path.is_file(): - raise RuntimeError("DINOv2 model file is missing") - _validate_dinov2_cache_metadata(model_path, model_name=model_name) - else: - _validate_downloaded_file( - candidate, - label="YOLOX pet detector model", - expected_sha256=( - str( - os.environ.get(PET_DETECTOR_MODEL_SHA256_ENV) - or DEFAULT_PET_DETECTOR_MODEL_SHA256 - ).lower() - ), - max_bytes=DEFAULT_PET_DETECTOR_MODEL_MAX_BYTES, - ) - return candidate - except (OSError, RuntimeError) as exc: - if override and root == Path(override).expanduser(): - raise RuntimeError( - f"Pet scanning unavailable: invalid model override artifact at {candidate}." - ) from exc - if root == user_cache: - model_path = candidate / f"{relative.name}.pt" if directory else candidate + def create_embedder(): try: - model_path.unlink(missing_ok=True) - if directory: - _dinov2_metadata_path(model_path).unlink(missing_ok=True) - except OSError: - _LOGGER.warning( - "Failed to quarantine invalid Pets model cache %s", - candidate, - exc_info=True, + return _DinoV2Embedder( + model_dir, + model_name=self._embedding_model_name, + allow_model_download=self._allow_model_download, ) - # Bundled artifacts are read-only. An invalid one must never become - # a download target and must not shadow a later valid artifact. - continue - return user_pet_model_cache_dir() / relative - - -def _download_file( - url: str, - destination: Path, - *, - label: str, - expected_sha256: str, - max_bytes: int, -) -> None: - destination = Path(destination) - if urlparse(url).scheme.lower() != "https": - raise RuntimeError(f"Pet scanning unavailable: {label} URL must use HTTPS.") - destination.parent.mkdir(parents=True, exist_ok=True) - try: - with tempfile.TemporaryDirectory( - prefix="iphoto-pet-model-", - dir=destination.parent, - ) as tmp_dir: - tmp_path = Path(tmp_dir) / destination.name - with ( - request.urlopen( # noqa: S310 - url, - timeout=_DOWNLOAD_TIMEOUT_SECONDS, - context=_download_ssl_context(url), - ) as response, - tmp_path.open("wb") as handle, - ): - total = 0 - while True: - chunk = response.read(_DOWNLOAD_CHUNK_SIZE) - if not chunk: - break - total += len(chunk) - if total > int(max_bytes): - raise RuntimeError(f"Downloaded {label} exceeds its size limit.") - handle.write(chunk) - _validate_downloaded_file( - tmp_path, - label=label, - expected_sha256=expected_sha256, - max_bytes=max_bytes, - ) - tmp_path.replace(destination) - except TimeoutError as exc: - raise RuntimeError( - f"Pet scanning unavailable: downloading {label} timed out. " - "Check your network connection or install the model manually." - ) from exc - except Exception as exc: - if isinstance(exc, RuntimeError) and str(exc).startswith("Pet scanning unavailable:"): - raise - raise RuntimeError( - f"Pet scanning unavailable: failed to download {label} from {url} " - f"({_error_reason(exc)}). Check your network connection, set " - f"{PET_DETECTOR_MODEL_URL_ENV}, or install the model manually." - ) from exc - - -def _validate_downloaded_file( - path: Path, - *, - label: str, - expected_sha256: str, - max_bytes: int, -) -> None: - size = Path(path).stat().st_size - if size <= 0: - raise RuntimeError(f"Downloaded {label} is empty.") - if size > int(max_bytes): - raise RuntimeError(f"Downloaded {label} exceeds its size limit.") - digest = _file_sha256(path) - if not expected_sha256 or digest != expected_sha256.lower(): - raise RuntimeError(f"Downloaded {label} failed SHA-256 verification.") - - -def _file_sha256(path: Path) -> str: - digest = hashlib.sha256() - with Path(path).open("rb") as handle: - for chunk in iter(lambda: handle.read(_DOWNLOAD_CHUNK_SIZE), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _dinov2_metadata_path(model_path: Path) -> Path: - return Path(model_path).with_suffix(f"{Path(model_path).suffix}.metadata.json") - - -def _validate_dinov2_cache_metadata(model_path: Path, *, model_name: str) -> None: - metadata_path = _dinov2_metadata_path(model_path) - if not metadata_path.is_file(): - raise RuntimeError( - "Pet scanning unavailable: DINOv2 TorchScript metadata is missing for " - f"{model_path}. Remove the incomplete cache so it can be rebuilt." - ) - try: - metadata = json.loads(metadata_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: - raise RuntimeError( - f"Pet scanning unavailable: invalid DINOv2 metadata at {metadata_path}." - ) from exc - expected = { - "model_name": model_name, - "source_repository": _EMBEDDER_MANIFEST["source_repository"], - "source_revision": _DINO_SOURCE_REVISION, - "torchscript_sha256": _EMBEDDER_MANIFEST["torchscript_sha256"], - "torchscript_size": _EMBEDDER_MANIFEST["torchscript_size"], - "input_shape": _EMBEDDER_MANIFEST["input_shape"], - "output_shape": _EMBEDDER_MANIFEST["output_shape"], - } - if any(metadata.get(key) != value for key, value in expected.items()): - raise RuntimeError( - f"Pet scanning unavailable: DINOv2 metadata contract mismatch at {metadata_path}." - ) - size_matches = int(metadata.get("torchscript_size") or -1) == model_path.stat().st_size - hash_matches = str(metadata.get("torchscript_sha256") or "").lower() == _file_sha256(model_path) - if not size_matches or not hash_matches: - raise RuntimeError( - f"Pet scanning unavailable: DINOv2 cache integrity check failed for {model_path}." - ) + except ( + _impl.PetRuntimeUnavailableError, + _impl.PetModelUnavailableError, + _impl.PetPipelineInvariantError, + ): + raise + except RuntimeError as exc: + raise _impl.PetModelUnavailableError(str(exc)) from exc + self._embedder = _LazyDinoV2Embedder(create_embedder) + return self._embedder -def _error_reason(exc: Exception) -> str: - return str(exc).strip() or exc.__class__.__name__ +def __getattr__(name: str): + """Expose legacy private helpers without replacing this module object.""" -def _download_ssl_context(url: str) -> ssl.SSLContext | None: - if not url.lower().startswith("https://"): - return None try: - import certifi - except ImportError: - return ssl.create_default_context() - _install_certifi_environment() - return ssl.create_default_context(cafile=certifi.where()) + return getattr(_impl, name) + except AttributeError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc -def _install_certifi_environment() -> None: - try: - import certifi - except ImportError: - return - cafile = certifi.where() - os.environ.setdefault("SSL_CERT_FILE", cafile) - os.environ.setdefault("REQUESTS_CA_BUNDLE", cafile) +def __dir__() -> list[str]: + return sorted(set(globals()) | set(dir(_impl))) diff --git a/src/iPhoto/pets/service.py b/src/iPhoto/pets/service.py index a10ece284..e527aba97 100644 --- a/src/iPhoto/pets/service.py +++ b/src/iPhoto/pets/service.py @@ -53,9 +53,9 @@ class _PetReadContext: def shared_pet_model_dir() -> Path: - from .pipeline import default_pet_model_dir + from .pipeline import default_pet_model_dir, pet_model_override_dir - return default_pet_model_dir() + return pet_model_override_dir() or default_pet_model_dir() def pet_library_paths(library_root: Path) -> PetLibraryPaths: diff --git a/tests/application/test_runtime_context.py b/tests/application/test_runtime_context.py index 3a4426c96..0d019e320 100644 --- a/tests/application/test_runtime_context.py +++ b/tests/application/test_runtime_context.py @@ -70,6 +70,7 @@ def __init__(self) -> None: self.bound_location_services: list[object | None] = [] self.asset_query_service_during_bind: object | None = None self.state_repository_during_bind: object | None = None + self.startup_recognition_requests = 0 def bind_path(self, root: Path) -> None: self.asset_query_service_during_bind = ( @@ -130,6 +131,9 @@ def start_scanning( ) -> None: self.scan_requests.append((root, list(include), list(exclude))) + def request_startup_recognition_after_idle(self) -> None: + self.startup_recognition_requests += 1 + def bind_scan_service(self, scan_service: object | None) -> None: self.bound_scan_services.append(scan_service) @@ -317,6 +321,21 @@ def test_resume_startup_tasks_can_defer_scan_until_gallery_opens( ] +def test_idle_startup_recognition_policy_can_be_disabled_for_controlled_ab( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + library_root = tmp_path / "library" + library_root.mkdir(parents=True) + context, library, _asset_runtime = _runtime_context(library_root) + context.resume_startup_tasks(defer_scan=True) + monkeypatch.setenv("IPHOTO_STARTUP_RECOGNITION_AUTO_START", "0") + + context.schedule_idle_startup_jobs() + + assert library.startup_recognition_requests == 0 + + def test_resume_startup_tasks_skips_scan_when_scope_complete(tmp_path: Path) -> None: library_root = tmp_path / "library" library_root.mkdir(parents=True) diff --git a/tests/contracts/test_pets_dino_source_contract.py b/tests/contracts/test_pets_dino_source_contract.py new file mode 100644 index 000000000..66049ea42 --- /dev/null +++ b/tests/contracts/test_pets_dino_source_contract.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + + +pytestmark = pytest.mark.skipif( + os.environ.get("IPHOTO_RUN_PETS_DINO_RELEASE_CONTRACT") != "1", + reason="real DINOv2 Release contract is opt-in", +) + + +def test_real_pinned_dinov2_release_downloads_and_loads(tmp_path: Path) -> None: + torch = pytest.importorskip("torch") + manifest = pet_pipeline._EMBEDDER_MANIFEST + target = tmp_path / "dinov2_vits14.pt" + pet_pipeline._download_file( + str(manifest["torchscript_url"]), + target, + label="DINOv2 TorchScript Release artifact", + expected_sha256=str(manifest["torchscript_sha256"]), + max_bytes=int(manifest["torchscript_size"]), + exact_size=int(manifest["torchscript_size"]), + ) + + model = torch.jit.load(str(target), map_location="cpu").eval() + example = torch.zeros(tuple(manifest["input_shape"]), dtype=torch.float32) + with torch.no_grad(): + output = model(example) + assert tuple(output.shape) == tuple(manifest["output_shape"]) diff --git a/tests/gui/test_main.py b/tests/gui/test_main.py index d7f5f56ff..ce0f80b3b 100644 --- a/tests/gui/test_main.py +++ b/tests/gui/test_main.py @@ -6,7 +6,8 @@ import pytest from PySide6.QtCore import QCoreApplication, QEvent, Qt -from PySide6.QtGui import QCloseEvent, QSurface +from PySide6.QtGui import QCloseEvent, QSurface, QWindow +from PySide6.QtWidgets import QDialog, QMainWindow, QMenu, QWidget from iPhoto.gui.main import ( _bootstrap_macos_external_tool_path, @@ -15,6 +16,7 @@ _prepare_top_level_rhi_surface, _startup_feature_plan, _startup_timing_plan, + _RecognitionIdleActivityFilter, _StartupInputGuard, ) from iPhoto.gui.ui import main_window as main_window_module @@ -399,6 +401,95 @@ def type(self): assert guard.eventFilter(child, _FakeEvent(QEvent.Type.MouseButtonPress)) is False +def test_recognition_idle_filter_resets_only_for_active_window_input() -> None: + installed_filters: list[object] = [] + removed_filters: list[object] = [] + activity: list[str] = [] + + class _FakeApp: + def installEventFilter(self, event_filter) -> None: # noqa: N802 + installed_filters.append(event_filter) + + def removeEventFilter(self, event_filter) -> None: # noqa: N802 + removed_filters.append(event_filter) + + class _FakeWindow: + def isAncestorOf(self, watched) -> bool: # noqa: N802 + return watched == "child" + + class _FakeEvent: + def __init__(self, event_type, buttons=Qt.MouseButton.NoButton) -> None: + self._event_type = event_type + self._buttons = buttons + + def type(self): + return self._event_type + + def buttons(self): + return self._buttons + + window = _FakeWindow() + activity_filter = _RecognitionIdleActivityFilter( + window, + _FakeApp(), + lambda: activity.append("active"), + ) + activity_filter.install() + + assert activity_filter.eventFilter("child", _FakeEvent(QEvent.Type.Wheel)) is False + assert activity == ["active"] + assert ( + activity_filter.eventFilter("child", _FakeEvent(QEvent.Type.MouseMove)) + is False + ) + assert activity == ["active"] + activity_filter.eventFilter( + "child", + _FakeEvent(QEvent.Type.MouseMove, Qt.MouseButton.LeftButton), + ) + assert activity == ["active", "active"] + activity_filter.eventFilter("external", _FakeEvent(QEvent.Type.KeyPress)) + assert activity == ["active", "active"] + + activity_filter.release() + assert installed_filters == [activity_filter] + assert removed_filters == [activity_filter] + + +def test_recognition_idle_filter_includes_owned_popup_modal_and_transient_windows(qapp) -> None: + activity: list[str] = [] + window = QMainWindow() + child = QWidget(window) + context_menu = QMenu(child) + modal = QDialog(window) + external = QDialog() + window.show() + qapp.processEvents() + transient = QWindow() + transient.setTransientParent(window.windowHandle()) + activity_filter = _RecognitionIdleActivityFilter( + window, + qapp, + lambda: activity.append("active"), + ) + activity_filter.install() + event = QEvent(QEvent.Type.KeyPress) + try: + assert activity_filter.eventFilter(context_menu, event) is False + assert activity_filter.eventFilter(modal, event) is False + assert activity_filter.eventFilter(transient, event) is False + assert activity == ["active", "active", "active"] + assert activity_filter.eventFilter(external, event) is False + assert activity == ["active", "active", "active"] + finally: + activity_filter.release() + transient.close() + external.close() + modal.close() + context_menu.close() + window.close() + + def test_settings_initialization_failure_emits_one_failed_terminal( monkeypatch, qapp, diff --git a/tests/test_library_manager.py b/tests/test_library_manager.py index 4b1cb78b8..e4ad2c5b1 100644 --- a/tests/test_library_manager.py +++ b/tests/test_library_manager.py @@ -5,7 +5,7 @@ import sys from pathlib import Path from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import Mock, call, patch import pytest @@ -606,6 +606,32 @@ def test_recognition_binding_does_not_start_ai_before_viewport_ready( start_ai.assert_called_once_with(root, startup=True) +def test_recognition_activation_waits_for_startup_metadata_scan( + tmp_path: Path, + qapp: QApplication, +) -> None: + root = tmp_path / "Library" + root.mkdir() + manager = LibraryRuntimeController() + manager._root = root + manager._recognition_services_root = root + manager._people_service = Mock() + manager._pet_service = Mock() + startup_worker = Mock() + startup_worker._defer_ai_workers_until_scan_finished = True + manager._current_scanner_worker = startup_worker + + with patch.object(manager, "_start_ai_scan_workers") as start_ai: + manager.activate_recognition_scans() + + start_ai.assert_not_called() + assert manager._startup_recognition_request == ( + root, + manager._recognition_generation, + ) + assert not manager._startup_recognition_timer.isActive() + + def test_upgraded_library_schedules_closed_input_pet_backfill_after_bind( tmp_path: Path, qapp: QApplication, @@ -755,7 +781,7 @@ def setObjectName(self, _name: str) -> None: def finish_input(self) -> None: self.input_closed = True - def start(self) -> None: + def start(self, _priority=None) -> None: self.started = True def cancel(self) -> None: @@ -829,7 +855,7 @@ def isRunning(self) -> bool: return False class _FaceWorker(_WorkerBase): - def start(self) -> None: + def start(self, _priority=None) -> None: face_starts.append(1) class _PetWorker(_WorkerBase): @@ -837,7 +863,7 @@ def __init__(self, *_args, **_kwargs) -> None: super().__init__(*_args, **_kwargs) pet_instances.append(self) - def start(self) -> None: + def start(self, _priority=None) -> None: if len(pet_instances) == 1: raise RuntimeError("injected Pet start failure") @@ -931,6 +957,7 @@ def test_startup_ai_workers_close_input_after_metadata_scan( manager = LibraryRuntimeController() manager.bind_path(root) created: list[object] = [] + runtime_prepared: list[str] = [] class _FakeSignal: def connect(self, _callback) -> None: @@ -938,6 +965,8 @@ def connect(self, _callback) -> None: class _FakeAiWorker: def __init__(self, *_args, **_kwargs) -> None: + if not created: + assert runtime_prepared == ["cv2"] self.statusChanged = _FakeSignal() self.finished = _FakeSignal() self.input_closed = False @@ -956,7 +985,7 @@ def wait(self, _timeout_ms: int) -> bool: def isRunning(self) -> bool: return False - def start(self) -> None: + def start(self, _priority=None) -> None: self.started = True monkeypatch.setitem( @@ -970,10 +999,33 @@ def start(self) -> None: SimpleNamespace(PetScanWorker=_FakeAiWorker), ) - with patch("iPhoto.library.scan_coordinator.mark") as profile_mark: + with ( + patch("iPhoto.library.scan_coordinator.mark") as profile_mark, + patch( + "iPhoto.library.scan_coordinator._prepare_face_runtime_imports", + side_effect=lambda: runtime_prepared.append("cv2"), + ) as prepare_runtime, + ): manager._start_ai_scan_workers(root, startup=True) - profile_mark.assert_called_once_with("startup_ai_scan.started", root=root) + profile_mark.assert_has_calls( + [ + call("startup_ai_scan.started", root=root), + call( + "recognition.worker.started", + generation=manager._recognition_generation, + worker="face", + startup=True, + ), + call( + "recognition.worker.started", + generation=manager._recognition_generation, + worker="pet", + startup=True, + ), + ] + ) + prepare_runtime.assert_called_once_with() assert len(created) == 2 assert all(worker.input_closed for worker in created) assert all(worker.started for worker in created) diff --git a/tests/test_model_release_provenance.py b/tests/test_model_release_provenance.py new file mode 100644 index 000000000..78fd79376 --- /dev/null +++ b/tests/test_model_release_provenance.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import pytest + +from tools.model_release_provenance import ProvenanceError, validate_release_provenance + + +HEAD = "a" * 40 +WORKFLOW_ID = 340207631 +WORKFLOW_PATH = ".github/workflows/pets-dino-source-contract.yml" +REPOSITORY = "OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager" + + +def _run(**overrides) -> dict: + value = { + "id": 123, + "workflow_id": WORKFLOW_ID, + "conclusion": "success", + "head_sha": HEAD, + "path": WORKFLOW_PATH, + "repository": {"full_name": REPOSITORY}, + } + value.update(overrides) + return value + + +def _build(**overrides) -> dict: + value = {"repository_commit": HEAD, "workflow_run_id": "123"} + value.update(overrides) + return value + + +def _validate(*, run=None, build=None, builder_commit=HEAD) -> str: + return validate_release_provenance( + run=run or _run(), + build=build or _build(), + builder_commit=builder_commit, + expected_repository=REPOSITORY, + expected_workflow_id=WORKFLOW_ID, + expected_workflow_path=WORKFLOW_PATH, + ) + + +def test_valid_release_provenance_returns_artifact_head() -> None: + assert _validate() == HEAD + + +@pytest.mark.parametrize( + ("run", "message"), + [ + (_run(head_sha=""), "head_sha"), + (_run(conclusion="failure"), "successfully"), + (_run(workflow_id=1), "different workflow"), + (_run(path=".github/workflows/other.yml"), "workflow path"), + (_run(repository={"full_name": "other/repo"}), "different repository"), + ], +) +def test_invalid_artifact_run_is_rejected(run: dict, message: str) -> None: + with pytest.raises(ProvenanceError, match=message): + _validate(run=run) + + +def test_builder_checkout_must_match_run_head() -> None: + with pytest.raises(ProvenanceError, match="builder checkout"): + _validate(builder_commit="b" * 40) + + +def test_build_manifest_commit_must_match_run_head() -> None: + with pytest.raises(ProvenanceError, match="repository_commit"): + _validate(build=_build(repository_commit="b" * 40)) + + +def test_build_manifest_run_id_must_match_artifact_run() -> None: + with pytest.raises(ProvenanceError, match="workflow_run_id"): + _validate(build=_build(workflow_run_id="999")) diff --git a/tests/test_pet_model_acquisition.py b/tests/test_pet_model_acquisition.py new file mode 100644 index 000000000..e04255488 --- /dev/null +++ b/tests/test_pet_model_acquisition.py @@ -0,0 +1,623 @@ +from __future__ import annotations + +import errno +import hashlib +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + +pet_impl = pet_pipeline._impl + +DETECTOR_RELATIVE = Path("detector") / "yolox_nano_coco.onnx" +EMBEDDER_RELATIVE = Path("embedding") / "dinov2_vits14" +_FAKE_DETECTOR_BYTES = b"detector" +_FAKE_DETECTOR_SHA256 = hashlib.sha256(_FAKE_DETECTOR_BYTES).hexdigest() + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(4096), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _valid_detector(path: Path) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(_FAKE_DETECTOR_BYTES) + return path + + +def _invalid_detector(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"invalid") + + +@pytest.fixture(autouse=True) +def _clear_override(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("IPHOTO_PET_MODEL_DIR", raising=False) + + +@pytest.fixture +def _fake_detector_manifest(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + pet_impl, + "DEFAULT_PET_DETECTOR_MODEL_SHA256", + _FAKE_DETECTOR_SHA256, + ) + + +class TestStoragePolicy: + def test_extension_first_lookup_without_writable_probe( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + _fake_detector_manifest: None, + ) -> None: + bundled = tmp_path / "extension" + cache = tmp_path / "cache" + bundled_model = _valid_detector(bundled / DETECTOR_RELATIVE) + _valid_detector(cache / DETECTOR_RELATIVE) + + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) + probed: list[Path] = [] + + def record_probe(path: Path) -> bool: + probed.append(path) + return True + + monkeypatch.setattr(pet_impl, "_directory_is_writable", record_probe) + resolved = pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) + assert resolved == bundled_model + assert not probed + + def test_cache_lookup_falls_back_after_invalid_bundled( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + _fake_detector_manifest: None, + ) -> None: + bundled = tmp_path / "extension" + cache = tmp_path / "cache" + _invalid_detector(bundled / DETECTOR_RELATIVE) + model = _valid_detector(cache / DETECTOR_RELATIVE) + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) + assert pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) == model + assert (bundled / DETECTOR_RELATIVE).exists() + + def test_missing_prefers_writable_extension( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + bundled = tmp_path / "extension" + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr( + pet_impl, + "_directory_is_writable", + lambda path: path == bundled, + ) + assert pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) == bundled / DETECTOR_RELATIVE + + def test_macos_app_installs_to_cache_without_probe( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + bundled = tmp_path / "iPhoto.app" / "Contents" / "Resources" / "models" + cache = tmp_path / "cache" + monkeypatch.setattr(pet_pipeline.sys, "platform", "darwin") + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) + + def fail_probe(_path: Path) -> bool: + raise AssertionError("packaged app bundles must not be probed") + + monkeypatch.setattr(pet_impl, "_directory_is_writable", fail_probe) + assert pet_pipeline.pet_model_install_root() == cache + assert pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) == cache / DETECTOR_RELATIVE + + @pytest.mark.parametrize("errno_value", [errno.EACCES, errno.EPERM, errno.EROFS]) + def test_permission_errors_are_typed_for_storage_fallback( + self, + errno_value: int, + tmp_path: Path, + ) -> None: + with pytest.raises(pet_pipeline._ModelStoragePermissionError): + pet_pipeline._raise_if_model_storage_error(OSError(errno_value, "denied"), tmp_path) + + @pytest.mark.parametrize("errno_value", [errno.ENOSPC, errno.EIO]) + def test_non_permission_errors_fail_without_retry( + self, + errno_value: int, + tmp_path: Path, + ) -> None: + with pytest.raises(OSError) as raised: + pet_pipeline._raise_if_model_storage_error(OSError(errno_value, "failed"), tmp_path) + assert raised.value.errno == errno_value + assert type(raised.value) is not pet_pipeline._ModelStoragePermissionError + + @pytest.mark.parametrize("errno_value", [errno.EACCES, errno.EPERM, errno.EROFS]) + def test_writability_probe_returns_false_only_for_permission_errors( + self, + errno_value: int, + ) -> None: + class FailingPath: + def mkdir(self, **_kwargs) -> None: + raise OSError(errno_value, "denied") + + assert pet_pipeline._directory_is_writable(FailingPath()) is False + + @pytest.mark.parametrize("errno_value", [errno.ENOSPC, errno.EIO]) + def test_writability_probe_surfaces_non_permission_errors( + self, + errno_value: int, + ) -> None: + class FailingPath: + def mkdir(self, **_kwargs) -> None: + raise OSError(errno_value, "failed") + + with pytest.raises(OSError) as raised: + pet_pipeline._directory_is_writable(FailingPath()) + assert raised.value.errno == errno_value + + def test_only_bundled_targets_have_user_cache_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + bundled = tmp_path / "extension" + cache = tmp_path / "cache" + target = bundled / EMBEDDER_RELATIVE / "dinov2_vits14.pt" + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) + assert pet_pipeline._model_storage_fallback_path(target) == ( + cache / EMBEDDER_RELATIVE / "dinov2_vits14.pt" + ) + assert pet_pipeline._model_storage_fallback_path(cache / "other.pt") is None + + +class TestOverrideAuthority: + def test_search_and_install_are_override_only(self, tmp_path, monkeypatch) -> None: + override = tmp_path / "override" + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(override)) + assert pet_pipeline.pet_model_search_roots() == (override,) + assert pet_pipeline.pet_model_install_root() == override + assert ( + pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) + == override / DETECTOR_RELATIVE + ) + assert pet_pipeline._model_storage_fallback_path(override / DETECTOR_RELATIVE) is None + + def test_pipeline_rejects_mismatched_model_root( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(tmp_path / "override")) + pipeline_instance = pet_pipeline.PetClusterPipeline(model_root=tmp_path / "models") + with pytest.raises(pet_pipeline.PetModelUnavailableError): + pipeline_instance._resolve_model_path(DETECTOR_RELATIVE) + + def test_invalid_dino_override_is_repaired_in_place_without_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + override = tmp_path / "override" + model_dir = override / EMBEDDER_RELATIVE + model_dir.mkdir(parents=True) + (model_dir / "dinov2_vits14.pt").write_bytes(b"legacy") + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(override)) + monkeypatch.setattr( + pet_impl, + "user_pet_model_cache_dir", + lambda: (_ for _ in ()).throw(AssertionError("override must not fall back")), + ) + + assert pet_pipeline.resolve_pet_model_path( + EMBEDDER_RELATIVE, + directory=True, + ) == model_dir + + def test_detector_override_permission_error_does_not_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + override = tmp_path / "override" + target = override / DETECTOR_RELATIVE + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(override)) + calls: list[Path] = [] + + def fail_download(_url: str, destination: Path, **_kwargs): + calls.append(Path(destination)) + raise pet_pipeline._ModelStoragePermissionError(errno.EACCES, "denied") + + monkeypatch.setattr(pet_impl, "_download_file", fail_download) + with pytest.raises(RuntimeError, match="not writable"): + pet_pipeline.ensure_pet_detector_model(target) + assert calls == [target] + + def test_invalid_detector_override_is_repaired_in_place_without_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + _fake_detector_manifest: None, + ) -> None: + override = tmp_path / "override" + target = override / DETECTOR_RELATIVE + _invalid_detector(target) + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(override)) + monkeypatch.setattr( + pet_impl, + "user_pet_model_cache_dir", + lambda: (_ for _ in ()).throw(AssertionError("override must not fall back")), + ) + downloads: list[Path] = [] + + def repair(_url: str, destination: Path, **_kwargs) -> Path: + downloads.append(Path(destination)) + Path(destination).write_bytes(_FAKE_DETECTOR_BYTES) + return Path(destination) + + monkeypatch.setattr(pet_impl, "_download_file", repair) + + resolved = pet_pipeline.resolve_pet_model_path(DETECTOR_RELATIVE) + assert resolved == target + assert pet_pipeline.ensure_pet_detector_model(resolved) == target + assert downloads == [target] + assert target.read_bytes() == _FAKE_DETECTOR_BYTES + + def test_invalid_detector_is_preserved_when_download_is_disabled( + self, + tmp_path: Path, + _fake_detector_manifest: None, + ) -> None: + target = tmp_path / "override" / DETECTOR_RELATIVE + _invalid_detector(target) + + with pytest.raises(RuntimeError, match="SHA-256"): + pet_pipeline.ensure_pet_detector_model( + target, + allow_model_download=False, + ) + + assert target.read_bytes() == b"invalid" + + def test_detector_repair_non_permission_unlink_error_does_not_retry( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + _fake_detector_manifest: None, + ) -> None: + target = tmp_path / "override" / DETECTOR_RELATIVE + _invalid_detector(target) + real_unlink = Path.unlink + + def fail_unlink(self: Path, *args, **kwargs): + if self == target: + raise OSError(errno.ENOSPC, "full") + return real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", fail_unlink) + monkeypatch.setattr( + pet_impl, + "_download_file", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("non-permission repair errors must not download") + ), + ) + + with pytest.raises(OSError) as raised: + pet_pipeline.ensure_pet_detector_model(target) + assert raised.value.errno == errno.ENOSPC + + def test_dino_override_permission_error_does_not_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + override = tmp_path / "override" + target = override / EMBEDDER_RELATIVE / "dinov2_vits14.pt" + monkeypatch.setenv("IPHOTO_PET_MODEL_DIR", str(override)) + embedder = pet_pipeline._DinoV2Embedder.__new__(pet_pipeline._DinoV2Embedder) + calls: list[Path] = [] + + def fail_build(path: Path): + calls.append(Path(path)) + raise pet_pipeline._ModelStoragePermissionError(errno.EACCES, "denied") + + embedder._build_dinov2_cache = fail_build + with pytest.raises(pet_pipeline.PetModelUnavailableError, match="not writable"): + embedder._download_dinov2_model(target) + assert calls == [target] + + +class TestDownloadsAndMetadata: + def test_embedder_rejects_an_unpinned_torch_runtime( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(__version__="2.11.0")) + with pytest.raises( + pet_pipeline.PetRuntimeUnavailableError, + match="requires torch==2.12.1", + ): + pet_pipeline._DinoV2Embedder( + tmp_path, + model_name="dinov2_vits14", + ) + + def test_exact_size_is_enforced(self, tmp_path: Path, monkeypatch) -> None: + target = tmp_path / "model.bin" + payload = b"exact-size" + + class Response: + def __init__(self, payload: bytes) -> None: + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def read(self, _size: int): + value, self.payload = self.payload, b"" + return value + + monkeypatch.setattr( + pet_pipeline.request, + "urlopen", + lambda *_args, **_kwargs: Response(payload), + ) + pet_pipeline._download_file( + "https://example.test/model.bin", + target, + label="test model", + expected_sha256=_sha256_target(payload), + max_bytes=100, + exact_size=len(payload), + ) + assert target.read_bytes() == payload + + short = payload[:-1] + monkeypatch.setattr( + pet_pipeline.request, + "urlopen", + lambda *_args, **_kwargs: ShortResponse(short), + ) + with pytest.raises(RuntimeError, match="wrong file size"): + pet_pipeline._download_file( + "https://example.test/model.bin", + target.with_suffix(".short"), + label="test model", + expected_sha256=_sha256_target(short), + max_bytes=100, + exact_size=len(payload), + ) + + def test_download_preserves_permission_error_for_storage_fallback( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + def fail_temp_dir(*_args, **_kwargs): + raise OSError(errno.EACCES, "denied") + + monkeypatch.setattr(pet_pipeline.tempfile, "TemporaryDirectory", fail_temp_dir) + with pytest.raises(pet_pipeline._ModelStoragePermissionError): + pet_pipeline._download_file( + "https://example.test/model.bin", + tmp_path / "model.bin", + label="test model", + expected_sha256="0" * 64, + max_bytes=100, + ) + + def test_manifest_dino_contract_uses_official_checkpoint(self) -> None: + manifest = pet_pipeline._EMBEDDER_MANIFEST + assert manifest["weights_url"] == ( + "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/" + "dinov2_vits14_pretrain.pth" + ) + assert manifest["weights_sha256"] == ( + "b938bf1bc15cd2ec0feacfe3a1bb553fe8ea9ca46a7e1d8d00217f29aef60cd9" + ) + assert manifest["weights_size"] == 88283115 + assert manifest["artifact_kind"] == "release_torchscript" + assert manifest["release_tag"] == "pet-models-v1" + assert manifest["cache_schema_version"] == 2 + assert manifest["producer_torch_version"] == "2.12.1" + assert manifest["producer_torchvision_version"] == "0.27.1" + assert manifest["torchscript_url"].endswith( + "/releases/download/pet-models-v1/dinov2_vits14.pt" + ) + + def test_legacy_derived_metadata_is_rejected(self, tmp_path: Path) -> None: + content = b"derived-torchscript" + metadata = { + "artifact_kind": "derived_checkpoint_cache", + "model_name": "dinov2_vits14", + "source_repository": pet_pipeline._EMBEDDER_MANIFEST["source_repository"], + "source_revision": pet_pipeline._DINO_SOURCE_REVISION, + "weights_sha256": pet_pipeline._DINO_WEIGHTS_SHA256, + "weights_size": pet_pipeline._DINO_WEIGHTS_SIZE, + "derived_torchscript_sha256": hashlib.sha256(content).hexdigest(), + "derived_torchscript_size": len(content), + "input_shape": [1, 3, 224, 224], + "output_shape": [1, 384], + } + model_path = tmp_path / "dinov2_vits14.pt" + model_path.write_bytes(content) + pet_pipeline._dinov2_metadata_path(model_path).write_text( + hashlib_json(metadata), encoding="utf-8" + ) + with pytest.raises(RuntimeError, match="metadata contract mismatch"): + pet_pipeline._validate_dinov2_cache_metadata( + model_path, + model_name="dinov2_vits14", + ) + + def test_dino_release_download_publishes_model_and_metadata_together( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + model_bytes = b"release-torchscript" + model_name = "dinov2_vits14" + model_path = tmp_path / EMBEDDER_RELATIVE / f"{model_name}.pt" + + class FakeTensor: + shape = (1, 384) + + class FakeModel: + def eval(self): + return self + + def to(self, _device): + return self + + def __call__(self, _example): + return FakeTensor() + + class FakeJit: + @staticmethod + def load(path: str, *, map_location): + assert Path(path).read_bytes() == model_bytes + assert map_location in {"cpu", "test-device"} + return FakeModel() + + class FakeNoGrad: + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + fake_torch = SimpleNamespace( + jit=FakeJit(), + float32=object(), + randn=lambda *_args, **_kwargs: object(), + no_grad=lambda: FakeNoGrad(), + ) + embedder = pet_pipeline._DinoV2Embedder.__new__(pet_pipeline._DinoV2Embedder) + embedder._torch = fake_torch + embedder._device = "test-device" + embedder._model_name = model_name + + digest = hashlib.sha256(model_bytes).hexdigest() + monkeypatch.setattr(pet_impl, "_DINO_TORCHSCRIPT_SHA256", digest) + monkeypatch.setattr(pet_impl, "_DINO_TORCHSCRIPT_SIZE", len(model_bytes)) + monkeypatch.setitem(pet_pipeline._EMBEDDER_MANIFEST, "torchscript_sha256", digest) + monkeypatch.setitem( + pet_pipeline._EMBEDDER_MANIFEST, + "torchscript_size", + len(model_bytes), + ) + + def fake_download(url: str, destination: Path, **kwargs) -> Path: + assert url == pet_pipeline._DINO_TORCHSCRIPT_URL + assert kwargs["expected_sha256"] == digest + assert kwargs["max_bytes"] == len(model_bytes) + assert kwargs["exact_size"] == len(model_bytes) + Path(destination).write_bytes(model_bytes) + return Path(destination) + + monkeypatch.setattr(pet_pipeline, "_install_certifi_environment", lambda: None) + monkeypatch.setattr(pet_pipeline, "_download_file", fake_download) + + loaded = embedder._build_dinov2_cache(model_path) + metadata_path = pet_pipeline._dinov2_metadata_path(model_path) + assert isinstance(loaded, FakeModel) + assert model_path.read_bytes() == model_bytes + assert metadata_path.is_file() + pet_pipeline._validate_dinov2_cache_metadata(model_path, model_name=model_name) + metadata = hashlib_json_load(metadata_path) + assert metadata["artifact_kind"] == "release_torchscript" + assert metadata["cache_schema_version"] == 2 + assert metadata["torchscript_sha256"] == digest + assert metadata["torchscript_size"] == len(model_bytes) + + def test_dino_publish_rolls_back_model_when_metadata_publish_is_denied( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + candidate = tmp_path / "candidate.pt" + candidate_metadata = tmp_path / "candidate.pt.metadata.json" + model_path = tmp_path / "published" / "dinov2_vits14.pt" + final_metadata = pet_pipeline._dinov2_metadata_path(model_path) + model_path.parent.mkdir(parents=True) + candidate.write_bytes(b"candidate") + candidate_metadata.write_text("{}", encoding="utf-8") + + real_replace = Path.replace + + def fail_metadata_replace(self: Path, target: Path): + if self == candidate_metadata: + raise OSError(errno.EACCES, "denied") + return real_replace(self, target) + + monkeypatch.setattr(Path, "replace", fail_metadata_replace) + + with pytest.raises(pet_pipeline._ModelStoragePermissionError): + pet_pipeline._publish_dinov2_cache_pair( + candidate, + candidate_metadata, + model_path, + ) + + assert not model_path.exists() + assert not final_metadata.exists() + + +class ShortResponse: + def __init__(self, payload: bytes) -> None: + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + def read(self, _size: int): + value, self.payload = self.payload, b"" + return value + + +def _sha256_target(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def hashlib_json(value: dict) -> str: + import json + + return json.dumps(value) + + +def hashlib_json_load(path: Path) -> dict: + import json + + return json.loads(path.read_text(encoding="utf-8")) + + +class TestLazyPipeline: + def test_empty_batch_does_not_initialize_models(self, tmp_path: Path) -> None: + pipeline_instance = pet_pipeline.PetClusterPipeline(model_root=tmp_path / "models") + + def fail_initialize(): + raise AssertionError("empty batches must not load or download models") + + pipeline_instance._ensure_detector = fail_initialize + pipeline_instance._ensure_embedder = fail_initialize + assert pipeline_instance.detect_pets_for_rows([], library_root=tmp_path, thumbnail_dir=tmp_path) == [] diff --git a/tests/test_pet_model_cache_integrity.py b/tests/test_pet_model_cache_integrity.py new file mode 100644 index 000000000..15406f805 --- /dev/null +++ b/tests/test_pet_model_cache_integrity.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + + +def _write_cache(tmp_path: Path, *, content: bytes, metadata: dict) -> Path: + model_path = tmp_path / "dinov2_vits14.pt" + metadata_path = pet_pipeline._dinov2_metadata_path(model_path) + model_path.parent.mkdir(parents=True, exist_ok=True) + model_path.write_bytes(content) + metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + return model_path + + +def _base_metadata() -> dict: + return pet_pipeline._dinov2_release_metadata(model_name="dinov2_vits14") + + +def test_prebuilt_torchscript_requires_manifest_integrity(tmp_path: Path) -> None: + manifest = pet_pipeline._EMBEDDER_MANIFEST + content = b"prebuilt-torchscript" + metadata = _base_metadata() + model_path = _write_cache(tmp_path, content=content, metadata=metadata) + + with pytest.raises(RuntimeError, match="integrity check failed"): + pet_pipeline._validate_dinov2_cache_metadata( + model_path, + model_name=manifest["model_name"], + ) + + +def test_legacy_derived_user_cache_requires_release_replacement( + tmp_path: Path, + monkeypatch, +) -> None: + monkeypatch.delenv("IPHOTO_PET_MODEL_DIR", raising=False) + cache_root = tmp_path / "cache" + monkeypatch.setattr(pet_pipeline, "user_pet_model_cache_dir", lambda: cache_root) + good = b"derived-torchscript" + corrupt = good[:-8] + metadata = { + "artifact_kind": "derived_checkpoint_cache", + "model_name": "dinov2_vits14", + "derived_torchscript_size": len(good), + "derived_torchscript_sha256": hashlib.sha256(good).hexdigest(), + } + model_path = _write_cache(tmp_path, content=corrupt, metadata=metadata) + + with pytest.raises(RuntimeError, match="metadata contract mismatch"): + pet_pipeline._validate_dinov2_cache_metadata( + model_path, + model_name=pet_pipeline._EMBEDDER_MANIFEST["model_name"], + ) diff --git a/tests/test_pet_model_concurrency.py b/tests/test_pet_model_concurrency.py new file mode 100644 index 000000000..28ce08238 --- /dev/null +++ b/tests/test_pet_model_concurrency.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import multiprocessing +import hashlib +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + + +_DETECTOR_BYTES = b"concurrent-detector" +_DETECTOR_SHA256 = hashlib.sha256(_DETECTOR_BYTES).hexdigest() + + +def _acquire_detector_once(model_path: str, start, downloads) -> None: + from iPhoto.pets import pipeline + + pipeline._impl.DEFAULT_PET_DETECTOR_MODEL_SHA256 = _DETECTOR_SHA256 + + def fake_download(_url: str, destination: Path, **_kwargs) -> Path: + with downloads.get_lock(): + downloads.value += 1 + time.sleep(0.2) + Path(destination).write_bytes(_DETECTOR_BYTES) + return Path(destination) + + pipeline._impl._download_file = fake_download + start.wait(10) + pipeline.ensure_pet_detector_model(Path(model_path)) + + +def _hold_dino_lock(model_path: str, ready, release) -> None: + from iPhoto.pets import pipeline + + with pipeline._dinov2_acquisition_lock(Path(model_path)): + ready.set() + release.wait(15) + + +def _wait_for_dino_lock(model_path: str, acquired) -> None: + from iPhoto.pets import pipeline + + with pipeline._dinov2_acquisition_lock(Path(model_path)): + acquired.set() + + +def test_dino_acquisition_lock_serializes_processes(tmp_path: Path) -> None: + model_path = tmp_path / "embedding" / "dinov2_vits14.pt" + context = multiprocessing.get_context("spawn") + ready = context.Event() + release = context.Event() + acquired = context.Event() + holder = context.Process( + target=_hold_dino_lock, + args=(str(model_path), ready, release), + ) + waiter = context.Process( + target=_wait_for_dino_lock, + args=(str(model_path), acquired), + ) + + try: + holder.start() + assert ready.wait(10), "first process did not acquire the DINOv2 cache lock" + waiter.start() + assert not acquired.wait(0.5), "second process bypassed the DINOv2 cache lock" + release.set() + assert acquired.wait(10), "second process did not acquire the released DINOv2 cache lock" + holder.join(10) + waiter.join(10) + assert holder.exitcode == 0 + assert waiter.exitcode == 0 + finally: + release.set() + for process in (holder, waiter): + if process.is_alive(): + process.terminate() + process.join(5) + + +def test_detector_acquisition_lock_downloads_once_across_processes(tmp_path: Path) -> None: + model_path = tmp_path / "detector" / "yolox_nano_coco.onnx" + context = multiprocessing.get_context("spawn") + start = context.Event() + downloads = context.Value("i", 0) + processes = [ + context.Process( + target=_acquire_detector_once, + args=(str(model_path), start, downloads), + ) + for _ in range(2) + ] + try: + for process in processes: + process.start() + start.set() + for process in processes: + process.join(15) + assert process.exitcode == 0 + assert downloads.value == 1 + assert model_path.read_bytes() == _DETECTOR_BYTES + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(5) + + +def test_waiting_builder_reuses_cache_published_by_first_builder( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_path = tmp_path / "embedding" / "dinov2_vits14.pt" + model_path.parent.mkdir(parents=True) + model_path.write_bytes(b"already-published") + pet_pipeline._dinov2_metadata_path(model_path).write_text("{}", encoding="utf-8") + + loaded = object() + load_calls: list[tuple[Path, str]] = [] + + def fake_load(path: str, *, map_location: str): + load_calls.append((Path(path), map_location)) + return loaded + + def fail_download(*_args, **_kwargs): + raise AssertionError("a waiting builder must not download after another builder publishes") + + monkeypatch.setattr(pet_pipeline, "_validate_dinov2_cache_metadata", lambda *_a, **_k: None) + monkeypatch.setattr(pet_pipeline, "_download_file", fail_download) + monkeypatch.setattr(pet_pipeline, "_install_certifi_environment", lambda: None) + + embedder = pet_pipeline._DinoV2Embedder.__new__(pet_pipeline._DinoV2Embedder) + embedder._model_name = "dinov2_vits14" + embedder._torch = SimpleNamespace(jit=SimpleNamespace(load=fake_load)) + + result = embedder._build_verified_dinov2_cpu_cache(model_path) + + assert result is loaded + assert load_calls == [(model_path, "cpu")] + + +def test_dino_publish_makes_model_visible_after_metadata( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + candidate = tmp_path / "candidate.pt" + candidate_metadata = tmp_path / "candidate.pt.metadata.json" + model_path = tmp_path / "published" / "dinov2_vits14.pt" + final_metadata = pet_pipeline._dinov2_metadata_path(model_path) + model_path.parent.mkdir(parents=True) + candidate.write_bytes(b"model") + candidate_metadata.write_text("{}", encoding="utf-8") + + real_replace = Path.replace + publish_order: list[Path] = [] + + def record_replace(self: Path, target: Path): + if self in {candidate, candidate_metadata}: + publish_order.append(self) + return real_replace(self, target) + + monkeypatch.setattr(Path, "replace", record_replace) + + pet_pipeline._publish_dinov2_cache_pair( + candidate, + candidate_metadata, + model_path, + ) + + assert publish_order == [candidate_metadata, candidate] + assert model_path.read_bytes() == b"model" + assert final_metadata.read_text(encoding="utf-8") == "{}" diff --git a/tests/test_pet_model_hardening.py b/tests/test_pet_model_hardening.py new file mode 100644 index 000000000..796c5a49b --- /dev/null +++ b/tests/test_pet_model_hardening.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + + +def test_dinov2_release_and_provenance_are_immutable() -> None: + manifest = pet_pipeline._EMBEDDER_MANIFEST + assert manifest["source_repository"] == "facebookresearch/dinov2" + assert manifest["source_revision"] == "7764ea0f912e53c92e82eb78a2a1631e92725fc8" + assert manifest["source_tree_sha1"] == "2a27257b79b0633b027a21014bc9360e3c1b3f43" + assert manifest["torchscript_url"] == ( + "https://github.com/OliverZhaohaibin/iPhotron-LocalPhotoAlbumManager/" + "releases/download/pet-models-v1/dinov2_vits14.pt" + ) + assert manifest["cache_schema_version"] == 2 + assert manifest["producer_torch_version"] == "2.12.1" + + +def test_production_pets_source_has_no_torch_hub_or_xformers_branch() -> None: + source_root = Path(pet_pipeline.__file__).parent + production_source = "\n".join( + path.read_text(encoding="utf-8") for path in source_root.glob("*.py") + ) + assert "torch.hub" not in production_source + assert "source=\"github\"" not in production_source + assert "XFORMERS" not in production_source + + +def test_embedder_construction_stays_lazy_until_first_crop( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[Path] = [] + + class ExplodingEmbedder: + def __init__(self, model_dir: Path, **_kwargs) -> None: + calls.append(Path(model_dir)) + raise AssertionError("real DINOv2 initialization should still be deferred") + + monkeypatch.setattr(pet_pipeline, "_DinoV2Embedder", ExplodingEmbedder) + pipeline = pet_pipeline.PetClusterPipeline(model_root=tmp_path / "models") + + lazy = pipeline._ensure_embedder() + assert calls == [] + + with pytest.raises(AssertionError, match="still be deferred"): + lazy.embed(None) + assert calls == [tmp_path / "models" / "embedding" / "dinov2_vits14"] + + +def test_device_failure_does_not_delete_published_dino_cache( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_path = tmp_path / "embedding" / "dinov2_vits14.pt" + metadata_path = pet_pipeline._dinov2_metadata_path(model_path) + + class FakeModel: + def eval(self): + return self + + def to(self, device): + if device == "test-device": + raise RuntimeError("device unavailable") + return self + + def fake_verified_cpu_build(_self, target: Path): + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"verified-cache") + pet_pipeline._dinov2_metadata_path(target).write_text("{}", encoding="utf-8") + return FakeModel() + + monkeypatch.setattr( + pet_pipeline._DinoV2Embedder, + "_build_verified_dinov2_cpu_cache", + fake_verified_cpu_build, + ) + embedder = pet_pipeline._DinoV2Embedder.__new__(pet_pipeline._DinoV2Embedder) + embedder._device = "test-device" + + with pytest.raises(pet_pipeline.PetModelUnavailableError, match="built and verified"): + embedder._build_dinov2_cache(model_path) + + assert model_path.read_bytes() == b"verified-cache" + assert metadata_path.is_file() + + +def test_dino_download_error_does_not_suggest_detector_override( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail_download(*_args, **_kwargs): + raise RuntimeError( + "Pet scanning unavailable: failed to download DINOv2 checkpoint from " + "https://example.test/model (network denied). Check your network connection, set " + f"{pet_pipeline.PET_DETECTOR_MODEL_URL_ENV}, or install the model manually." + ) + + monkeypatch.setattr(pet_pipeline, "_original_download_file", fail_download) + + with pytest.raises(RuntimeError) as raised: + pet_pipeline._download_file( + "https://example.test/model", + tmp_path / "model.bin", + label="DINOv2 checkpoint", + expected_sha256="0" * 64, + max_bytes=100, + ) + + message = str(raised.value) + assert pet_pipeline.PET_DETECTOR_MODEL_URL_ENV not in message + assert "Check your network connection or install the model manually." in message diff --git a/tests/test_pet_model_network_error_policy.py b/tests/test_pet_model_network_error_policy.py new file mode 100644 index 000000000..3e2fabfe6 --- /dev/null +++ b/tests/test_pet_model_network_error_policy.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import errno +from pathlib import Path + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + + +def test_network_permission_error_does_not_trigger_storage_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundled = tmp_path / "extension" + cache = tmp_path / "cache" + target = bundled / "detector" / "yolox_nano_coco.onnx" + monkeypatch.delenv("IPHOTO_PET_MODEL_DIR", raising=False) + monkeypatch.setattr(pet_pipeline, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_pipeline, "user_pet_model_cache_dir", lambda: cache) + calls: list[str] = [] + + def fail_urlopen(url: str, **_kwargs): + calls.append(url) + raise PermissionError(errno.EACCES, "network denied") + + monkeypatch.setattr(pet_pipeline.request, "urlopen", fail_urlopen) + + with pytest.raises(RuntimeError, match="failed to download") as raised: + pet_pipeline.ensure_pet_detector_model(target) + + assert not isinstance(raised.value, pet_pipeline._ModelStoragePermissionError) + assert len(calls) == 1 + assert not cache.exists() diff --git a/tests/test_pet_model_resolver_concurrency.py b/tests/test_pet_model_resolver_concurrency.py new file mode 100644 index 000000000..ba555ec5d --- /dev/null +++ b/tests/test_pet_model_resolver_concurrency.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from iPhoto.pets import pipeline as pet_pipeline + +pet_impl = pet_pipeline._impl + + +def test_resolver_does_not_delete_metadata_during_dino_publish( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + relative = Path("embedding") / "dinov2_vits14" + bundled = tmp_path / "bundled" + user_cache = tmp_path / "cache" + final_dir = user_cache / relative + final_dir.mkdir(parents=True) + model_path = final_dir / "dinov2_vits14.pt" + final_metadata = pet_pipeline._dinov2_metadata_path(model_path) + + candidate = tmp_path / "candidate.pt" + candidate_metadata = tmp_path / "candidate.pt.metadata.json" + candidate.write_bytes(b"model") + candidate_metadata.write_text("{}", encoding="utf-8") + + monkeypatch.setattr(pet_impl, "pet_model_override_dir", lambda: None) + monkeypatch.setattr(pet_impl, "bundled_pet_model_dir", lambda: bundled) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: user_cache) + monkeypatch.setattr( + pet_impl, + "pet_model_search_roots", + lambda: (bundled, user_cache), + ) + monkeypatch.setattr(pet_impl, "pet_model_install_root", lambda: user_cache) + + real_replace = Path.replace + resolver_results: list[Path] = [] + + def interleaved_replace(self: Path, target: Path): + result = real_replace(self, target) + if self == candidate_metadata: + assert final_metadata.is_file() + assert not model_path.exists() + resolver_results.append( + pet_pipeline.resolve_pet_model_path(relative, directory=True) + ) + assert final_metadata.is_file() + assert not model_path.exists() + return result + + monkeypatch.setattr(Path, "replace", interleaved_replace) + + pet_pipeline._publish_dinov2_cache_pair( + candidate, + candidate_metadata, + model_path, + ) + + assert resolver_results == [final_dir] + assert model_path.read_bytes() == b"model" + assert final_metadata.read_text(encoding="utf-8") == "{}" diff --git a/tests/test_pet_remediation_contracts.py b/tests/test_pet_remediation_contracts.py index 2ac03e792..d534e37d4 100644 --- a/tests/test_pet_remediation_contracts.py +++ b/tests/test_pet_remediation_contracts.py @@ -32,6 +32,8 @@ from iPhoto.pets.status import is_pet_scan_candidate from iPhoto.utils.pathutils import LibraryAssetPathError, resolve_library_asset_path +pet_impl = pet_pipeline._impl + class _AssetStore: def __init__(self, asset_id: str) -> None: @@ -396,7 +398,7 @@ def save_then_fail(image, bbox, output_path, *, padding_ratio): Path(output_path).parent.mkdir(parents=True, exist_ok=True) Path(output_path).write_bytes(b"thumbnail") - monkeypatch.setattr(pet_pipeline, "save_pet_thumbnail", save_then_fail) + monkeypatch.setattr(pet_impl, "save_pet_thumbnail", save_then_fail) thumbnail_dir = tmp_path / ".iPhoto" / "pets" / "thumbnails" / ".staging" / "op" results = pipeline.detect_pets_for_rows( [{"id": "asset-a", "rel": "album/a.jpg"}], @@ -1101,8 +1103,6 @@ def test_incremental_rescan_preserves_manual_cross_species_identity( assert coordinator._journal.unfinished() == () assert f"Preserving mixed-species Pet identity {dog_identity}" in caplog.text - # The non-dominant species member must also keep the durable identity. - # This is the regression case that previously fell through the species gate. caplog.clear() rescanned_cat = replace( _detection("cat-rescan", asset_id="asset-cat"), @@ -1500,23 +1500,25 @@ def test_model_resolver_skips_empty_cache_for_complete_bundled_embedder( hashlib.sha256(model_path.read_bytes()).hexdigest(), ) monkeypatch.setitem(manifest, "torchscript_size", model_path.stat().st_size) + monkeypatch.setattr( + pet_impl, + "_DINO_TORCHSCRIPT_SHA256", + hashlib.sha256(model_path.read_bytes()).hexdigest(), + ) + monkeypatch.setattr( + pet_impl, + "_DINO_TORCHSCRIPT_SIZE", + model_path.stat().st_size, + ) model_path.with_suffix(".pt.metadata.json").write_text( json.dumps( - { - "model_name": "dinov2_vits14", - "source_repository": manifest["source_repository"], - "source_revision": manifest["source_revision"], - "torchscript_sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(), - "torchscript_size": model_path.stat().st_size, - "input_shape": manifest["input_shape"], - "output_shape": manifest["output_shape"], - } + pet_pipeline._dinov2_release_metadata(model_name="dinov2_vits14") ), encoding="utf-8", ) - monkeypatch.setattr(pet_pipeline, "user_pet_model_cache_dir", lambda: cache) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) monkeypatch.setattr( - pet_pipeline, + pet_impl, "pet_model_search_roots", lambda: (cache, bundled), ) @@ -1524,7 +1526,7 @@ def test_model_resolver_skips_empty_cache_for_complete_bundled_embedder( assert pet_pipeline.resolve_pet_model_path(relative, directory=True) == bundled_dir -def test_model_resolver_removes_corrupt_user_cache_and_uses_bundled( +def test_model_resolver_preserves_corrupt_user_cache_and_uses_bundled( tmp_path: Path, monkeypatch, ) -> None: @@ -1543,16 +1545,16 @@ def validate(path, **kwargs): if Path(path) == cached_model: raise RuntimeError("bad hash") - monkeypatch.setattr(pet_pipeline, "_validate_downloaded_file", validate) - monkeypatch.setattr(pet_pipeline, "user_pet_model_cache_dir", lambda: cache) + monkeypatch.setattr(pet_impl, "_validate_downloaded_file", validate) + monkeypatch.setattr(pet_impl, "user_pet_model_cache_dir", lambda: cache) monkeypatch.setattr( - pet_pipeline, + pet_impl, "pet_model_search_roots", lambda: (cache, bundled), ) assert pet_pipeline.resolve_pet_model_path(relative) == bundled_model - assert not cached_model.exists() + assert cached_model.exists() def test_thumbnail_publish_compensates_when_later_replace_fails( diff --git a/tests/test_pet_service.py b/tests/test_pet_service.py index a35b7549f..7b5bc4973 100644 --- a/tests/test_pet_service.py +++ b/tests/test_pet_service.py @@ -34,7 +34,6 @@ _decode_yolox_predictions, _dedupe_supported_species_boxes, _DetectedPetBox, - _DinoV2Embedder, _map_yolox_box_to_source, _pet_box_overlaps_people_boxes, _preprocess_yolox, @@ -2447,59 +2446,16 @@ def recording_temporary_directory(*args, **kwargs): assert not temporary_directories[0].exists() -def test_dinov2_runtime_downloads_only_the_fixed_torchscript_artifact( - tmp_path: Path, - monkeypatch, -) -> None: - class FakeModel: - def eval(self): - return self - - def to(self, _device): - return self - - class FakeJit: - @staticmethod - def load(path: str, *, map_location: str): - assert Path(path).read_bytes() == b"fixed-torchscript" - assert map_location == "cpu" - return FakeModel() - - embedder = _DinoV2Embedder.__new__(_DinoV2Embedder) - embedder._torch = SimpleNamespace(jit=FakeJit()) - embedder._device = "cpu" - embedder._model_name = "dinov2_vits14" - monkeypatch.setattr(pet_pipeline, "_install_certifi_environment", lambda: None) - monkeypatch.setitem( - pet_pipeline._EMBEDDER_MANIFEST, - "torchscript_url", - "https://models.example.test/dinov2_vits14.pt", - ) - monkeypatch.setitem( - pet_pipeline._EMBEDDER_MANIFEST, - "torchscript_sha256", - hashlib.sha256(b"fixed-torchscript").hexdigest(), +def test_dinov2_manifest_declares_official_checkpoint_source() -> None: + embedder = PET_MODEL_MANIFEST["embedder"] + assert embedder["weights_url"] == ( + "https://dl.fbaipublicfiles.com/dinov2/dinov2_vits14/" + "dinov2_vits14_pretrain.pth" ) - monkeypatch.setitem( - pet_pipeline._EMBEDDER_MANIFEST, - "torchscript_size", - len(b"fixed-torchscript"), + assert embedder["weights_sha256"] == ( + "b938bf1bc15cd2ec0feacfe3a1bb553fe8ea9ca46a7e1d8d00217f29aef60cd9" ) - - def fake_download(url, path, **kwargs): - assert url == "https://models.example.test/dinov2_vits14.pt" - assert kwargs["expected_sha256"] == hashlib.sha256(b"fixed-torchscript").hexdigest() - assert kwargs["max_bytes"] == len(b"fixed-torchscript") - Path(path).parent.mkdir(parents=True, exist_ok=True) - Path(path).write_bytes(b"fixed-torchscript") - - monkeypatch.setattr(pet_pipeline, "_download_file", fake_download) - model_path = tmp_path / "dinov2_vits14.pt" - model = embedder._download_dinov2_model(model_path) - - assert isinstance(model, FakeModel) - assert model_path.read_bytes() == b"fixed-torchscript" - assert pet_pipeline._dinov2_metadata_path(model_path).is_file() + assert embedder["weights_size"] == 88283115 def test_pet_embedding_source_is_pinned_to_commit() -> None: diff --git a/tests/test_startup_benchmark.py b/tests/test_startup_benchmark.py index 58d57d87c..e4c3202eb 100644 --- a/tests/test_startup_benchmark.py +++ b/tests/test_startup_benchmark.py @@ -29,6 +29,7 @@ def _write_profile( revision: str = "candidate", artifact_sha256: str = "a" * 64, build_environment_fingerprint: str = "f" * 64, + scenario_env_names: str = "", ) -> None: context = { "run_id": path.stem, @@ -42,6 +43,7 @@ def _write_profile( "cache_controlled": controlled, "cache_eviction_method": "purge" if controlled else "uncontrolled", "scenario": scenario, + "scenario_env_names": scenario_env_names, "build_environment_fingerprint": build_environment_fingerprint, "artifact_sha256": artifact_sha256, "manifest_source_revision": revision, @@ -110,6 +112,59 @@ def test_analyse_and_summarize_valid_profiles(tmp_path) -> None: assert summary["metrics"]["first_usable_thumbnail_ms"]["p95"] == 768.0 +def test_recognition_resource_snapshots_are_correlated_to_activation(tmp_path) -> None: + path = tmp_path / "resources.jsonl" + _write_profile(path) + payloads = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + payloads = [item for item in payloads if item["stage"] != "launcher.process_finished"] + payloads.extend( + [ + { + "stage": "recognition.startup.activated", + "elapsed_ms": 700.0, + "wall_time": 1000.7, + "pid": 123, + "details": {"generation": 1}, + }, + *[ + { + "stage": "launcher.resource_sample", + "elapsed_ms": elapsed, + "wall_time": 1000.0 + elapsed / 1000.0, + "pid": 123, + "details": { + "cpu_ms": elapsed / 10.0, + "rss_bytes": int(elapsed * 1000), + "read_bytes": int(elapsed * 10), + "write_bytes": int(elapsed * 5), + }, + } + for elapsed in (300.0, 700.0, 2200.0, 5700.0) + ], + { + "stage": "launcher.process_finished", + "elapsed_ms": 6000.0, + "wall_time": 1006.0, + "pid": 123, + "details": {"return_code": 0, "timed_out": False}, + }, + ] + ) + payloads.sort(key=lambda item: float(item["elapsed_ms"])) + path.write_text("".join(json.dumps(item) + "\n" for item in payloads), encoding="utf-8") + + run = analyse_run(path) + summary = summarize_profiles([path]) + + assert run["metrics"]["interactive_recognition_activation_ms"] == 400.0 + assert run["metrics"]["max_post_recognition_gui_stall_ms"] == 0.0 + assert run["resource_snapshots"]["interactive"]["rss_bytes"] == 300000 + assert run["resource_snapshots"]["recognition_activation"]["rss_bytes"] == 700000 + assert run["resource_snapshots"]["recognition_plus_1500ms"]["rss_bytes"] == 2200000 + assert run["resource_snapshots"]["recognition_plus_5000ms"]["rss_bytes"] == 5700000 + assert summary["resources"]["recognition_plus_5000ms"]["cpu_ms"]["p95"] == 570.0 + + def test_missing_terminal_and_path_leak_are_rejected(tmp_path) -> None: path = tmp_path / "broken.jsonl" _write_profile(path, include_terminal=False) @@ -159,6 +214,41 @@ def test_nonzero_or_timed_out_process_is_rejected(tmp_path) -> None: assert any("timed out" in error for error in timed_out_run["errors"]) +def test_quick_close_rejects_recognition_worker_start(tmp_path) -> None: + path = tmp_path / "quick-close.jsonl" + _write_profile(path, scenario="recognition-quick-close") + payloads = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + payloads.append( + { + "stage": "recognition.worker.started", + "elapsed_ms": 610.0, + "wall_time": 1000.61, + "pid": 123, + "details": {"generation": 1, "worker": "face"}, + } + ) + payloads.sort(key=lambda item: float(item["elapsed_ms"])) + path.write_text("".join(json.dumps(item) + "\n" for item in payloads), encoding="utf-8") + + run = analyse_run(path) + + assert run["valid"] is False + assert "quick-close started a recognition worker" in run["errors"] + + +def test_feature_scoped_ab_arm_does_not_require_auto_activation(tmp_path) -> None: + path = tmp_path / "feature-scoped.jsonl" + _write_profile( + path, + scenario="recognition-auto-models-present", + scenario_env_names="IPHOTO_STARTUP_RECOGNITION_AUTO_START", + ) + + run = analyse_run(path) + + assert run["valid"] is True + + def test_mixed_batch_contexts_are_rejected(tmp_path) -> None: local = tmp_path / "local.jsonl" network = tmp_path / "network.jsonl" @@ -212,6 +302,7 @@ def test_collect_isolates_profile_and_aggregates_subprocess(tmp_path) -> None: from pathlib import Path assert os.environ["IPHOTO_STARTUP_BENCHMARK_AUTO_EXIT_MS"] == "25" +assert os.environ["IPHOTO_TEST_SCENARIO"] == "enabled" path = Path(os.environ["IPHOTO_STARTUP_PROFILE_PATH"]) context = { "run_id": os.environ["IPHOTO_STARTUP_RUN_ID"], @@ -292,6 +383,8 @@ def test_collect_isolates_profile_and_aggregates_subprocess(tmp_path) -> None: "1", "--auto-exit-delay-ms", "25", + "--set-env", + "IPHOTO_TEST_SCENARIO=enabled", "--output-dir", str(output), "--", @@ -336,6 +429,44 @@ def test_packaged_collect_requires_matching_build_manifest(tmp_path) -> None: assert result == 2 +def test_template_restore_is_confined_to_output_active_library(tmp_path) -> None: + template = tmp_path / "template" + template.mkdir() + outside = tmp_path / "outside" + output = tmp_path / "output" + + result = benchmark_main( + [ + "collect", + "--revision", + "candidate", + "--scenario", + "recognition-auto-models-present", + "--library", + str(outside), + "--library-template", + str(template), + "--confirm-dedicated-library", + "--confirm-template-restore", + "--runtime", + "source", + "--cache-state", + "hot", + "--samples", + "1", + "--output-dir", + str(output), + "--", + sys.executable, + "-c", + "pass", + ] + ) + + assert result == 2 + assert not outside.exists() + + def test_comparison_rejects_environment_mismatch(tmp_path) -> None: baseline_path = tmp_path / "baseline.jsonl" candidate_path = tmp_path / "candidate.jsonl" diff --git a/tests/test_startup_pet_backfill.py b/tests/test_startup_pet_backfill.py new file mode 100644 index 000000000..82d160d6b --- /dev/null +++ b/tests/test_startup_pet_backfill.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +pytest.importorskip("PySide6", reason="PySide6 is required for library tests", exc_type=ImportError) +pytest.importorskip("PySide6.QtWidgets", reason="Qt widgets not available", exc_type=ImportError) + +from PySide6.QtWidgets import QApplication + +from iPhoto.library.runtime_controller import LibraryRuntimeController + + +@pytest.fixture(scope="module") +def qapp() -> QApplication: + os.environ.setdefault("QT_QPA_PLATFORM", "offscreen") + app = QApplication.instance() + if app is None: + app = QApplication([]) + yield app + + +def _startup_worker() -> Mock: + worker = Mock(cancelled=False, failed=False) + worker.scan_service = Mock() + worker.scan_started_at_ms = 1 + worker.scan_job_id = "startup-test" + worker._defer_ai_workers_until_scan_finished = True + return worker + + +def test_successful_startup_scan_arms_people_and_pets_idle_gate( + tmp_path: Path, + qapp: QApplication, +) -> None: + root = tmp_path / "Library" + root.mkdir() + manager = LibraryRuntimeController() + manager.bind_path(root) + worker = _startup_worker() + manager._current_scanner_worker = worker + manager._live_scan_root = root + manager.request_startup_recognition_after_idle() + + with ( + patch.object(manager, "_arm_startup_recognition_idle_timer") as arm_idle, + patch.object(manager._scan_thread_pool, "start") as start_pool, + ): + manager._on_scan_finished(worker, root, [{"rel": "pet.jpg"}]) + qapp.processEvents() + + worker.scan_service.finalize_scan_result.assert_called_once() + arm_idle.assert_called_once_with(root, manager._recognition_generation) + start_pool.assert_called_once() + + +def test_startup_recognition_idle_gate_is_invalidated_by_shutdown_generation( + tmp_path: Path, + qapp: QApplication, +) -> None: + root = tmp_path / "Library" + root.mkdir() + manager = LibraryRuntimeController() + manager.bind_path(root) + worker = _startup_worker() + manager._current_scanner_worker = worker + manager._live_scan_root = root + manager.request_startup_recognition_after_idle() + + def invalidate_recognition_generation(*_args) -> None: + manager._recognition_generation += 1 + + manager.scanFinished.connect(invalidate_recognition_generation) + + with ( + patch.object(manager, "_arm_startup_recognition_idle_timer") as arm_idle, + patch.object(manager._scan_thread_pool, "start") as start_pool, + ): + manager._on_scan_finished(worker, root, [{"rel": "pet.jpg"}]) + qapp.processEvents() + + arm_idle.assert_not_called() + start_pool.assert_called_once() + + +def test_idle_timeout_lazily_binds_and_starts_both_recognition_services( + tmp_path: Path, + qapp: QApplication, +) -> None: + root = tmp_path / "Library" + root.mkdir() + manager = LibraryRuntimeController() + manager._root = root + people_service = object() + pet_service = object() + manager._library_session = SimpleNamespace( + library_root=root, + people=people_service, + pets=pet_service, + ) + + with ( + patch.object(manager, "bind_recognition_services") as bind_services, + patch.object(manager, "activate_recognition_scans") as activate, + ): + manager.request_startup_recognition_after_idle() + assert manager._startup_recognition_timer.isActive() + assert manager._startup_recognition_timer.interval() == 1500 + manager.notify_user_activity() + assert manager._startup_recognition_timer.isActive() + manager._startup_recognition_timer.stop() + manager._activate_startup_recognition_after_idle() + + bind_services.assert_called_once_with(people_service, pet_service) + activate.assert_called_once_with() + + +def test_failed_startup_scan_cancels_pending_idle_activation( + tmp_path: Path, + qapp: QApplication, +) -> None: + root = tmp_path / "Library" + root.mkdir() + manager = LibraryRuntimeController() + manager.bind_path(root) + worker = _startup_worker() + worker.failed = True + manager._current_scanner_worker = worker + manager._live_scan_root = root + manager.request_startup_recognition_after_idle() + + manager._on_scan_finished(worker, root, []) + + assert manager._startup_recognition_request is None + assert not manager._startup_recognition_timer.isActive() diff --git a/tools/convert_dinov2_torchscript.py b/tools/convert_dinov2_torchscript.py old mode 100644 new mode 100755 index 44a49019c..ca8aa4a98 --- a/tools/convert_dinov2_torchscript.py +++ b/tools/convert_dinov2_torchscript.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Build and verify the fixed DINOv2 TorchScript release artifact. -This is a development/release tool. Production code must never execute Torch Hub. +This is a release-only tool. Production never executes DINOv2 source code. """ from __future__ import annotations @@ -9,11 +9,15 @@ import argparse import hashlib import json +import os +import platform +import subprocess import sys import tempfile from pathlib import Path import torch +import torchvision REPOSITORY_ROOT = Path(__file__).resolve().parents[1] MANIFEST_PATH = REPOSITORY_ROOT / "src" / "iPhoto" / "pets" / "model_manifest.json" @@ -30,6 +34,10 @@ def _sha256(path: Path) -> str: def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("output", type=Path, help="TorchScript artifact destination") + parser.add_argument("--source-dir", required=True, type=Path) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--runtime-metadata", type=Path) + parser.add_argument("--build-manifest", type=Path) parser.add_argument( "--require-manifest-identity", action="store_true", @@ -44,18 +52,45 @@ def main() -> int: repository = str(manifest["source_repository"]) revision = str(manifest["source_revision"]) model_name = str(manifest["model_name"]) - source = f"{repository}:{revision}" - + source_dir = args.source_dir.resolve() + checkpoint = args.checkpoint.resolve() + if _sha256(checkpoint) != str(manifest["weights_sha256"]): + raise RuntimeError("DINOv2 checkpoint SHA-256 does not match the manifest") + if checkpoint.stat().st_size != int(manifest["weights_size"]): + raise RuntimeError("DINOv2 checkpoint size does not match the manifest") + source_commit = _git_value(source_dir, "HEAD") + source_tree = _git_value(source_dir, "HEAD^{tree}") + if source_commit != revision or source_tree != str(manifest["source_tree_sha1"]): + raise RuntimeError("DINOv2 source checkout does not match the pinned commit/tree") + + os.environ.setdefault("XFORMERS_DISABLED", "1") torch.manual_seed(0) example = torch.randn(tuple(manifest["input_shape"]), dtype=torch.float32) - model = torch.hub.load(source, model_name, source="github", trust_repo=True).eval().cpu() + model = ( + torch.hub.load( + str(source_dir), + model_name, + source="local", + trust_repo=True, + pretrained=False, + ) + .eval() + .cpu() + ) + state_dict = torch.load(str(checkpoint), map_location="cpu", weights_only=True) + model.load_state_dict(state_dict, strict=True) args.output.parent.mkdir(parents=True, exist_ok=True) with tempfile.TemporaryDirectory(prefix="iphoto-dinov2-release-") as temp_dir: candidate = Path(temp_dir) / args.output.name with torch.no_grad(): eager_output = model(example) - traced = torch.jit.trace(model, example, strict=False) + traced = torch.jit.trace( + model, + example, + strict=False, + check_trace=False, + ) traced.save(str(candidate)) scripted = torch.jit.load(str(candidate), map_location="cpu").eval() scripted_output = scripted(example) @@ -68,7 +103,7 @@ def main() -> int: raise RuntimeError( f"output shape mismatch: {tuple(scripted_output.shape)} != {expected_shape}" ) - torch.testing.assert_close(scripted_output, eager_output, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(scripted_output, eager_output, rtol=1e-3, atol=3e-5) artifact_sha256 = _sha256(candidate) artifact_size = candidate.stat().st_size if args.require_manifest_identity and ( @@ -81,6 +116,43 @@ def main() -> int: ) candidate.replace(args.output) + runtime_metadata = { + "artifact_kind": "release_torchscript", + "cache_schema_version": int(manifest["cache_schema_version"]), + "model_name": model_name, + "release_tag": manifest["release_tag"], + "source_repository": repository, + "source_revision": revision, + "source_tree_sha1": source_tree, + "weights_sha256": manifest["weights_sha256"], + "weights_size": int(manifest["weights_size"]), + "producer_python_version": manifest["producer_python_version"], + "producer_torch_version": manifest["producer_torch_version"], + "producer_torchvision_version": manifest["producer_torchvision_version"], + "torchscript_url": manifest["torchscript_url"], + "torchscript_sha256": artifact_sha256, + "torchscript_size": artifact_size, + "input_shape": manifest["input_shape"], + "output_shape": manifest["output_shape"], + } + if args.runtime_metadata is not None: + _write_json(args.runtime_metadata, runtime_metadata) + + build_manifest = { + **runtime_metadata, + "repository_commit": os.environ.get("IPHOTO_MODEL_BUILD_COMMIT") + or os.environ.get("GITHUB_SHA"), + "workflow_run_id": os.environ.get("GITHUB_RUN_ID"), + "python_runtime": platform.python_version(), + "torch_runtime": torch.__version__, + "torchvision_runtime": torchvision.__version__, + "numeric_equivalence": True, + "numeric_equivalence_rtol": 1e-3, + "numeric_equivalence_atol": 3e-5, + } + if args.build_manifest is not None: + _write_json(args.build_manifest, build_manifest) + print( json.dumps( { @@ -96,5 +168,17 @@ def main() -> int: return 0 +def _git_value(repository: Path, revision: str) -> str: + return subprocess.check_output( + ["git", "-C", str(repository), "rev-parse", revision], + text=True, + ).strip() + + +def _write_json(path: Path, value: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + if __name__ == "__main__": sys.exit(main()) diff --git a/tools/model_release_provenance.py b/tools/model_release_provenance.py new file mode 100644 index 000000000..f55aef776 --- /dev/null +++ b/tools/model_release_provenance.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Validate that a model artifact, workflow run, and Release tag share one builder.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + + +class ProvenanceError(ValueError): + pass + + +def validate_release_provenance( + *, + run: dict[str, Any], + build: dict[str, Any], + builder_commit: str, + expected_repository: str, + expected_workflow_id: int, + expected_workflow_path: str, +) -> str: + try: + run_id = int(run["id"]) + workflow_id = int(run["workflow_id"]) + repository = str(run["repository"]["full_name"]) + except (KeyError, TypeError, ValueError) as exc: + raise ProvenanceError("artifact run metadata is incomplete") from exc + head_sha = str(run.get("head_sha") or "").lower() + if re.fullmatch(r"[0-9a-f]{40}", head_sha) is None: + raise ProvenanceError("artifact run head_sha is invalid") + if str(run.get("conclusion") or "") != "success": + raise ProvenanceError("artifact run did not complete successfully") + if workflow_id != int(expected_workflow_id): + raise ProvenanceError("artifact run belongs to a different workflow") + if str(run.get("path") or "") != expected_workflow_path: + raise ProvenanceError("artifact run workflow path is invalid") + if repository != expected_repository: + raise ProvenanceError("artifact run belongs to a different repository") + if str(builder_commit).lower() != head_sha: + raise ProvenanceError("builder checkout does not match artifact run head_sha") + if str(build.get("repository_commit") or "").lower() != head_sha: + raise ProvenanceError("build manifest repository_commit does not match run head_sha") + if str(build.get("workflow_run_id") or "") != str(run_id): + raise ProvenanceError("build manifest workflow_run_id does not match artifact run") + return head_sha + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-json", required=True, type=Path) + parser.add_argument("--build-manifest", required=True, type=Path) + parser.add_argument("--builder-commit", required=True) + parser.add_argument("--expected-repository", required=True) + parser.add_argument("--expected-workflow-id", required=True, type=int) + parser.add_argument("--expected-workflow-path", required=True) + args = parser.parse_args(argv) + try: + head_sha = validate_release_provenance( + run=json.loads(args.run_json.read_text(encoding="utf-8")), + build=json.loads(args.build_manifest.read_text(encoding="utf-8")), + builder_commit=args.builder_commit, + expected_repository=args.expected_repository, + expected_workflow_id=args.expected_workflow_id, + expected_workflow_path=args.expected_workflow_path, + ) + except (OSError, json.JSONDecodeError, ProvenanceError) as exc: + print(f"model release provenance error: {exc}", file=sys.stderr) + return 2 + print(head_sha) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/startup_benchmark.py b/tools/startup_benchmark.py index 51260bff2..29c147dc0 100755 --- a/tools/startup_benchmark.py +++ b/tools/startup_benchmark.py @@ -56,6 +56,14 @@ "probe_ms", "max_gui_job_ms", "max_post_interactive_gui_stall_ms", + "max_post_recognition_gui_stall_ms", + "interactive_recognition_activation_ms", +) +RESOURCE_SNAPSHOT_NAMES = ( + "interactive", + "recognition_activation", + "recognition_plus_1500ms", + "recognition_plus_5000ms", ) BATCH_CONTEXT_KEYS = ( "revision", @@ -68,6 +76,8 @@ "cache_controlled", "cache_eviction_method", "scenario", + "scenario_env_names", + "library_restored_per_sample", "build_environment_fingerprint", "artifact_sha256", "manifest_source_revision", @@ -216,6 +226,44 @@ def _duration(start: dict[str, Any] | None, end: dict[str, Any] | None) -> float return round(max(0.0, float(end["elapsed_ms"]) - float(start["elapsed_ms"])), 3) +def _nearest_resource_sample( + samples: Sequence[dict[str, Any]], + target_wall_time: float, +) -> dict[str, Any] | None: + if not samples: + return None + return min( + samples, + key=lambda item: abs(float(item["wall_time"]) - float(target_wall_time)), + ) + + +def _resource_snapshots(events: Sequence[dict[str, Any]]) -> dict[str, dict[str, Any] | None]: + samples = [event for event in events if event["stage"] == "launcher.resource_sample"] + interactive = _event(events, "interactive") + activation = next( + (event for event in events if event["stage"] == "recognition.startup.activated"), + None, + ) + targets = { + "interactive": float(interactive["wall_time"]) if interactive is not None else None, + "recognition_activation": ( + float(activation["wall_time"]) if activation is not None else None + ), + "recognition_plus_1500ms": ( + float(activation["wall_time"]) + 1.5 if activation is not None else None + ), + "recognition_plus_5000ms": ( + float(activation["wall_time"]) + 5.0 if activation is not None else None + ), + } + snapshots: dict[str, dict[str, Any] | None] = {} + for name, target in targets.items(): + sample = _nearest_resource_sample(samples, target) if target is not None else None + snapshots[name] = dict(_details(sample)) if sample is not None else None + return snapshots + + def _contains_path(value: Any) -> bool: if isinstance(value, str): return bool(_ABSOLUTE_PATH.search(value)) @@ -262,6 +310,8 @@ def analyse_run(path: Path, *, require_gallery: bool = True) -> dict[str, Any]: previous_elapsed = -1.0 for event in events: + if event["stage"] == "launcher.resource_sample": + continue elapsed = float(event["elapsed_ms"]) if elapsed < previous_elapsed: errors.append("event elapsed_ms is out of order") @@ -350,6 +400,10 @@ def analyse_run(path: Path, *, require_gallery: bool = True) -> dict[str, Any]: probe_finished = next( (event for event in events if event["stage"] == "startup.probe.finished"), None ) + recognition_activated = next( + (event for event in events if event["stage"] == "recognition.startup.activated"), + None, + ) context = next( (event.get("context") for event in events if isinstance(event.get("context"), dict)), @@ -360,6 +414,35 @@ def analyse_run(path: Path, *, require_gallery: bool = True) -> dict[str, Any]: eligible = not errors and (cache_state != "cold" or cache_controlled) if cache_state == "cold" and not cache_controlled: errors.append("cold cache was not controlled; excluded from formal statistics") + scenario = str(context.get("scenario", "")) + if scenario == "recognition-quick-close" and any( + event["stage"] == "recognition.worker.started" for event in events + ): + errors.append("quick-close started a recognition worker") + feature_scoped_ab = ( + "IPHOTO_STARTUP_RECOGNITION_AUTO_START" + in str(context.get("scenario_env_names", "")).split(",") + ) + if scenario in { + "recognition-auto-models-present", + "recognition-auto-missing-models", + "recognition-auto-50k-pending", + } and not feature_scoped_ab: + if recognition_activated is None: + errors.append("recognition activation event is missing") + snapshots = _resource_snapshots(events) + missing_snapshots = [name for name, value in snapshots.items() if value is None] + if missing_snapshots: + errors.append( + "recognition resource snapshots are missing: " + + ", ".join(missing_snapshots) + ) + if feature_scoped_ab and any( + event["stage"] == "recognition.worker.started" + and bool(_details(event).get("startup", False)) + for event in events + ): + errors.append("feature-scoped A/B arm started a startup recognition worker") metrics = { "process_start_app_created_ms": process_app_ms, @@ -393,6 +476,26 @@ def analyse_run(path: Path, *, require_gallery: bool = True) -> dict[str, Any]: "max_post_interactive_gui_stall_ms": ( round(max(post_interactive_durations), 3) if post_interactive_durations else 0.0 ), + "max_post_recognition_gui_stall_ms": ( + round( + max( + ( + float(_details(event).get("duration_ms", 0.0)) + for event in events + if event["stage"] == "startup.gui_job.finished" + and recognition_activated is not None + and float(event["elapsed_ms"]) + >= float(recognition_activated["elapsed_ms"]) + ), + default=0.0, + ), + 3, + ) + ), + "interactive_recognition_activation_ms": _duration( + interactive, + recognition_activated, + ), } terminal_details = _details(terminal) if terminal is not None else {} return { @@ -404,6 +507,7 @@ def analyse_run(path: Path, *, require_gallery: bool = True) -> dict[str, Any]: "error_code": terminal_details.get("code"), "context": context, "metrics": metrics, + "resource_snapshots": _resource_snapshots(events), } @@ -437,6 +541,25 @@ def summarize_profiles(paths: Iterable[Path], *, require_gallery: bool = True) - context = ( eligible[0].get("context", {}) if eligible else (runs[0].get("context", {}) if runs else {}) ) + resource_summary: dict[str, dict[str, dict[str, float | int | None]]] = {} + for snapshot_name in RESOURCE_SNAPSHOT_NAMES: + resource_summary[snapshot_name] = {} + for field in ("cpu_ms", "rss_bytes", "read_bytes", "write_bytes"): + values = [ + float(snapshot[field]) + for run in eligible + if isinstance( + snapshot := run.get("resource_snapshots", {}).get(snapshot_name), + dict, + ) + and snapshot.get(field) is not None + ] + resource_summary[snapshot_name][field] = { + "count": len(values), + "p50": nearest_rank(values, 50), + "p95": nearest_rank(values, 95), + "max": round(max(values), 3) if values else None, + } return { "schema_version": 1, "context": context, @@ -456,6 +579,7 @@ def summarize_profiles(paths: Iterable[Path], *, require_gallery: bool = True) - "terminal_counts": dict(sorted(terminal_counts.items())), "error_codes": dict(sorted(error_codes.items())), "metrics": metric_summary, + "resources": resource_summary, "runs": runs, } @@ -475,7 +599,13 @@ def add(name: str, passed: bool, actual: Any, limit: Any) -> None: matching_context_keys = tuple( key for key in BATCH_CONTEXT_KEYS - if key not in {"revision", "artifact_sha256", "manifest_source_revision"} + if key + not in { + "revision", + "artifact_sha256", + "manifest_source_revision", + "scenario_env_names", + } ) mismatched_context = { key: (baseline_context.get(key), candidate_context.get(key)) @@ -535,6 +665,19 @@ def add(name: str, passed: bool, actual: Any, limit: Any) -> None: mismatched_context or "matched", "same platform/backend/scenario/cache/build context", ) + if str(candidate_context.get("scenario", "")).startswith("recognition-"): + add( + "recognition_policy_arms_are_distinct", + "IPHOTO_STARTUP_RECOGNITION_AUTO_START" + in str(baseline_context.get("scenario_env_names", "")).split(",") + and "IPHOTO_STARTUP_RECOGNITION_AUTO_START" + not in str(candidate_context.get("scenario_env_names", "")).split(","), + ( + baseline_context.get("scenario_env_names"), + candidate_context.get("scenario_env_names"), + ), + "feature-scoped baseline vs automatic candidate", + ) add( "baseline_has_30_eligible_samples", baseline.get("eligible_count", 0) >= 30, @@ -657,6 +800,23 @@ def _summary_markdown(summary: dict[str, Any]) -> str: f"| `{metric}` | {stats.get('count', 0)} | {stats.get('p50')} | " f"{stats.get('p95')} | {stats.get('max')} |" ) + lines.extend( + ( + "", + "## Recognition resource snapshots", + "", + "| Snapshot | Resource | Count | P50 | P95 | Max |", + "|---|---|---:|---:|---:|---:|", + ) + ) + for snapshot_name in RESOURCE_SNAPSHOT_NAMES: + snapshot = summary.get("resources", {}).get(snapshot_name, {}) + for field in ("cpu_ms", "rss_bytes", "read_bytes", "write_bytes"): + stats = snapshot.get(field, {}) + lines.append( + f"| `{snapshot_name}` | `{field}` | {stats.get('count', 0)} | " + f"{stats.get('p50')} | {stats.get('p95')} | {stats.get('max')} |" + ) pending = [] platform_name = str(context.get("platform", "")) architecture = str(context.get("architecture", "")) @@ -700,6 +860,67 @@ def _append_event(path: Path, event: dict[str, Any]) -> None: stream.write(json.dumps(event, ensure_ascii=False) + "\n") +def _wait_with_resource_sampling( + process: subprocess.Popen, + *, + profile_path: Path, + launched_wall: float, + timeout_seconds: float, + interval_ms: int, +) -> tuple[int, bool]: + try: + import psutil + except ImportError as exc: + raise ProfileError("resource sampling requires psutil") from exc + observed = psutil.Process(process.pid) + deadline = time.monotonic() + float(timeout_seconds) + timed_out = False + while process.poll() is None: + try: + memory = observed.memory_info() + cpu = observed.cpu_times() + io_getter = getattr(observed, "io_counters", None) + io = io_getter() if callable(io_getter) else None + details = { + "cpu_ms": round((float(cpu.user) + float(cpu.system)) * 1000.0, 3), + "rss_bytes": int(memory.rss), + "read_bytes": ( + int(getattr(io, "read_bytes")) + if io is not None and hasattr(io, "read_bytes") + else None + ), + "write_bytes": ( + int(getattr(io, "write_bytes")) + if io is not None and hasattr(io, "write_bytes") + else None + ), + } + _append_event( + profile_path, + { + "stage": "launcher.resource_sample", + "elapsed_ms": round((time.time() - launched_wall) * 1000.0, 3), + "wall_time": time.time(), + "pid": process.pid, + "details": details, + }, + ) + except (psutil.Error, OSError): + pass + if time.monotonic() >= deadline: + timed_out = True + process.terminate() + break + time.sleep(max(10, int(interval_ms)) / 1000.0) + if timed_out: + try: + return process.wait(timeout=2.0), True + except subprocess.TimeoutExpired: + process.kill() + return process.wait(timeout=2.0), True + return int(process.wait()), False + + def collect(args: argparse.Namespace) -> int: command = list(args.command) if command and command[0] == "--": @@ -710,6 +931,12 @@ def collect(args: argparse.Namespace) -> int: raise ProfileError("samples must be positive") if not args.confirm_dedicated_library: raise ProfileError("refusing to benchmark without --confirm-dedicated-library") + environment_overrides: dict[str, str] = {} + for declaration in args.set_env: + name, separator, value = str(declaration).partition("=") + if not separator or not name or not re.fullmatch(r"[A-Z][A-Z0-9_]*", name): + raise ProfileError(f"invalid --set-env declaration: {declaration}") + environment_overrides[name] = value build_identity: dict[str, Any] = {} if args.runtime == "packaged": if args.build_manifest is None: @@ -721,8 +948,6 @@ def collect(args: argparse.Namespace) -> int: cwd=args.cwd, ) benchmark_library = args.library.expanduser().resolve() - if not benchmark_library.is_dir(): - raise ProfileError(f"benchmark library is not a directory: {benchmark_library}") controlled = args.cache_state != "cold" or bool(args.confirm_controlled_cold_cache) method = args.cache_eviction_method.strip() or "uncontrolled" if args.cache_state == "cold" and (not controlled or method == "uncontrolled"): @@ -730,7 +955,30 @@ def collect(args: argparse.Namespace) -> int: output_dir = args.output_dir.resolve() output_dir.mkdir(parents=True, exist_ok=True) + library_template = ( + args.library_template.expanduser().resolve() + if args.library_template is not None + else None + ) + if library_template is None: + if not benchmark_library.is_dir(): + raise ProfileError(f"benchmark library is not a directory: {benchmark_library}") + else: + if not args.confirm_template_restore: + raise ProfileError("library template restore requires --confirm-template-restore") + if not library_template.is_dir(): + raise ProfileError(f"library template is not a directory: {library_template}") + if benchmark_library != output_dir / "active-library": + raise ProfileError( + "restored benchmark library must be OUTPUT_DIR/active-library" + ) + if benchmark_library == library_template: + raise ProfileError("benchmark library and template must be different directories") for index in range(1, args.samples + 1): + if library_template is not None: + if benchmark_library.exists(): + shutil.rmtree(benchmark_library) + shutil.copytree(library_template, benchmark_library) run_id = ( f"{args.revision}-{args.scenario}-{args.cache_state}-{index:03d}-{uuid.uuid4().hex[:8]}" ) @@ -754,6 +1002,8 @@ def collect(args: argparse.Namespace) -> int: "cache_controlled": controlled, "cache_eviction_method": method, "scenario": args.scenario, + "scenario_env_names": ",".join(sorted(environment_overrides)), + "library_restored_per_sample": library_template is not None, **build_identity, } launched_wall = time.time() @@ -794,6 +1044,7 @@ def collect(args: argparse.Namespace) -> int: "IPHOTO_SETTINGS_PATH": str(settings_path), } ) + environment.update(environment_overrides) settings_path.write_text( json.dumps({"basic_library_path": str(benchmark_library)}, ensure_ascii=False) + "\n", encoding="utf-8", @@ -811,16 +1062,25 @@ def collect(args: argparse.Namespace) -> int: stdout=stdout, stderr=stderr, ) - try: - return_code = process.wait(timeout=args.timeout_seconds) - except subprocess.TimeoutExpired: - timed_out = True - process.terminate() + if args.sample_resources: + return_code, timed_out = _wait_with_resource_sampling( + process, + profile_path=profile_path, + launched_wall=launched_wall, + timeout_seconds=args.timeout_seconds, + interval_ms=args.resource_sample_interval_ms, + ) + else: try: - return_code = process.wait(timeout=2.0) + return_code = process.wait(timeout=args.timeout_seconds) except subprocess.TimeoutExpired: - process.kill() - return_code = process.wait(timeout=2.0) + timed_out = True + process.terminate() + try: + return_code = process.wait(timeout=2.0) + except subprocess.TimeoutExpired: + process.kill() + return_code = process.wait(timeout=2.0) try: last_elapsed_ms = max(float(event["elapsed_ms"]) for event in load_events(profile_path)) except ProfileError: @@ -854,6 +1114,8 @@ def build_parser() -> argparse.ArgumentParser: collect_parser.add_argument("--revision", required=True) collect_parser.add_argument("--scenario", required=True) collect_parser.add_argument("--library", type=Path, required=True) + collect_parser.add_argument("--library-template", type=Path) + collect_parser.add_argument("--confirm-template-restore", action="store_true") collect_parser.add_argument("--confirm-dedicated-library", action="store_true") collect_parser.add_argument("--runtime", choices=("source", "packaged"), required=True) collect_parser.add_argument("--build-manifest", type=Path) @@ -865,6 +1127,15 @@ def build_parser() -> argparse.ArgumentParser: collect_parser.add_argument("--samples", type=int, default=30) collect_parser.add_argument("--timeout-seconds", type=float, default=30.0) collect_parser.add_argument("--auto-exit-delay-ms", type=int, default=250) + collect_parser.add_argument("--sample-resources", action="store_true") + collect_parser.add_argument("--resource-sample-interval-ms", type=int, default=100) + collect_parser.add_argument( + "--set-env", + action="append", + default=[], + metavar="NAME=VALUE", + help="set a non-secret scenario environment variable in the child process", + ) collect_parser.add_argument("--allow-degraded", action="store_true") collect_parser.add_argument("--cwd", type=Path, default=Path.cwd()) collect_parser.add_argument("--output-dir", type=Path, required=True) diff --git a/tools/validate_dinov2_torchscript.py b/tools/validate_dinov2_torchscript.py new file mode 100644 index 000000000..c176154d6 --- /dev/null +++ b/tools/validate_dinov2_torchscript.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Validate one fixed DINOv2 TorchScript artifact without source code.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("artifact", type=Path) + parser.add_argument("metadata", type=Path) + parser.add_argument("--manifest", type=Path) + parser.add_argument("--metadata-only", action="store_true") + args = parser.parse_args() + + metadata = json.loads(args.metadata.read_text(encoding="utf-8")) + digest = _sha256(args.artifact) + size = args.artifact.stat().st_size + if digest != metadata["torchscript_sha256"] or size != metadata["torchscript_size"]: + raise RuntimeError("TorchScript artifact does not match its release metadata") + if args.manifest is not None: + manifest = json.loads(args.manifest.read_text(encoding="utf-8"))["embedder"] + for key in ( + "artifact_kind", + "cache_schema_version", + "model_name", + "release_tag", + "producer_python_version", + "producer_torch_version", + "producer_torchvision_version", + "torchscript_url", + "torchscript_sha256", + "torchscript_size", + "input_shape", + "output_shape", + ): + if metadata.get(key) != manifest.get(key): + raise RuntimeError(f"release metadata does not match manifest field {key}") + if args.metadata_only: + print(json.dumps({"artifact_sha256": digest, "artifact_size": size}, sort_keys=True)) + return 0 + + import torch + + runtime_version = str(torch.__version__).split("+", 1)[0] + if runtime_version != metadata["producer_torch_version"]: + raise RuntimeError( + f"torch runtime mismatch: {runtime_version} != {metadata['producer_torch_version']}" + ) + + model = torch.jit.load(str(args.artifact), map_location="cpu").eval() + example = torch.zeros(tuple(metadata["input_shape"]), dtype=torch.float32) + with torch.no_grad(): + output = model(example) + if isinstance(output, (list, tuple)): + output = output[0] + if tuple(output.shape) != tuple(metadata["output_shape"]): + raise RuntimeError( + f"output shape mismatch: {tuple(output.shape)} != {tuple(metadata['output_shape'])}" + ) + output_digest = hashlib.sha256(output.detach().cpu().contiguous().numpy().tobytes()).hexdigest() + print( + json.dumps( + { + "artifact_sha256": digest, + "artifact_size": size, + "output_sha256": output_digest, + "output_shape": list(output.shape), + "torch_version": runtime_version, + }, + sort_keys=True, + ) + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main())