Add README with file type guidelines and snake_case naming rule; update .gitignore typo
This commit is contained in:
DatTT127
@@ -0,0 +1,46 @@
|
||||
# Backend Specification
|
||||
|
||||
## Purpose
|
||||
Orchestrates API routers, role checks, Socratic circuit-breaker state evaluations, and coordinates ML inference, telemetry collection, and data persistence.
|
||||
|
||||
## Owner
|
||||
Core Backend Team
|
||||
|
||||
## Boundary
|
||||
FastAPI server, API routers, authentication middleware, circuit breaker engine, report generator, RAG coordinator, ledger logger, and connections to Postgres, S3, Redis, Triton, Qdrant, ladybugDB.
|
||||
|
||||
## Internal Design
|
||||
- Built with FastAPI (Python) and Uvicorn for async HTTP server.
|
||||
- Authentication middleware validates JWT tokens and enforces RBAC (roles: RO_RADIOLOGIST, RO_THERAPIST).
|
||||
- Socratic circuit-breaker engine monitors interaction telemetry (hover duration, decision time, override magnitude) and triggers safety dialogs.
|
||||
- Clinical Report Engine uses ReportLab to generate bilingual PDF reports per Circular 46/2018/TT-BYT.
|
||||
- RAG Coordinator orchestrates Retrieval-Augmented Generation: dense vector lookup in Qdrant, graph traversal in ladybugDB, prompt enrichment, LLM generation on Triton (PhoGPT/MedGemma), and hallucination guarding.
|
||||
- Ledger Logger appends immutable, cryptographically chained audit logs to Postgres via triggers preventing UPDATE/DELETE.
|
||||
- Connections: Postgres (via SQLAlchemy), S3 (via boto3), Redis (via redis-py), Triton (via gRPC), Qdrant (via gRPC/HTTP), ladybugDB (via in-process C++ bindings).
|
||||
- Model weights loaded at startup from internal registry; cached in memory.
|
||||
- API endpoints layered: public clinical (sessions, analysis, reports, feedback) and internal/local safety (explanations, safety, drift, RAG, activations, annotations, ground-truth, escalation, morphology, telemetry).
|
||||
|
||||
## Interface Contract
|
||||
See `bento/backend/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
- frontend
|
||||
|
||||
## Breaking-change Policy
|
||||
See `bento/backend/spec/interface-contract.md`.
|
||||
|
||||
## References
|
||||
- NFR-7 (Real-Time UI Screen Refresh ≤200ms)
|
||||
- NFR-10 (Generative Safety Guardrails)
|
||||
- NFR-11 (Frontline Usability & Training)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25776 (Generate GradCAM & CoT Explanation Panel)
|
||||
- UC-02423 (Log High-Trust Concur Block)
|
||||
- UC_Q2_* (All Quadrant 2 safety workflows)
|
||||
- UC_Q3_* (All Quadrant 3 subservience workflows)
|
||||
- UC_Q4_* (All Quadrant 4 double-blind workflows)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Sections 1.2, 2.1-2.6)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Sections 2.1-2.6)
|
||||
- DATA_ENGINEERING_SPEC.md (Sections 4-12 for domain objects)
|
||||
- CI_CD_DEPLOYMENT_PIPELINE.md (Section 9.2 for docker-compose)
|
||||
@@ -0,0 +1,46 @@
|
||||
# Backend Interface Contract
|
||||
|
||||
## Purpose
|
||||
Orchestrates API routers, role checks, Socratic circuit-breaker state evaluations, and coordinates ML inference, telemetry collection, and data persistence.
|
||||
|
||||
## Owner
|
||||
Core Backend Team
|
||||
|
||||
## Provides
|
||||
- api endpoints (session management, frame upload, analysis jobs, reporting, feedback, safety endpoints)
|
||||
- model inference orchestration (dispatches to Triton, aggregates results)
|
||||
- telemetry collection (edge-based behavioral summaries, audit logs)
|
||||
- data persistence coordination (writes to Postgres, S3, Redis)
|
||||
|
||||
## Consumes
|
||||
- data:storage-spec (Postgres DB, S3 object store, Redis cache)
|
||||
- ml:inference-spec (Triton server for angle, inflammation, segmentation, severity)
|
||||
- knowledge:guideline-spec (Qdrant vector DB, ladybugDB graph DB for grounded explanations)
|
||||
|
||||
## Consumers
|
||||
- frontend
|
||||
|
||||
## Not Directly Consumable
|
||||
- data internals (Postgres tables, S3 object layout, Redis keys)
|
||||
- ml internals (Triton model details, GPU kernels)
|
||||
- knowledge internals (Qdrant vectors, ladybugDB graph)
|
||||
|
||||
## Breaking-change Policy
|
||||
- API versioning via path (e.g., /api/v1/).
|
||||
- Backward compatibility maintained for one minor version.
|
||||
- Deprecation notices issued in release notes.
|
||||
- Model interface changes (input/output tensors) require version bump.
|
||||
|
||||
## References
|
||||
- NFR-7 (Real-Time UI Screen Refresh ≤200ms)
|
||||
- NFR-10 (Generative Safety Guardrails)
|
||||
- NFR-11 (Frontline Usability & Training)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25776 (Generate GradCAM & CoT Explanation Panel)
|
||||
- UC-02423 (Log High-Trust Concur Block)
|
||||
- UC_Q2_* (All Quadrant 2 safety workflows)
|
||||
- UC_Q3_* (All Quadrant 3 subservience workflows)
|
||||
- UC_Q4_* (All Quadrant 4 double-blind workflows)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Sections 1.2, 2.1-2.6)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Sections 2.1-2.6)
|
||||
@@ -0,0 +1,52 @@
|
||||
# Data Model Specification
|
||||
|
||||
## Purpose
|
||||
Manages persistent storage, caching, and object storage for clinical data including patient records, imaging studies, analysis results, and audit trails using PostgreSQL, Redis, and S3 with calibration services.
|
||||
|
||||
## Owner
|
||||
Data / Domain Team
|
||||
|
||||
## Boundary
|
||||
PostgreSQL database clusters, Redis cache instances, S3 buckets/object storage, database connection pools, cache invalidation strategies, and storage lifecycle management policies.
|
||||
|
||||
## Internal Design
|
||||
- PostgreSQL 15 with TimescaleDB extension for temporal data
|
||||
- Schema organization: patient, study, session, analysis, audit, calibration namespaces
|
||||
- Connection pooling via PgBouncer for efficient database access
|
||||
- Redis 7 for session caching, rate limiting, and temporary computation results
|
||||
- S3 bucket structure: raw-imagery, processed-results, exports, backups with lifecycle policies
|
||||
- Encryption-at-rest for sensitive data (PHI) using AES-256-GCM
|
||||
- Automated backups with point-in-time recovery (PITR) capabilities
|
||||
- Read replicas for query distribution and reporting workloads
|
||||
- Calibration service: pixel-to-mm conversion factors stored per device/protocol
|
||||
- Data retention policies: active data (2 years), archived data (7 years), purged data (>7 years)
|
||||
- Migration system using Flyway for schema versioning
|
||||
- Monitoring: query performance, connection pool utilization, replication lag
|
||||
|
||||
## Interface Contract
|
||||
See `bento/data/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
- backend:api-spec (for CRUD operations on patient/study/session data)
|
||||
- backend:ledger-spec (for audit trail storage)
|
||||
- ml:engine-spec (for model weights and training data)
|
||||
- knowledge:spec (for exporting vectorizable content)
|
||||
|
||||
## Breaking-change Policy
|
||||
- Database schema versioning via semantic versioning aligned with API versions
|
||||
- Table/column additions: backward compatible (MINOR version)
|
||||
- Table/column removals or type changes: require MAJOR version with migration path
|
||||
- API contract changes follow backend interface contract policies
|
||||
- Data migration scripts provided for breaking changes
|
||||
- Deprecation notices for schema changes 60 days in advance
|
||||
|
||||
## References
|
||||
- NFR-7 (Data Durability: 99.999999999% annual)
|
||||
- NFR-8 (Recovery Time Objective ≤4 hours)
|
||||
- NFR-9 (Storage Cost Efficiency)
|
||||
- NFR-13 (Audit Trail Immutability)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-92006 (Save Analysis Results)
|
||||
- UC-01580 (Export Study Package)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 2.3)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Section 5.2)
|
||||
@@ -0,0 +1,51 @@
|
||||
# Data Model Interface Contract
|
||||
|
||||
## Purpose
|
||||
Manages persistent storage, caching, and object storage for clinical data including patient records, imaging studies, analysis results, and audit trails using PostgreSQL, Redis, and S3 with calibration services.
|
||||
|
||||
## Owner
|
||||
Data / Domain Team
|
||||
|
||||
## Provides
|
||||
- persistent-storage (ACID-compliant patient/study/session records)
|
||||
- object-storage (binary imagery, masks, overlays, exported reports)
|
||||
- caching-layer (session state, rate limiting counters, temp computation results)
|
||||
- calibration-service (pixel-to-mm conversion factors per device/protocol)
|
||||
- backup-and-recovery (point-in-time recovery, cross-region replication)
|
||||
- data-export-functionality (DICOM, PDF, CSV formats)
|
||||
|
||||
## Consumes
|
||||
- (None - foundational storage layer)
|
||||
|
||||
## Consumers
|
||||
- backend:api-spec (patient/study/session CRUD operations, search)
|
||||
- backend:ledger-spec (audit trail append-only storage)
|
||||
- ml:engine-spec (model artifacts storage/retrieval, training datasets)
|
||||
- knowledge:spec (guideline documents, vectorization source material)
|
||||
- infra:spec (shared PostgreSQL/Redis instances for platform services)
|
||||
|
||||
## Not Directly Consumable
|
||||
- internal table structures beyond published schemas
|
||||
- connection pool tuning parameters
|
||||
- Redis key naming conventions for internal caching
|
||||
- S3 bucket policies beyond published access patterns
|
||||
- backup encryption key management
|
||||
- vacuum/analyze maintenance schedules
|
||||
|
||||
## Breaking-change Policy
|
||||
- Database schema versioning via semantic versioning (MAJOR.MINOR.PATCH aligned with API)
|
||||
- Additive changes (new tables/columns): backward compatible (MINOR version)
|
||||
- Breaking changes (removed columns, type alterations): require MAJOR version
|
||||
- Migration scripts provided for all breaking changes with rollback procedures
|
||||
- Deprecation notices for schema changes issued 90 days in advance
|
||||
- Storage interface changes (S3 prefixes, Redis keys) follow same versioning
|
||||
- Consumers must opt-in to breaking changes via feature flags
|
||||
|
||||
## References
|
||||
- NFR-7 (Data Durability: 99.999999999% annual)
|
||||
- NFR-8 (Recovery Time Objective ≤4 hours)
|
||||
- NFR-9 (Storage Cost Efficiency)
|
||||
- NFR-13 (Audit Trail Immutability)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-92006 (Save Analysis Results)
|
||||
- UC-01580 (Export Study Package)
|
||||
@@ -0,0 +1,41 @@
|
||||
# Frontend Specification
|
||||
|
||||
## Purpose
|
||||
Provides interactive clinical workspace, runs edge models (LiteRT, MediaPipe), handles client-side encryption (WebCrypto), offline sync via IndexedDB/Service Worker, and renders UI viewport with graphics adapter fallback.
|
||||
|
||||
## Owner
|
||||
Web Experience Team
|
||||
|
||||
## Boundary
|
||||
PWA service worker, graphics adapter layer (IGraphicsViewport), local browser storage (IndexedDB), and UI components (React, Zustand).
|
||||
|
||||
## Internal Design
|
||||
- Built as a Single Page Application (SPA) using React with TypeScript.
|
||||
- State managed via Zustand store.
|
||||
- Client-side encryption via WebCrypto API (AES-256-GCM) before local storage.
|
||||
- Offline synchronization via Dexie.js (IndexedDB) and Service Worker that queues actions and retries on reconnection.
|
||||
- Graphics rendering abstracted via IGraphicsViewport interface with WebGLThreeAdapter (Three.js) and CPUSpriteAdapter fallback.
|
||||
- Edge ML executed in Web Workers: DICOM parser (cornerstone-core), LiteRT angle classifier (MobileNetV4), MediaPipe ROI pre-cropper.
|
||||
- UI components render multi-layered canvas, workspace controls, diagnostic ribbons, and explanation panels.
|
||||
- Communication with backend via HTTPS to NGINX gateway, JWT-based authentication, role-based access control (RBAC).
|
||||
|
||||
## Interface Contract
|
||||
See `bento/frontend/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
(None)
|
||||
|
||||
## Breaking-change Policy
|
||||
See `bento/frontend/spec/interface-contract.md`.
|
||||
|
||||
## References
|
||||
- NFR-1 (Collaborative Rendering Speed ≤3s)
|
||||
- NFR-4 (Client Memory Footprint ≤150MB)
|
||||
- NFR-14 (Legacy Hardware Compatibility)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25637 (Expose Pixel-Level Activation Logic)
|
||||
- UC-60739 (Isolate Visual Noise/Artifacts)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Sections 2.4, 3.1, 3.2, 3.3)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 2.4)
|
||||
- PROJECT_VIS.md (Section 3.1, 3.2)
|
||||
@@ -0,0 +1,39 @@
|
||||
# Frontend Interface Contract
|
||||
|
||||
## Purpose
|
||||
Provides interactive clinical workspace, runs edge models (LiteRT, MediaPipe), handles client-side encryption (WebCrypto), offline sync via IndexedDB/Service Worker, and renders UI viewport with graphics adapter fallback.
|
||||
|
||||
## Owner
|
||||
Web Experience Team
|
||||
|
||||
## Provides
|
||||
- ui viewport (interactive canvas, overlays, controls)
|
||||
- edge ml (angle classification, inflammation detection, ROI pre-cropping)
|
||||
- offline sync (local cache, sync queue, background synchronization)
|
||||
|
||||
## Consumes
|
||||
- backend:api-spec (REST API endpoints for session, analysis, reporting, feedback)
|
||||
- knowledge:guideline-spec (GraphRAG pipeline for grounded explanations, evidence arbitration)
|
||||
|
||||
## Consumers
|
||||
(None)
|
||||
|
||||
## Not Directly Consumable
|
||||
- backend internals (e.g., FastAPI route implementations, Triton model details)
|
||||
- knowledge internals (Qdrant vectors, ladybugDB graph structure)
|
||||
- data internals (Postgres schema, S3 object layout)
|
||||
|
||||
## Breaking-change Policy
|
||||
- API versioning via path (e.g., /api/v1/).
|
||||
- Backward compatibility maintained for one minor version.
|
||||
- Deprecation notices issued in release notes.
|
||||
- Breaking changes require major version bump.
|
||||
|
||||
## References
|
||||
- NFR-1 (Collaborative Rendering Speed ≤3s)
|
||||
- NFR-4 (Client Memory Footprint ≤150MB)
|
||||
- NFR-14 (Legacy Hardware Compatibility)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25637 (Expose Pixel-Level Activation Logic)
|
||||
- UC-60739 (Isolate Visual Noise/Artifacts)
|
||||
@@ -0,0 +1,42 @@
|
||||
# Offline Interface Contract
|
||||
|
||||
## Purpose
|
||||
Provides offline caching and synchronization for the PWA frontend, enabling continued operation during network interruptions and seamless sync upon reconnection.
|
||||
|
||||
## Owner
|
||||
Web Experience Team (same as frontend)
|
||||
|
||||
## Parent
|
||||
frontend
|
||||
|
||||
Boundary
|
||||
IndexedDB database via Dexie.js, Service Worker for background sync, and local queue for offline actions.
|
||||
|
||||
## Provides
|
||||
- offline cache (IndexedDB storage of encrypted patient sessions, DICOM frames, annotation vectors)
|
||||
- sync queue (Service Worker-intercepted pending actions)
|
||||
- local persistence (survives browser reloads, network drops)
|
||||
|
||||
## Consumes
|
||||
(None)
|
||||
|
||||
## Consumers
|
||||
- frontend
|
||||
|
||||
## Not Directly Consumable
|
||||
- frontend internals (e.g., React components, Zustand store)
|
||||
- backend internals
|
||||
|
||||
## Breaking-change Policy
|
||||
- Changes to IndexedDB schema require version migration scripts.
|
||||
- Sync protocol changes are backward-compatible; old clients can still sync via tombstone markers.
|
||||
|
||||
## References
|
||||
- NFR-8 (Local Network Fault Tolerance)
|
||||
- NFR-4 (Client Memory Footprint)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25637 (Expose Pixel-Level Activation Logic)
|
||||
- UC-60739 (Isolate Visual Noise/Artifacts)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Section 3.2)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 3.2)
|
||||
@@ -0,0 +1,40 @@
|
||||
# Offline Specification
|
||||
|
||||
## Purpose
|
||||
Provides offline caching and synchronization for the PWA frontend, enabling continued operation during network interruptions and seamless sync upon reconnection.
|
||||
|
||||
## Owner
|
||||
Web Experience Team (same as frontend)
|
||||
|
||||
## Parent
|
||||
frontend
|
||||
|
||||
Boundary
|
||||
IndexedDB database via Dexie.js, Service Worker for background sync, and local queue for offline actions.
|
||||
|
||||
## Internal Design
|
||||
- Dexie.js wrapper around IndexedDB with encrypted stores for patient sessions, frames, and annotation layers.
|
||||
- Service Worker intercepts network requests (fetch, XMLHttpRequest) and caches responses; queues POST/PUT/PATCH actions when offline.
|
||||
- On reconnection, Service Worker processes queued actions in order, with idempotent retries.
|
||||
- Data stored in IndexedDB is encrypted via WebCrypto before writing; decrypted on read.
|
||||
- Schema includes tables: sessions, frames, annotations, audit logs, calibration data.
|
||||
- Versioning handled via Dexie.js version upgrades with migration scripts.
|
||||
|
||||
## Interface Contract
|
||||
See `bento/frontend/subprojects/offline/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
- frontend
|
||||
|
||||
## Breaking-change Policy
|
||||
See `bento/frontend/subprojects/offline/spec/interface-contract.md`.
|
||||
|
||||
## References
|
||||
- NFR-8 (Local Network Fault Tolerance)
|
||||
- NFR-4 (Client Memory Footprint)
|
||||
- UC-48376 (Load Patient Scan Session)
|
||||
- UC-47988 (Review Suggested Synovitis Grade)
|
||||
- UC-25637 (Expose Pixel-Level Activation Logic)
|
||||
- UC-60739 (Isolate Visual Noise/Artifacts)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Section 3.2)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 3.2)
|
||||
@@ -0,0 +1,31 @@
|
||||
# Install Docker image
|
||||
docker network create jenkins
|
||||
|
||||
# install docker-integratable image
|
||||
docker run --name jenkins-docker --rm --detach \
|
||||
-p 8080:8080 -p 50000:50000 \
|
||||
--restart=on-failure \
|
||||
--privileged --network jenkins --network-alias docker \
|
||||
--env DOCKER_TLS_CERTDIR=/certs \
|
||||
--volume jenkins-docker-certs:/certs/client \
|
||||
--volume jenkins-data:/var/jenkins_home \
|
||||
--publish 2376:2376 \
|
||||
docker:dind --storage-driver overlay2 \
|
||||
-v jenkins_home:/var/jenkins_home jenkins/jenkins:lts
|
||||
|
||||
# install the docker images
|
||||
docker run \
|
||||
--name jenkins-blueocean \
|
||||
--restart=on-failure \
|
||||
--detach \
|
||||
--network jenkins \
|
||||
--env DOCKER_HOST=tcp://docker:2376 \
|
||||
--env DOCKER_CERT_PATH=/certs/client \
|
||||
--env DOCKER_TLS_VERIFY=1 \
|
||||
--publish 8080:8080 \
|
||||
--publish 50000:50000 \
|
||||
--volume jenkins-data:/var/jenkins_home \
|
||||
--volume jenkins-docker-certs:/certs/client:ro \
|
||||
jenkins/jenkins:lts
|
||||
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
prometheus:
|
||||
image: prom/prometheus:latest
|
||||
volumes:
|
||||
- ./prometheus.yml:/etc/prometheus/prometheus.yml
|
||||
ports:
|
||||
- "9090:9090"
|
||||
|
||||
grafana:
|
||||
image: grafana/grafana:latest
|
||||
ports:
|
||||
- "3000:3000"
|
||||
environment:
|
||||
- GF_SECURITY_ADMIN_PASSWORD=admin
|
||||
@@ -0,0 +1,12 @@
|
||||
global:
|
||||
scrape_interval: 30s # Poll every 30 seconds instead of hammering it every 5s
|
||||
scrape_timeout: 25s # Give it 25 full seconds to respond before timing out
|
||||
|
||||
scrape_configs:
|
||||
- job_name: 'triton'
|
||||
metrics_path: '/metrics'
|
||||
scheme: 'https'
|
||||
tls_config:
|
||||
insecure_skip_verify: true
|
||||
static_configs:
|
||||
- targets: ['dtj-tran--triton-s3-service-unified-triton-server.modal.run']
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "angle_classify_convnext_tiny"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "angle_classify_densenet"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "angle_classify_efficientnet"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "angle_classify_resnet50"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "angle_classify_swin_v2_s"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,412 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate a Triton ensemble model repository for the VKIST vision pipeline.
|
||||
|
||||
The VKIST design docs define the logical vision flow as:
|
||||
|
||||
1. Angle classification selects the scan view.
|
||||
2. Inflammation detection checks whether synovitis/effusion is present.
|
||||
3. Segmentation models produce anatomical masks for supported view branches.
|
||||
|
||||
The 11 resident Triton models in this repository all consume the same raw image
|
||||
tensor (`input_image`) and return `logits`. Triton ensemble scheduling moves
|
||||
those tensors through the ensemble graph internally, so this generator maps the
|
||||
external client input once and exposes every component model's logits as a
|
||||
terminal ensemble output.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
|
||||
ENSEMBLE_NAME = "my_vision_pipeline_ensemble"
|
||||
MODEL_ROOT_DEFAULT = Path(__file__).resolve().parent
|
||||
OUTPUT_DIR_DEFAULT = Path(__file__).resolve().parent / ENSEMBLE_NAME
|
||||
VERSION_DIR = "1"
|
||||
|
||||
# Ordered by the architecture documents: angle classification -> inflammation
|
||||
# detection -> segmentation. These are the 11 component models already resident
|
||||
# under s3://vkist-ml-model/.
|
||||
ANGLE_CLASSIFICATION_MODELS = [
|
||||
"angle_classify_convnext_tiny",
|
||||
"angle_classify_resnet50",
|
||||
"angle_classify_swin_v2_s",
|
||||
"angle_classify_densenet",
|
||||
"angle_classify_efficientnet",
|
||||
]
|
||||
|
||||
INFLAMMATION_MODELS = [
|
||||
"inflammation_model_efficientnet_b0_ultrasound_2_cls",
|
||||
]
|
||||
|
||||
SEGMENTATION_MODELS = [
|
||||
"segmentation_model_unet_resnet101",
|
||||
"segmentation_model_unet3plus_att",
|
||||
"segmentation_model_post_deeplabv3_resnet101",
|
||||
"segmentation_model_post_deeplabv3",
|
||||
"segmentation_model_post_efficientfeedback",
|
||||
]
|
||||
|
||||
ALL_MODEL_NAMES = [
|
||||
*ANGLE_CLASSIFICATION_MODELS,
|
||||
*INFLAMMATION_MODELS,
|
||||
*SEGMENTATION_MODELS,
|
||||
]
|
||||
|
||||
DEFAULT_INPUT_NAME = "input"
|
||||
DEFAULT_INPUT_DATA_TYPE = "TYPE_UINT8"
|
||||
DEFAULT_INPUT_DIMS = [-1, -1, -1, 3]
|
||||
DEFAULT_MODEL_INPUT_NAME = "input_image"
|
||||
DEFAULT_MODEL_OUTPUT_NAME = "logits"
|
||||
DEFAULT_MODEL_OUTPUT_DATA_TYPE = "TYPE_FP32"
|
||||
DEFAULT_MAX_BATCH_SIZE = 8
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelConfig:
|
||||
"""Parsed Triton config for one resident component model."""
|
||||
|
||||
name: str
|
||||
platform: str
|
||||
max_batch_size: int
|
||||
input_name: str
|
||||
input_data_type: str
|
||||
input_dims: list[int]
|
||||
output_name: str
|
||||
output_data_type: str
|
||||
output_dims: list[int]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EnsembleTensor:
|
||||
"""One ensemble output and its matching internal model-output tensor."""
|
||||
|
||||
model_name: str
|
||||
internal_name: str
|
||||
output_name: str
|
||||
data_type: str
|
||||
dims: list[int]
|
||||
|
||||
|
||||
def parse_int_list(text: str) -> list[int]:
|
||||
"""Parse a Triton dims list such as '[ -1, 7, 512, 512 ]'."""
|
||||
|
||||
return [int(item) for item in re.findall(r"-?\d+", text)]
|
||||
|
||||
|
||||
def parse_scalar(text: str, key: str, default: str | None = None) -> str | None:
|
||||
"""Parse a quoted scalar from pbtxt text."""
|
||||
|
||||
match = re.search(rf"^\s*{re.escape(key)}:\s*\"([^\"]+)\"", text, flags=re.MULTILINE)
|
||||
if match:
|
||||
return match.group(1)
|
||||
return default
|
||||
|
||||
|
||||
def parse_block_fields(block_text: str) -> dict[str, str]:
|
||||
"""Parse the simple scalar fields used by the existing model configs."""
|
||||
|
||||
fields: dict[str, str] = {}
|
||||
for key in ("name", "platform", "data_type"):
|
||||
value = parse_scalar(block_text, key)
|
||||
if value is not None:
|
||||
fields[key] = value
|
||||
return fields
|
||||
|
||||
|
||||
def parse_first_model_config(config_path: Path, fallback_name: str) -> ModelConfig:
|
||||
"""Parse a component model config.pbtxt and fall back to known defaults."""
|
||||
|
||||
text = config_path.read_text(encoding="utf-8")
|
||||
name = parse_scalar(text, "name", fallback_name) or fallback_name
|
||||
platform = parse_scalar(text, "platform", "unknown") or "unknown"
|
||||
max_batch_match = re.search(r"^\s*max_batch_size:\s*(-?\d+)", text, flags=re.MULTILINE)
|
||||
max_batch_size = int(max_batch_match.group(1)) if max_batch_match else DEFAULT_MAX_BATCH_SIZE
|
||||
|
||||
input_match = re.search(r"input\s*\[(.*?)\]\s*output", text, flags=re.DOTALL)
|
||||
output_match = re.search(r"output\s*\[(.*?)\]\s*$", text, flags=re.DOTALL)
|
||||
|
||||
input_fields = parse_block_fields(input_match.group(1)) if input_match else {}
|
||||
output_fields = parse_block_fields(output_match.group(1)) if output_match else {}
|
||||
|
||||
input_dims_match = re.search(r"dims:\s*\[(.*?)\]", input_match.group(1), flags=re.DOTALL) if input_match else None
|
||||
output_dims_match = re.search(r"dims:\s*\[(.*?)\]", output_match.group(1), flags=re.DOTALL) if output_match else None
|
||||
|
||||
input_dims = parse_int_list(input_dims_match.group(1)) if input_dims_match else DEFAULT_INPUT_DIMS
|
||||
output_dims = parse_int_list(output_dims_match.group(1)) if output_dims_match else [4]
|
||||
|
||||
return ModelConfig(
|
||||
name=name,
|
||||
platform=platform,
|
||||
max_batch_size=max_batch_size,
|
||||
input_name=input_fields.get("name", DEFAULT_MODEL_INPUT_NAME),
|
||||
input_data_type=input_fields.get("data_type", DEFAULT_INPUT_DATA_TYPE),
|
||||
input_dims=input_dims,
|
||||
output_name=output_fields.get("name", DEFAULT_MODEL_OUTPUT_NAME),
|
||||
output_data_type=output_fields.get("data_type", DEFAULT_MODEL_OUTPUT_DATA_TYPE),
|
||||
output_dims=output_dims,
|
||||
)
|
||||
|
||||
|
||||
def load_model_configs(model_root: Path, model_names: Iterable[str]) -> dict[str, ModelConfig]:
|
||||
"""Load component configs from the local S3 mirror used by Triton."""
|
||||
|
||||
configs: dict[str, ModelConfig] = {}
|
||||
for model_name in model_names:
|
||||
config_path = model_root / model_name / "config.pbtxt"
|
||||
configs[model_name] = parse_first_model_config(config_path, model_name)
|
||||
return configs
|
||||
|
||||
|
||||
def tensor_name_for_model(model_name: str) -> str:
|
||||
"""Create a stable internal ensemble tensor name for one model."""
|
||||
|
||||
return f"{model_name}_logits"
|
||||
|
||||
|
||||
def output_name_for_model(model_name: str) -> str:
|
||||
"""Create the public ensemble output name for one component model."""
|
||||
|
||||
return model_name.upper()
|
||||
|
||||
|
||||
def ensemble_output_dims(model_config: ModelConfig, max_batch_size: int) -> list[int]:
|
||||
"""Return non-batch output dims for an ensemble output block.
|
||||
|
||||
Existing component configs include a leading `-1` output dim for dynamic
|
||||
batching. Triton ensemble output blocks should declare only non-batch dims
|
||||
when `max_batch_size` is positive.
|
||||
"""
|
||||
|
||||
dims = list(model_config.output_dims)
|
||||
if max_batch_size > 0 and dims and dims[0] == -1:
|
||||
return dims[1:]
|
||||
return dims
|
||||
|
||||
|
||||
def format_dims(dims: Iterable[int]) -> str:
|
||||
"""Format Triton dims as `[ 7, 512, 512 ]`."""
|
||||
|
||||
return "[ " + ", ".join(str(dim) for dim in dims) + " ]"
|
||||
|
||||
|
||||
def build_ensemble_tensors(
|
||||
model_configs: dict[str, ModelConfig],
|
||||
max_batch_size: int,
|
||||
) -> list[EnsembleTensor]:
|
||||
"""Build ordered public outputs for the ensemble config."""
|
||||
|
||||
tensors: list[EnsembleTensor] = []
|
||||
for model_name in ALL_MODEL_NAMES:
|
||||
model_config = model_configs[model_name]
|
||||
tensors.append(
|
||||
EnsembleTensor(
|
||||
model_name=model_name,
|
||||
internal_name=tensor_name_for_model(model_name),
|
||||
output_name=output_name_for_model(model_name),
|
||||
data_type=model_config.output_data_type,
|
||||
dims=ensemble_output_dims(model_config, max_batch_size),
|
||||
)
|
||||
)
|
||||
return tensors
|
||||
|
||||
|
||||
def format_input_block(input_name: str, input_data_type: str, input_dims: list[int]) -> str:
|
||||
"""Render the ensemble input block."""
|
||||
|
||||
return f""" {{
|
||||
name: "{input_name}"
|
||||
data_type: {input_data_type}
|
||||
dims: {format_dims(input_dims)}
|
||||
}}"""
|
||||
|
||||
|
||||
def format_output_block(tensor: EnsembleTensor) -> str:
|
||||
"""Render one ensemble output block."""
|
||||
|
||||
return f""" {{
|
||||
name: "{tensor.output_name}"
|
||||
data_type: {tensor.data_type}
|
||||
dims: {format_dims(tensor.dims)}
|
||||
}}"""
|
||||
|
||||
|
||||
def format_step_block(model_name: str, model_input_name: str, input_value: str, model_output_name: str, output_value: str) -> str:
|
||||
"""Render one Triton ensemble_scheduling step."""
|
||||
|
||||
return f""" {{
|
||||
model_name: "{model_name}"
|
||||
model_version: -1
|
||||
input_map {{
|
||||
key: "{model_input_name}"
|
||||
value: "{input_value}"
|
||||
}}
|
||||
output_map {{
|
||||
key: "{model_output_name}"
|
||||
value: "{output_value}"
|
||||
}}
|
||||
}}"""
|
||||
|
||||
|
||||
def build_config_pbtxt(
|
||||
ensemble_name: str,
|
||||
max_batch_size: int,
|
||||
input_name: str,
|
||||
input_data_type: str,
|
||||
input_dims: list[int],
|
||||
tensors: list[EnsembleTensor],
|
||||
model_configs: dict[str, ModelConfig],
|
||||
) -> str:
|
||||
"""Build the complete Triton ensemble config.pbtxt string."""
|
||||
|
||||
input_blocks = [format_input_block(input_name, input_data_type, input_dims)]
|
||||
output_blocks = [format_output_block(tensor) for tensor in tensors]
|
||||
|
||||
steps: list[str] = []
|
||||
for model_name in ALL_MODEL_NAMES:
|
||||
model_config = model_configs[model_name]
|
||||
tensor = next(item for item in tensors if item.model_name == model_name)
|
||||
steps.append(
|
||||
format_step_block(
|
||||
model_name=model_name,
|
||||
model_input_name=model_config.input_name,
|
||||
input_value=input_name,
|
||||
model_output_name=model_config.output_name,
|
||||
output_value=tensor.internal_name,
|
||||
)
|
||||
)
|
||||
|
||||
sections = [
|
||||
f"name: \"{ensemble_name}\"",
|
||||
"platform: \"ensemble\"",
|
||||
f"max_batch_size: {max_batch_size}",
|
||||
"input [\n" + ",\n".join(input_blocks) + "\n]",
|
||||
"output [\n" + ",\n".join(output_blocks) + "\n]",
|
||||
"ensemble_scheduling {\n step [\n" + ",\n".join(steps) + "\n ]\n}",
|
||||
"",
|
||||
]
|
||||
return "\n".join(sections)
|
||||
|
||||
|
||||
def parse_dims_arg(value: str) -> list[int]:
|
||||
"""Parse a comma-separated dims CLI value."""
|
||||
|
||||
return [int(item.strip()) for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def prepare_output_dir(output_dir: Path, clean: bool) -> None:
|
||||
"""Create or refresh the local Triton model repository directory."""
|
||||
|
||||
if clean and output_dir.exists():
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(output_dir)
|
||||
version_dir = output_dir / VERSION_DIR
|
||||
version_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def generate_ensemble(
|
||||
model_root: Path,
|
||||
output_dir: Path,
|
||||
ensemble_name: str,
|
||||
max_batch_size: int,
|
||||
input_name: str,
|
||||
input_data_type: str,
|
||||
input_dims: list[int],
|
||||
clean: bool,
|
||||
) -> Path:
|
||||
"""Generate the ensemble repository and return the written config path."""
|
||||
|
||||
model_configs = load_model_configs(model_root, ALL_MODEL_NAMES)
|
||||
tensors = build_ensemble_tensors(model_configs, max_batch_size)
|
||||
config_text = build_config_pbtxt(
|
||||
ensemble_name=ensemble_name,
|
||||
max_batch_size=max_batch_size,
|
||||
input_name=input_name,
|
||||
input_data_type=input_data_type,
|
||||
input_dims=input_dims,
|
||||
tensors=tensors,
|
||||
model_configs=model_configs,
|
||||
)
|
||||
|
||||
prepare_output_dir(output_dir, clean=clean)
|
||||
config_path = output_dir / VERSION_DIR / "config.pbtxt"
|
||||
config_path.write_text(config_text, encoding="utf-8")
|
||||
return config_path
|
||||
|
||||
|
||||
def build_arg_parser() -> argparse.ArgumentParser:
|
||||
"""Build CLI arguments for local generation."""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate Triton ensemble config for the VKIST 11-model vision pipeline."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-root",
|
||||
type=Path,
|
||||
default=MODEL_ROOT_DEFAULT,
|
||||
help="Local directory containing the 11 resident model config.pbtxt files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=Path,
|
||||
default=OUTPUT_DIR_DEFAULT,
|
||||
help="Local Triton model repository directory to create.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ensemble-name",
|
||||
default=ENSEMBLE_NAME,
|
||||
help="Triton ensemble model name and S3 top-level folder name.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-batch-size",
|
||||
type=int,
|
||||
default=DEFAULT_MAX_BATCH_SIZE,
|
||||
help="Triton max_batch_size for the ensemble.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-name",
|
||||
default=DEFAULT_INPUT_NAME,
|
||||
help="External client input tensor name.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-data-type",
|
||||
default=DEFAULT_INPUT_DATA_TYPE,
|
||||
help="External client input tensor data type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-dims",
|
||||
default=",".join(str(item) for item in DEFAULT_INPUT_DIMS),
|
||||
help="Comma-separated external input dims, excluding batch dimension.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-clean",
|
||||
action="store_true",
|
||||
help="Do not remove an existing output directory before generation.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""CLI entry point."""
|
||||
|
||||
args = build_arg_parser().parse_args()
|
||||
config_path = generate_ensemble(
|
||||
model_root=args.model_root,
|
||||
output_dir=args.output_dir,
|
||||
ensemble_name=args.ensemble_name,
|
||||
max_batch_size=args.max_batch_size,
|
||||
input_name=args.input_name,
|
||||
input_data_type=args.input_data_type,
|
||||
input_dims=parse_dims_arg(args.input_dims),
|
||||
clean=not args.no_clean,
|
||||
)
|
||||
print(f"Wrote Triton ensemble config: {config_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "inflammation_model_efficientnet_b0_ultrasound_2_cls"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 2 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,215 @@
|
||||
name: "msk_vision_pipeline_ensemble"
|
||||
platform: "ensemble"
|
||||
backend: "ensemble"
|
||||
max_batch_size: 8
|
||||
|
||||
# MODIFY HERE: Declare 2 separate input ports with fixed dimensions; no longer using the flexible -1 dimension
|
||||
input [
|
||||
{
|
||||
name: "input_224"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 224, 224 ]
|
||||
},
|
||||
{
|
||||
name: "input_512"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
|
||||
output [
|
||||
{
|
||||
name: "angle_classify_convnext_tiny_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
},
|
||||
{
|
||||
name: "angle_classify_resnet50_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
},
|
||||
{
|
||||
name: "angle_classify_swin_v2_s_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
},
|
||||
{
|
||||
name: "angle_classify_densenet_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
},
|
||||
{
|
||||
name: "angle_classify_efficientnet_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 4 ]
|
||||
},
|
||||
{
|
||||
name: "inflammation_model_efficientnet_b0_ultrasound_2_cls_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 2 ]
|
||||
},
|
||||
{
|
||||
name: "segmentation_model_unet_resnet101_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
},
|
||||
{
|
||||
name: "segmentation_model_unet3plus_att_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
},
|
||||
{
|
||||
name: "segmentation_model_post_deeplabv3_resnet101_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
},
|
||||
{
|
||||
name: "segmentation_model_post_deeplabv3_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
},
|
||||
{
|
||||
name: "segmentation_model_post_efficientfeedback_logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
|
||||
ensemble_scheduling {
|
||||
step [
|
||||
# ---- 224x224 MODEL GROUP (Processes input data from the input_224 variable) ----
|
||||
{
|
||||
model_name: "angle_classify_convnext_tiny"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "angle_classify_convnext_tiny_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "angle_classify_resnet50"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "angle_classify_resnet50_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "angle_classify_swin_v2_s"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "angle_classify_swin_v2_s_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "angle_classify_densenet"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "angle_classify_densenet_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "angle_classify_efficientnet"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "angle_classify_efficientnet_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "inflammation_model_efficientnet_b0_ultrasound_2_cls"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_224"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "inflammation_model_efficientnet_b0_ultrasound_2_cls_logits"
|
||||
}
|
||||
},
|
||||
# ---- 512x512 MODEL GROUP (Processes input data from the input_512 variable) ----
|
||||
{
|
||||
model_name: "segmentation_model_unet_resnet101"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_512"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "segmentation_model_unet_resnet101_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "segmentation_model_unet3plus_att"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_512"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "segmentation_model_unet3plus_att_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "segmentation_model_post_deeplabv3_resnet101"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_512"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "segmentation_model_post_deeplabv3_resnet101_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "segmentation_model_post_deeplabv3"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_512"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "segmentation_model_post_deeplabv3_logits"
|
||||
}
|
||||
},
|
||||
{
|
||||
model_name: "segmentation_model_post_efficientfeedback"
|
||||
model_version: -1
|
||||
input_map {
|
||||
key: "input_image"
|
||||
value: "input_512"
|
||||
}
|
||||
output_map {
|
||||
key: "logits"
|
||||
value: "segmentation_model_post_efficientfeedback_logits"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "segmentation_model_post_deeplabv3"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "segmentation_model_post_deeplabv3_resnet101"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "segmentation_model_post_efficientfeedback"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "segmentation_model_unet3plus_att"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
name: "segmentation_model_unet_resnet101"
|
||||
platform: "pytorch_libtorch"
|
||||
max_batch_size: 8
|
||||
input [
|
||||
{
|
||||
name: "input_image"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 3, 512, 512 ]
|
||||
}
|
||||
]
|
||||
output [
|
||||
{
|
||||
name: "logits"
|
||||
data_type: TYPE_FP32
|
||||
dims: [ 7, 512, 512 ]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,346 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Upload the generated Triton ensemble repository to AWS S3.
|
||||
|
||||
This script mirrors the local Triton model repository folder into the active
|
||||
VKIST model bucket root:
|
||||
|
||||
s3://vkist-ml-model/my_vision_pipeline_ensemble/
|
||||
|
||||
It uploads every file and also creates zero-byte directory marker objects so the
|
||||
S3 prefix reflects the same nested structure Triton expects in a model
|
||||
repository.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import mimetypes
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
try:
|
||||
import boto3
|
||||
from boto3.s3.transfer import TransferConfig
|
||||
except ImportError: # pragma: no cover - exercised only when boto3 is absent.
|
||||
boto3 = None
|
||||
TransferConfig = None
|
||||
|
||||
|
||||
ENSEMBLE_NAME = "my_vision_pipeline_ensemble"
|
||||
DEFAULT_SOURCE_DIR = Path(__file__).resolve().parent / ENSEMBLE_NAME
|
||||
DEFAULT_BUCKET_URI = "s3://vkist-ml-model/"
|
||||
DEFAULT_TRANSFER_CONFIG = None
|
||||
|
||||
|
||||
def require_env(name: str) -> str:
|
||||
"""Read a required AWS credential/environment value."""
|
||||
|
||||
value = os.environ.get(name)
|
||||
if not value:
|
||||
raise RuntimeError(f"Missing required environment variable: {name}")
|
||||
return value
|
||||
|
||||
|
||||
def parse_s3_uri(uri: str) -> tuple[str, str]:
|
||||
"""Parse an S3 URI into bucket and prefix."""
|
||||
|
||||
if not uri.startswith("s3://"):
|
||||
raise ValueError(f"S3 URI must start with 's3://': {uri}")
|
||||
|
||||
body = uri.removeprefix("s3://").strip()
|
||||
if not body:
|
||||
raise ValueError("S3 URI must include a bucket name.")
|
||||
|
||||
parts = body.split("/", 1)
|
||||
bucket = parts[0]
|
||||
prefix = parts[1] if len(parts) > 1 else ""
|
||||
return bucket, normalize_prefix(prefix)
|
||||
|
||||
|
||||
def normalize_prefix(prefix: str) -> str:
|
||||
"""Ensure an S3 prefix ends with '/' when non-empty."""
|
||||
|
||||
prefix = prefix.strip().replace("\\", "/")
|
||||
if prefix and not prefix.endswith("/"):
|
||||
return prefix + "/"
|
||||
return prefix
|
||||
|
||||
|
||||
def relpath_for(path: Path, root: Path) -> Path:
|
||||
"""Return a POSIX-style relative path under a root directory."""
|
||||
|
||||
return path.relative_to(root).as_posix()
|
||||
|
||||
|
||||
def collect_local_tree(source_dir: Path) -> tuple[list[Path], list[Path]]:
|
||||
"""Collect local files and directories to mirror into S3."""
|
||||
|
||||
if not source_dir.exists():
|
||||
raise FileNotFoundError(f"Source directory does not exist: {source_dir}")
|
||||
if not source_dir.is_dir():
|
||||
raise NotADirectoryError(f"Source path is not a directory: {source_dir}")
|
||||
|
||||
files = [path for path in source_dir.rglob("*") if path.is_file()]
|
||||
directories = [path for path in source_dir.rglob("*") if path.is_dir()]
|
||||
directories.append(source_dir)
|
||||
return sorted(files), sorted(directories, reverse=True)
|
||||
|
||||
|
||||
def file_key_for(source_dir: Path, file_path: Path, prefix: str) -> str:
|
||||
"""Build the S3 object key for a local file."""
|
||||
|
||||
return prefix + relpath_for(file_path, source_dir).replace("\\", "/")
|
||||
|
||||
|
||||
def directory_marker_key_for(source_dir: Path, directory_path: Path, prefix: str) -> str:
|
||||
"""Build the S3 directory marker key for a local directory."""
|
||||
|
||||
if directory_path == source_dir:
|
||||
return prefix
|
||||
return prefix + relpath_for(directory_path, source_dir).replace("\\", "/") + "/"
|
||||
|
||||
|
||||
def content_type_for(path: Path) -> str:
|
||||
"""Guess a safe MIME type for an S3 object."""
|
||||
|
||||
guessed_type, _ = mimetypes.guess_type(str(path))
|
||||
return guessed_type or "application/octet-stream"
|
||||
|
||||
|
||||
def create_s3_client() -> object:
|
||||
"""Create a Boto3 S3 client from local AWS environment variables."""
|
||||
|
||||
if boto3 is None:
|
||||
raise RuntimeError("boto3 is required. Install it with: pip install boto3")
|
||||
|
||||
access_key = require_env("AWS_ACCESS_KEY_ID")
|
||||
secret_key = require_env("AWS_SECRET_ACCESS_KEY")
|
||||
|
||||
client_kwargs = {
|
||||
"aws_access_key_id": access_key,
|
||||
"aws_secret_access_key": secret_key,
|
||||
}
|
||||
|
||||
session_token = os.environ.get("AWS_SESSION_TOKEN")
|
||||
if session_token:
|
||||
client_kwargs["aws_session_token"] = session_token
|
||||
|
||||
region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
|
||||
if region:
|
||||
client_kwargs["region_name"] = region
|
||||
|
||||
endpoint_url = os.environ.get("AWS_ENDPOINT_URL")
|
||||
if endpoint_url:
|
||||
client_kwargs["endpoint_url"] = endpoint_url
|
||||
|
||||
return boto3.client("s3", **client_kwargs)
|
||||
|
||||
|
||||
def put_directory_marker(
|
||||
s3_client: object,
|
||||
bucket: str,
|
||||
key: str,
|
||||
dry_run: bool = False,
|
||||
) -> None:
|
||||
"""Create a zero-byte S3 marker object for one directory prefix."""
|
||||
|
||||
if dry_run:
|
||||
print(f"[dry-run] marker s3://{bucket}/{key}")
|
||||
return
|
||||
|
||||
s3_client.put_object(
|
||||
Bucket=bucket,
|
||||
Key=key,
|
||||
Body=b"",
|
||||
ContentType="application/x-directory",
|
||||
)
|
||||
|
||||
|
||||
def upload_file(
|
||||
s3_client: object,
|
||||
bucket: str,
|
||||
local_path: Path,
|
||||
key: str,
|
||||
transfer_config: TransferConfig | None,
|
||||
dry_run: bool = False,
|
||||
) -> None:
|
||||
"""Upload one local file to S3."""
|
||||
|
||||
if dry_run:
|
||||
print(f"[dry-run] upload {local_path} -> s3://{bucket}/{key}")
|
||||
return
|
||||
|
||||
if transfer_config is None:
|
||||
if TransferConfig is None:
|
||||
raise RuntimeError("boto3 is required. Install it with: pip install boto3")
|
||||
transfer_config = TransferConfig(
|
||||
multipart_threshold=64 * 1024 * 1024,
|
||||
multipart_chunksize=16 * 1024 * 1024,
|
||||
max_concurrency=8,
|
||||
use_threads=True,
|
||||
)
|
||||
|
||||
s3_client.upload_file(
|
||||
Filename=str(local_path),
|
||||
Bucket=bucket,
|
||||
Key=key,
|
||||
ExtraArgs={
|
||||
"ContentType": content_type_for(local_path),
|
||||
"CacheControl": "no-store",
|
||||
},
|
||||
Config=transfer_config,
|
||||
)
|
||||
|
||||
|
||||
def list_existing_keys(s3_client: object, bucket: str, prefix: str) -> set[str]:
|
||||
"""List all existing object keys below an S3 prefix."""
|
||||
|
||||
paginator = s3_client.get_paginator("list_objects_v2")
|
||||
keys: set[str] = set()
|
||||
|
||||
for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
|
||||
for item in page.get("Contents", []):
|
||||
keys.add(item["Key"])
|
||||
|
||||
return keys
|
||||
|
||||
|
||||
def delete_stale_keys(
|
||||
s3_client: object,
|
||||
bucket: str,
|
||||
prefix: str,
|
||||
desired_keys: Iterable[str],
|
||||
dry_run: bool = False,
|
||||
) -> int:
|
||||
"""Delete objects under the prefix that are not present locally."""
|
||||
|
||||
desired = set(desired_keys)
|
||||
existing = list_existing_keys(s3_client, bucket, prefix)
|
||||
stale = sorted(existing - desired)
|
||||
|
||||
if not stale:
|
||||
return 0
|
||||
|
||||
if dry_run:
|
||||
for key in stale:
|
||||
print(f"[dry-run] delete s3://{bucket}/{key}")
|
||||
return len(stale)
|
||||
|
||||
for key in stale:
|
||||
s3_client.delete_object(Bucket=bucket, Key=key)
|
||||
|
||||
return len(stale)
|
||||
|
||||
|
||||
def mirror_to_s3(
|
||||
source_dir: Path,
|
||||
bucket_uri: str,
|
||||
prefix: str | None = None,
|
||||
delete: bool = False,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Mirror a local Triton model repository directory into S3."""
|
||||
|
||||
bucket, bucket_prefix = parse_s3_uri(bucket_uri)
|
||||
effective_prefix = normalize_prefix(prefix if prefix is not None else source_dir.name + "/")
|
||||
s3_prefix = bucket_prefix + effective_prefix
|
||||
|
||||
files, directories = collect_local_tree(source_dir)
|
||||
s3_client = create_s3_client() if (not dry_run) or delete else None
|
||||
|
||||
desired_keys: set[str] = set()
|
||||
|
||||
for directory_path in directories:
|
||||
key = directory_marker_key_for(source_dir, directory_path, s3_prefix)
|
||||
desired_keys.add(key)
|
||||
put_directory_marker(s3_client, bucket, key, dry_run=dry_run)
|
||||
|
||||
for file_path in files:
|
||||
key = file_key_for(source_dir, file_path, s3_prefix)
|
||||
desired_keys.add(key)
|
||||
upload_file(
|
||||
s3_client=s3_client,
|
||||
bucket=bucket,
|
||||
local_path=file_path,
|
||||
key=key,
|
||||
transfer_config=DEFAULT_TRANSFER_CONFIG,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
|
||||
deleted = 0
|
||||
if delete:
|
||||
deleted = delete_stale_keys(
|
||||
s3_client=s3_client,
|
||||
bucket=bucket,
|
||||
prefix=s3_prefix,
|
||||
desired_keys=desired_keys,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
|
||||
return {
|
||||
"files_uploaded": len(files),
|
||||
"directories_marked": len(directories),
|
||||
"keys_desired": len(desired_keys),
|
||||
"keys_deleted": deleted,
|
||||
}
|
||||
|
||||
|
||||
def build_arg_parser() -> argparse.ArgumentParser:
|
||||
"""Build CLI arguments for S3 upload."""
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Mirror a generated Triton ensemble repository to AWS S3."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source-dir",
|
||||
type=Path,
|
||||
default=DEFAULT_SOURCE_DIR,
|
||||
help="Local generated Triton model repository directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bucket-uri",
|
||||
default=DEFAULT_BUCKET_URI,
|
||||
help="S3 model bucket root URI, for example s3://vkist-ml-model/.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefix",
|
||||
default=None,
|
||||
help="Optional S3 prefix under the bucket. Defaults to the source folder name.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--delete",
|
||||
action="store_true",
|
||||
help="Delete stale objects under the target prefix that are missing locally.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Print S3 actions without writing or deleting objects.",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""CLI entry point."""
|
||||
|
||||
args = build_arg_parser().parse_args()
|
||||
summary = mirror_to_s3(
|
||||
source_dir=args.source_dir,
|
||||
bucket_uri=args.bucket_uri,
|
||||
prefix=args.prefix,
|
||||
delete=args.delete,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
|
||||
print(
|
||||
"Uploaded ensemble mirror: "
|
||||
f"files={summary['files_uploaded']}, "
|
||||
f"directories={summary['directories_marked']}, "
|
||||
f"desired_keys={summary['keys_desired']}, "
|
||||
f"deleted={summary['keys_deleted']}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,29 @@
|
||||
# Create folder
|
||||
|
||||
# aws s3api put-object --bucket vkist-ml-model --key angle_classify_densenet/
|
||||
aws s3api put-object --bucket vkist-ml-model --key angle_classify_efficientnet/ # best_efficientnet_b2.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key angle_classify_resnet50/ # best_resnet50.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key angle_classify_swin_v2_s/ # best_swin_v2_s.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key segmentation_model_post_deeplabv3_resnet101/ # best_model_deeplabv3_resnet101_seed_16.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key segmentation_model_post_deeplabv3/ # best_model_Deeplav3.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key segmentation_model_post_efficientfeedback/ # efficientfeedback.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key segmentation_model_unet_resnet101/ # unet_resnet101.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key segmentation_model_unet3plus_att/ # unet3plus_att.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key inflammation_model_efficientnet_b0_ultrasound_2_cls/ # efficientnet_b0_ultrasound_2_class.pth
|
||||
aws s3api put-object --bucket vkist-ml-model --key msk_vision_pipeline_ensemble
|
||||
|
||||
# upload model
|
||||
aws s3 mv s3://vkist-ml-model/best_densenet.pth s3://vkist-ml-model/angle_classify_densenet/best_densenet.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_efficientnet_b2.pth s3://vkist-ml-model/angle_classify_efficientnet/best_efficientnet_b2.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_model_deeplabv3_resnet101_seed_16.pth s3://vkist-ml-model/segmentation_model_post_deeplabv3_resnet101/best_model_deeplabv3_resnet101_seed_16.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_model_Deeplav3.pth s3://vkist-ml-model/segmentation_model_post_deeplabv3/best_model_Deeplav3.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_resnet50.pth s3://vkist-ml-model/angle_classify_resnet50/best_resnet50.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_swin_v2_s.pth s3://vkist-ml-model/angle_classify_swin_v2_s/best_swin_v2_s.pth
|
||||
aws s3 mv s3://vkist-ml-model/efficientfeedback.pth s3://vkist-ml-model/segmentation_model_post_efficientfeedback/efficientfeedback.pth
|
||||
aws s3 mv s3://vkist-ml-model/efficientnet_b0_ultrasound_2_class.pth s3://vkist-ml-model/inflammation_model_efficientnet_b0_ultrasound_2_cls/efficientnet_b0_ultrasound_2_class.pth
|
||||
aws s3 mv s3://vkist-ml-model/unet_resnet101.pth s3://vkist-ml-model/segmentation_model_unet_resnet101/unet_resnet101.pth
|
||||
aws s3 mv s3://vkist-ml-model/unet3plus_att.pth s3://vkist-ml-model/segmentation_model_unet3plus_att/unet3plus_att.pth
|
||||
aws s3 mv s3://vkist-ml-model/best_convnext_tiny.pth s3://vkist-ml-model/angle_classify_convnext_tiny/best_convnext_tiny.pth
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
# Classification Models
|
||||
aws s3 mv s3://vkist-ml-model/angle_classify_densenet/best_densenet.pth s3://vkist-ml-model/angle_classify_densenet/1/best_densenet.pth
|
||||
aws s3 mv s3://vkist-ml-model/angle_classify_efficientnet/best_efficientnet_b2.pth s3://vkist-ml-model/angle_classify_efficientnet/1/best_efficientnet_b2.pth
|
||||
aws s3 mv s3://vkist-ml-model/angle_classify_resnet50/best_resnet50.pth s3://vkist-ml-model/angle_classify_resnet50/1/best_resnet50.pth
|
||||
aws s3 mv s3://vkist-ml-model/angle_classify_swin_v2_s/best_swin_v2_s.pth s3://vkist-ml-model/angle_classify_swin_v2_s/1/best_swin_v2_s.pth
|
||||
aws s3 mv s3://vkist-ml-model/angle_classify_convnext_tiny/best_convnext_tiny.pth s3://vkist-ml-model/angle_classify_convnext_tiny/1/best_convnext_tiny.pth
|
||||
aws s3 mv s3://vkist-ml-model/inflammation_model_efficientnet_b0_ultrasound_2_cls/efficientnet_b0_ultrasound_2_class.pth s3://vkist-ml-model/inflammation_model_efficientnet_b0_ultrasound_2_cls/1/efficientnet_b0_ultrasound_2_class.pth
|
||||
|
||||
# Segmentation Models
|
||||
aws s3 mv s3://vkist-ml-model/segmentation_model_post_deeplabv3_resnet101/best_model_deeplabv3_resnet101_seed_16.pth s3://vkist-ml-model/segmentation_model_post_deeplabv3_resnet101/1/best_model_deeplabv3_resnet101_seed_16.pth
|
||||
aws s3 mv s3://vkist-ml-model/segmentation_model_post_deeplabv3/best_model_Deeplav3.pth s3://vkist-ml-model/segmentation_model_post_deeplabv3/1/best_model_Deeplav3.pth
|
||||
aws s3 mv s3://vkist-ml-model/segmentation_model_post_efficientfeedback/efficientfeedback.pth s3://vkist-ml-model/segmentation_model_post_efficientfeedback/1/efficientfeedback.pth
|
||||
aws s3 mv s3://vkist-ml-model/segmentation_model_unet_resnet101/unet_resnet101.pth s3://vkist-ml-model/segmentation_model_unet_resnet101/1/unet_resnet101.pth
|
||||
aws s3 mv s3://vkist-ml-model/segmentation_model_unet3plus_att/unet3plus_att.pth s3://vkist-ml-model/segmentation_model_unet3plus_att/1/unet3plus_att.pth
|
||||
@@ -0,0 +1,29 @@
|
||||
#!/bin/bash
|
||||
|
||||
S3_BUCKET="s3://vkist-ml-model"
|
||||
|
||||
# Set the exact local baseline path where your updated config.pbtxt models reside
|
||||
LOCAL_CONFIG_DIR="/Users/davestran/Downloads/vkist_internship/PILOT_PROJECT/workspace/sprint_1_2/codebase/infra/implementation/s3"
|
||||
|
||||
echo "📤 Syncing local config.pbtxt modifications up to S3 bucket repository..."
|
||||
|
||||
# Classification Configs
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/angle_classify_convnext_tiny/config.pbtxt" "$S3_BUCKET/angle_classify_convnext_tiny/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/angle_classify_densenet/config.pbtxt" "$S3_BUCKET/angle_classify_densenet/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/angle_classify_efficientnet/config.pbtxt" "$S3_BUCKET/angle_classify_efficientnet/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/angle_classify_resnet50/config.pbtxt" "$S3_BUCKET/angle_classify_resnet50/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/angle_classify_swin_v2_s/config.pbtxt" "$S3_BUCKET/angle_classify_swin_v2_s/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/inflammation_model_efficientnet_b0_ultrasound_2_cls/config.pbtxt" "$S3_BUCKET/inflammation_model_efficientnet_b0_ultrasound_2_cls/config.pbtxt"
|
||||
|
||||
# Segmentation Configs
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/segmentation_model_post_deeplabv3/config.pbtxt" "$S3_BUCKET/segmentation_model_post_deeplabv3/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/segmentation_model_post_deeplabv3_resnet101/config.pbtxt" "$S3_BUCKET/segmentation_model_post_deeplabv3_resnet101/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/segmentation_model_post_efficientfeedback/config.pbtxt" "$S3_BUCKET/segmentation_model_post_efficientfeedback/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/segmentation_model_unet3plus_att/config.pbtxt" "$S3_BUCKET/segmentation_model_unet3plus_att/config.pbtxt"
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/segmentation_model_unet_resnet101/config.pbtxt" "$S3_BUCKET/segmentation_model_unet_resnet101/config.pbtxt"
|
||||
|
||||
|
||||
# Ensemble Config
|
||||
aws s3 cp "$LOCAL_CONFIG_DIR/msk_vision_pipeline_ensemble/config.pbtxt" "$S3_BUCKET/msk_vision_pipeline_ensemble/config.pbtxt"
|
||||
|
||||
echo "✅ Configuration files successfully synced to S3 backend!"
|
||||
@@ -0,0 +1,32 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Define the absolute local path to your compiled TorchScript folder
|
||||
LOCAL_DIR="PILOT_PROJECT/workspace/sprint_1_2/codebase/infra/implementation/MODEL_ZIP_PILOT_LT"
|
||||
S3_BUCKET="s3://vkist-ml-model"
|
||||
|
||||
echo "📤 Starting upload workflow of LibTorch binaries to S3 bucket layout..."
|
||||
|
||||
# ==========================================
|
||||
# 1. Classification Models
|
||||
# ==========================================
|
||||
|
||||
echo "🔄 Uploading: Classification models..."
|
||||
aws s3 cp "$LOCAL_DIR/best_densenet.pth" "$S3_BUCKET/angle_classify_densenet/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/best_efficientnet_b2.pth" "$S3_BUCKET/angle_classify_efficientnet/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/best_resnet50.pth" "$S3_BUCKET/angle_classify_resnet50/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/best_swin_v2_s.pth" "$S3_BUCKET/angle_classify_swin_v2_s/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/best_convnext_tiny.pth" "$S3_BUCKET/angle_classify_convnext_tiny/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/efficientnet_b0_ultrasound_2_class.pth" "$S3_BUCKET/inflammation_model_efficientnet_b0_ultrasound_2_cls/1/model.pt"
|
||||
|
||||
# ==========================================
|
||||
# 2. Segmentation Models
|
||||
# ==========================================
|
||||
|
||||
echo "🔄 Uploading: Segmentation models..."
|
||||
aws s3 cp "$LOCAL_DIR/best_model_deeplabv3_resnet101_seed_16.pth" "$S3_BUCKET/segmentation_model_post_deeplabv3_resnet101/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/best_model_Deeplav3.pth" "$S3_BUCKET/segmentation_model_post_deeplabv3/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/efficientfeedback.pth" "$S3_BUCKET/segmentation_model_post_efficientfeedback/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/unet_resnet101.pth" "$S3_BUCKET/segmentation_model_unet_resnet101/1/model.pt"
|
||||
aws s3 cp "$LOCAL_DIR/unet3plus_att.pth" "$S3_BUCKET/segmentation_model_unet3plus_att/1/model.pt"
|
||||
|
||||
echo "🎉 All local LibTorch models compiled down and synchronized with S3 Triton targets successfully!"
|
||||
@@ -0,0 +1,216 @@
|
||||
import subprocess
|
||||
import modal
|
||||
import time
|
||||
|
||||
# run from root: vkist_internship (will manaage later)
|
||||
|
||||
triton_image = (
|
||||
modal.Image.from_registry(
|
||||
tag="nvcr.io/nvidia/tritonserver:24.02-py3",
|
||||
add_python="3.12"
|
||||
)
|
||||
# Step B: Install minimal system dependencies (replacing your apt-get RUN command)
|
||||
.apt_install(
|
||||
"libgl1",
|
||||
"libglib2.0-0" # Crucial runtime hook for OpenCV / Ultralytics
|
||||
)
|
||||
# Step C: Install PyTorch pinned strictly to CUDA 12.1 wheel indices
|
||||
.run_commands(
|
||||
"python3 -m pip install --upgrade pip setuptools",
|
||||
"python3 -m pip install torch==2.5.0 torchaudio==2.5.0 torchvision==0.20.0 --index-url https://download.pytorch.org/whl/cu121",
|
||||
"python3 -m pip install transformers==4.57.3 timm==1.0.22 ultralytics==8.3.0 opencv-python grpcio protobuf",
|
||||
"python3 -m pip install fastapi[standard]",
|
||||
"python3 -m pip install tritonclient[http,cuda]"
|
||||
)
|
||||
)
|
||||
|
||||
app = modal.App("triton-s3-service", image=triton_image)
|
||||
from fastapi import FastAPI, Response, Request,HTTPException
|
||||
from fastapi.responses import StreamingResponse # 👈 ADD THIS IMPORT
|
||||
import httpx
|
||||
web_app = FastAPI()
|
||||
# -------------------------------------------------------------
|
||||
# FASTAPI PROXY ROUTING (Living inside the container)
|
||||
# -------------------------------------------------------------
|
||||
|
||||
@web_app.get("/v2/health/ready")
|
||||
async def forward_health():
|
||||
"""Proxies external HTTP REST calls straight to Triton's internal inference engine"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get("http://127.0.0.1:8000/v2/health/ready")
|
||||
return Response(content=response.content, status_code=response.status_code, media_type=response.headers.get("content-type"))
|
||||
except Exception as e:
|
||||
return Response(content=f"Triton booting models from S3... Error: {str(e)}", status_code=503)
|
||||
|
||||
@web_app.get("/metrics")
|
||||
@web_app.get("/")
|
||||
async def forward_metrics():
|
||||
"""Proxies external metric calls straight to Triton's internal metrics engine"""
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get("http://127.0.0.1:8002/metrics")
|
||||
return Response(content=response.content, status_code=response.status_code, media_type=response.headers.get("content-type"))
|
||||
except Exception as e:
|
||||
return Response(content=f"Waiting for metrics channel... Error: {str(e)}", status_code=503)
|
||||
|
||||
# 👇 ADD THIS CATCH-ALL ROUTE HERE 👇
|
||||
@web_app.api_route("/v2/{path:path}", methods=["GET", "POST"])
|
||||
async def proxy_all_triton_request(path: str, request: Request):
|
||||
import tritonclient.grpc.aio as grpcclient
|
||||
from tritonclient.grpc import service_pb2, service_pb2_grpc
|
||||
from tritonclient.grpc import _utils as grpc_utils #InferenceServerClient
|
||||
import grpc
|
||||
import numpy as np
|
||||
# 1. Keep HTTP proxy ONLY for metadata/health checks
|
||||
if "infer" not in path:
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
url = f"http://127.0.0.1:8000/v2/{path}"
|
||||
headers = dict(request.headers)
|
||||
headers.pop("host", None)
|
||||
triton_response = await client.request(
|
||||
method=request.method, url=url, headers=headers, content=await request.body()
|
||||
)
|
||||
return Response(
|
||||
content=triton_response.content,
|
||||
status_code=triton_response.status_code,
|
||||
headers=dict(triton_response.headers)
|
||||
)
|
||||
|
||||
# 2. 🚀 FOR INFERENCE: Convert the incoming HTTP raw body into a gRPC call
|
||||
if "infer" in path:
|
||||
try:
|
||||
# Extract model name from the route path (e.g., v2/models/MODEL_NAME/infer)
|
||||
parts = path.split("/")
|
||||
model_name = parts[parts.index("models") + 1]
|
||||
|
||||
# Read incoming raw binary HTTP payload
|
||||
raw_http_body = await request.body()
|
||||
|
||||
header_length_str = request.headers.get("Inference-Header-Content-Length")
|
||||
if not header_length_str:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Missing 'Inference-Header-Content-Length' header required for binary Triton transcoding."
|
||||
)
|
||||
|
||||
header_length = int(header_length_str)
|
||||
|
||||
# --- 💥 KSERVE V2 MULTI-PART BODY PARSING ---
|
||||
# Extract the front JSON metadata and the trailing raw binary tensors
|
||||
import json
|
||||
json_bytes = raw_http_body[:header_length]
|
||||
binary_data = raw_http_body[header_length:]
|
||||
request_metadata = json.loads(json_bytes.decode('utf-8'))
|
||||
|
||||
# Setup async gRPC connection
|
||||
triton_url = "127.0.0.1:8001"
|
||||
|
||||
# Configure channels to accept large payload returns (100MB limit override)
|
||||
max_msg_length = 100 * 1024 * 1024
|
||||
channel_options = [
|
||||
('grpc.max_receive_message_length', max_msg_length),
|
||||
('grpc.max_send_message_length', max_msg_length),
|
||||
]
|
||||
async with grpc.aio.insecure_channel(triton_url, options=channel_options) as channel:
|
||||
stub = service_pb2_grpc.GRPCInferenceServiceStub(channel=channel)
|
||||
|
||||
# Construct the native ModelInferRequest protobuf
|
||||
grpc_request = service_pb2.ModelInferRequest()
|
||||
grpc_request.model_name = model_name
|
||||
grpc_request.model_version = ""
|
||||
|
||||
# Populate inputs dynamically from incoming KServe metadata
|
||||
binary_offset = 0
|
||||
for input_tensor in request_metadata.get("inputs", []):
|
||||
# Correct Protobuf repeated field instantiation via .add()
|
||||
infer_input = grpc_request.inputs.add()
|
||||
infer_input.name = input_tensor["name"]
|
||||
infer_input.datatype = input_tensor["datatype"]
|
||||
infer_input.shape.extend(input_tensor["shape"]) # Explicit clean integers!
|
||||
|
||||
# Extract the binary slice matching this tensor out of the raw payload block
|
||||
if "parameters" in input_tensor and "binary_data_size" in input_tensor["parameters"]:
|
||||
data_size = input_tensor["parameters"]["binary_data_size"]
|
||||
grpc_request.raw_input_contents.append(
|
||||
binary_data[binary_offset : binary_offset + data_size]
|
||||
)
|
||||
binary_offset += data_size
|
||||
|
||||
# Request output tensor mappings dynamically based on what the client requested
|
||||
for output_tensor in request_metadata.get("outputs", []):
|
||||
infer_output = grpc_request.outputs.add()
|
||||
infer_output.name = output_tensor["name"]
|
||||
# Signal Triton to return output via raw binary buffers
|
||||
infer_output.parameters["binary_data"].bool_param = True
|
||||
|
||||
# ✅ Send the transcoding payload straight into Triton over internal gRPC loop
|
||||
grpc_response = await stub.ModelInfer(request=grpc_request, timeout=None)
|
||||
|
||||
# --- 💥 TRANSCODE gRPC RESPONSE BACK TO MULTI-PART KSERVE HTTP ---
|
||||
response_metadata = {
|
||||
"model_name": grpc_response.model_name,
|
||||
"model_version": grpc_response.model_version,
|
||||
"outputs": []
|
||||
}
|
||||
|
||||
response_binary_body = b""
|
||||
for i, output in enumerate(grpc_response.outputs):
|
||||
out_desc = {
|
||||
"name": output.name,
|
||||
"datatype": output.datatype,
|
||||
"shape": list(output.shape),
|
||||
"parameters": {
|
||||
"binary_data_size": len(grpc_response.raw_output_contents[i])
|
||||
}
|
||||
}
|
||||
response_metadata["outputs"].append(out_desc)
|
||||
response_binary_body += grpc_response.raw_output_contents[i]
|
||||
|
||||
# Re-bundle into [JSON metadata] + [Raw binary output chunks]
|
||||
response_json_bytes = json.dumps(response_metadata).encode('utf-8')
|
||||
output_http_body = response_json_bytes + response_binary_body
|
||||
|
||||
return Response(
|
||||
content=output_http_body,
|
||||
status_code=200,
|
||||
headers={
|
||||
"Content-Type": "application/octet-stream",
|
||||
"Inference-Header-Content-Length": str(len(response_json_bytes))
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print(f"CRITICAL TRANSLATION EXCEPTION: {traceback.format_exc()}")
|
||||
return Response(
|
||||
content=f"Internal gRPC Pipeline Multiplex Error: {str(e)}",
|
||||
status_code=502
|
||||
)
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# THE UNIFIED SERVICE FUNCTION (1 Container, 1 GPU, 1 Triton Process)
|
||||
# -------------------------------------------------------------
|
||||
|
||||
@app.function(
|
||||
gpu="T4", # for the expense
|
||||
timeout=3600,
|
||||
max_containers=3, # Strict production capping
|
||||
min_containers=1, # for keeping warm and prevention,
|
||||
buffer_containers=2, # Number of additional idle containers to maintain under active load.
|
||||
scaledown_window=30, # Max time (in seconds) a container can remain idle while scaling down.
|
||||
secrets=[modal.Secret.from_name("aws-secrets")]
|
||||
)
|
||||
@modal.asgi_app()
|
||||
def unified_triton_server():
|
||||
print("🚀 Booting ONE Triton Instance inside ONE A100 Container...")
|
||||
|
||||
# Spawns Triton in the background. It will automatically read
|
||||
# your "aws-secrets" environment keys to mount s3://vkist-ml-model/
|
||||
cmd = ["tritonserver", "--model-repository=s3://vkist-ml-model/"]
|
||||
subprocess.Popen(cmd)
|
||||
|
||||
print("📋 Triton background process delegated. Handing routing control over to FastAPI.")
|
||||
|
||||
# Returns immediately! FastAPI now takes over the container lifecycle
|
||||
return web_app
|
||||
@@ -0,0 +1 @@
|
||||
modal deploy PILOT_PROJECT/workspace/sprint_1_2/codebase/infra/implementation/triton_run/modal_triton.py
|
||||
@@ -0,0 +1,41 @@
|
||||
# Infrastructure Specification
|
||||
|
||||
## Purpose
|
||||
Provides platform-level infrastructure services including network routing, high availability, reverse proxy, and foundational resource management to ensure secure, reliable, and observable operation of the VKIST MSK system.
|
||||
|
||||
## Owner
|
||||
Platform Engineering Team
|
||||
|
||||
## Boundary
|
||||
- NGINX reverse proxy (TLS termination, request routing, rate limiting)
|
||||
- Keepalived for VRRP-based high availability and failover
|
||||
- Shared PostgreSQL and Redis instances (coordinated with Data room for provisioning)
|
||||
- System-level monitoring, logging, and alerting foundations
|
||||
- Network segmentation, firewall rules, and VPN/gateway configuration
|
||||
- Infrastructure-as-code (Terraform) for provisioning and lifecycle management
|
||||
|
||||
## Internal Design
|
||||
- NGINX configured as ingress controller with SSL/TLS termination, path-based routing to backend services, and WebSocket support for real-time features.
|
||||
- Keepalived deployed in active-passive mode across cluster nodes, assigning a virtual IP (VIP) for seamless failover.
|
||||
- PostgreSQL and Redis instances are provisioned and managed via Terraform; connection details are exposed as environment variables to consuming rooms.
|
||||
- Foundational logging: structured JSON logs shipped to centralized observability stack (outside scope of this spec).
|
||||
- Security: network policies restrict inter-room communication to declared interfaces; NGINX enforces authentication headers and rate limits.
|
||||
- Observability: exposes Prometheus metrics endpoints; health checks for liveness and readiness.
|
||||
|
||||
## Interface Contract
|
||||
See `infra/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
- All rooms (frontend, backend, ml, data, knowledge) consume networking and availability guarantees.
|
||||
- Backend and ML rooms consume reverse proxy for external API exposure.
|
||||
- Data room consumes shared storage provisioning (Postgres, Redis) for stateful services.
|
||||
|
||||
## Breaking-change Policy
|
||||
See `infra/spec/interface-contract.md`.
|
||||
|
||||
## References
|
||||
- NFR-2 (System Availability ≥99.9% Monthly)
|
||||
- NFR-8 (Network Latency ≤50ms Inter-Region)
|
||||
- NFR-12 (Infrastructure as Code & Immutable Deployments)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Sections 3.1-3.4)
|
||||
- DATA_ENGINEERING_SPEC.md (Section 2 for storage provisioning)
|
||||
@@ -0,0 +1,43 @@
|
||||
# Infrastructure Interface Contract
|
||||
|
||||
## Purpose
|
||||
Provides platform-level infrastructure services including network routing, high availability, reverse proxy, and foundational resource management to ensure secure, reliable, and observable operation of the VKIST MSK system.
|
||||
|
||||
## Owner
|
||||
Platform Engineering Team
|
||||
|
||||
## Provides
|
||||
- network routing and security (NGINX reverse proxy with TLS termination, request routing, rate limiting, WAF)
|
||||
- high availability and failover (Keepalived VRRP, virtual IP, health checks)
|
||||
- foundational resource provisioning (PostgreSQL and Redis connection details via environment variables)
|
||||
- infrastructure observability (Prometheus metrics endpoints, health check endpoints)
|
||||
- foundational logging and monitoring (structured logs, alerting foundations)
|
||||
|
||||
## Consumes
|
||||
- (none) – Infra room provides foundational services; it does not consume other rooms' interfaces for its core purpose.
|
||||
Note: Infra relies on underlying cloud/provider services (VMs, networking, storage) which are outside the scope of this interface contract.
|
||||
|
||||
## Consumers
|
||||
- frontend: consumes network routing and availability for accessing the application.
|
||||
- backend: consumes reverse proxy for external API exposure, HA for service continuity.
|
||||
- ml: consumes reverse proxy for model serving endpoints (Triton), HA for inference reliability.
|
||||
- data: consumes foundational resource provisioning (Postgres, Redis) for stateful services; consumes HA for storage durability.
|
||||
- knowledge: consumes reverse proxy for external access to knowledge services (if exposed), HA for service durability.
|
||||
|
||||
## Not Directly Consumable
|
||||
- internal NGINX configuration details (upstreams, SSL certificates)
|
||||
- Keepalived VRRP configuration and scripts
|
||||
- Terraform state and provider specifics
|
||||
- underlying VM/hardware details
|
||||
|
||||
## Breaking-change Policy
|
||||
- Changes to provided network endpoints (e.g., port, path prefixes) require version bump and backward compatibility period of one release.
|
||||
- Deprecation of any provided interface will be communicated with at least one release notice.
|
||||
- Resource provisioning interface (environment variable names) is considered stable; changes will be backward compatible where possible.
|
||||
|
||||
## References
|
||||
- NFR-2 (System Availability ≥99.9% Monthly)
|
||||
- NFR-8 (Network Latency ≤50ms Inter-Region)
|
||||
- NFR-12 (Infrastructure as Code & Immutable Deployments)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Sections 3.1-3.4)
|
||||
- DATA_ENGINEERING_SPEC.md (Section 2 for storage provisioning)
|
||||
@@ -0,0 +1,50 @@
|
||||
# Knowledge Stack Interface Contract
|
||||
|
||||
## Purpose
|
||||
Provides semantic and graph-based retrieval augmented generation (RAG) for clinical guideline explanations, evidence arbitration, and diagnostic reasoning using vector embeddings (Qdrant) and ontology relationships (ladybugDB) with LLM grounding.
|
||||
|
||||
## Owner
|
||||
Knowledge Engineering Team
|
||||
|
||||
## Provides
|
||||
- guideline embeddings and vector similarity search (Qdrant)
|
||||
- ontology relationships and graph traversal (ladybugDB)
|
||||
- grounded explanation generation (retrieval + LLM + grounding)
|
||||
- evidence arbitration and belief propagation for conflicting evidence
|
||||
- knowledge versioning and temporal validity tracking
|
||||
- hallucination detection and policy enforcement
|
||||
|
||||
## Consumes
|
||||
- (none) – Knowledge stack provides foundational AI services; it does not consume other rooms' interfaces for its core purpose.
|
||||
Note: Knowledge consumes internal storage (Qdrant, ladybugDB) and embedding/LLM services which are part of its boundary.
|
||||
|
||||
## Consumers
|
||||
- frontend: consumes grounded explanations for UI display via `frontend:guideline-spec`.
|
||||
- backend: consumes explanation generation for analysis reporting via `backend:api-spec`.
|
||||
- ml: consumes model activation explanations via `ml:engine-spec`.
|
||||
|
||||
## Not Directly Consumable
|
||||
- internal Qdrant collection names, vector dimensions, and indexing parameters
|
||||
- ladybugDB schema details (predicate names, ontology version)
|
||||
- embedding model specifics (model ID, tensor shapes)
|
||||
- LLM prompt templates and decoding parameters
|
||||
- knowledge curation pipeline details (source ingestion, validation)
|
||||
|
||||
## Breaking-change Policy
|
||||
- Knowledge schema versioning via semantic versioning (MAJOR.MINOR.PATCH)
|
||||
- Embedding model changes: require MAJOR version if dimension or architecture changes
|
||||
- Ontology updates: backward compatible additions (MINOR), breaking changes require MAJOR
|
||||
- LLM interface changes: versioned endpoints with deprecation windows
|
||||
- Knowledge consumers must validate compatibility with new versions
|
||||
- Deprecation notices for breaking changes 60 days in advance
|
||||
- Automated migration tools for knowledge base version upgrades
|
||||
|
||||
## References
|
||||
- NFR-3 (Explanation Latency ≤2s @ 95th percentile)
|
||||
- NFR-6 (Guideline Coverage ≥95% of common synovitis queries)
|
||||
- NFR-10 (Explanation Factuality Score ≥0.9)
|
||||
- UC-25776 (Generate Grounded Explanation for Analysis)
|
||||
- UC-65473 (Resolve Conflicting Evidence)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 3.3)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Section 4.3)
|
||||
- GUIDELINE_SOURCES.md (Appendix B)
|
||||
@@ -0,0 +1,50 @@
|
||||
# Knowledge Stack Specification
|
||||
|
||||
## Purpose
|
||||
Provides semantic and graph-based retrieval augmented generation (RAG) for clinical guideline explanations, evidence arbitration, and diagnostic reasoning using vector embeddings (Qdrant) and ontology relationships (ladybugDB) with LLM grounding.
|
||||
|
||||
## Owner
|
||||
Knowledge Engineering Team
|
||||
|
||||
## Boundary
|
||||
Qdrant vector database instances, ladybugDB graph database instances, embedding model servers, LLM inference endpoints, knowledge curation pipelines, and validation/verification workflows.
|
||||
|
||||
## Internal Design
|
||||
- Hybrid knowledge architecture: vector similarity search + graph traversal
|
||||
- Qdrant: stores guideline section embeddings (BioClinicalBERT, PubMedBERT) with payload metadata
|
||||
- ladybugDB: stores ontology concepts (SNOMED-CT, LOINC, RadLex) and relational axioms
|
||||
- EmbeddingGemma: generates 768-dimension vectors for text chunks
|
||||
- PhoGPT/MedGPT: LLM for answer generation with constrained decoding
|
||||
- Retrieval pipeline: hybrid search (vector + BM25) → graph expansion → reranking
|
||||
- Grounding module: verifies LLM outputs against source guidelines with citation extraction
|
||||
- Arbitration engine: resolves conflicting evidence using belief propagation
|
||||
- Continuous integration: automated guideline ingestion from trusted sources (NIH, CDC, radiology societies)
|
||||
- Versioned knowledge bases with temporal validity tracking
|
||||
- Monitoring: retrieval relevance, grounding accuracy, latency SLOs
|
||||
|
||||
## Interface Contract
|
||||
See `bento/knowledge/spec/interface-contract.md`.
|
||||
|
||||
## Consumers
|
||||
- frontend:guideline-spec (for displaying grounded explanations in UI)
|
||||
- backend:api-spec (for analysis explanation generation)
|
||||
- ml:engine-spec (for generating model activation explanations)
|
||||
|
||||
## Breaking-change Policy
|
||||
- Knowledge schema versioning via semantic versioning
|
||||
- Embedding model changes: require MAJOR version if dimension or architecture changes
|
||||
- Ontology updates: backward compatible additions (MINOR), breaking changes require MAJOR
|
||||
- LLM interface changes: versioned endpoints with deprecation windows
|
||||
- Knowledge consumers must validate compatibility with new versions
|
||||
- Deprecation notices for breaking changes 60 days in advance
|
||||
- Automated migration tools for knowledge base version upgrades
|
||||
|
||||
## References
|
||||
- NFR-3 (Explanation Latency ≤2s @ 95th percentile)
|
||||
- NFR-6 (Guideline Coverage ≥95% of common synovitis queries)
|
||||
- NFR-10 (Explanation Factuality Score ≥0.9)
|
||||
- UC-25776 (Generate Grounded Explanation for Analysis)
|
||||
- UC-65473 (Resolve Conflicting Evidence)
|
||||
- SOLUTION_ARCHITECTURE_SPEC.md (Section 3.3)
|
||||
- SOFTWARE_SYSTEM_DESIGN_FR_25.md (Section 4.3)
|
||||
- GUIDELINE_SOURCES.md (Appendix B)
|
||||
@@ -0,0 +1,232 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
import timm
|
||||
import torch.nn.functional as F
|
||||
from torchinfo import summary
|
||||
from torch.nn import Softmax, Parameter
|
||||
|
||||
|
||||
def convblock(in_channels,out_channels,kernel_size=3,stride=1,dilation=1,padding=1):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=kernel_size,stride=stride,dilation=dilation,padding=padding),
|
||||
nn.BatchNorm2d(out_channels),
|
||||
nn.ReLU()
|
||||
)
|
||||
|
||||
|
||||
class DecoderBlock(nn.Module):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
skip_channels,
|
||||
out_channels,):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
convblock(in_channels=in_channels + skip_channels,out_channels=out_channels,kernel_size=1,padding=0),
|
||||
convblock(in_channels=out_channels,out_channels=out_channels,kernel_size=3,padding=1)
|
||||
)
|
||||
self.conv2 = convblock(in_channels=out_channels,out_channels=out_channels,kernel_size=3,padding=1)
|
||||
|
||||
def forward(self,x,skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear")
|
||||
if skip is not None:
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
class ASPP_module(nn.Module):
|
||||
def __init__(self, inplanes, planes, rate):
|
||||
super(ASPP_module, self).__init__()
|
||||
if rate == 1:
|
||||
kernel_size = 1
|
||||
padding = 0
|
||||
else:
|
||||
kernel_size = 3
|
||||
padding = rate
|
||||
self.atrous_convolution = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
|
||||
stride=1, padding=padding, dilation=rate, bias=False)
|
||||
self.bn = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
self._init_weight()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.atrous_convolution(x)
|
||||
x = self.bn(x)
|
||||
|
||||
return self.relu(x)
|
||||
|
||||
def _init_weight(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(m.weight)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
m.weight.data.fill_(1)
|
||||
m.bias.data.zero_()
|
||||
|
||||
class CAM_Module(nn.Module):
|
||||
def __init__(self):
|
||||
super(CAM_Module, self).__init__()
|
||||
self.gamma = Parameter(torch.zeros(1))
|
||||
self.softmax = Softmax(dim=-1)
|
||||
|
||||
def forward(self, x, fb):
|
||||
batch_size, chnnels, width, height = x.shape
|
||||
proj_query = fb.view(batch_size, chnnels, -1)
|
||||
proj_key = fb.view(batch_size, chnnels, -1).permute(0, 2, 1)
|
||||
energy = torch.bmm(proj_query, proj_key)
|
||||
energy_new = torch.max(energy, -1, keepdim=True)[0].expand_as(energy) - energy
|
||||
attention = self.softmax(energy_new)
|
||||
proj_value = fb.view(batch_size, chnnels, -1)
|
||||
|
||||
out = torch.bmm(attention, proj_value)
|
||||
out = out.view(batch_size, chnnels, height, width)
|
||||
|
||||
return x + self.gamma * out
|
||||
|
||||
def INF(B, H, W):
|
||||
return -torch.diag(torch.tensor(float("inf")).to('cuda').repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)
|
||||
|
||||
class S_Module(nn.Module):
|
||||
def __init__(self, in_dim):
|
||||
super(S_Module, self).__init__()
|
||||
self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
|
||||
self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1)
|
||||
self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
|
||||
self.softmax = Softmax(dim=3)
|
||||
self.INF = INF
|
||||
self.gamma = nn.Parameter(torch.zeros(1))
|
||||
|
||||
def forward(self, x):
|
||||
m_batchsize, _, height, width = x.size()
|
||||
proj_query = self.query_conv(x)
|
||||
proj_query_H = proj_query.permute(0, 3, 1, 2).contiguous().view(m_batchsize * width, -1, height).permute(0, 2,
|
||||
1)
|
||||
proj_query_W = proj_query.permute(0, 2, 1, 3).contiguous().view(m_batchsize * height, -1, width).permute(0, 2,
|
||||
1)
|
||||
proj_key = self.key_conv(x)
|
||||
proj_key_H = proj_key.permute(0, 3, 1, 2).contiguous().view(m_batchsize * width, -1, height)
|
||||
proj_key_W = proj_key.permute(0, 2, 1, 3).contiguous().view(m_batchsize * height, -1, width)
|
||||
proj_value = self.value_conv(x) #D
|
||||
proj_value_H = proj_value.permute(0, 3, 1, 2).contiguous().view(m_batchsize * width, -1, height)
|
||||
proj_value_W = proj_value.permute(0, 2, 1, 3).contiguous().view(m_batchsize * height, -1, width)
|
||||
energy_H = (torch.bmm(proj_query_H, proj_key_H) + self.INF(m_batchsize, height, width)).view(m_batchsize, width,
|
||||
height,
|
||||
height).permute(0,
|
||||
2,
|
||||
1,
|
||||
3)
|
||||
energy_W = torch.bmm(proj_query_W, proj_key_W).view(m_batchsize, height, width, width)
|
||||
concate = self.softmax(torch.cat([energy_H, energy_W], 3))
|
||||
|
||||
self.pau_attention = concate
|
||||
att_H = concate[:, :, :, 0:height].permute(0, 2, 1, 3).contiguous().view(m_batchsize * width, height, height)
|
||||
att_W = concate[:, :, :, height:height + width].contiguous().view(m_batchsize * height, width, width)
|
||||
out_H = torch.bmm(proj_value_H, att_H.permute(0, 2, 1)).view(m_batchsize, width, -1, height).permute(0, 2, 3, 1)
|
||||
out_W = torch.bmm(proj_value_W, att_W.permute(0, 2, 1)).view(m_batchsize, height, -1, width).permute(0, 2, 1, 3)
|
||||
|
||||
return self.gamma * (out_H + out_W) + x
|
||||
|
||||
class FeedbackSpatialAttention(nn.Module):
|
||||
def __init__(self, in_channel,feedback=False) :
|
||||
super().__init__()
|
||||
self.x_att = S_Module(in_dim=in_channel)
|
||||
self.fb_att = S_Module(in_dim=in_channel)
|
||||
self.feedback = feedback
|
||||
self.gamma = nn.Parameter(torch.zeros(1))
|
||||
self.gamma2 = nn.Parameter(torch.zeros(1))
|
||||
|
||||
def forward(self,x,fb=None):
|
||||
x_att = self.gamma*self.x_att(x)
|
||||
if fb!=None:
|
||||
fb_att = self.fb_att(fb)
|
||||
x_att = x_att + self.gamma2*fb_att
|
||||
|
||||
output = x + x_att
|
||||
return output
|
||||
|
||||
class StageAttentionwCAM(nn.Module):
|
||||
def __init__(self,in_channel,out_channel,cbam=False):
|
||||
super().__init__()
|
||||
self.down = nn.MaxPool2d(kernel_size=2,stride=2)
|
||||
self.oneconv = convblock(in_channels=in_channel,out_channels=out_channel,kernel_size=1,padding=0)
|
||||
self.pam = FeedbackSpatialAttention(in_channel=out_channel)
|
||||
self.cam = CAM_Module()
|
||||
|
||||
def forward(self,x,fb=None):
|
||||
if fb!=None:
|
||||
fb = self.down(fb)
|
||||
fb = self.oneconv(fb)
|
||||
out = self.pam(x,fb)
|
||||
out2 = self.cam(x,fb)
|
||||
return out+out2
|
||||
|
||||
class EfficientFeedbackNetwork(nn.Module):
|
||||
def __init__(self, in_channels=3,num_class=3,feedback=False):
|
||||
super().__init__()
|
||||
self.encoder = timm.create_model(model_name='efficientnet_b0',pretrained=True,features_only=True)
|
||||
channel_size = [24,40,112,320]
|
||||
skip_channel = [16,24,40,112]
|
||||
out_channel = [16,24,40,112]
|
||||
|
||||
sa_input = [num_class,16,24,40,112]
|
||||
sa_out = [16,24,40,112,320]
|
||||
|
||||
blocks = [
|
||||
DecoderBlock(in_ch, skip_ch, out_ch)
|
||||
for in_ch, skip_ch, out_ch in zip(channel_size, skip_channel, out_channel)
|
||||
]
|
||||
self.decoder = nn.ModuleList(blocks)
|
||||
attblocks = [
|
||||
StageAttentionwCAM(in_channel=in_c,out_channel=out_c) for in_c,out_c in zip(sa_input,sa_out)
|
||||
]
|
||||
self.attblocks = nn.ModuleList(attblocks)
|
||||
|
||||
rates = [2, 4, 6, 8]
|
||||
self.aspp1 = ASPP_module(320, 100, rate=rates[0])
|
||||
self.aspp2 = ASPP_module(320, 100, rate=rates[1])
|
||||
self.aspp3 = ASPP_module(320, 100, rate=rates[2])
|
||||
self.aspp4 = ASPP_module(320, 100, rate=rates[3])
|
||||
self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
|
||||
nn.Conv2d(320, 100, 1, stride=1, bias=False),
|
||||
nn.BatchNorm2d(100),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.ref_aspp = nn.Conv2d(500, 320, 1, bias=False)
|
||||
|
||||
self.head = nn.Sequential(
|
||||
nn.Upsample(scale_factor=2,mode='bilinear',align_corners=True),
|
||||
nn.Conv2d(in_channels=16,out_channels=num_class,kernel_size=1)
|
||||
)
|
||||
self.feedback = feedback
|
||||
if feedback:
|
||||
strides = [2,4,8,16]
|
||||
|
||||
def forward(self,x,fb=None):
|
||||
encoder = self.encoder(x)
|
||||
if fb!=None:
|
||||
for i in range(len(encoder)-1):
|
||||
if i==0:
|
||||
encoder[i] = self.attblocks[i](encoder[i],fb)
|
||||
else:
|
||||
encoder[i] = self.attblocks[i](encoder[i],encoder[i-1])
|
||||
aspp1 = self.aspp1(encoder[-1])
|
||||
aspp2 = self.aspp2(encoder[-1])
|
||||
aspp3 = self.aspp3(encoder[-1])
|
||||
aspp4 = self.aspp4(encoder[-1])
|
||||
aspp5 = self.global_avg_pool(encoder[-1])
|
||||
aspp5 = F.interpolate(aspp5, size=aspp4.size()[2:], mode='bilinear', align_corners=True)
|
||||
aspp_all = torch.cat([aspp1,aspp2,aspp3,aspp4,aspp5],dim=1)
|
||||
|
||||
dec0 = self.ref_aspp(aspp_all)
|
||||
if fb!=None:
|
||||
dec0 = self.attblocks[-1](dec0,encoder[-2])
|
||||
for i in range(len(encoder)-2,-1,-1):
|
||||
dec0 = self.decoder[i](dec0,encoder[i])
|
||||
head = self.head(dec0)
|
||||
return head
|
||||
|
||||
|
||||
# if __name__=='__main__':
|
||||
# model = EfficientFeedbackNetwork(num_class=2)
|
||||
# summary(model,((1,3,512,512),(1,2,512,512)))
|
||||
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .build_sam import (
|
||||
build_sam,
|
||||
build_sam_vit_h,
|
||||
build_sam_vit_l,
|
||||
build_sam_vit_b,
|
||||
sam_model_registry,
|
||||
)
|
||||
from .predictor import SamPredictor
|
||||
from .automatic_mask_generator import SamAutomaticMaskGenerator
|
||||
@@ -0,0 +1,383 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.ops.boxes import batched_nms, box_area # type: ignore
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .modeling import Sam
|
||||
from .predictor import SamPredictor
|
||||
from .utils.amg import (
|
||||
MaskData,
|
||||
area_from_rle,
|
||||
batch_iterator,
|
||||
batched_mask_to_box,
|
||||
box_xyxy_to_xywh,
|
||||
build_all_layer_point_grids,
|
||||
calculate_stability_score,
|
||||
coco_encode_rle,
|
||||
generate_crop_boxes,
|
||||
is_box_near_crop_edge,
|
||||
mask_to_rle_pytorch,
|
||||
remove_small_regions,
|
||||
rle_to_mask,
|
||||
uncrop_boxes_xyxy,
|
||||
uncrop_masks,
|
||||
uncrop_points,
|
||||
)
|
||||
|
||||
|
||||
class SamAutomaticMaskGenerator:
|
||||
def __init__(
|
||||
self,
|
||||
model: Sam,
|
||||
points_per_side: Optional[int] = 32,
|
||||
points_per_batch: int = 64,
|
||||
pred_iou_thresh: float = 0.88,
|
||||
stability_score_thresh: float = 0.95,
|
||||
stability_score_offset: float = 1.0,
|
||||
box_nms_thresh: float = 0.7,
|
||||
crop_n_layers: int = 0,
|
||||
crop_nms_thresh: float = 0.7,
|
||||
crop_overlap_ratio: float = 512 / 1500,
|
||||
crop_n_points_downscale_factor: int = 1,
|
||||
point_grids: Optional[List[np.ndarray]] = None,
|
||||
min_mask_region_area: int = 0,
|
||||
output_mode: str = "binary_mask",
|
||||
) -> None:
|
||||
"""
|
||||
Using a SAM model, generates masks for the entire image.
|
||||
Generates a grid of point prompts over the image, then filters
|
||||
low quality and duplicate masks. The default settings are chosen
|
||||
for SAM with a ViT-H backbone.
|
||||
|
||||
Arguments:
|
||||
model (Sam): The SAM model to use for mask prediction.
|
||||
points_per_side (int or None): The number of points to be sampled
|
||||
along one side of the image. The total number of points is
|
||||
points_per_side**2. If None, 'point_grids' must provide explicit
|
||||
point sampling.
|
||||
points_per_batch (int): Sets the number of points run simultaneously
|
||||
by the model. Higher numbers may be faster but use more GPU memory.
|
||||
pred_iou_thresh (float): A filtering threshold in [0,1], using the
|
||||
model's predicted mask quality.
|
||||
stability_score_thresh (float): A filtering threshold in [0,1], using
|
||||
the stability of the mask under changes to the cutoff used to binarize
|
||||
the model's mask predictions.
|
||||
stability_score_offset (float): The amount to shift the cutoff when
|
||||
calculated the stability score.
|
||||
box_nms_thresh (float): The box IoU cutoff used by non-maximal
|
||||
suppression to filter duplicate masks.
|
||||
crop_n_layers (int): If >0, mask prediction will be run again on
|
||||
crops of the image. Sets the number of layers to run, where each
|
||||
layer has 2**i_layer number of image crops.
|
||||
crop_nms_thresh (float): The box IoU cutoff used by non-maximal
|
||||
suppression to filter duplicate masks between different crops.
|
||||
crop_overlap_ratio (float): Sets the degree to which crops overlap.
|
||||
In the first crop layer, crops will overlap by this fraction of
|
||||
the image length. Later layers with more crops scale down this overlap.
|
||||
crop_n_points_downscale_factor (int): The number of points-per-side
|
||||
sampled in layer n is scaled down by crop_n_points_downscale_factor**n.
|
||||
point_grids (list(np.ndarray) or None): A list over explicit grids
|
||||
of points used for sampling, normalized to [0,1]. The nth grid in the
|
||||
list is used in the nth crop layer. Exclusive with points_per_side.
|
||||
min_mask_region_area (int): If >0, postprocessing will be applied
|
||||
to remove disconnected regions and holes in masks with area smaller
|
||||
than min_mask_region_area. Requires opencv.
|
||||
output_mode (str): The form masks are returned in. Can be 'binary_mask',
|
||||
'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.
|
||||
For large resolutions, 'binary_mask' may consume large amounts of
|
||||
memory.
|
||||
"""
|
||||
|
||||
assert (points_per_side is None) != (
|
||||
point_grids is None
|
||||
), "Exactly one of points_per_side or point_grid must be provided."
|
||||
if points_per_side is not None:
|
||||
self.point_grids = build_all_layer_point_grids(
|
||||
points_per_side,
|
||||
crop_n_layers,
|
||||
crop_n_points_downscale_factor,
|
||||
)
|
||||
elif point_grids is not None:
|
||||
self.point_grids = point_grids
|
||||
else:
|
||||
raise ValueError("Can't have both points_per_side and point_grid be None.")
|
||||
|
||||
assert output_mode in [
|
||||
"binary_mask",
|
||||
"uncompressed_rle",
|
||||
"coco_rle",
|
||||
], f"Unknown output_mode {output_mode}."
|
||||
if output_mode == "coco_rle":
|
||||
from pycocotools import mask as mask_utils # type: ignore # noqa: F401
|
||||
|
||||
if min_mask_region_area > 0:
|
||||
import cv2 # type: ignore # noqa: F401
|
||||
|
||||
self.predictor = SamPredictor(model)
|
||||
self.points_per_batch = points_per_batch
|
||||
self.pred_iou_thresh = pred_iou_thresh
|
||||
self.stability_score_thresh = stability_score_thresh
|
||||
self.stability_score_offset = stability_score_offset
|
||||
self.box_nms_thresh = box_nms_thresh
|
||||
self.crop_n_layers = crop_n_layers
|
||||
self.crop_nms_thresh = crop_nms_thresh
|
||||
self.crop_overlap_ratio = crop_overlap_ratio
|
||||
self.crop_n_points_downscale_factor = crop_n_points_downscale_factor
|
||||
self.min_mask_region_area = min_mask_region_area
|
||||
self.output_mode = output_mode
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generates masks for the given image.
|
||||
|
||||
Arguments:
|
||||
image (np.ndarray): The image to generate masks for, in HWC uint8 format.
|
||||
|
||||
Returns:
|
||||
list(dict(str, any)): A list over records for masks. Each record is
|
||||
a dict containing the following keys:
|
||||
segmentation (dict(str, any) or np.ndarray): The mask. If
|
||||
output_mode='binary_mask', is an array of shape HW. Otherwise,
|
||||
is a dictionary containing the RLE.
|
||||
bbox (list(float)): The box around the mask, in XYWH format.
|
||||
area (int): The area in pixels of the mask.
|
||||
predicted_iou (float): The model's own prediction of the mask's
|
||||
quality. This is filtered by the pred_iou_thresh parameter.
|
||||
point_coords (list(list(float))): The point coordinates input
|
||||
to the model to generate this mask.
|
||||
stability_score (float): A measure of the mask's quality. This
|
||||
is filtered on using the stability_score_thresh parameter.
|
||||
crop_box (list(float)): The crop of the image used to generate
|
||||
the mask, given in XYWH format.
|
||||
"""
|
||||
|
||||
# Generate masks
|
||||
mask_data = self._generate_masks(image)
|
||||
|
||||
# Filter small disconnected regions and holes in masks
|
||||
if self.min_mask_region_area > 0:
|
||||
mask_data = self.postprocess_small_regions(
|
||||
mask_data,
|
||||
self.min_mask_region_area,
|
||||
max(self.box_nms_thresh, self.crop_nms_thresh),
|
||||
)
|
||||
|
||||
# Encode masks
|
||||
if self.output_mode == "coco_rle":
|
||||
mask_data["segmentations"] = [
|
||||
coco_encode_rle(rle) for rle in mask_data["rles"]
|
||||
]
|
||||
elif self.output_mode == "binary_mask":
|
||||
mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]]
|
||||
else:
|
||||
mask_data["segmentations"] = mask_data["rles"]
|
||||
|
||||
# Write mask records
|
||||
curr_anns = []
|
||||
for idx in range(len(mask_data["segmentations"])):
|
||||
ann = {
|
||||
"segmentation": mask_data["segmentations"][idx],
|
||||
"area": area_from_rle(mask_data["rles"][idx]),
|
||||
"bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),
|
||||
"predicted_iou": mask_data["iou_preds"][idx].item(),
|
||||
"point_coords": [mask_data["points"][idx].tolist()],
|
||||
"stability_score": mask_data["stability_score"][idx].item(),
|
||||
"crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),
|
||||
}
|
||||
curr_anns.append(ann)
|
||||
|
||||
return curr_anns
|
||||
|
||||
def _generate_masks(self, image: np.ndarray) -> MaskData:
|
||||
orig_size = image.shape[:2]
|
||||
crop_boxes, layer_idxs = generate_crop_boxes(
|
||||
orig_size, self.crop_n_layers, self.crop_overlap_ratio
|
||||
)
|
||||
|
||||
# Iterate over image crops
|
||||
data = MaskData()
|
||||
for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
|
||||
crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
|
||||
data.cat(crop_data)
|
||||
|
||||
# Remove duplicate masks between crops
|
||||
if len(crop_boxes) > 1:
|
||||
# Prefer masks from smaller crops
|
||||
scores = 1 / box_area(data["crop_boxes"])
|
||||
scores = scores.to(data["boxes"].device)
|
||||
keep_by_nms = batched_nms(
|
||||
data["boxes"].float(),
|
||||
scores,
|
||||
torch.zeros_like(data["boxes"][:, 0]), # categories
|
||||
iou_threshold=self.crop_nms_thresh,
|
||||
)
|
||||
data.filter(keep_by_nms)
|
||||
|
||||
data.to_numpy()
|
||||
return data
|
||||
|
||||
def _process_crop(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
crop_box: List[int],
|
||||
crop_layer_idx: int,
|
||||
orig_size: Tuple[int, ...],
|
||||
) -> MaskData:
|
||||
# Crop the image and calculate embeddings
|
||||
x0, y0, x1, y1 = crop_box
|
||||
cropped_im = image[y0:y1, x0:x1, :]
|
||||
cropped_im_size = cropped_im.shape[:2]
|
||||
self.predictor.set_image(cropped_im)
|
||||
|
||||
# Get points for this crop
|
||||
points_scale = np.array(cropped_im_size)[None, ::-1]
|
||||
points_for_image = self.point_grids[crop_layer_idx] * points_scale
|
||||
|
||||
# Generate masks for this crop in batches
|
||||
data = MaskData()
|
||||
for (points,) in batch_iterator(self.points_per_batch, points_for_image):
|
||||
batch_data = self._process_batch(
|
||||
points, cropped_im_size, crop_box, orig_size
|
||||
)
|
||||
data.cat(batch_data)
|
||||
del batch_data
|
||||
self.predictor.reset_image()
|
||||
|
||||
# Remove duplicates within this crop.
|
||||
keep_by_nms = batched_nms(
|
||||
data["boxes"].float(),
|
||||
data["iou_preds"],
|
||||
torch.zeros_like(data["boxes"][:, 0]), # categories
|
||||
iou_threshold=self.box_nms_thresh,
|
||||
)
|
||||
data.filter(keep_by_nms)
|
||||
|
||||
# Return to the original image frame
|
||||
data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)
|
||||
data["points"] = uncrop_points(data["points"], crop_box)
|
||||
data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])
|
||||
|
||||
return data
|
||||
|
||||
def _process_batch(
|
||||
self,
|
||||
points: np.ndarray,
|
||||
im_size: Tuple[int, ...],
|
||||
crop_box: List[int],
|
||||
orig_size: Tuple[int, ...],
|
||||
) -> MaskData:
|
||||
orig_h, orig_w = orig_size
|
||||
|
||||
# Run model on this batch
|
||||
transformed_points = self.predictor.transform.apply_coords(points, im_size)
|
||||
in_points = torch.as_tensor(transformed_points, device=self.predictor.device)
|
||||
in_labels = torch.ones(
|
||||
in_points.shape[0], dtype=torch.int, device=in_points.device
|
||||
)
|
||||
masks, iou_preds, _ = self.predictor.predict_torch(
|
||||
in_points[:, None, :],
|
||||
in_labels[:, None],
|
||||
multimask_output=True,
|
||||
return_logits=True,
|
||||
)
|
||||
|
||||
# Serialize predictions and store in MaskData
|
||||
data = MaskData(
|
||||
masks=masks.flatten(0, 1),
|
||||
iou_preds=iou_preds.flatten(0, 1),
|
||||
points=torch.as_tensor(points.repeat(masks.shape[1], axis=0)),
|
||||
)
|
||||
del masks
|
||||
|
||||
# Filter by predicted IoU
|
||||
if self.pred_iou_thresh > 0.0:
|
||||
keep_mask = data["iou_preds"] > self.pred_iou_thresh
|
||||
data.filter(keep_mask)
|
||||
|
||||
# Calculate stability score
|
||||
data["stability_score"] = calculate_stability_score(
|
||||
data["masks"],
|
||||
self.predictor.model.mask_threshold,
|
||||
self.stability_score_offset,
|
||||
)
|
||||
if self.stability_score_thresh > 0.0:
|
||||
keep_mask = data["stability_score"] >= self.stability_score_thresh
|
||||
data.filter(keep_mask)
|
||||
|
||||
# Threshold masks and calculate boxes
|
||||
data["masks"] = data["masks"] > self.predictor.model.mask_threshold
|
||||
data["boxes"] = batched_mask_to_box(data["masks"])
|
||||
|
||||
# Filter boxes that touch crop boundaries
|
||||
keep_mask = ~is_box_near_crop_edge(
|
||||
data["boxes"], crop_box, [0, 0, orig_w, orig_h]
|
||||
)
|
||||
if not torch.all(keep_mask):
|
||||
data.filter(keep_mask)
|
||||
|
||||
# Compress to RLE
|
||||
data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w)
|
||||
data["rles"] = mask_to_rle_pytorch(data["masks"])
|
||||
del data["masks"]
|
||||
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def postprocess_small_regions(
|
||||
mask_data: MaskData, min_area: int, nms_thresh: float
|
||||
) -> MaskData:
|
||||
"""
|
||||
Removes small disconnected regions and holes in masks, then reruns
|
||||
box NMS to remove any new duplicates.
|
||||
|
||||
Edits mask_data in place.
|
||||
|
||||
Requires open-cv as a dependency.
|
||||
"""
|
||||
if len(mask_data["rles"]) == 0:
|
||||
return mask_data
|
||||
|
||||
# Filter small disconnected regions and holes
|
||||
new_masks = []
|
||||
scores = []
|
||||
for rle in mask_data["rles"]:
|
||||
mask = rle_to_mask(rle)
|
||||
|
||||
mask, changed = remove_small_regions(mask, min_area, mode="holes")
|
||||
unchanged = not changed
|
||||
mask, changed = remove_small_regions(mask, min_area, mode="islands")
|
||||
unchanged = unchanged and not changed
|
||||
|
||||
new_masks.append(torch.as_tensor(mask).unsqueeze(0))
|
||||
# Give score=0 to changed masks and score=1 to unchanged masks
|
||||
# so NMS will prefer ones that didn't need postprocessing
|
||||
scores.append(float(unchanged))
|
||||
|
||||
# Recalculate boxes and remove any new duplicates
|
||||
masks = torch.cat(new_masks, dim=0)
|
||||
boxes = batched_mask_to_box(masks)
|
||||
keep_by_nms = batched_nms(
|
||||
boxes.float(),
|
||||
torch.as_tensor(scores),
|
||||
torch.zeros_like(boxes[:, 0]), # categories
|
||||
iou_threshold=nms_thresh,
|
||||
)
|
||||
|
||||
# Only recalculate RLEs for masks that have changed
|
||||
for i_mask in keep_by_nms:
|
||||
if scores[i_mask] == 0.0:
|
||||
mask_torch = masks[i_mask].unsqueeze(0)
|
||||
mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]
|
||||
mask_data["boxes"][i_mask] = boxes[i_mask] # update res directly
|
||||
mask_data.filter(keep_by_nms)
|
||||
|
||||
return mask_data
|
||||
@@ -0,0 +1,149 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
import urllib.request
|
||||
import torch
|
||||
|
||||
from .modeling import (
|
||||
ImageEncoderViT,
|
||||
MaskDecoder,
|
||||
PromptEncoder,
|
||||
Sam,
|
||||
TwoWayTransformer,
|
||||
)
|
||||
|
||||
|
||||
def build_sam_vit_h(checkpoint=None):
|
||||
return _build_sam(
|
||||
encoder_embed_dim=1280,
|
||||
encoder_depth=32,
|
||||
encoder_num_heads=16,
|
||||
encoder_global_attn_indexes=[7, 15, 23, 31],
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
|
||||
build_sam = build_sam_vit_h
|
||||
|
||||
|
||||
def build_sam_vit_l(checkpoint=None):
|
||||
return _build_sam(
|
||||
encoder_embed_dim=1024,
|
||||
encoder_depth=24,
|
||||
encoder_num_heads=16,
|
||||
encoder_global_attn_indexes=[5, 11, 17, 23],
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
|
||||
def build_sam_vit_b(checkpoint=None):
|
||||
return _build_sam(
|
||||
encoder_embed_dim=768,
|
||||
encoder_depth=12,
|
||||
encoder_num_heads=12,
|
||||
encoder_global_attn_indexes=[2, 5, 8, 11],
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
|
||||
sam_model_registry = {
|
||||
"default": build_sam_vit_h,
|
||||
"vit_h": build_sam_vit_h,
|
||||
"vit_l": build_sam_vit_l,
|
||||
"vit_b": build_sam_vit_b,
|
||||
}
|
||||
|
||||
|
||||
def _build_sam(
|
||||
encoder_embed_dim,
|
||||
encoder_depth,
|
||||
encoder_num_heads,
|
||||
encoder_global_attn_indexes,
|
||||
checkpoint=None,
|
||||
):
|
||||
prompt_embed_dim = 256
|
||||
image_size = 1024
|
||||
vit_patch_size = 16
|
||||
image_embedding_size = image_size // vit_patch_size
|
||||
sam = Sam(
|
||||
image_encoder=ImageEncoderViT(
|
||||
depth=encoder_depth,
|
||||
embed_dim=encoder_embed_dim,
|
||||
img_size=image_size,
|
||||
mlp_ratio=4,
|
||||
norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
|
||||
num_heads=encoder_num_heads,
|
||||
patch_size=vit_patch_size,
|
||||
qkv_bias=True,
|
||||
use_rel_pos=True,
|
||||
global_attn_indexes=encoder_global_attn_indexes,
|
||||
window_size=14,
|
||||
out_chans=prompt_embed_dim,
|
||||
),
|
||||
prompt_encoder=PromptEncoder(
|
||||
embed_dim=prompt_embed_dim,
|
||||
image_embedding_size=(image_embedding_size, image_embedding_size),
|
||||
input_image_size=(image_size, image_size),
|
||||
mask_in_chans=16,
|
||||
),
|
||||
mask_decoder=MaskDecoder(
|
||||
num_multimask_outputs=3,
|
||||
transformer=TwoWayTransformer(
|
||||
depth=2,
|
||||
embedding_dim=prompt_embed_dim,
|
||||
mlp_dim=2048,
|
||||
num_heads=8,
|
||||
),
|
||||
transformer_dim=prompt_embed_dim,
|
||||
iou_head_depth=3,
|
||||
iou_head_hidden_dim=256,
|
||||
),
|
||||
pixel_mean=[123.675, 116.28, 103.53],
|
||||
pixel_std=[58.395, 57.12, 57.375],
|
||||
)
|
||||
sam.eval()
|
||||
checkpoint = Path(checkpoint)
|
||||
if checkpoint.name == "sam_vit_b_01ec64.pth" and not checkpoint.exists():
|
||||
cmd = input("Download sam_vit_b_01ec64.pth from facebook AI? [y]/n: ")
|
||||
if len(cmd) == 0 or cmd.lower() == "y":
|
||||
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
||||
print("Downloading SAM ViT-B checkpoint...")
|
||||
urllib.request.urlretrieve(
|
||||
"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
|
||||
checkpoint,
|
||||
)
|
||||
print(checkpoint.name, " is downloaded!")
|
||||
elif checkpoint.name == "sam_vit_h_4b8939.pth" and not checkpoint.exists():
|
||||
cmd = input("Download sam_vit_h_4b8939.pth from facebook AI? [y]/n: ")
|
||||
if len(cmd) == 0 or cmd.lower() == "y":
|
||||
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
||||
print("Downloading SAM ViT-H checkpoint...")
|
||||
urllib.request.urlretrieve(
|
||||
"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
|
||||
checkpoint,
|
||||
)
|
||||
print(checkpoint.name, " is downloaded!")
|
||||
elif checkpoint.name == "sam_vit_l_0b3195.pth" and not checkpoint.exists():
|
||||
cmd = input("Download sam_vit_l_0b3195.pth from facebook AI? [y]/n: ")
|
||||
if len(cmd) == 0 or cmd.lower() == "y":
|
||||
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
||||
print("Downloading SAM ViT-L checkpoint...")
|
||||
urllib.request.urlretrieve(
|
||||
"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
|
||||
checkpoint,
|
||||
)
|
||||
print(checkpoint.name, " is downloaded!")
|
||||
|
||||
if checkpoint is not None:
|
||||
with open(checkpoint, "rb") as f:
|
||||
# state_dict = torch.load(f["model"], map_location=torch.device('cpu'))
|
||||
checkpoint = torch.load(f, map_location=torch.device('cpu'))
|
||||
state_dict = checkpoint["model"]
|
||||
|
||||
sam.load_state_dict(state_dict)
|
||||
return sam
|
||||
@@ -0,0 +1,12 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .sam import Sam
|
||||
from .image_encoder import ImageEncoderViT
|
||||
from .mask_decoder import MaskDecoder
|
||||
from .prompt_encoder import PromptEncoder
|
||||
from .transformer import TwoWayTransformer
|
||||
@@ -0,0 +1,44 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from typing import Type
|
||||
|
||||
|
||||
class MLPBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
mlp_dim: int,
|
||||
act: Type[nn.Module] = nn.GELU,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.lin1 = nn.Linear(embedding_dim, mlp_dim)
|
||||
self.lin2 = nn.Linear(mlp_dim, embedding_dim)
|
||||
self.act = act()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.lin2(self.act(self.lin1(x)))
|
||||
|
||||
|
||||
# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
|
||||
# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa
|
||||
class LayerNorm2d(nn.Module):
|
||||
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(num_channels))
|
||||
self.bias = nn.Parameter(torch.zeros(num_channels))
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None, None] * x + self.bias[:, None, None]
|
||||
return x
|
||||
@@ -0,0 +1,420 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from typing import Optional, Tuple, Type
|
||||
|
||||
from .common import LayerNorm2d, MLPBlock
|
||||
|
||||
|
||||
# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa
|
||||
class ImageEncoderViT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
img_size: int = 1024,
|
||||
patch_size: int = 16,
|
||||
in_chans: int = 3,
|
||||
embed_dim: int = 768,
|
||||
depth: int = 12,
|
||||
num_heads: int = 12,
|
||||
mlp_ratio: float = 4.0,
|
||||
out_chans: int = 256,
|
||||
qkv_bias: bool = True,
|
||||
norm_layer: Type[nn.Module] = nn.LayerNorm,
|
||||
act_layer: Type[nn.Module] = nn.GELU,
|
||||
use_abs_pos: bool = True,
|
||||
use_rel_pos: bool = False,
|
||||
rel_pos_zero_init: bool = True,
|
||||
window_size: int = 0,
|
||||
global_attn_indexes: Tuple[int, ...] = (),
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
img_size (int): Input image size.
|
||||
patch_size (int): Patch size.
|
||||
in_chans (int): Number of input image channels.
|
||||
embed_dim (int): Patch embedding dimension.
|
||||
depth (int): Depth of ViT.
|
||||
num_heads (int): Number of attention heads in each ViT block.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
||||
qkv_bias (bool): If True, add a learnable bias to query, key, value.
|
||||
norm_layer (nn.Module): Normalization layer.
|
||||
act_layer (nn.Module): Activation layer.
|
||||
use_abs_pos (bool): If True, use absolute positional embeddings.
|
||||
use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
|
||||
rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
|
||||
window_size (int): Window size for window attention blocks.
|
||||
global_attn_indexes (list): Indexes for blocks using global attention.
|
||||
"""
|
||||
super().__init__()
|
||||
self.img_size = img_size
|
||||
|
||||
self.patch_embed = PatchEmbed(
|
||||
kernel_size=(patch_size, patch_size),
|
||||
stride=(patch_size, patch_size),
|
||||
in_chans=in_chans,
|
||||
embed_dim=embed_dim,
|
||||
)
|
||||
|
||||
self.pos_embed: Optional[nn.Parameter] = None
|
||||
if use_abs_pos:
|
||||
# Initialize absolute positional embedding with pretrain image size.
|
||||
self.pos_embed = nn.Parameter(
|
||||
torch.zeros(
|
||||
1, img_size // patch_size, img_size // patch_size, embed_dim
|
||||
)
|
||||
)
|
||||
|
||||
self.blocks = nn.ModuleList()
|
||||
for i in range(depth):
|
||||
block = Block(
|
||||
dim=embed_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_ratio=mlp_ratio,
|
||||
qkv_bias=qkv_bias,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
use_rel_pos=use_rel_pos,
|
||||
rel_pos_zero_init=rel_pos_zero_init,
|
||||
window_size=window_size if i not in global_attn_indexes else 0,
|
||||
input_size=(img_size // patch_size, img_size // patch_size),
|
||||
)
|
||||
self.blocks.append(block)
|
||||
|
||||
self.neck = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
embed_dim,
|
||||
out_chans,
|
||||
kernel_size=1,
|
||||
bias=False,
|
||||
),
|
||||
LayerNorm2d(out_chans),
|
||||
nn.Conv2d(
|
||||
out_chans,
|
||||
out_chans,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
bias=False,
|
||||
),
|
||||
LayerNorm2d(out_chans),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.patch_embed(x)
|
||||
if self.pos_embed is not None:
|
||||
x = x + self.pos_embed
|
||||
|
||||
for blk in self.blocks:
|
||||
x = blk(x)
|
||||
|
||||
x = self.neck(x.permute(0, 3, 1, 2))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
"""Transformer blocks with support of window attention and residual propagation blocks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qkv_bias: bool = True,
|
||||
norm_layer: Type[nn.Module] = nn.LayerNorm,
|
||||
act_layer: Type[nn.Module] = nn.GELU,
|
||||
use_rel_pos: bool = False,
|
||||
rel_pos_zero_init: bool = True,
|
||||
window_size: int = 0,
|
||||
input_size: Optional[Tuple[int, int]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads in each ViT block.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
|
||||
qkv_bias (bool): If True, add a learnable bias to query, key, value.
|
||||
norm_layer (nn.Module): Normalization layer.
|
||||
act_layer (nn.Module): Activation layer.
|
||||
use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
|
||||
rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
|
||||
window_size (int): Window size for window attention blocks. If it equals 0, then
|
||||
use global attention.
|
||||
input_size (tuple(int, int) or None): Input resolution for calculating the relative
|
||||
positional parameter size.
|
||||
"""
|
||||
super().__init__()
|
||||
self.norm1 = norm_layer(dim)
|
||||
self.attn = Attention(
|
||||
dim,
|
||||
num_heads=num_heads,
|
||||
qkv_bias=qkv_bias,
|
||||
use_rel_pos=use_rel_pos,
|
||||
rel_pos_zero_init=rel_pos_zero_init,
|
||||
input_size=input_size if window_size == 0 else (window_size, window_size),
|
||||
)
|
||||
|
||||
self.norm2 = norm_layer(dim)
|
||||
self.mlp = MLPBlock(
|
||||
embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer
|
||||
)
|
||||
|
||||
self.window_size = window_size
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
shortcut = x
|
||||
x = self.norm1(x)
|
||||
# Window partition
|
||||
if self.window_size > 0:
|
||||
H, W = x.shape[1], x.shape[2]
|
||||
x, pad_hw = window_partition(x, self.window_size)
|
||||
|
||||
x = self.attn(x)
|
||||
# Reverse window partition
|
||||
if self.window_size > 0:
|
||||
x = window_unpartition(x, self.window_size, pad_hw, (H, W))
|
||||
|
||||
x = shortcut + x
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""Multi-head Attention block with relative position embeddings."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_heads: int = 8,
|
||||
qkv_bias: bool = True,
|
||||
use_rel_pos: bool = False,
|
||||
rel_pos_zero_init: bool = True,
|
||||
input_size: Optional[Tuple[int, int]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
num_heads (int): Number of attention heads.
|
||||
qkv_bias (bool): If True, add a learnable bias to query, key, value.
|
||||
rel_pos (bool): If True, add relative positional embeddings to the attention map.
|
||||
rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
|
||||
input_size (tuple(int, int) or None): Input resolution for calculating the relative
|
||||
positional parameter size.
|
||||
"""
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
self.scale = head_dim**-0.5
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
self.use_rel_pos = use_rel_pos
|
||||
if self.use_rel_pos:
|
||||
assert (
|
||||
input_size is not None
|
||||
), "Input size must be provided if using relative positional encoding."
|
||||
# initialize relative positional embeddings
|
||||
self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
|
||||
self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
B, H, W, _ = x.shape
|
||||
# qkv with shape (3, B, nHead, H * W, C)
|
||||
qkv = (
|
||||
self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
|
||||
)
|
||||
# q, k, v with shape (B * nHead, H * W, C)
|
||||
q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
|
||||
|
||||
attn = (q * self.scale) @ k.transpose(-2, -1)
|
||||
|
||||
if self.use_rel_pos:
|
||||
attn = add_decomposed_rel_pos(
|
||||
attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)
|
||||
)
|
||||
|
||||
attn = attn.softmax(dim=-1)
|
||||
x = (
|
||||
(attn @ v)
|
||||
.view(B, self.num_heads, H, W, -1)
|
||||
.permute(0, 2, 3, 1, 4)
|
||||
.reshape(B, H, W, -1)
|
||||
)
|
||||
x = self.proj(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def window_partition(
|
||||
x: torch.Tensor, window_size: int
|
||||
) -> Tuple[torch.Tensor, Tuple[int, int]]:
|
||||
"""
|
||||
Partition into non-overlapping windows with padding if needed.
|
||||
Args:
|
||||
x (tensor): input tokens with [B, H, W, C].
|
||||
window_size (int): window size.
|
||||
|
||||
Returns:
|
||||
windows: windows after partition with [B * num_windows, window_size, window_size, C].
|
||||
(Hp, Wp): padded height and width before partition
|
||||
"""
|
||||
B, H, W, C = x.shape
|
||||
|
||||
pad_h = (window_size - H % window_size) % window_size
|
||||
pad_w = (window_size - W % window_size) % window_size
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
|
||||
Hp, Wp = H + pad_h, W + pad_w
|
||||
|
||||
x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
|
||||
windows = (
|
||||
x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
|
||||
)
|
||||
return windows, (Hp, Wp)
|
||||
|
||||
|
||||
def window_unpartition(
|
||||
windows: torch.Tensor,
|
||||
window_size: int,
|
||||
pad_hw: Tuple[int, int],
|
||||
hw: Tuple[int, int],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Window unpartition into original sequences and removing padding.
|
||||
Args:
|
||||
windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].
|
||||
window_size (int): window size.
|
||||
pad_hw (Tuple): padded height and width (Hp, Wp).
|
||||
hw (Tuple): original height and width (H, W) before padding.
|
||||
|
||||
Returns:
|
||||
x: unpartitioned sequences with [B, H, W, C].
|
||||
"""
|
||||
Hp, Wp = pad_hw
|
||||
H, W = hw
|
||||
B = windows.shape[0] // (Hp * Wp // window_size // window_size)
|
||||
x = windows.view(
|
||||
B, Hp // window_size, Wp // window_size, window_size, window_size, -1
|
||||
)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
|
||||
|
||||
if Hp > H or Wp > W:
|
||||
x = x[:, :H, :W, :].contiguous()
|
||||
return x
|
||||
|
||||
|
||||
def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Get relative positional embeddings according to the relative positions of
|
||||
query and key sizes.
|
||||
Args:
|
||||
q_size (int): size of query q.
|
||||
k_size (int): size of key k.
|
||||
rel_pos (Tensor): relative position embeddings (L, C).
|
||||
|
||||
Returns:
|
||||
Extracted positional embeddings according to relative positions.
|
||||
"""
|
||||
max_rel_dist = int(2 * max(q_size, k_size) - 1)
|
||||
# Interpolate rel pos if needed.
|
||||
if rel_pos.shape[0] != max_rel_dist:
|
||||
# Interpolate rel pos.
|
||||
rel_pos_resized = F.interpolate(
|
||||
rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
|
||||
size=max_rel_dist,
|
||||
mode="linear",
|
||||
)
|
||||
rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
|
||||
else:
|
||||
rel_pos_resized = rel_pos
|
||||
|
||||
# Scale the coords with short length if shapes for q and k are different.
|
||||
q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
|
||||
k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
|
||||
relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
|
||||
|
||||
return rel_pos_resized[relative_coords.long()]
|
||||
|
||||
|
||||
def add_decomposed_rel_pos(
|
||||
attn: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
rel_pos_h: torch.Tensor,
|
||||
rel_pos_w: torch.Tensor,
|
||||
q_size: Tuple[int, int],
|
||||
k_size: Tuple[int, int],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
|
||||
https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950
|
||||
Args:
|
||||
attn (Tensor): attention map.
|
||||
q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
|
||||
rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
|
||||
rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
|
||||
q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
|
||||
k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
|
||||
|
||||
Returns:
|
||||
attn (Tensor): attention map with added relative positional embeddings.
|
||||
"""
|
||||
q_h, q_w = q_size
|
||||
k_h, k_w = k_size
|
||||
Rh = get_rel_pos(q_h, k_h, rel_pos_h)
|
||||
Rw = get_rel_pos(q_w, k_w, rel_pos_w)
|
||||
|
||||
B, _, dim = q.shape
|
||||
r_q = q.reshape(B, q_h, q_w, dim)
|
||||
rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
|
||||
rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
|
||||
|
||||
attn = (
|
||||
attn.view(B, q_h, q_w, k_h, k_w)
|
||||
+ rel_h[:, :, :, :, None]
|
||||
+ rel_w[:, :, :, None, :]
|
||||
).view(B, q_h * q_w, k_h * k_w)
|
||||
|
||||
return attn
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""
|
||||
Image to Patch Embedding.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_size: Tuple[int, int] = (16, 16),
|
||||
stride: Tuple[int, int] = (16, 16),
|
||||
padding: Tuple[int, int] = (0, 0),
|
||||
in_chans: int = 3,
|
||||
embed_dim: int = 768,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
kernel_size (Tuple): kernel size of the projection layer.
|
||||
stride (Tuple): stride of the projection layer.
|
||||
padding (Tuple): padding size of the projection layer.
|
||||
in_chans (int): Number of input image channels.
|
||||
embed_dim (int): Patch embedding dimension.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.proj = nn.Conv2d(
|
||||
in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.proj(x)
|
||||
# B C H W -> B H W C
|
||||
x = x.permute(0, 2, 3, 1)
|
||||
return x
|
||||
@@ -0,0 +1,190 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from typing import List, Tuple, Type
|
||||
|
||||
from .common import LayerNorm2d
|
||||
|
||||
|
||||
class MaskDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
transformer_dim: int,
|
||||
transformer: nn.Module,
|
||||
num_multimask_outputs: int = 3,
|
||||
activation: Type[nn.Module] = nn.GELU,
|
||||
iou_head_depth: int = 3,
|
||||
iou_head_hidden_dim: int = 256,
|
||||
) -> None:
|
||||
"""
|
||||
Predicts masks given an image and prompt embeddings, using a
|
||||
transformer architecture.
|
||||
|
||||
Arguments:
|
||||
transformer_dim (int): the channel dimension of the transformer
|
||||
transformer (nn.Module): the transformer used to predict masks
|
||||
num_multimask_outputs (int): the number of masks to predict
|
||||
when disambiguating masks
|
||||
activation (nn.Module): the type of activation to use when
|
||||
upscaling masks
|
||||
iou_head_depth (int): the depth of the MLP used to predict
|
||||
mask quality
|
||||
iou_head_hidden_dim (int): the hidden dimension of the MLP
|
||||
used to predict mask quality
|
||||
"""
|
||||
super().__init__()
|
||||
self.transformer_dim = transformer_dim
|
||||
self.transformer = transformer
|
||||
|
||||
self.num_multimask_outputs = num_multimask_outputs
|
||||
|
||||
self.iou_token = nn.Embedding(1, transformer_dim)
|
||||
self.num_mask_tokens = num_multimask_outputs + 1
|
||||
self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
|
||||
|
||||
self.output_upscaling = nn.Sequential(
|
||||
nn.ConvTranspose2d(
|
||||
transformer_dim, transformer_dim // 4, kernel_size=2, stride=2
|
||||
),
|
||||
LayerNorm2d(transformer_dim // 4),
|
||||
activation(),
|
||||
nn.ConvTranspose2d(
|
||||
transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2
|
||||
),
|
||||
activation(),
|
||||
)
|
||||
self.output_hypernetworks_mlps = nn.ModuleList(
|
||||
[
|
||||
MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
|
||||
for i in range(self.num_mask_tokens)
|
||||
]
|
||||
)
|
||||
|
||||
self.iou_prediction_head = MLP(
|
||||
transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_embeddings: torch.Tensor,
|
||||
image_pe: torch.Tensor,
|
||||
sparse_prompt_embeddings: torch.Tensor,
|
||||
dense_prompt_embeddings: torch.Tensor,
|
||||
multimask_output: bool,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Predict masks given image and prompt embeddings.
|
||||
|
||||
Arguments:
|
||||
image_embeddings (torch.Tensor): the embeddings from the image encoder
|
||||
image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
|
||||
sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
|
||||
dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
|
||||
multimask_output (bool): Whether to return multiple masks or a single
|
||||
mask.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: batched predicted masks
|
||||
torch.Tensor: batched predictions of mask quality
|
||||
"""
|
||||
masks, iou_pred = self.predict_masks(
|
||||
image_embeddings=image_embeddings,
|
||||
image_pe=image_pe,
|
||||
sparse_prompt_embeddings=sparse_prompt_embeddings,
|
||||
dense_prompt_embeddings=dense_prompt_embeddings,
|
||||
)
|
||||
|
||||
# Select the correct mask or masks for output
|
||||
if multimask_output:
|
||||
mask_slice = slice(1, None)
|
||||
else:
|
||||
mask_slice = slice(0, 1)
|
||||
masks = masks[:, mask_slice, :, :]
|
||||
iou_pred = iou_pred[:, mask_slice]
|
||||
|
||||
# Prepare output
|
||||
return masks, iou_pred
|
||||
|
||||
def predict_masks(
|
||||
self,
|
||||
image_embeddings: torch.Tensor,
|
||||
image_pe: torch.Tensor,
|
||||
sparse_prompt_embeddings: torch.Tensor,
|
||||
dense_prompt_embeddings: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Predicts masks. See 'forward' for more details."""
|
||||
# Concatenate output tokens
|
||||
output_tokens = torch.cat(
|
||||
[self.iou_token.weight, self.mask_tokens.weight], dim=0
|
||||
)
|
||||
output_tokens = output_tokens.unsqueeze(0).expand(
|
||||
sparse_prompt_embeddings.size(0), -1, -1
|
||||
)
|
||||
tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
|
||||
|
||||
# Expand per-image data in batch direction to be per-mask
|
||||
if image_embeddings.shape[0] != tokens.shape[0]:
|
||||
src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
|
||||
else:
|
||||
src = image_embeddings
|
||||
src = src + dense_prompt_embeddings
|
||||
pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
|
||||
b, c, h, w = src.shape
|
||||
|
||||
# Run the transformer
|
||||
hs, src = self.transformer(src, pos_src, tokens)
|
||||
iou_token_out = hs[:, 0, :]
|
||||
mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
|
||||
|
||||
# Upscale mask embeddings and predict masks using the mask tokens
|
||||
src = src.transpose(1, 2).view(b, c, h, w)
|
||||
upscaled_embedding = self.output_upscaling(src)
|
||||
hyper_in_list: List[torch.Tensor] = []
|
||||
for i in range(self.num_mask_tokens):
|
||||
hyper_in_list.append(
|
||||
self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
|
||||
)
|
||||
hyper_in = torch.stack(hyper_in_list, dim=1)
|
||||
b, c, h, w = upscaled_embedding.shape
|
||||
masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
|
||||
|
||||
# Generate mask quality predictions
|
||||
iou_pred = self.iou_prediction_head(iou_token_out)
|
||||
|
||||
return masks, iou_pred
|
||||
|
||||
|
||||
# Lightly adapted from
|
||||
# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
|
||||
class MLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
output_dim: int,
|
||||
num_layers: int,
|
||||
sigmoid_output: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
h = [hidden_dim] * (num_layers - 1)
|
||||
self.layers = nn.ModuleList(
|
||||
nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
|
||||
)
|
||||
self.sigmoid_output = sigmoid_output
|
||||
|
||||
def forward(self, x):
|
||||
for i, layer in enumerate(self.layers):
|
||||
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
|
||||
if self.sigmoid_output:
|
||||
x = F.sigmoid(x)
|
||||
return x
|
||||
@@ -0,0 +1,226 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from typing import Any, Optional, Tuple, Type
|
||||
|
||||
from .common import LayerNorm2d
|
||||
|
||||
|
||||
class PromptEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
image_embedding_size: Tuple[int, int],
|
||||
input_image_size: Tuple[int, int],
|
||||
mask_in_chans: int,
|
||||
activation: Type[nn.Module] = nn.GELU,
|
||||
) -> None:
|
||||
"""
|
||||
Encodes prompts for input to SAM's mask decoder.
|
||||
|
||||
Arguments:
|
||||
embed_dim (int): The prompts' embedding dimension
|
||||
image_embedding_size (tuple(int, int)): The spatial size of the
|
||||
image embedding, as (H, W).
|
||||
input_image_size (int): The padded size of the image as input
|
||||
to the image encoder, as (H, W).
|
||||
mask_in_chans (int): The number of hidden channels used for
|
||||
encoding input masks.
|
||||
activation (nn.Module): The activation to use when encoding
|
||||
input masks.
|
||||
"""
|
||||
super().__init__()
|
||||
self.embed_dim = embed_dim
|
||||
self.input_image_size = input_image_size
|
||||
self.image_embedding_size = image_embedding_size
|
||||
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
|
||||
|
||||
self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners
|
||||
point_embeddings = [
|
||||
nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)
|
||||
]
|
||||
self.point_embeddings = nn.ModuleList(point_embeddings)
|
||||
self.not_a_point_embed = nn.Embedding(1, embed_dim)
|
||||
|
||||
self.mask_input_size = (
|
||||
4 * image_embedding_size[0],
|
||||
4 * image_embedding_size[1],
|
||||
)
|
||||
self.mask_downscaling = nn.Sequential(
|
||||
nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
|
||||
LayerNorm2d(mask_in_chans // 4),
|
||||
activation(),
|
||||
nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
|
||||
LayerNorm2d(mask_in_chans),
|
||||
activation(),
|
||||
nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
|
||||
)
|
||||
self.no_mask_embed = nn.Embedding(1, embed_dim)
|
||||
|
||||
def get_dense_pe(self) -> torch.Tensor:
|
||||
"""
|
||||
Returns the positional encoding used to encode point prompts,
|
||||
applied to a dense set of points the shape of the image encoding.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Positional encoding with shape
|
||||
1x(embed_dim)x(embedding_h)x(embedding_w)
|
||||
"""
|
||||
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
|
||||
|
||||
def _embed_points(
|
||||
self,
|
||||
points: torch.Tensor,
|
||||
labels: torch.Tensor,
|
||||
pad: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Embeds point prompts."""
|
||||
points = points + 0.5 # Shift to center of pixel
|
||||
if pad:
|
||||
padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
|
||||
padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
|
||||
points = torch.cat([points, padding_point], dim=1)
|
||||
labels = torch.cat([labels, padding_label], dim=1)
|
||||
point_embedding = self.pe_layer.forward_with_coords(
|
||||
points, self.input_image_size
|
||||
)
|
||||
point_embedding[labels == -1] = 0.0
|
||||
point_embedding[labels == -1] += self.not_a_point_embed.weight
|
||||
point_embedding[labels == 0] += self.point_embeddings[0].weight
|
||||
point_embedding[labels == 1] += self.point_embeddings[1].weight
|
||||
return point_embedding
|
||||
|
||||
def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
|
||||
"""Embeds box prompts."""
|
||||
boxes = boxes + 0.5 # Shift to center of pixel
|
||||
coords = boxes.reshape(-1, 2, 2)
|
||||
corner_embedding = self.pe_layer.forward_with_coords(
|
||||
coords, self.input_image_size
|
||||
)
|
||||
corner_embedding[:, 0, :] += self.point_embeddings[2].weight
|
||||
corner_embedding[:, 1, :] += self.point_embeddings[3].weight
|
||||
return corner_embedding
|
||||
|
||||
def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
|
||||
"""Embeds mask inputs."""
|
||||
mask_embedding = self.mask_downscaling(masks)
|
||||
return mask_embedding
|
||||
|
||||
def _get_batch_size(
|
||||
self,
|
||||
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
|
||||
boxes: Optional[torch.Tensor],
|
||||
masks: Optional[torch.Tensor],
|
||||
) -> int:
|
||||
"""
|
||||
Gets the batch size of the output given the batch size of the input prompts.
|
||||
"""
|
||||
if points is not None:
|
||||
return points[0].shape[0]
|
||||
elif boxes is not None:
|
||||
return boxes.shape[0]
|
||||
elif masks is not None:
|
||||
return masks.shape[0]
|
||||
else:
|
||||
return 1
|
||||
|
||||
def _get_device(self) -> torch.device:
|
||||
return self.point_embeddings[0].weight.device
|
||||
|
||||
def forward(
|
||||
self,
|
||||
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
|
||||
boxes: Optional[torch.Tensor],
|
||||
masks: Optional[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Embeds different types of prompts, returning both sparse and dense
|
||||
embeddings.
|
||||
|
||||
Arguments:
|
||||
points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
|
||||
and labels to embed.
|
||||
boxes (torch.Tensor or none): boxes to embed
|
||||
masks (torch.Tensor or none): masks to embed
|
||||
|
||||
Returns:
|
||||
torch.Tensor: sparse embeddings for the points and boxes, with shape
|
||||
BxNx(embed_dim), where N is determined by the number of input points
|
||||
and boxes.
|
||||
torch.Tensor: dense embeddings for the masks, in the shape
|
||||
Bx(embed_dim)x(embed_H)x(embed_W)
|
||||
"""
|
||||
bs = self._get_batch_size(points, boxes, masks)
|
||||
sparse_embeddings = torch.empty(
|
||||
(bs, 0, self.embed_dim), device=self._get_device()
|
||||
)
|
||||
if points is not None:
|
||||
coords, labels = points
|
||||
point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
|
||||
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
|
||||
if boxes is not None:
|
||||
box_embeddings = self._embed_boxes(boxes)
|
||||
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
|
||||
|
||||
if masks is not None:
|
||||
dense_embeddings = self._embed_masks(masks)
|
||||
else:
|
||||
dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
|
||||
bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
|
||||
)
|
||||
|
||||
return sparse_embeddings, dense_embeddings
|
||||
|
||||
|
||||
class PositionEmbeddingRandom(nn.Module):
|
||||
"""
|
||||
Positional encoding using random spatial frequencies.
|
||||
"""
|
||||
|
||||
def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
|
||||
super().__init__()
|
||||
if scale is None or scale <= 0.0:
|
||||
scale = 1.0
|
||||
self.register_buffer(
|
||||
"positional_encoding_gaussian_matrix",
|
||||
scale * torch.randn((2, num_pos_feats)),
|
||||
)
|
||||
|
||||
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
|
||||
"""Positionally encode points that are normalized to [0,1]."""
|
||||
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
|
||||
coords = 2 * coords - 1
|
||||
coords = coords @ self.positional_encoding_gaussian_matrix
|
||||
coords = 2 * np.pi * coords
|
||||
# outputs d_1 x ... x d_n x C shape
|
||||
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
|
||||
|
||||
def forward(self, size: Tuple[int, int]) -> torch.Tensor:
|
||||
"""Generate positional encoding for a grid of the specified size."""
|
||||
h, w = size
|
||||
device: Any = self.positional_encoding_gaussian_matrix.device
|
||||
grid = torch.ones((h, w), device=device, dtype=torch.float32)
|
||||
y_embed = grid.cumsum(dim=0) - 0.5
|
||||
x_embed = grid.cumsum(dim=1) - 0.5
|
||||
y_embed = y_embed / h
|
||||
x_embed = x_embed / w
|
||||
|
||||
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
|
||||
return pe.permute(2, 0, 1) # C x H x W
|
||||
|
||||
def forward_with_coords(
|
||||
self, coords_input: torch.Tensor, image_size: Tuple[int, int]
|
||||
) -> torch.Tensor:
|
||||
"""Positionally encode points that are not normalized to [0,1]."""
|
||||
coords = coords_input.clone()
|
||||
coords[:, :, 0] = coords[:, :, 0] / image_size[1]
|
||||
coords[:, :, 1] = coords[:, :, 1] / image_size[0]
|
||||
return self._pe_encoding(coords.to(torch.float)) # B x N x C
|
||||
@@ -0,0 +1,181 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from .image_encoder import ImageEncoderViT
|
||||
from .mask_decoder import MaskDecoder
|
||||
from .prompt_encoder import PromptEncoder
|
||||
|
||||
|
||||
class Sam(nn.Module):
|
||||
mask_threshold: float = 0.0
|
||||
image_format: str = "RGB"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image_encoder: ImageEncoderViT,
|
||||
prompt_encoder: PromptEncoder,
|
||||
mask_decoder: MaskDecoder,
|
||||
pixel_mean: List[float] = [123.675, 116.28, 103.53],
|
||||
pixel_std: List[float] = [58.395, 57.12, 57.375],
|
||||
) -> None:
|
||||
"""
|
||||
SAM predicts object masks from an image and input prompts.
|
||||
|
||||
Arguments:
|
||||
image_encoder (ImageEncoderViT): The backbone used to encode the
|
||||
image into image embeddings that allow for efficient mask prediction.
|
||||
prompt_encoder (PromptEncoder): Encodes various types of input prompts.
|
||||
mask_decoder (MaskDecoder): Predicts masks from the image embeddings
|
||||
and encoded prompts.
|
||||
pixel_mean (list(float)): Mean values for normalizing pixels in the input image.
|
||||
pixel_std (list(float)): Std values for normalizing pixels in the input image.
|
||||
"""
|
||||
super().__init__()
|
||||
self.image_encoder = image_encoder
|
||||
self.prompt_encoder = prompt_encoder
|
||||
self.mask_decoder = mask_decoder
|
||||
self.register_buffer(
|
||||
"pixel_mean", torch.Tensor(pixel_mean).view(-1, 1, 1), False
|
||||
)
|
||||
self.register_buffer("pixel_std", torch.Tensor(pixel_std).view(-1, 1, 1), False)
|
||||
|
||||
@property
|
||||
def device(self) -> Any:
|
||||
return self.pixel_mean.device
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
batched_input: List[Dict[str, Any]],
|
||||
multimask_output: bool,
|
||||
) -> List[Dict[str, torch.Tensor]]:
|
||||
"""
|
||||
Predicts masks end-to-end from provided images and prompts.
|
||||
If prompts are not known in advance, using SamPredictor is
|
||||
recommended over calling the model directly.
|
||||
|
||||
Arguments:
|
||||
batched_input (list(dict)): A list over input images, each a
|
||||
dictionary with the following keys. A prompt key can be
|
||||
excluded if it is not present.
|
||||
'image': The image as a torch tensor in 3xHxW format,
|
||||
already transformed for input to the model.
|
||||
'original_size': (tuple(int, int)) The original size of
|
||||
the image before transformation, as (H, W).
|
||||
'point_coords': (torch.Tensor) Batched point prompts for
|
||||
this image, with shape BxNx2. Already transformed to the
|
||||
input frame of the model.
|
||||
'point_labels': (torch.Tensor) Batched labels for point prompts,
|
||||
with shape BxN.
|
||||
'boxes': (torch.Tensor) Batched box inputs, with shape Bx4.
|
||||
Already transformed to the input frame of the model.
|
||||
'mask_inputs': (torch.Tensor) Batched mask inputs to the model,
|
||||
in the form Bx1xHxW.
|
||||
multimask_output (bool): Whether the model should predict multiple
|
||||
disambiguating masks, or return a single mask.
|
||||
|
||||
Returns:
|
||||
(list(dict)): A list over input images, where each element is
|
||||
as dictionary with the following keys.
|
||||
'masks': (torch.Tensor) Batched binary mask predictions,
|
||||
with shape BxCxHxW, where B is the number of input prompts,
|
||||
C is determined by multimask_output, and (H, W) is the
|
||||
original size of the image.
|
||||
'iou_predictions': (torch.Tensor) The model's predictions
|
||||
of mask quality, in shape BxC.
|
||||
'low_res_logits': (torch.Tensor) Low resolution logits with
|
||||
shape BxCxHxW, where H=W=256. Can be passed as mask input
|
||||
to subsequent iterations of prediction.
|
||||
"""
|
||||
input_images = torch.stack(
|
||||
[self.preprocess(x["image"]) for x in batched_input], dim=0
|
||||
)
|
||||
image_embeddings = self.image_encoder(input_images)
|
||||
|
||||
outputs = []
|
||||
for image_record, curr_embedding in zip(batched_input, image_embeddings):
|
||||
if "point_coords" in image_record:
|
||||
points = (image_record["point_coords"], image_record["point_labels"])
|
||||
else:
|
||||
points = None
|
||||
sparse_embeddings, dense_embeddings = self.prompt_encoder(
|
||||
points=points,
|
||||
boxes=image_record.get("boxes", None),
|
||||
masks=image_record.get("mask_inputs", None),
|
||||
)
|
||||
low_res_masks, iou_predictions = self.mask_decoder(
|
||||
image_embeddings=curr_embedding.unsqueeze(0),
|
||||
image_pe=self.prompt_encoder.get_dense_pe(),
|
||||
sparse_prompt_embeddings=sparse_embeddings,
|
||||
dense_prompt_embeddings=dense_embeddings,
|
||||
multimask_output=multimask_output,
|
||||
)
|
||||
masks = self.postprocess_masks(
|
||||
low_res_masks,
|
||||
input_size=image_record["image"].shape[-2:],
|
||||
original_size=image_record["original_size"],
|
||||
)
|
||||
masks = masks > self.mask_threshold
|
||||
outputs.append(
|
||||
{
|
||||
"masks": masks,
|
||||
"iou_predictions": iou_predictions,
|
||||
"low_res_logits": low_res_masks,
|
||||
}
|
||||
)
|
||||
return outputs
|
||||
|
||||
def postprocess_masks(
|
||||
self,
|
||||
masks: torch.Tensor,
|
||||
input_size: Tuple[int, ...],
|
||||
original_size: Tuple[int, ...],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Remove padding and upscale masks to the original image size.
|
||||
|
||||
Arguments:
|
||||
masks (torch.Tensor): Batched masks from the mask_decoder,
|
||||
in BxCxHxW format.
|
||||
input_size (tuple(int, int)): The size of the image input to the
|
||||
model, in (H, W) format. Used to remove padding.
|
||||
original_size (tuple(int, int)): The original size of the image
|
||||
before resizing for input to the model, in (H, W) format.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): Batched masks in BxCxHxW format, where (H, W)
|
||||
is given by original_size.
|
||||
"""
|
||||
masks = F.interpolate(
|
||||
masks,
|
||||
(self.image_encoder.img_size, self.image_encoder.img_size),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
masks = masks[..., : input_size[0], : input_size[1]]
|
||||
masks = F.interpolate(
|
||||
masks, original_size, mode="bilinear", align_corners=False
|
||||
)
|
||||
return masks
|
||||
|
||||
def preprocess(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Normalize pixel values and pad to a square input."""
|
||||
# Normalize colors
|
||||
x = (x - self.pixel_mean) / self.pixel_std
|
||||
|
||||
# Pad
|
||||
h, w = x.shape[-2:]
|
||||
padh = self.image_encoder.img_size - h
|
||||
padw = self.image_encoder.img_size - w
|
||||
x = F.pad(x, (0, padw, 0, padh))
|
||||
return x
|
||||
@@ -0,0 +1,243 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
import math
|
||||
from typing import Tuple, Type
|
||||
|
||||
from .common import MLPBlock
|
||||
|
||||
|
||||
class TwoWayTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
depth: int,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer decoder that attends to an input image using
|
||||
queries whose positional embedding is supplied.
|
||||
|
||||
Args:
|
||||
depth (int): number of layers in the transformer
|
||||
embedding_dim (int): the channel dimension for the input embeddings
|
||||
num_heads (int): the number of heads for multihead attention. Must
|
||||
divide embedding_dim
|
||||
mlp_dim (int): the channel dimension internal to the MLP block
|
||||
activation (nn.Module): the activation to use in the MLP block
|
||||
"""
|
||||
super().__init__()
|
||||
self.depth = depth
|
||||
self.embedding_dim = embedding_dim
|
||||
self.num_heads = num_heads
|
||||
self.mlp_dim = mlp_dim
|
||||
self.layers = nn.ModuleList()
|
||||
|
||||
for i in range(depth):
|
||||
self.layers.append(
|
||||
TwoWayAttentionBlock(
|
||||
embedding_dim=embedding_dim,
|
||||
num_heads=num_heads,
|
||||
mlp_dim=mlp_dim,
|
||||
activation=activation,
|
||||
attention_downsample_rate=attention_downsample_rate,
|
||||
skip_first_layer_pe=(i == 0),
|
||||
)
|
||||
)
|
||||
|
||||
self.final_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm_final_attn = nn.LayerNorm(embedding_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
image_embedding: Tensor,
|
||||
image_pe: Tensor,
|
||||
point_embedding: Tensor,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Args:
|
||||
image_embedding (torch.Tensor): image to attend to. Should be shape
|
||||
B x embedding_dim x h x w for any h and w.
|
||||
image_pe (torch.Tensor): the positional encoding to add to the image. Must
|
||||
have the same shape as image_embedding.
|
||||
point_embedding (torch.Tensor): the embedding to add to the query points.
|
||||
Must have shape B x N_points x embedding_dim for any N_points.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the processed point_embedding
|
||||
torch.Tensor: the processed image_embedding
|
||||
"""
|
||||
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
|
||||
bs, c, h, w = image_embedding.shape
|
||||
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
|
||||
image_pe = image_pe.flatten(2).permute(0, 2, 1)
|
||||
|
||||
# Prepare queries
|
||||
queries = point_embedding
|
||||
keys = image_embedding
|
||||
|
||||
# Apply transformer blocks and final layernorm
|
||||
for layer in self.layers:
|
||||
queries, keys = layer(
|
||||
queries=queries,
|
||||
keys=keys,
|
||||
query_pe=point_embedding,
|
||||
key_pe=image_pe,
|
||||
)
|
||||
|
||||
# Apply the final attention layer from the points to the image
|
||||
q = queries + point_embedding
|
||||
k = keys + image_pe
|
||||
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm_final_attn(queries)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class TwoWayAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
mlp_dim: int = 2048,
|
||||
activation: Type[nn.Module] = nn.ReLU,
|
||||
attention_downsample_rate: int = 2,
|
||||
skip_first_layer_pe: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
A transformer block with four layers: (1) self-attention of sparse
|
||||
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
|
||||
block on sparse inputs, and (4) cross attention of dense inputs to sparse
|
||||
inputs.
|
||||
|
||||
Arguments:
|
||||
embedding_dim (int): the channel dimension of the embeddings
|
||||
num_heads (int): the number of heads in the attention layers
|
||||
mlp_dim (int): the hidden dimension of the mlp block
|
||||
activation (nn.Module): the activation of the mlp block
|
||||
skip_first_layer_pe (bool): skip the PE on the first layer
|
||||
"""
|
||||
super().__init__()
|
||||
self.self_attn = Attention(embedding_dim, num_heads)
|
||||
self.norm1 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.cross_attn_token_to_image = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
self.norm2 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
|
||||
self.norm3 = nn.LayerNorm(embedding_dim)
|
||||
|
||||
self.norm4 = nn.LayerNorm(embedding_dim)
|
||||
self.cross_attn_image_to_token = Attention(
|
||||
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
|
||||
)
|
||||
|
||||
self.skip_first_layer_pe = skip_first_layer_pe
|
||||
|
||||
def forward(
|
||||
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
# Self attention block
|
||||
if self.skip_first_layer_pe:
|
||||
queries = self.self_attn(q=queries, k=queries, v=queries)
|
||||
else:
|
||||
q = queries + query_pe
|
||||
attn_out = self.self_attn(q=q, k=q, v=queries)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm1(queries)
|
||||
|
||||
# Cross attention block, tokens attending to image embedding
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
|
||||
queries = queries + attn_out
|
||||
queries = self.norm2(queries)
|
||||
|
||||
# MLP block
|
||||
mlp_out = self.mlp(queries)
|
||||
queries = queries + mlp_out
|
||||
queries = self.norm3(queries)
|
||||
|
||||
# Cross attention block, image embedding attending to tokens
|
||||
q = queries + query_pe
|
||||
k = keys + key_pe
|
||||
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
|
||||
keys = keys + attn_out
|
||||
keys = self.norm4(keys)
|
||||
|
||||
return queries, keys
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
"""
|
||||
An attention layer that allows for downscaling the size of the embedding
|
||||
after projection to queries, keys, and values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedding_dim: int,
|
||||
num_heads: int,
|
||||
downsample_rate: int = 1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
self.internal_dim = embedding_dim // downsample_rate
|
||||
self.num_heads = num_heads
|
||||
assert (
|
||||
self.internal_dim % num_heads == 0
|
||||
), "num_heads must divide embedding_dim."
|
||||
|
||||
self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
|
||||
self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
|
||||
self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
|
||||
self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
|
||||
|
||||
def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
|
||||
b, n, c = x.shape
|
||||
x = x.reshape(b, n, num_heads, c // num_heads)
|
||||
return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head
|
||||
|
||||
def _recombine_heads(self, x: Tensor) -> Tensor:
|
||||
b, n_heads, n_tokens, c_per_head = x.shape
|
||||
x = x.transpose(1, 2)
|
||||
return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
|
||||
# Input projections
|
||||
q = self.q_proj(q)
|
||||
k = self.k_proj(k)
|
||||
v = self.v_proj(v)
|
||||
|
||||
# Separate into heads
|
||||
q = self._separate_heads(q, self.num_heads)
|
||||
k = self._separate_heads(k, self.num_heads)
|
||||
v = self._separate_heads(v, self.num_heads)
|
||||
|
||||
# Attention
|
||||
_, _, _, c_per_head = q.shape
|
||||
attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens
|
||||
attn = attn / math.sqrt(c_per_head)
|
||||
attn = torch.softmax(attn, dim=-1)
|
||||
|
||||
# Get output
|
||||
out = attn @ v
|
||||
out = self._recombine_heads(out)
|
||||
out = self.out_proj(out)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,286 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from PILOT_PROJECT.workspace.sprint_1_2.Design_Material.codebase.ml.implementation.cv.arch.segment_anything.modeling import Sam
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from .utils.transforms import ResizeLongestSide
|
||||
|
||||
|
||||
class SamPredictor:
|
||||
def __init__(
|
||||
self,
|
||||
sam_model: Sam,
|
||||
) -> None:
|
||||
"""
|
||||
Uses SAM to calculate the image embedding for an image, and then
|
||||
allow repeated, efficient mask prediction given prompts.
|
||||
|
||||
Arguments:
|
||||
sam_model (Sam): The model to use for mask prediction.
|
||||
"""
|
||||
super().__init__()
|
||||
self.model = sam_model
|
||||
self.transform = ResizeLongestSide(sam_model.image_encoder.img_size)
|
||||
self.reset_image()
|
||||
|
||||
def set_image(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
image_format: str = "RGB",
|
||||
) -> None:
|
||||
"""
|
||||
Calculates the image embeddings for the provided image, allowing
|
||||
masks to be predicted with the 'predict' method.
|
||||
|
||||
Arguments:
|
||||
image (np.ndarray): The image for calculating masks. Expects an
|
||||
image in HWC uint8 format, with pixel values in [0, 255].
|
||||
image_format (str): The color format of the image, in ['RGB', 'BGR'].
|
||||
"""
|
||||
assert image_format in [
|
||||
"RGB",
|
||||
"BGR",
|
||||
], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
|
||||
if image_format != self.model.image_format:
|
||||
image = image[..., ::-1]
|
||||
|
||||
# Transform the image to the form expected by the model
|
||||
input_image = self.transform.apply_image(image)
|
||||
input_image_torch = torch.as_tensor(input_image, device=self.device)
|
||||
input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[
|
||||
None, :, :, :
|
||||
]
|
||||
|
||||
self.set_torch_image(input_image_torch, image.shape[:2])
|
||||
|
||||
@torch.no_grad()
|
||||
def set_torch_image(
|
||||
self,
|
||||
transformed_image: torch.Tensor,
|
||||
original_image_size: Tuple[int, ...],
|
||||
) -> None:
|
||||
"""
|
||||
Calculates the image embeddings for the provided image, allowing
|
||||
masks to be predicted with the 'predict' method. Expects the input
|
||||
image to be already transformed to the format expected by the model.
|
||||
|
||||
Arguments:
|
||||
transformed_image (torch.Tensor): The input image, with shape
|
||||
1x3xHxW, which has been transformed with ResizeLongestSide.
|
||||
original_image_size (tuple(int, int)): The size of the image
|
||||
before transformation, in (H, W) format.
|
||||
"""
|
||||
assert (
|
||||
len(transformed_image.shape) == 4
|
||||
and transformed_image.shape[1] == 3
|
||||
and max(*transformed_image.shape[2:]) == self.model.image_encoder.img_size
|
||||
), f"set_torch_image input must be BCHW with long side {self.model.image_encoder.img_size}."
|
||||
self.reset_image()
|
||||
|
||||
self.original_size = original_image_size
|
||||
self.input_size = tuple(transformed_image.shape[-2:])
|
||||
input_image = self.model.preprocess(transformed_image)
|
||||
self.features = self.model.image_encoder(input_image)
|
||||
self.is_image_set = True
|
||||
|
||||
def predict(
|
||||
self,
|
||||
point_coords: Optional[np.ndarray] = None,
|
||||
point_labels: Optional[np.ndarray] = None,
|
||||
box: Optional[np.ndarray] = None,
|
||||
mask_input: Optional[np.ndarray] = None,
|
||||
multimask_output: bool = True,
|
||||
return_logits: bool = False,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Predict masks for the given input prompts, using the currently set image.
|
||||
|
||||
Arguments:
|
||||
point_coords (np.ndarray or None): A Nx2 array of point prompts to the
|
||||
model. Each point is in (X,Y) in pixels.
|
||||
point_labels (np.ndarray or None): A length N array of labels for the
|
||||
point prompts. 1 indicates a foreground point and 0 indicates a
|
||||
background point.
|
||||
box (np.ndarray or None): A length 4 array given a box prompt to the
|
||||
model, in XYXY format.
|
||||
mask_input (np.ndarray): A low resolution mask input to the model, typically
|
||||
coming from a previous prediction iteration. Has form 1xHxW, where
|
||||
for SAM, H=W=256.
|
||||
multimask_output (bool): If true, the model will return three masks.
|
||||
For ambiguous input prompts (such as a single click), this will often
|
||||
produce better masks than a single prediction. If only a single
|
||||
mask is needed, the model's predicted quality score can be used
|
||||
to select the best mask. For non-ambiguous prompts, such as multiple
|
||||
input prompts, multimask_output=False can give better results.
|
||||
return_logits (bool): If true, returns un-thresholded masks logits
|
||||
instead of a binary mask.
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The output masks in CxHxW format, where C is the
|
||||
number of masks, and (H, W) is the original image size.
|
||||
(np.ndarray): An array of length C containing the model's
|
||||
predictions for the quality of each mask.
|
||||
(np.ndarray): An array of shape CxHxW, where C is the number
|
||||
of masks and H=W=256. These low resolution logits can be passed to
|
||||
a subsequent iteration as mask input.
|
||||
"""
|
||||
if not self.is_image_set:
|
||||
raise RuntimeError(
|
||||
"An image must be set with .set_image(...) before mask prediction."
|
||||
)
|
||||
|
||||
# Transform input prompts
|
||||
coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
|
||||
if point_coords is not None:
|
||||
assert (
|
||||
point_labels is not None
|
||||
), "point_labels must be supplied if point_coords is supplied."
|
||||
point_coords = self.transform.apply_coords(point_coords, self.original_size)
|
||||
coords_torch = torch.as_tensor(
|
||||
point_coords, dtype=torch.float, device=self.device
|
||||
)
|
||||
labels_torch = torch.as_tensor(
|
||||
point_labels, dtype=torch.int, device=self.device
|
||||
)
|
||||
coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
|
||||
if box is not None:
|
||||
box = self.transform.apply_boxes(box, self.original_size)
|
||||
box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)
|
||||
box_torch = box_torch[None, :]
|
||||
if mask_input is not None:
|
||||
mask_input_torch = torch.as_tensor(
|
||||
mask_input, dtype=torch.float, device=self.device
|
||||
)
|
||||
mask_input_torch = mask_input_torch[None, :, :, :]
|
||||
|
||||
masks, iou_predictions, low_res_masks = self.predict_torch(
|
||||
coords_torch,
|
||||
labels_torch,
|
||||
box_torch,
|
||||
mask_input_torch,
|
||||
multimask_output,
|
||||
return_logits=return_logits,
|
||||
)
|
||||
|
||||
masks_np = masks[0].detach().cpu().numpy()
|
||||
iou_predictions_np = iou_predictions[0].detach().cpu().numpy()
|
||||
low_res_masks_np = low_res_masks[0].detach().cpu().numpy()
|
||||
return masks_np, iou_predictions_np, low_res_masks_np
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_torch(
|
||||
self,
|
||||
point_coords: Optional[torch.Tensor],
|
||||
point_labels: Optional[torch.Tensor],
|
||||
boxes: Optional[torch.Tensor] = None,
|
||||
mask_input: Optional[torch.Tensor] = None,
|
||||
multimask_output: bool = True,
|
||||
return_logits: bool = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Predict masks for the given input prompts, using the currently set image.
|
||||
Input prompts are batched torch tensors and are expected to already be
|
||||
transformed to the input frame using ResizeLongestSide.
|
||||
|
||||
Arguments:
|
||||
point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
|
||||
model. Each point is in (X,Y) in pixels.
|
||||
point_labels (torch.Tensor or None): A BxN array of labels for the
|
||||
point prompts. 1 indicates a foreground point and 0 indicates a
|
||||
background point.
|
||||
boxes (np.ndarray or None): A Bx4 array given a box prompt to the
|
||||
model, in XYXY format.
|
||||
mask_input (np.ndarray): A low resolution mask input to the model, typically
|
||||
coming from a previous prediction iteration. Has form Bx1xHxW, where
|
||||
for SAM, H=W=256. Masks returned by a previous iteration of the
|
||||
predict method do not need further transformation.
|
||||
multimask_output (bool): If true, the model will return three masks.
|
||||
For ambiguous input prompts (such as a single click), this will often
|
||||
produce better masks than a single prediction. If only a single
|
||||
mask is needed, the model's predicted quality score can be used
|
||||
to select the best mask. For non-ambiguous prompts, such as multiple
|
||||
input prompts, multimask_output=False can give better results.
|
||||
return_logits (bool): If true, returns un-thresholded masks logits
|
||||
instead of a binary mask.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor): The output masks in BxCxHxW format, where C is the
|
||||
number of masks, and (H, W) is the original image size.
|
||||
(torch.Tensor): An array of shape BxC containing the model's
|
||||
predictions for the quality of each mask.
|
||||
(torch.Tensor): An array of shape BxCxHxW, where C is the number
|
||||
of masks and H=W=256. These low res logits can be passed to
|
||||
a subsequent iteration as mask input.
|
||||
"""
|
||||
if not self.is_image_set:
|
||||
raise RuntimeError(
|
||||
"An image must be set with .set_image(...) before mask prediction."
|
||||
)
|
||||
|
||||
if point_coords is not None:
|
||||
points = (point_coords, point_labels)
|
||||
else:
|
||||
points = None
|
||||
|
||||
# Embed prompts
|
||||
sparse_embeddings, dense_embeddings = self.model.prompt_encoder(
|
||||
points=points,
|
||||
boxes=boxes,
|
||||
masks=mask_input,
|
||||
)
|
||||
|
||||
# Predict masks
|
||||
low_res_masks, iou_predictions = self.model.mask_decoder(
|
||||
image_embeddings=self.features,
|
||||
image_pe=self.model.prompt_encoder.get_dense_pe(),
|
||||
sparse_prompt_embeddings=sparse_embeddings,
|
||||
dense_prompt_embeddings=dense_embeddings,
|
||||
multimask_output=multimask_output,
|
||||
)
|
||||
|
||||
# Upscale the masks to the original image resolution
|
||||
masks = self.model.postprocess_masks(
|
||||
low_res_masks, self.input_size, self.original_size
|
||||
)
|
||||
|
||||
if not return_logits:
|
||||
masks = masks > self.model.mask_threshold
|
||||
|
||||
return masks, iou_predictions, low_res_masks
|
||||
|
||||
def get_image_embedding(self) -> torch.Tensor:
|
||||
"""
|
||||
Returns the image embeddings for the currently set image, with
|
||||
shape 1xCxHxW, where C is the embedding dimension and (H,W) are
|
||||
the embedding spatial dimension of SAM (typically C=256, H=W=64).
|
||||
"""
|
||||
if not self.is_image_set:
|
||||
raise RuntimeError(
|
||||
"An image must be set with .set_image(...) to generate an embedding."
|
||||
)
|
||||
assert (
|
||||
self.features is not None
|
||||
), "Features must exist if an image has been set."
|
||||
return self.features
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
return self.model.device
|
||||
|
||||
def reset_image(self) -> None:
|
||||
"""Resets the currently set image."""
|
||||
self.is_image_set = False
|
||||
self.features = None
|
||||
self.orig_h = None
|
||||
self.orig_w = None
|
||||
self.input_h = None
|
||||
self.input_w = None
|
||||
@@ -0,0 +1,6 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
@@ -0,0 +1,347 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import math
|
||||
from copy import deepcopy
|
||||
from itertools import product
|
||||
from typing import Any, Dict, Generator, ItemsView, List, Tuple
|
||||
|
||||
|
||||
class MaskData:
|
||||
"""
|
||||
A structure for storing masks and their related data in batched format.
|
||||
Implements basic filtering and concatenation.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
for v in kwargs.values():
|
||||
assert isinstance(
|
||||
v, (list, np.ndarray, torch.Tensor)
|
||||
), "MaskData only supports list, numpy arrays, and torch tensors."
|
||||
self._stats = dict(**kwargs)
|
||||
|
||||
def __setitem__(self, key: str, item: Any) -> None:
|
||||
assert isinstance(
|
||||
item, (list, np.ndarray, torch.Tensor)
|
||||
), "MaskData only supports list, numpy arrays, and torch tensors."
|
||||
self._stats[key] = item
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
del self._stats[key]
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
return self._stats[key]
|
||||
|
||||
def items(self) -> ItemsView[str, Any]:
|
||||
return self._stats.items()
|
||||
|
||||
def filter(self, keep: torch.Tensor) -> None:
|
||||
for k, v in self._stats.items():
|
||||
if v is None:
|
||||
self._stats[k] = None
|
||||
elif isinstance(v, torch.Tensor):
|
||||
self._stats[k] = v[torch.as_tensor(keep, device=v.device)]
|
||||
elif isinstance(v, np.ndarray):
|
||||
self._stats[k] = v[keep.detach().cpu().numpy()]
|
||||
elif isinstance(v, list) and keep.dtype == torch.bool:
|
||||
self._stats[k] = [a for i, a in enumerate(v) if keep[i]]
|
||||
elif isinstance(v, list):
|
||||
self._stats[k] = [v[i] for i in keep]
|
||||
else:
|
||||
raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
|
||||
|
||||
def cat(self, new_stats: "MaskData") -> None:
|
||||
for k, v in new_stats.items():
|
||||
if k not in self._stats or self._stats[k] is None:
|
||||
self._stats[k] = deepcopy(v)
|
||||
elif isinstance(v, torch.Tensor):
|
||||
self._stats[k] = torch.cat([self._stats[k], v], dim=0)
|
||||
elif isinstance(v, np.ndarray):
|
||||
self._stats[k] = np.concatenate([self._stats[k], v], axis=0)
|
||||
elif isinstance(v, list):
|
||||
self._stats[k] = self._stats[k] + deepcopy(v)
|
||||
else:
|
||||
raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
|
||||
|
||||
def to_numpy(self) -> None:
|
||||
for k, v in self._stats.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
self._stats[k] = v.detach().cpu().numpy()
|
||||
|
||||
|
||||
def is_box_near_crop_edge(
|
||||
boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0
|
||||
) -> torch.Tensor:
|
||||
"""Filter masks at the edge of a crop, but not at the edge of the original image."""
|
||||
crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)
|
||||
orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)
|
||||
boxes = uncrop_boxes_xyxy(boxes, crop_box).float()
|
||||
near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)
|
||||
near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)
|
||||
near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)
|
||||
return torch.any(near_crop_edge, dim=1)
|
||||
|
||||
|
||||
def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:
|
||||
box_xywh = deepcopy(box_xyxy)
|
||||
box_xywh[2] = box_xywh[2] - box_xywh[0]
|
||||
box_xywh[3] = box_xywh[3] - box_xywh[1]
|
||||
return box_xywh
|
||||
|
||||
|
||||
def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
|
||||
assert len(args) > 0 and all(
|
||||
len(a) == len(args[0]) for a in args
|
||||
), "Batched iteration must have inputs of all the same size."
|
||||
n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
|
||||
for b in range(n_batches):
|
||||
yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]
|
||||
|
||||
|
||||
def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Encodes masks to an uncompressed RLE, in the format expected by
|
||||
pycoco tools.
|
||||
"""
|
||||
# Put in fortran order and flatten h,w
|
||||
b, h, w = tensor.shape
|
||||
tensor = tensor.permute(0, 2, 1).flatten(1)
|
||||
|
||||
# Compute change indices
|
||||
diff = tensor[:, 1:] ^ tensor[:, :-1]
|
||||
change_indices = diff.nonzero()
|
||||
|
||||
# Encode run length
|
||||
out = []
|
||||
for i in range(b):
|
||||
cur_idxs = change_indices[change_indices[:, 0] == i, 1]
|
||||
cur_idxs = torch.cat(
|
||||
[
|
||||
torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),
|
||||
cur_idxs + 1,
|
||||
torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),
|
||||
]
|
||||
)
|
||||
btw_idxs = cur_idxs[1:] - cur_idxs[:-1]
|
||||
counts = [] if tensor[i, 0] == 0 else [0]
|
||||
counts.extend(btw_idxs.detach().cpu().tolist())
|
||||
out.append({"size": [h, w], "counts": counts})
|
||||
return out
|
||||
|
||||
|
||||
def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:
|
||||
"""Compute a binary mask from an uncompressed RLE."""
|
||||
h, w = rle["size"]
|
||||
mask = np.empty(h * w, dtype=bool)
|
||||
idx = 0
|
||||
parity = False
|
||||
for count in rle["counts"]:
|
||||
mask[idx : idx + count] = parity
|
||||
idx += count
|
||||
parity ^= True
|
||||
mask = mask.reshape(w, h)
|
||||
return mask.transpose() # Put in C order
|
||||
|
||||
|
||||
def area_from_rle(rle: Dict[str, Any]) -> int:
|
||||
return sum(rle["counts"][1::2])
|
||||
|
||||
|
||||
def calculate_stability_score(
|
||||
masks: torch.Tensor, mask_threshold: float, threshold_offset: float
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Computes the stability score for a batch of masks. The stability
|
||||
score is the IoU between the binary masks obtained by thresholding
|
||||
the predicted mask logits at high and low values.
|
||||
"""
|
||||
# One mask is always contained inside the other.
|
||||
# Save memory by preventing unnecessary cast to torch.int64
|
||||
intersections = (
|
||||
(masks > (mask_threshold + threshold_offset))
|
||||
.sum(-1, dtype=torch.int16)
|
||||
.sum(-1, dtype=torch.int32)
|
||||
)
|
||||
unions = (
|
||||
(masks > (mask_threshold - threshold_offset))
|
||||
.sum(-1, dtype=torch.int16)
|
||||
.sum(-1, dtype=torch.int32)
|
||||
)
|
||||
return intersections / unions
|
||||
|
||||
|
||||
def build_point_grid(n_per_side: int) -> np.ndarray:
|
||||
"""Generates a 2D grid of points evenly spaced in [0,1]x[0,1]."""
|
||||
offset = 1 / (2 * n_per_side)
|
||||
points_one_side = np.linspace(offset, 1 - offset, n_per_side)
|
||||
points_x = np.tile(points_one_side[None, :], (n_per_side, 1))
|
||||
points_y = np.tile(points_one_side[:, None], (1, n_per_side))
|
||||
points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)
|
||||
return points
|
||||
|
||||
|
||||
def build_all_layer_point_grids(
|
||||
n_per_side: int, n_layers: int, scale_per_layer: int
|
||||
) -> List[np.ndarray]:
|
||||
"""Generates point grids for all crop layers."""
|
||||
points_by_layer = []
|
||||
for i in range(n_layers + 1):
|
||||
n_points = int(n_per_side / (scale_per_layer**i))
|
||||
points_by_layer.append(build_point_grid(n_points))
|
||||
return points_by_layer
|
||||
|
||||
|
||||
def generate_crop_boxes(
|
||||
im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float
|
||||
) -> Tuple[List[List[int]], List[int]]:
|
||||
"""
|
||||
Generates a list of crop boxes of different sizes. Each layer
|
||||
has (2**i)**2 boxes for the ith layer.
|
||||
"""
|
||||
crop_boxes, layer_idxs = [], []
|
||||
im_h, im_w = im_size
|
||||
short_side = min(im_h, im_w)
|
||||
|
||||
# Original image
|
||||
crop_boxes.append([0, 0, im_w, im_h])
|
||||
layer_idxs.append(0)
|
||||
|
||||
def crop_len(orig_len, n_crops, overlap):
|
||||
return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))
|
||||
|
||||
for i_layer in range(n_layers):
|
||||
n_crops_per_side = 2 ** (i_layer + 1)
|
||||
overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))
|
||||
|
||||
crop_w = crop_len(im_w, n_crops_per_side, overlap)
|
||||
crop_h = crop_len(im_h, n_crops_per_side, overlap)
|
||||
|
||||
crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]
|
||||
crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]
|
||||
|
||||
# Crops in XYWH format
|
||||
for x0, y0 in product(crop_box_x0, crop_box_y0):
|
||||
box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]
|
||||
crop_boxes.append(box)
|
||||
layer_idxs.append(i_layer + 1)
|
||||
|
||||
return crop_boxes, layer_idxs
|
||||
|
||||
|
||||
def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
|
||||
x0, y0, _, _ = crop_box
|
||||
offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)
|
||||
# Check if boxes has a channel dimension
|
||||
if len(boxes.shape) == 3:
|
||||
offset = offset.unsqueeze(1)
|
||||
return boxes + offset
|
||||
|
||||
|
||||
def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
|
||||
x0, y0, _, _ = crop_box
|
||||
offset = torch.tensor([[x0, y0]], device=points.device)
|
||||
# Check if points has a channel dimension
|
||||
if len(points.shape) == 3:
|
||||
offset = offset.unsqueeze(1)
|
||||
return points + offset
|
||||
|
||||
|
||||
def uncrop_masks(
|
||||
masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int
|
||||
) -> torch.Tensor:
|
||||
x0, y0, x1, y1 = crop_box
|
||||
if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:
|
||||
return masks
|
||||
# Coordinate transform masks
|
||||
pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)
|
||||
pad = (x0, pad_x - x0, y0, pad_y - y0)
|
||||
return torch.nn.functional.pad(masks, pad, value=0)
|
||||
|
||||
|
||||
def remove_small_regions(
|
||||
mask: np.ndarray, area_thresh: float, mode: str
|
||||
) -> Tuple[np.ndarray, bool]:
|
||||
"""
|
||||
Removes small disconnected regions and holes in a mask. Returns the
|
||||
mask and an indicator of if the mask has been modified.
|
||||
"""
|
||||
import cv2 # type: ignore
|
||||
|
||||
assert mode in ["holes", "islands"]
|
||||
correct_holes = mode == "holes"
|
||||
working_mask = (correct_holes ^ mask).astype(np.uint8)
|
||||
n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)
|
||||
sizes = stats[:, -1][1:] # Row 0 is background label
|
||||
small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]
|
||||
if len(small_regions) == 0:
|
||||
return mask, False
|
||||
fill_labels = [0] + small_regions
|
||||
if not correct_holes:
|
||||
fill_labels = [i for i in range(n_labels) if i not in fill_labels]
|
||||
# If every region is below threshold, keep largest
|
||||
if len(fill_labels) == 0:
|
||||
fill_labels = [int(np.argmax(sizes)) + 1]
|
||||
mask = np.isin(regions, fill_labels)
|
||||
return mask, True
|
||||
|
||||
|
||||
def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:
|
||||
from pycocotools import mask as mask_utils # type: ignore
|
||||
|
||||
h, w = uncompressed_rle["size"]
|
||||
rle = mask_utils.frPyObjects(uncompressed_rle, h, w)
|
||||
rle["counts"] = rle["counts"].decode("utf-8") # Necessary to serialize with json
|
||||
return rle
|
||||
|
||||
|
||||
def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Calculates boxes in XYXY format around masks. Return [0,0,0,0] for
|
||||
an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
|
||||
"""
|
||||
# torch.max below raises an error on empty inputs, just skip in this case
|
||||
if torch.numel(masks) == 0:
|
||||
return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
|
||||
|
||||
# Normalize shape to CxHxW
|
||||
shape = masks.shape
|
||||
h, w = shape[-2:]
|
||||
if len(shape) > 2:
|
||||
masks = masks.flatten(0, -3)
|
||||
else:
|
||||
masks = masks.unsqueeze(0)
|
||||
|
||||
# Get top and bottom edges
|
||||
in_height, _ = torch.max(masks, dim=-1)
|
||||
in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]
|
||||
bottom_edges, _ = torch.max(in_height_coords, dim=-1)
|
||||
in_height_coords = in_height_coords + h * (~in_height)
|
||||
top_edges, _ = torch.min(in_height_coords, dim=-1)
|
||||
|
||||
# Get left and right edges
|
||||
in_width, _ = torch.max(masks, dim=-2)
|
||||
in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]
|
||||
right_edges, _ = torch.max(in_width_coords, dim=-1)
|
||||
in_width_coords = in_width_coords + w * (~in_width)
|
||||
left_edges, _ = torch.min(in_width_coords, dim=-1)
|
||||
|
||||
# If the mask is empty the right edge will be to the left of the left edge.
|
||||
# Replace these boxes with [0, 0, 0, 0]
|
||||
empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
|
||||
out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
|
||||
out = out * (~empty_filter).unsqueeze(-1)
|
||||
|
||||
# Return to original shape
|
||||
if len(shape) > 2:
|
||||
out = out.reshape(*shape[:-2], 4)
|
||||
else:
|
||||
out = out[0]
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,158 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
from ..modeling import Sam
|
||||
from .amg import calculate_stability_score
|
||||
|
||||
|
||||
class SamOnnxModel(nn.Module):
|
||||
"""
|
||||
This model should not be called directly, but is used in ONNX export.
|
||||
It combines the prompt encoder, mask decoder, and mask postprocessing of Sam,
|
||||
with some functions modified to enable model tracing. Also supports extra
|
||||
options controlling what information. See the ONNX export script for details.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Sam,
|
||||
return_single_mask: bool,
|
||||
use_stability_score: bool = False,
|
||||
return_extra_metrics: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.mask_decoder = model.mask_decoder
|
||||
self.model = model
|
||||
self.img_size = model.image_encoder.img_size
|
||||
self.return_single_mask = return_single_mask
|
||||
self.use_stability_score = use_stability_score
|
||||
self.stability_score_offset = 1.0
|
||||
self.return_extra_metrics = return_extra_metrics
|
||||
|
||||
@staticmethod
|
||||
def resize_longest_image_size(
|
||||
input_image_size: torch.Tensor, longest_side: int
|
||||
) -> torch.Tensor:
|
||||
input_image_size = input_image_size.to(torch.float32)
|
||||
scale = longest_side / torch.max(input_image_size)
|
||||
transformed_size = scale * input_image_size
|
||||
transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)
|
||||
return transformed_size
|
||||
|
||||
def _embed_points(
|
||||
self, point_coords: torch.Tensor, point_labels: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
point_coords = point_coords + 0.5
|
||||
point_coords = point_coords / self.img_size
|
||||
point_embedding = self.model.prompt_encoder.pe_layer._pe_encoding(point_coords)
|
||||
point_labels = point_labels.unsqueeze(-1).expand_as(point_embedding)
|
||||
|
||||
point_embedding = point_embedding * (point_labels != -1)
|
||||
point_embedding = (
|
||||
point_embedding
|
||||
+ self.model.prompt_encoder.not_a_point_embed.weight * (point_labels == -1)
|
||||
)
|
||||
|
||||
for i in range(self.model.prompt_encoder.num_point_embeddings):
|
||||
point_embedding = (
|
||||
point_embedding
|
||||
+ self.model.prompt_encoder.point_embeddings[i].weight
|
||||
* (point_labels == i)
|
||||
)
|
||||
|
||||
return point_embedding
|
||||
|
||||
def _embed_masks(
|
||||
self, input_mask: torch.Tensor, has_mask_input: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
mask_embedding = has_mask_input * self.model.prompt_encoder.mask_downscaling(
|
||||
input_mask
|
||||
)
|
||||
mask_embedding = mask_embedding + (
|
||||
1 - has_mask_input
|
||||
) * self.model.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
|
||||
return mask_embedding
|
||||
|
||||
def mask_postprocessing(
|
||||
self, masks: torch.Tensor, orig_im_size: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
masks = F.interpolate(
|
||||
masks,
|
||||
size=(self.img_size, self.img_size),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
prepadded_size = self.resize_longest_image_size(orig_im_size, self.img_size).to(
|
||||
torch.int64
|
||||
)
|
||||
masks = masks[..., : prepadded_size[0], : prepadded_size[1]] # type: ignore
|
||||
|
||||
orig_im_size = orig_im_size.to(torch.int64)
|
||||
h, w = orig_im_size[0], orig_im_size[1]
|
||||
masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False)
|
||||
return masks
|
||||
|
||||
def select_masks(
|
||||
self, masks: torch.Tensor, iou_preds: torch.Tensor, num_points: int
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
# Determine if we should return the multiclick mask or not from the number of points.
|
||||
# The reweighting is used to avoid control flow.
|
||||
score_reweight = torch.tensor(
|
||||
[[1000] + [0] * (self.model.mask_decoder.num_mask_tokens - 1)]
|
||||
).to(iou_preds.device)
|
||||
score = iou_preds + (num_points - 2.5) * score_reweight
|
||||
best_idx = torch.argmax(score, dim=1)
|
||||
masks = masks[torch.arange(masks.shape[0]), best_idx, :, :].unsqueeze(1)
|
||||
iou_preds = iou_preds[torch.arange(masks.shape[0]), best_idx].unsqueeze(1)
|
||||
|
||||
return masks, iou_preds
|
||||
|
||||
@torch.no_grad()
|
||||
def forward(
|
||||
self,
|
||||
image_embeddings: torch.Tensor,
|
||||
point_coords: torch.Tensor,
|
||||
point_labels: torch.Tensor,
|
||||
mask_input: torch.Tensor,
|
||||
has_mask_input: torch.Tensor,
|
||||
orig_im_size: torch.Tensor,
|
||||
):
|
||||
sparse_embedding = self._embed_points(point_coords, point_labels)
|
||||
dense_embedding = self._embed_masks(mask_input, has_mask_input)
|
||||
|
||||
masks, scores = self.model.mask_decoder.predict_masks(
|
||||
image_embeddings=image_embeddings,
|
||||
image_pe=self.model.prompt_encoder.get_dense_pe(),
|
||||
sparse_prompt_embeddings=sparse_embedding,
|
||||
dense_prompt_embeddings=dense_embedding,
|
||||
)
|
||||
|
||||
if self.use_stability_score:
|
||||
scores = calculate_stability_score(
|
||||
masks, self.model.mask_threshold, self.stability_score_offset
|
||||
)
|
||||
|
||||
if self.return_single_mask:
|
||||
masks, scores = self.select_masks(masks, scores, point_coords.shape[1])
|
||||
|
||||
upscaled_masks = self.mask_postprocessing(masks, orig_im_size)
|
||||
|
||||
if self.return_extra_metrics:
|
||||
stability_scores = calculate_stability_score(
|
||||
upscaled_masks, self.model.mask_threshold, self.stability_score_offset
|
||||
)
|
||||
areas = (upscaled_masks > self.model.mask_threshold).sum(-1).sum(-1)
|
||||
return upscaled_masks, scores, stability_scores, areas, masks
|
||||
|
||||
return upscaled_masks, scores, masks
|
||||
@@ -0,0 +1,111 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from torchvision.transforms.functional import resize, to_pil_image # type: ignore
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
class ResizeLongestSide:
|
||||
"""
|
||||
Resizes images to the longest side 'target_length', as well as provides
|
||||
methods for resizing coordinates and boxes. Provides methods for
|
||||
transforming both numpy array and batched torch tensors.
|
||||
"""
|
||||
|
||||
def __init__(self, target_length: int) -> None:
|
||||
self.target_length = target_length
|
||||
|
||||
def apply_image(self, image: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Expects a numpy array with shape HxWxC in uint8 format.
|
||||
"""
|
||||
target_size = self.get_preprocess_shape(
|
||||
image.shape[0], image.shape[1], self.target_length
|
||||
)
|
||||
return np.array(resize(to_pil_image(image), target_size))
|
||||
|
||||
def apply_coords(
|
||||
self, coords: np.ndarray, original_size: Tuple[int, ...]
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Expects a numpy array of length 2 in the final dimension. Requires the
|
||||
original image size in (H, W) format.
|
||||
"""
|
||||
old_h, old_w = original_size
|
||||
new_h, new_w = self.get_preprocess_shape(old_h, old_w, self.target_length)
|
||||
new_coords = np.empty_like(coords)
|
||||
new_coords[..., 0] = coords[..., 0] * (new_w / old_w)
|
||||
new_coords[..., 1] = coords[..., 1] * (new_h / old_h)
|
||||
return new_coords
|
||||
|
||||
def apply_boxes(
|
||||
self, boxes: np.ndarray, original_size: Tuple[int, ...]
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Expects a numpy array shape Bx4. Requires the original image size
|
||||
in (H, W) format.
|
||||
"""
|
||||
boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)
|
||||
return boxes.reshape(-1, 4)
|
||||
|
||||
def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Expects batched images with shape BxCxHxW and float format. This
|
||||
transformation may not exactly match apply_image. apply_image is
|
||||
the transformation expected by the model.
|
||||
"""
|
||||
# Expects an image in BCHW format. May not exactly match apply_image.
|
||||
target_size = self.get_preprocess_shape(
|
||||
image.shape[2], image.shape[3], self.target_length
|
||||
)
|
||||
return F.interpolate(
|
||||
image, target_size, mode="bilinear", align_corners=False, antialias=True
|
||||
)
|
||||
|
||||
def apply_coords_torch(
|
||||
self, coords: torch.Tensor, original_size: Tuple[int, ...]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Expects a torch tensor with length 2 in the last dimension. Requires the
|
||||
original image size in (H, W) format.
|
||||
"""
|
||||
old_h, old_w = original_size
|
||||
new_h, new_w = self.get_preprocess_shape(
|
||||
original_size[0], original_size[1], self.target_length
|
||||
)
|
||||
coords = deepcopy(coords).to(torch.float)
|
||||
coords[..., 0] = coords[..., 0] * (new_w / old_w)
|
||||
coords[..., 1] = coords[..., 1] * (new_h / old_h)
|
||||
return coords
|
||||
|
||||
def apply_boxes_torch(
|
||||
self, boxes: torch.Tensor, original_size: Tuple[int, ...]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Expects a torch tensor with shape Bx4. Requires the original image
|
||||
size in (H, W) format.
|
||||
"""
|
||||
boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)
|
||||
return boxes.reshape(-1, 4)
|
||||
|
||||
@staticmethod
|
||||
def get_preprocess_shape(
|
||||
oldh: int, oldw: int, long_side_length: int
|
||||
) -> Tuple[int, int]:
|
||||
"""
|
||||
Compute the output size given input size and target long side length.
|
||||
"""
|
||||
scale = long_side_length * 1.0 / max(oldh, oldw)
|
||||
newh, neww = oldh * scale, oldw * scale
|
||||
neww = int(neww + 0.5)
|
||||
newh = int(newh + 0.5)
|
||||
return (newh, neww)
|
||||
@@ -0,0 +1,269 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
# ====================== Self-Attention Block ======================
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, in_dim):
|
||||
super(SelfAttention, self).__init__()
|
||||
self.query_conv = nn.Conv2d(in_dim, max(1, in_dim // 8), kernel_size=1)
|
||||
self.key_conv = nn.Conv2d(in_dim, max(1, in_dim // 8), kernel_size=1)
|
||||
self.value_conv = nn.Conv2d(in_dim, in_dim, kernel_size=1)
|
||||
self.gamma = nn.Parameter(torch.zeros(1))
|
||||
|
||||
def forward(self, x):
|
||||
B, C, H, W = x.size()
|
||||
query = self.query_conv(x).view(B, -1, H * W).permute(0, 2, 1) # B, N, C'
|
||||
key = self.key_conv(x).view(B, -1, H * W) # B, C', N
|
||||
energy = torch.bmm(query, key) # B, N, N
|
||||
attention = F.softmax(energy, dim=-1)
|
||||
value = self.value_conv(x).view(B, -1, H * W) # B, C, N
|
||||
out = torch.bmm(value, attention.permute(0, 2, 1)) # B, C, N
|
||||
out = out.view(B, C, H, W)
|
||||
out = self.gamma * out + x
|
||||
return out
|
||||
|
||||
|
||||
# ====================== Attention Gate ======================
|
||||
class AttentionGate(nn.Module):
|
||||
def __init__(self, F_g, F_l, F_int):
|
||||
super(AttentionGate, self).__init__()
|
||||
self.W_g = nn.Sequential(
|
||||
nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
|
||||
nn.BatchNorm2d(F_int)
|
||||
)
|
||||
|
||||
self.W_x = nn.Sequential(
|
||||
nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
|
||||
nn.BatchNorm2d(F_int)
|
||||
)
|
||||
|
||||
self.psi = nn.Sequential(
|
||||
nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
|
||||
nn.BatchNorm2d(1),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, g, x):
|
||||
# g: gating signal (deeper) ; x: skip connection (shallower)
|
||||
g1 = self.W_g(g)
|
||||
x1 = self.W_x(x)
|
||||
|
||||
# ensure same spatial size
|
||||
# if g1.size()[2:] != x1.size()[2:]:
|
||||
g1 = F.interpolate(g1, size=x1.size()[2:], mode='bilinear', align_corners=True)
|
||||
# g1 = F.interpolate(g1, size=(int(x1.size(2)), int(x1.size(3))), mode='bilinear', align_corners=True)
|
||||
|
||||
psi = self.relu(g1 + x1)
|
||||
psi = self.psi(psi)
|
||||
|
||||
# ensure psi matches x spatially
|
||||
# if psi.size()[2:] != x.size()[2:]:
|
||||
psi = F.interpolate(psi, size=x.size()[2:], mode='bilinear', align_corners=True)
|
||||
|
||||
return x * psi
|
||||
|
||||
|
||||
# ====================== Basic Conv Block ======================
|
||||
class conv_block(nn.Module):
|
||||
def __init__(self, in_ch, out_ch):
|
||||
super(conv_block, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
|
||||
nn.BatchNorm2d(out_ch),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
# ====================== UNet3+ with Attention (fixed channels) ======================
|
||||
class UNet3Plus_Attention(nn.Module):
|
||||
def __init__(self, in_channels=3, num_classes=7, filters=[64, 128, 256, 512, 1024]):
|
||||
super(UNet3Plus_Attention, self).__init__()
|
||||
self.num_classes = num_classes
|
||||
F1, F2, F3, F4, F5 = filters # e.g. 64,128,256,512,1024
|
||||
|
||||
# ---------------- Encoder ----------------
|
||||
self.encoder1 = conv_block(in_channels, F1)
|
||||
self.pool1 = nn.MaxPool2d(2)
|
||||
self.encoder2 = conv_block(F1, F2)
|
||||
self.pool2 = nn.MaxPool2d(2)
|
||||
self.encoder3 = conv_block(F2, F3)
|
||||
self.pool3 = nn.MaxPool2d(2)
|
||||
self.encoder4 = conv_block(F3, F4)
|
||||
self.pool4 = nn.MaxPool2d(2)
|
||||
self.encoder5 = conv_block(F4, F5)
|
||||
|
||||
# ---------------- Self-Attention at bottleneck ----------------
|
||||
self.self_att = SelfAttention(F5)
|
||||
|
||||
# ---------------- 1x1 conv for multi-scale fusion ----------------
|
||||
# These compress encoder features to F1 channels for easier fusion where used.
|
||||
self.conv_1x1_1 = nn.Conv2d(F1, F1, kernel_size=1)
|
||||
self.conv_1x1_2 = nn.Conv2d(F2, F1, kernel_size=1)
|
||||
self.conv_1x1_3 = nn.Conv2d(F3, F1, kernel_size=1)
|
||||
self.conv_1x1_4 = nn.Conv2d(F4, F1, kernel_size=1)
|
||||
self.conv_1x1_5 = nn.Conv2d(F5, F1, kernel_size=1)
|
||||
|
||||
# ---------------- Decoder Blocks ----------------
|
||||
# IMPORTANT: set in_channels = actual concatenation channels at each level
|
||||
# Decoder4 concatenates upsampled conv_1x1 outputs (all compressed to F1) => 5*F1
|
||||
self.decoder4_cat_conv = conv_block(F1 * 5, F4) # output channels F4
|
||||
|
||||
# Decoder3 concatenation channels:
|
||||
# d1_up (conv_1x1_1) : F1
|
||||
# d2_up (conv_1x1_2) : F1
|
||||
# d3_up (conv_1x1_3) : F1
|
||||
# d4_up (decoder4 output) : F4
|
||||
# d5_up (conv_1x1_5) : F1
|
||||
dec3_in = F1 + F1 + F1 + F4 + F1 # = 4*F1 + F4
|
||||
self.decoder3_cat_conv = conv_block(dec3_in, F3)
|
||||
|
||||
# Decoder2 concatenation channels:
|
||||
# d1_up: F1
|
||||
# d2_up: F1
|
||||
# d3_up: decoder3 output F3
|
||||
# d4_up: decoder4 output F4
|
||||
# d5_up: F1
|
||||
dec2_in = F1 + F1 + F3 + F4 + F1
|
||||
self.decoder2_cat_conv = conv_block(dec2_in, F2)
|
||||
|
||||
# Decoder1 concatenation channels:
|
||||
# d1_up: F1
|
||||
# d2_up: decoder2 output F2
|
||||
# d3_up: decoder3 output F3
|
||||
# d4_up: decoder4 output F4
|
||||
# d5_up: F1
|
||||
dec1_in = F1 + F2 + F3 + F4 + F1
|
||||
self.decoder1_cat_conv = conv_block(dec1_in, F1)
|
||||
|
||||
# ---------------- Attention Gates ----------------
|
||||
self.att4 = AttentionGate(F_g=F5, F_l=F4, F_int=max(1, F4 // 2))
|
||||
self.att3 = AttentionGate(F_g=F4, F_l=F3, F_int=max(1, F3 // 2))
|
||||
self.att2 = AttentionGate(F_g=F3, F_l=F2, F_int=max(1, F2 // 2))
|
||||
self.att1 = AttentionGate(F_g=F2, F_l=F1, F_int=max(1, F1 // 2))
|
||||
|
||||
# ---------------- Final Fusion ----------------
|
||||
self.final_reduce = nn.Sequential(
|
||||
nn.Conv2d(F1 + F2 + F3 + F4 + F1, F1, kernel_size=1),
|
||||
nn.BatchNorm2d(F1),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
# Note: final_reduce input channels use same sum as dec1 input (we fuse d1 + up(d2) + up(d3) + up(d4) + up(x5)
|
||||
# but to keep consistent we will recompute in forward and pass through this module (in_ch matches dec1_in)
|
||||
# For simplicity, we use a 1x1 to F1.
|
||||
|
||||
self.final_conv = nn.Conv2d(F1, self.num_classes, kernel_size=1)
|
||||
|
||||
self._init_weights()
|
||||
|
||||
def _init_weights(self):
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
||||
if getattr(m, "bias", None) is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# ---------------- Forward ----------------
|
||||
def forward(self, x):
|
||||
# Encoder
|
||||
x1 = self.encoder1(x) # [B, F1, H, W]
|
||||
x2 = self.encoder2(self.pool1(x1)) # [B, F2, H/2, W/2]
|
||||
x3 = self.encoder3(self.pool2(x2)) # [B, F3, H/4, W/4]
|
||||
x4 = self.encoder4(self.pool3(x3)) # [B, F4, H/8, W/8]
|
||||
x5 = self.encoder5(self.pool4(x4)) # [B, F5, H/16, W/16]
|
||||
|
||||
# Bottleneck self-attention
|
||||
x5 = self.self_att(x5)
|
||||
|
||||
# Attention gates on skip connections (ensure spatial alignment inside gate)
|
||||
x4 = self.att4(x5, x4)
|
||||
x3 = self.att3(x4, x3)
|
||||
x2 = self.att2(x3, x2)
|
||||
x1 = self.att1(x2, x1)
|
||||
|
||||
# ---- Prepare compressed feature maps (1x1) ----
|
||||
# These are used for multi-scale fusion where appropriate
|
||||
e1c = self.conv_1x1_1(x1) # [B, F1, H, W]
|
||||
e2c = self.conv_1x1_2(x2) # [B, F1, H/2, W/2]
|
||||
e3c = self.conv_1x1_3(x3) # [B, F1, H/4, W/4]
|
||||
e4c = self.conv_1x1_4(x4) # [B, F1, H/8, W/8]
|
||||
e5c = self.conv_1x1_5(x5) # [B, F1, H/16, W/16]
|
||||
|
||||
# ---------------- Decoder 4 (produce d4 with out channels F4) ----------------
|
||||
target4 = x4.size()[2:] # spatial of level 4
|
||||
d1_u4 = F.interpolate(e1c, size=target4, mode='bilinear', align_corners=True) # F1
|
||||
d2_u4 = F.interpolate(e2c, size=target4, mode='bilinear', align_corners=True) # F1
|
||||
d3_u4 = F.interpolate(e3c, size=target4, mode='bilinear', align_corners=True) # F1
|
||||
d4_c = e4c # F1
|
||||
d5_u4 = F.interpolate(e5c, size=target4, mode='bilinear', align_corners=True) # F1
|
||||
|
||||
d4_in = torch.cat([d1_u4, d2_u4, d3_u4, d4_c, d5_u4], dim=1) # 5*F1
|
||||
d4 = self.decoder4_cat_conv(d4_in) # out channels = F4
|
||||
|
||||
# ---------------- Decoder 3 (out channels F3) ----------------
|
||||
target3 = x3.size()[2:]
|
||||
d1_u3 = F.interpolate(e1c, size=target3, mode='bilinear', align_corners=True) # F1
|
||||
d2_u3 = F.interpolate(e2c, size=target3, mode='bilinear', align_corners=True) # F1
|
||||
d3_c = e3c # F1
|
||||
d4_u3 = F.interpolate(d4, size=target3, mode='bilinear', align_corners=True) # F4
|
||||
d5_u3 = F.interpolate(e5c, size=target3, mode='bilinear', align_corners=True) # F1
|
||||
|
||||
d3_in = torch.cat([d1_u3, d2_u3, d3_c, d4_u3, d5_u3], dim=1) # channels = 4*F1 + F4
|
||||
d3 = self.decoder3_cat_conv(d3_in) # out channels = F3
|
||||
|
||||
# ---------------- Decoder 2 (out channels F2) ----------------
|
||||
target2 = x2.size()[2:]
|
||||
d1_u2 = F.interpolate(e1c, size=target2, mode='bilinear', align_corners=True) # F1
|
||||
d2_c = e2c # F1
|
||||
d3_u2 = F.interpolate(d3, size=target2, mode='bilinear', align_corners=True) # F3
|
||||
d4_u2 = F.interpolate(d4, size=target2, mode='bilinear', align_corners=True) # F4
|
||||
d5_u2 = F.interpolate(e5c, size=target2, mode='bilinear', align_corners=True) # F1
|
||||
|
||||
d2_in = torch.cat([d1_u2, d2_c, d3_u2, d4_u2, d5_u2], dim=1) # channels = F1 + F1 + F3 + F4 + F1
|
||||
d2 = self.decoder2_cat_conv(d2_in) # out channels = F2
|
||||
|
||||
# ---------------- Decoder 1 (out channels F1) ----------------
|
||||
target1 = x1.size()[2:]
|
||||
d1_c = e1c # F1
|
||||
d2_u1 = F.interpolate(d2, size=target1, mode='bilinear', align_corners=True) # F2
|
||||
d3_u1 = F.interpolate(d3, size=target1, mode='bilinear', align_corners=True) # F3
|
||||
d4_u1 = F.interpolate(d4, size=target1, mode='bilinear', align_corners=True) # F4
|
||||
d5_u1 = F.interpolate(e5c, size=target1, mode='bilinear', align_corners=True) # F1
|
||||
|
||||
d1_in = torch.cat([d1_c, d2_u1, d3_u1, d4_u1, d5_u1], dim=1) # channels = F1 + F2 + F3 + F4 + F1
|
||||
d1 = self.decoder1_cat_conv(d1_in) # out channels = F1
|
||||
|
||||
# ---------------- Final feature fusion ----------------
|
||||
# Fuse d1 and upsampled coarser decoders + upsampled bottleneck
|
||||
# We will form a concat consistent with final_reduce 1x1 (which maps dec1_in -> F1)
|
||||
f_d2 = F.interpolate(d2, size=d1.size()[2:], mode='bilinear', align_corners=True)
|
||||
f_d3 = F.interpolate(d3, size=d1.size()[2:], mode='bilinear', align_corners=True)
|
||||
f_d4 = F.interpolate(d4, size=d1.size()[2:], mode='bilinear', align_corners=True)
|
||||
f_x5 = F.interpolate(e5c, size=d1.size()[2:], mode='bilinear', align_corners=True)
|
||||
|
||||
final_concat = torch.cat([d1, f_d2, f_d3, f_d4, f_x5], dim=1) # channels = F1 + F2 + F3 + F4 + F1
|
||||
# Reduce to F1
|
||||
final_feat = self.final_reduce(final_concat)
|
||||
out = self.final_conv(final_feat) # [B, num_classes, H, W]
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# # quick sanity test (only run if file executed directly)
|
||||
# if __name__ == "__main__":
|
||||
# model = UNet3Plus_Attention(in_channels=3, num_classes=7).cuda()
|
||||
# x = torch.randn(2, 3, 512, 512).cuda()
|
||||
# y = model(x)
|
||||
# print("Output:", y.shape) # expect [2, 7, 512, 512]
|
||||
Reference in New Issue
Block a user