update
This commit is contained in:
DatTT127
Submodule workspace/sprint_1_2/Design_Material/LEGACY/VKIST_ML/codebase-vkist-ultrasound-legacy deleted from 8b2758928a
@@ -1,578 +0,0 @@
|
||||
# KNEE ULTRASOUND ANALYSIS API: TÀI LIỆU KỸ THUẬT & HƯỚNG DẪN SỬ DỤNG
|
||||
|
||||
**Phiên bản:** 1.0 | **Ngày:** 03/2026
|
||||
**Framework:** FastAPI + PyTorch | **Python:** 3.10+
|
||||
|
||||
---
|
||||
|
||||
## 1. Tổng quan hệ thống
|
||||
|
||||
Knee Ultrasound Analysis API được xây dựng bằng FastAPI, chuyên phục vụ tác vụ phân tích ảnh siêu âm đầu gối. Hệ thống tích hợp nhiều mô hình Học máy sâu (Deep Learning) nhằm tự động hóa quy trình phân loại mặt cắt, phát hiện dấu hiệu viêm, phân đoạn cấu trúc giải phẫu và đo lường kích thước tổn thương định lượng.
|
||||
|
||||
### 1.1 Các chức năng chính
|
||||
|
||||
* **Phân loại góc chụp siêu âm (Angle Classification):** Tự động nhận dạng mặt cắt/góc chụp từ ảnh đầu vào bao gồm các nhãn: `med-lat`, `post-trans`, `sup-trans-flex`, và `sup-up-long`.
|
||||
|
||||
* **Phát hiện viêm (Inflammation Detection):** Xác định sự hiện diện của tình trạng viêm khớp gối qua hai góc chụp chính là `sup-up-long` và `post-trans`.
|
||||
|
||||
* **Phân đoạn ảnh ngữ nghĩa (Segmentation):** Tách biệt các cấu trúc giải phẫu đích (dịch khớp, gân, xương, màng hoạt dịch...) thành các phân vùng mặt nạ màu riêng biệt.
|
||||
|
||||
* **Đo độ dày tự động (Thickness Measurement):** Tự động tính toán khoảng cách hình học theo đơn vị milimét ($mm$) giữa các phân vùng mô mềm đã được phân đoạn (chỉ áp dụng đối với mặt cắt góc `sup-up-long`).
|
||||
|
||||
* **Đánh giá mức độ viêm (Severity Analysis):** Xếp hạng thang điểm mức độ nghiêm trọng của viêm từ cấp độ 0 (Rất nhẹ) đến cấp độ 3 (Nặng) dựa trên tỷ lệ diện tích dịch khớp và sự tăng sinh màng hoạt dịch.
|
||||
|
||||
|
||||
|
||||
### 1.2 Luồng xử lý dữ liệu tổng thể (Pipeline Processing)
|
||||
|
||||
Khi nhận tập tin ảnh từ máy trạm (Client), hệ thống sẽ thực hiện phân nhánh xử lý logic động dựa trên kết quả của khối phân loại góc chụp:
|
||||
|
||||
| Góc chụp phát hiện | Quy trình xử lý chi tiết trong Backend pipeline |
|
||||
| --- | --- |
|
||||
| **`post-trans`** | Phân loại góc $\rightarrow$ Phát hiện viêm $\rightarrow$ Phân đoạn ảnh POST $\rightarrow$ Trả kết quả JSON & Mask.|
|
||||
| **`sup-up-long`** | Phân loại góc $\rightarrow$ Phát hiện viêm $\rightarrow$ Phân đoạn ảnh SUP $\rightarrow$ Đo độ dày mô $\rightarrow$ Đánh giá mức độ nặng $\rightarrow$ Trả kết quả.|
|
||||
| **`med-lat`** or **`sup-trans-flex`** | Chỉ thực hiện phân loại góc $\rightarrow$ Trả kết quả trực tiếp (Bỏ qua nhánh phân đoạn & đo lường).|
|
||||
|
||||
```plantuml
|
||||
@startuml
|
||||
skinparam PackageStyle rect
|
||||
skinparam BackgroundColor #FFFFFF
|
||||
skinparam ArrowColor #2C3E50
|
||||
|
||||
title KIẾN TRÚC ĐIỀU HƯỚNG PIPELINE BACKEND (FASTAPI)
|
||||
|
||||
start
|
||||
|
||||
:Nhận ảnh siêu âm thô (Multipart HTTP Post);
|
||||
:Chạy mô hình Phân loại Góc chụp (Angle Classification);
|
||||
|
||||
if (Góc chụp nhận diện được là gì?) then (post-trans)
|
||||
:Chạy mô hình Phát hiện Viêm;
|
||||
:Chạy mô hình Phân đoạn góc POST (Deeplabv3 ResNet101);
|
||||
:Kết xuất Ảnh Overlay Mask & Metadata;
|
||||
elseif (sup-up-long) then (sup-up-long)
|
||||
:Chạy mô hình Phát hiện Viêm;
|
||||
:Chạy mô hình Phân đoạn góc SUP (Tùy chọn: Deeplabv3 / UNet3+ / vv.);
|
||||
:Tính toán độ dày tự động (Thickness Measurement);
|
||||
:Tính toán điểm Severity (Severity Score Analysis);
|
||||
:Kết xuất Ảnh Overlay Mask & Thống kê định lượng;
|
||||
else (med-lat / sup-trans-flex)
|
||||
:Ghi nhận nhãn góc chụp;
|
||||
note right: Không cấu hình phân đoạn\ncho các mặt cắt này
|
||||
endif
|
||||
|
||||
:Đóng gói Response JSON (Success=True);
|
||||
:Trả kết quả về cho Client Application;
|
||||
stop
|
||||
@enduml
|
||||
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Hướng dẫn cài đặt & Triển khai môi trường
|
||||
|
||||
### 2.1 Yêu cầu hệ thống tối thiểu
|
||||
|
||||
| Thành phần cấu phần | Thông số kỹ thuật yêu cầu tối thiểu |
|
||||
| --- | --- |
|
||||
| **Hệ điều hành** | Ubuntu 20.04+ / Windows 10+ / macOS 12+ |
|
||||
| **Môi trường Python** | Phiên bản 3.10 cố định |
|
||||
| **Bộ nhớ RAM** | 16 GB trở lên |
|
||||
| **Bộ xử lý đồ họa (GPU)** | NVIDIA GPU hỗ trợ nền tảng CUDA 12.4 (Khuyến nghị để tối ưu tốc độ) |
|
||||
| **Dung lượng VRAM** | Tối thiểu 8 GB (Khuyến nghị 16 GB nếu chạy song song đồng thời nhiều mô hình)|
|
||||
| **Ổ cứng lưu trữ** | Tối thiểu 15 GB dung lượng trống (Dành cho bộ cài đặt và file weights `.pth`) |
|
||||
| **Bộ công cụ bổ trợ** | CUDA Toolkit 12.4 & cuDNN 9.x tương thích tương ứng|
|
||||
|
||||
### 2.2 Khởi tạo môi trường ảo
|
||||
|
||||
Tải mã nguồn dự án từ kho lưu trữ Git:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/itvkist/vkist-ultrasound.git
|
||||
cd vkist-ultrasound
|
||||
|
||||
```
|
||||
|
||||
Khuyến nghị thiết lập môi trường bằng **Conda** nhằm quản lý và cô lập các gói phụ thuộc:
|
||||
|
||||
```bash
|
||||
conda create -n vkist-ultrasound python=3.10 -y
|
||||
conda activate vkist-ultrasound
|
||||
|
||||
```
|
||||
|
||||
*Hoặc khởi tạo nhanh bằng mô-đun thư viện chuẩn `venv` nếu hệ thống chưa cài đặt Anaconda*:
|
||||
|
||||
```bash
|
||||
# Trên nền tảng hệ điều hành Linux / macOS
|
||||
python3.10 -m venv venv
|
||||
source venv/bin/activate
|
||||
|
||||
# Trên nền tảng hệ điều hành Windows
|
||||
venv\Scripts\activate
|
||||
|
||||
```
|
||||
|
||||
### 2.3 Cài đặt các gói thư viện phụ thuộc (Dependencies)
|
||||
|
||||
Thực hiện cài đặt các thư viện lõi quy định trong tệp cấu hình:
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
|
||||
```
|
||||
|
||||
Trong trường hợp nhân của khung phần mềm PyTorch không nhận diện được phần cứng CUDA, tiến hành ghi đè cài đặt thủ công phiên bản biên dịch GPU:
|
||||
|
||||
```bash
|
||||
pip install torch==2.5.0+cu124 torchvision==0.20.0+cu124 --index-url https://download.pytorch.org/whl/cu124
|
||||
|
||||
```
|
||||
|
||||
> ⚠️ **LƯU Ý QUAN TRỌNG VỀ PACKAGE NATTEN:**
|
||||
> Dòng cấu hình cài đặt gói `natten==0.17.3+torch250cu124` mặc định đã bị gắn chú thích (`#` comment out) trong tệp `requirements.txt`. Nếu bạn sử dụng các kiến trúc mạng Transformer nâng cao yêu cầu gói này, bắt buộc cài đặt thủ công qua liên kết phân phối bánh xe (wheels) chính thức:
|
||||
> `pip install natten==0.17.3+torch250cu124 -f https://shi-labs.com/natten/wheels/`
|
||||
>
|
||||
>
|
||||
|
||||
### 2.4 Cấu trúc cây thư mục dự án chuẩn
|
||||
|
||||
Để đảm bảo hệ thống FastAPI khởi chạy chính xác và tự động nạp các tệp trọng số mô hình, cấu trúc cây thư mục dự án cần được sắp xếp như sau:
|
||||
|
||||
```text
|
||||
project/
|
||||
├── app.py # Trình khởi chạy máy chủ chính - FastAPI Server
|
||||
├── requirements.txt # Danh sách các gói thư viện phụ thuộc hệ thống
|
||||
├── arch/ # Thư mục mã nguồn chứa định nghĩa kiến trúc mạng Custom
|
||||
│ ├── efficientfeedback.py # Định nghĩa mạng EfficientFeedbackNetwork
|
||||
│ └── unet3plus_att.py # Định nghĩa mạng phân đoạn UNet3Plus_Attention
|
||||
├── models/ # Thư mục lưu trữ tệp trọng số nhị phân (.pth)
|
||||
│ ├── best_convnext_tiny.pth
|
||||
│ ├── best_densenet.pth
|
||||
│ ├── best_resnet50.pth
|
||||
│ ├── best_efficientnet_b2.pth
|
||||
│ ├── best_swin_v2_s.pth
|
||||
│ ├── efficientnet_b0_ultrasound_2_class.pth
|
||||
│ ├── best_model_Deeplav3.pth
|
||||
│ ├── unet_resnet101.pth
|
||||
│ ├── efficientfeedback.pth
|
||||
│ ├── unet3plus_att.pth
|
||||
│ └── best_model_deeplabv3_resnet101_seed_16.pth
|
||||
├── templates/ # Các tài nguyên phục vụ giao diện Web tích hợp
|
||||
│ ├── index.html
|
||||
│ ├── css/
|
||||
│ └── js/
|
||||
├── uploads/ # Thư mục lưu trữ tạm ảnh thô đầu vào (Tự động khởi tạo)
|
||||
└── results/ # Thư mục lưu trữ ảnh kết quả phân đoạn (Tự động khởi tạo)
|
||||
|
||||
```
|
||||
|
||||
### 2.5 Khởi động máy chủ dịch vụ
|
||||
|
||||
Thực thi lệnh chạy máy chủ tại thư mục gốc:
|
||||
|
||||
```bash
|
||||
python app.py
|
||||
|
||||
```
|
||||
|
||||
**Mô phỏng nhật ký màn hình console khi máy chủ khởi chạy thành công:**
|
||||
|
||||
```độc_thoại
|
||||
[INFO] Loading deep learning weights safely to memory...
|
||||
[INFO] Device successfully mapped: cuda (NVIDIA GeForce RTX 4090)
|
||||
[INFO] FastAPI Application core initialized.
|
||||
[INFO] Uvicorn server running on http://localhost:8000 (Press CTRL+C to quit)
|
||||
|
||||
```
|
||||
|
||||
* **Giao diện Web UI kiểm thử trực quan:** `http://localhost:8000`
|
||||
|
||||
|
||||
* **Tài liệu API tương tác tự động (Swagger UI):** `http://localhost:8000/docs`
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 3. Tài liệu đặc tả API (API Reference)
|
||||
|
||||
### 3.1 Trạng thái hoạt động (Health Check)
|
||||
|
||||
* **Endpoint:** `GET /api/health`
|
||||
|
||||
|
||||
* **Chức năng:** Kiểm tra tính sẵn sàng phục vụ của cụm dịch vụ API Backend.
|
||||
|
||||
|
||||
* **Định dạng dữ liệu phản hồi (Response JSON):**
|
||||
```json
|
||||
{
|
||||
"status": "healthy"
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
### 3.2 Phân tích ảnh siêu âm chính - `POST /api/analyze`
|
||||
|
||||
Mã hóa dữ liệu đầu vào dưới dạng tệp tin `multipart/form-data` để gửi tới mạng mạng nơ-ron xử lý.
|
||||
|
||||
#### Các tham số yêu cầu (Request Parameters)
|
||||
|
||||
| Tham số cấu hình | Phương thức truyền | Kiểu dữ liệu | Giá trị mặc định | Định nghĩa chức năng chi tiết |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **`image`** | Multipart Form | Binary File | *Bắt buộc* | Tệp tin ảnh siêu âm đầu gối cần xử lý (Hỗ trợ mở rộng định dạng: `.jpg`, `.png`, `.bmp`).|
|
||||
| **`angle_model`** | Query String | String | `convnext` | Tên định danh mô hình đảm nhận tác vụ phân loại góc chụp.|
|
||||
| **`inflammation_model`** | Query String | String | `efficientnet_b0` | Mô hình phát hiện tình trạng viêm (Hiện tại cố định cấu hình mạng).|
|
||||
| **`segment_model_sup`** | Query String | String | `deeplabv3` | Mô hình phân đoạn cấu trúc giải phẫu dành cho mặt cắt góc `sup-up-long`.|
|
||||
| **`segment_model_post`** | Query String | String | `deeplabv3_resnet101` | Mô hình phân đoạn cấu trúc giải phẫu dành cho mặt cắt góc `post-trans`.|
|
||||
|
||||
#### Danh sách định danh mô hình khả dụng trong hệ thống
|
||||
|
||||
| Phân nhóm Task | Tên tham số truyền vào | Kiến trúc mạng nơ-ron gốc | Mô tả đặc tính đầu ra |
|
||||
| --- | --- | --- | --- |
|
||||
| **Phân loại Góc chụp** | `convnext` | ConvNeXt Tiny | Phân cấp phân loại ra 4 lớp nhãn đầu ra.|
|
||||
| | `densenet` | DenseNet-121 | Mạng kết nối dày đặc.|
|
||||
| | `resnet50` | ResNet-50 | Kiến trúc mạng dư thừa tiêu chuẩn.|
|
||||
| | `efficientnet_b2` | EfficientNet-B2 | Tối ưu hóa đa quy mô tài nguyên mạng.|
|
||||
| | `swin` | Swin Transformer V2-S | Kiến trúc Attention cửa sổ dịch chuyển.|
|
||||
| **Phân đoạn góc SUP** | `deeplabv3` | DeepLabV3 ResNet-50 | Trích xuất đặc trưng đa tỷ lệ với 7 lớp đầu ra.|
|
||||
| | `unet_resnet101` | UNet + ResNet-101 | Kiến trúc Encoder-Decoder kết hợp ResNet.|
|
||||
| | `efficientfeedback` | EfficientFeedbackNetwork | Thiết kế tùy biến riêng có liên kết phản hồi dữ liệu.|
|
||||
| | `unet3plus` | UNet3+ with Attention | Cơ chế Attention kết hợp kết nối toàn diện Full-scale.|
|
||||
| **Phân đoạn góc POST** | `deeplabv3_resnet101` | DeepLabV3 ResNet-101 | Cấu trúc chuyên sâu phân đoạn góc nhìn mặt sau.|
|
||||
|
||||
#### Cấu trúc dữ liệu JSON phản hồi (Response Body Schema)
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"filename": "d3b07384-d113-4ec8-a141-8664b07fb8d3.jpg",
|
||||
"images": {
|
||||
"original": "/uploads/d3b07384-d113-4ec8-a141-8664b07fb8d3.jpg",
|
||||
"segmented": "/results/seg_d3b07384-d113-4ec8-a141-8664b07fb8d3.jpg"
|
||||
},
|
||||
"models_used": {
|
||||
"angle_model": "convnext",
|
||||
"inflammation_model": "efficientnet_b0",
|
||||
"segmentation_model": "deeplabv3"
|
||||
},
|
||||
"angle": {
|
||||
"class": "sup-up-long",
|
||||
"confidence": 98.45
|
||||
},
|
||||
"inflammation": {
|
||||
"detected": true,
|
||||
"confidence": 94.20
|
||||
},
|
||||
"segmentation": {
|
||||
"angle_type": "sup",
|
||||
"classes_detected": ["background", "effusion", "fat", "femur", "synovium", "tendon"],
|
||||
"color_legend": {
|
||||
"background": [0, 0, 0],
|
||||
"effusion": [255, 0, 0],
|
||||
"synovium": [255, 0, 255]
|
||||
}
|
||||
},
|
||||
"measurement": {
|
||||
"thickness_mm": 6.87,
|
||||
"thickness_px": 100,
|
||||
"location_x": 256
|
||||
},
|
||||
"severity": {
|
||||
"level": 3,
|
||||
"severity": "Nặng",
|
||||
"combined_score": 18.4,
|
||||
"effusion": {
|
||||
"pixels": 25400,
|
||||
"ratio": 0.096,
|
||||
"thickness": 6.87
|
||||
},
|
||||
"synovium": {
|
||||
"pixels": 12500,
|
||||
"ratio": 0.047
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Ví dụ mã triển khai gọi dịch vụ (Client Invocations)
|
||||
|
||||
* **Sử dụng lệnh Client cURL CLI:**
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/api/analyze?angle_model=convnext&segment_model_sup=deeplabv3" \
|
||||
-F "image=@/data/medical/knee_sample.jpg"
|
||||
|
||||
```
|
||||
|
||||
* **Triển khai ứng dụng gọi qua script Python (Requests):**
|
||||
|
||||
```python
|
||||
import requests
|
||||
|
||||
url = "http://localhost:8000/api/analyze"
|
||||
query_parameters = {
|
||||
"angle_model": "swin",
|
||||
"segment_model_sup": "unet3plus"
|
||||
}
|
||||
|
||||
target_image_path = "path/to/clinical_knee.jpg"
|
||||
with open(target_image_path, "rb") as image_file:
|
||||
payload = {"image": image_file}
|
||||
api_response = requests.post(url, params=query_parameters, files=payload)
|
||||
|
||||
parsed_result = api_response.json()
|
||||
print("Phân loại góc:", parsed_result["angle"]["class"])
|
||||
print("Số liệu đo lường hình học:", parsed_result.get("measurement"))
|
||||
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. Các thông số cấu hình lõi hệ thống
|
||||
|
||||
### 4.1 Hằng số hệ thống trong `app.py`
|
||||
|
||||
| Tên định danh hằng số | Giá trị mặc định | Diễn giải chức năng kỹ thuật |
|
||||
| --- | --- | --- |
|
||||
| `UPLOAD_FOLDER` | `'uploads'` | Đường dẫn cục bộ lưu trữ file ảnh thô nhận từ máy trạm.|
|
||||
| `RESULTS_FOLDER` | `'results'` | Đường dẫn lưu ảnh màu sau phân đoạn (Color Mask Overlayed).|
|
||||
| `TEMPLATES_FOLDER` | `'templates'` | Thư mục chứa mã nguồn giao diện phân tích Web UI.|
|
||||
| `PIXEL_TO_MM` | $\frac{45.0}{655.0} \approx 0.0687$ | Hệ số chuyển đổi từ độ phân giải pixel sang kích thước thực tế ($mm$). Phụ thuộc cố định vào cấu hình đầu ra của phần cứng máy quét siêu âm.|
|
||||
| `DEFAULT_MEASURE_IDS` | `[1, 5]` | Danh sách mảng chứa ID nhãn lớp cấu trúc giải phẫu kích hoạt thuật toán đo độ dày: `1 = effusion` (Dịch khớp), `5 = synovium` (Màng hoạt dịch).|
|
||||
| `device` | `cuda` hoặc `cpu` | Khối phần cứng thực thi tính toán đồ họa (Tự động thiết lập dựa trên tính khả dụng của driver NVIDIA).|
|
||||
|
||||
### 4.2 Cấu hình Pipeline tiền xử lý và biến đổi ma trận ảnh (Transforms)
|
||||
|
||||
Hệ thống phân tách ảnh đầu vào thành các luồng biến đổi riêng biệt trước khi nạp vào tensor mô hình tùy thuộc vào mục tiêu xử lý chuyên biệt:
|
||||
|
||||
| Luồng xử lý Pipeline ảnh | Kích thước chuyển đổi (Resize) | Quy định chuẩn hóa phân phối ma trận (Normalization) |
|
||||
| --- | --- | --- |
|
||||
| **Phân loại góc & Phát hiện viêm** | <br>$224 \times 224$ pixel | Áp dụng phân phối phân cấp:<br> $\text{mean} = [0.485, 0.456, 0.406]$, <br> $\text{std} = [0.229, 0.224, 0.225]$ |
|
||||
| **Phân đoạn cấu trúc (Segmentation)** | <br>$512 \times 512$ pixel | Không áp dụng chuẩn hóa phân phối (Chỉ thực thi hàm chuyển đổi tensor `ToTensor()`) |
|
||||
|
||||
---
|
||||
|
||||
## 5. Ràng buộc kỹ thuật & Quy tắc thiết kế hệ thống
|
||||
|
||||
### 5.1 Quản lý và giải phóng tài nguyên bộ nhớ GPU (VRAM Leak Warning)
|
||||
|
||||
Trong phiên bản hiện tại, logic xử lý nội tại của API kích hoạt các hàm `load_angle_model()`, `load_inflammation_model()`, và `load_segmentation_model_*()` trực tiếp bên trong vòng đời của mỗi phiên request nhận về. Hành vi này ép buộc GPU liên tục nạp lại dữ liệu tệp `.pth` vào VRAM cho mỗi giao dịch HTTP, sinh ra độ trễ (Overhead) I/O lớn và tiềm ẩn nguy cơ tràn bộ nhớ hệ thống. Khi triển khai môi trường Production, bắt buộc phải tái cấu trúc chuyển các hàm này thành Singleton dịch vụ (Tải một lần duy nhất lúc khởi động tiến trình Web Server).
|
||||
|
||||
### 5.2 Ràng buộc phi tuyến tính của tham số vật lý `PIXEL_TO_MM`
|
||||
|
||||
Hằng số quy đổi $\text{PIXEL\_TO\_MM} = \frac{45.0}{655.0}$ là một giá trị được cấu hình cứng (Hardcoded) trong mã nguồn, đặc trưng duy nhất cho một dòng máy siêu âm lâm sàng có tỷ lệ hiển thị $45mm$ tương đương với độ phân giải vùng quét $655\text{ px}$. Khi hệ thống thu thập ảnh siêu âm từ các thiết bị chuẩn đoán hình ảnh khác, hoặc thay đổi độ phân giải ảnh xuất ra, số liệu đo khoảng cách tổn thương sẽ sai lệch nghiêm trọng nếu hằng số này không được hiệu chuẩn lại thông qua ma trận nội quan của máy quét mới.
|
||||
|
||||
### 5.3 Quy tắc ánh xạ phân lớp (Class Remapping Matrix) đối với mô hình Custom
|
||||
|
||||
Hai mô hình tùy biến sâu phục vụ mặt cắt góc nhìn phía trên bánh chè (`UNet3+` và `EfficientFeedback Network`) được huấn luyện trên tập dữ liệu đặc thù sở hữu thứ tự cấu trúc mảng nhãn đầu ra lệch pha hoàn toàn so với kiến trúc phân cấp chuẩn của hệ thống. Để thống nhất dữ liệu trả về cho Client, khối Backend API thực hiện cơ chế tự động chuyển đổi chỉ mục mảng (Index Remapping) theo bảng đặc tả logic dưới đây:
|
||||
|
||||
| Chỉ mục Mô hình gốc (Output Model Index) | Chỉ mục chuẩn hóa hệ thống (Standard System Index) | Tên nhãn lớp giải phẫu tương ứng (Anatomical Label Class) |
|
||||
| --- | --- | --- |
|
||||
| `0` | `0` | **`background`** (Nền ảnh không chứa cấu trúc) |
|
||||
| `1` | `2` | **`fat`** (Lớp mô mỡ dưới da) |
|
||||
| `2` | `6` | **`tendon`** (Cấu trúc gân cơ) |
|
||||
| `3` | `1` | **`effusion`** (Vùng tụ dịch khớp gối ổ viêm) |
|
||||
| `4` | `4` | **`femur`** (Ranh giới cấu trúc xương đùi) |
|
||||
| `5` | `5` | **`synovium`** (Màng hoạt dịch bao quanh khớp) |
|
||||
| `6` | `3` | **`fat-pat`** (Tổ chức mỡ Hoffa) |
|
||||
|
||||
### 5.4 Cơ chế tự động dọn dẹp tập tin tồn đọng (Garbage Collection Task)
|
||||
|
||||
Các tập tin ảnh thô tải lên thư mục `uploads/` và ảnh xử lý nhị phân kết xuất lưu trong `results/` được ghi dưới dạng định danh chuỗi không trùng lặp UUID và lưu trữ vô thời hạn trên đĩa cứng hệ thống. Hệ thống lõi của ứng dụng không tích hợp cơ chế tự động giải phóng (Auto-deletion) các tệp cũ. Khi chạy vận hành dài hạn trong hệ thống y tế thực tế, bắt buộc phải cấu hình thêm Background Task (sử dụng thư viện Asyncio) hoặc thiết lập dịch vụ Cronjob của hệ điều hành để dọn dẹp định kỳ tránh cạn kiệt dung lượng ổ đĩa lưu trữ.
|
||||
|
||||
---
|
||||
|
||||
## 6. Giải pháp mở rộng tính năng mã nguồn (Backend Optimization Guide)
|
||||
|
||||
### 6.1 Tăng tốc độ phản hồi bằng Cơ chế Caching Mô hình Toàn cục
|
||||
|
||||
Thay thế kiến trúc nạp tải mô hình cũ bằng một kho lưu trữ Cache tĩnh trong bộ nhớ RAM, tối ưu hóa thời gian xử lý request từ mức giây xuống mức mili-giây:
|
||||
|
||||
```python
|
||||
# Cấu hình biến lưu trữ toàn cục ở đầu tệp app.py
|
||||
global_model_cache = {}
|
||||
|
||||
def get_cached_angle_model(selected_model_name: str):
|
||||
cache_lookup_key = f"angle_classification_{selected_model_name}"
|
||||
if cache_lookup_key not in global_model_cache:
|
||||
# Thực hiện nạp trọng số mô hình từ đĩa cứng lần đầu tiên
|
||||
global_model_cache[cache_lookup_key] = load_angle_model(selected_model_name)
|
||||
return global_model_cache[cache_lookup_key]
|
||||
|
||||
# Áp dụng logic tương tự cho get_inflammation_model(), get_seg_model_sup(), get_seg_model_post()
|
||||
|
||||
```
|
||||
|
||||
### 6.2 Thêm mới một kiến trúc phân loại góc chụp (Ví dụ: Vision Transformer - ViT)
|
||||
|
||||
Để tích hợp một mạng nơ-ron mới vào hệ thống xử lý, tuân thủ nghiêm ngặt quy trình 3 bước sau:
|
||||
|
||||
* **Bước 1:** Bổ sung khối xử lý điều kiện rẽ nhánh logic vào hàm khởi tạo mô hình `load_angle_model()`:
|
||||
|
||||
```python
|
||||
elif model_name == 'vit':
|
||||
from torchvision.models import vit_b_16
|
||||
# Khởi tạo mạng mạng ViT không nạp trọng số mặc định ImageNet
|
||||
model = vit_b_16(weights=None)
|
||||
# Tái cấu trúc tầng phân loại tuyến tính cuối cùng tương thích với 4 lớp nhãn đầu gối
|
||||
model.heads[0] = nn.Linear(model.heads[0].in_features, 4)
|
||||
# Tải tệp trọng số huấn luyện cục bộ từ thư mục mô hình
|
||||
checkpoint_tensor = torch.load('models/best_vit_b16.pth', map_location=device, weights_only=False)
|
||||
model.load_state_dict(checkpoint_tensor)
|
||||
|
||||
```
|
||||
|
||||
* **Bước 2:** Di chuyển tệp trọng số huấn luyện nhị phân của mạng (`best_vit_b16.pth`) vào chính xác không gian lưu trữ của thư mục `/models/`.
|
||||
|
||||
* **Bước 3:** Ứng dụng phía Client có thể kích hoạt mạng mới bằng cách truyền giá trị định danh qua tham số URL: `/api/analyze?angle_model=vit`.
|
||||
|
||||
|
||||
|
||||
### 6.3 Xử lý luồng dữ liệu phân đoạn song song đồng thời (Batch Processing API)
|
||||
|
||||
Tối ưu hóa năng lực phục vụ của Server đối với bài toán nhận diện hàng loạt ảnh cùng lúc từ phòng khám bằng endpoint xử lý không đồng bộ:
|
||||
|
||||
```python
|
||||
from fastapi import errors, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from typing import List
|
||||
|
||||
@app.post('/api/analyze_batch', status_code=status.HTTP_200_OK)
|
||||
async def analyze_batch_images(images: List[UploadFile] = File(...)):
|
||||
import asyncio
|
||||
|
||||
# Đóng gói các tác vụ phân tích ảnh đơn lẻ vào một danh sách hàng đợi task không đồng bộ
|
||||
async_tasks_queue = [analyze_single_image_pipeline(img) for img in images]
|
||||
|
||||
# Kích hoạt thực thi đồng thời trên luồng phần cứng thông qua gather cơ chế
|
||||
compiled_batch_results = await asyncio.gather(*async_tasks_queue)
|
||||
|
||||
return JSONResponse({
|
||||
"results": compiled_batch_results,
|
||||
"processed_count": len(compiled_batch_results)
|
||||
})
|
||||
|
||||
```
|
||||
|
||||
### 6.4 Bản đóng gói container hóa ứng dụng (Production Dockerfile)
|
||||
|
||||
Đóng gói toàn bộ ML Stack bao gồm trình điều khiển GPU NVIDIA CUDA để triển khai đồng bộ trên các hạ tầng Cloud hoặc máy chủ On-Premise của bệnh viện:
|
||||
|
||||
```dockerfile
|
||||
# Sử dụng Base Image chứa sẵn môi trường CUDA 12.4 và cuDNN 9 của NVIDIA
|
||||
FROM nvidia/cuda:12.4.0-cudnn9-runtime-ubuntu22.04
|
||||
|
||||
# Cài đặt môi trường Python 3.10 và các gói hệ thống cốt lõi
|
||||
RUN apt-get update && apt-get install -y \
|
||||
python3.10 \
|
||||
python3-pip \
|
||||
git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Thiết lập không gian làm việc nội bộ bên trong container
|
||||
WORKDIR /app
|
||||
|
||||
# Sao chép và cài đặt danh sách các thư viện Python
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Sao chép toàn bộ mã nguồn ứng dụng vào Container
|
||||
COPY . .
|
||||
|
||||
# Lắng nghe và kích hoạt ứng dụng Web Server FastAPI
|
||||
CMD ["python", "app.py"]
|
||||
|
||||
```
|
||||
|
||||
* **Lệnh khởi dựng Image hệ thống:** `docker build -t medical-api-service .`
|
||||
|
||||
* **Lệnh kích hoạt Container chia sẻ tài nguyên phần cứng GPU vật lý:**
|
||||
|
||||
```bash
|
||||
docker run --gpus all -p 8000:8000 -v $(pwd)/models:/app/models medical-api-service
|
||||
|
||||
```
|
||||
|
||||
### 6.5 Bộ chuyển đổi tiếp nhận trực tiếp luồng dữ liệu ảnh y tế chuẩn DICOM
|
||||
|
||||
Mở rộng chức năng cho phép hệ thống API đọc trực tiếp tệp tin ảnh gốc dạng `.dcm` trích xuất trực tiếp từ các thiết bị siêu âm chuẩn lâm sàng trong bệnh viện mà không cần qua bước chuyển đổi định dạng thủ công:
|
||||
|
||||
```python
|
||||
import pydicom
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
@app.post('/api/analyze_dicom')
|
||||
async def analyze_dicom_file(file: UploadFile = File(...)):
|
||||
# Đọc luồng byte nhị phân trực tiếp từ tệp DICOM tải lên
|
||||
dicom_dataset = pydicom.dcmread(file.file)
|
||||
|
||||
# Trích xuất ma trận pixel thô từ thẻ DICOM Pixel Data
|
||||
raw_pixel_array = dicom_dataset.pixel_array
|
||||
|
||||
# Chuyển đổi ma trận mảng Numpy sang định dạng ảnh PIL tương thích với Pipeline biến đổi
|
||||
converted_image_pil = Image.fromarray(raw_pixel_array).convert('RGB')
|
||||
|
||||
# Chuyển tiếp ảnh đã đổi định dạng vào luồng xử lý tự động nội bộ của API
|
||||
analysis_output_json = await execute_core_analysis_pipeline(converted_image_pil)
|
||||
return analysis_output_json
|
||||
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phụ lục: Đặc tả Dữ liệu định lượng lâm sàng
|
||||
|
||||
### Phụ lục A: Bảng phân định mã màu mặt nạ phân đoạn ngữ nghĩa (Color Map Legend)
|
||||
|
||||
1. Cấu trúc Mặt cắt mặt trên bánh chè - Góc SUP (`sup-up-long`)
|
||||
|
||||
Góc SUP tập trung khoanh vùng các lớp mô mềm phía trước đầu gối phục vụ thuật toán tính toán độ dày dịch tụ.
|
||||
|
||||
| Từ khóa nhãn (Key) | Cấu trúc giải phẫu đích | Mã màu hiển thị (RGB) | Trực quan hóa màu sắc |
|
||||
| --- | --- | --- | --- |
|
||||
| `background` | Nền ảnh siêu âm | `[0, 0, 0]` | ⬛ Đen (Không chứa dữ liệu) |
|
||||
| `effusion` | Vùng dịch khớp tụ ổ viêm | `[255, 0, 0]` | 🟥 Đỏ |
|
||||
| `fat` | Tổ chức mô mỡ dưới da | `[255, 255, 0]` | 🟨 Vàng |
|
||||
| `fat-pat` | Khối mỡ Hoffa | `[0, 255, 255]` | 🟦 Lam sáng |
|
||||
| `femur` | Cấu trúc bề mặt xương đùi | `[0, 255, 0]` | 🟩 Xanh lá |
|
||||
| `synovium` | Lớp màng hoạt dịch tăng sinh | `[255, 0, 255]` | 🟪 Tím |
|
||||
| `tendon` | Vùng bó gân cơ | `[0, 0, 255]` | 🟦 Xanh dương |
|
||||
|
||||
> 🔄 **QUY TẮC CHUYỂN ĐỔI CHUYỂN GÓC (SUP $\rightarrow$ POST):**
|
||||
> Khi hệ thống chuyển đổi trạng thái phân tích sang mặt cắt phía sau khớp gối (Góc `POST`), ma trận thuật toán phân đoạn sẽ tự động tái cấu trúc màu sắc ngữ nghĩa: Vùng tổn thương chứa **`effusion`** (màu đỏ) sẽ chuyển trạng thái biểu diễn thành **`baker's cyst`** (Kén Baker), và tổ chức cấu trúc vùng **`fat-pat`** (màu lam sáng) sẽ hoán đổi ý nghĩa thành vùng **`muscle`** (Cơ bắp vùng khoeo).
|
||||
>
|
||||
>
|
||||
|
||||
2. Cấu trúc Mặt cắt mặt sau vùng khoeo chân - Góc POST (`post-trans`)
|
||||
|
||||
| Từ khóa nhãn (Key) | Cấu trúc giải phẫu đích | Mã màu hiển thị (RGB) | Trực quan hóa màu sắc |
|
||||
| --- | --- | --- | --- |
|
||||
| `background` | Nền ảnh siêu âm | `[0, 0, 0]` | ⬛ Đen |
|
||||
| `baker's cyst` | Tổ chức kén hoạt dịch vùng khoeo (Baker) | `[255, 0, 0]` | 🟥 Đỏ |
|
||||
| `fat` | Lớp mô mỡ | `[255, 255, 0]` | 🟨 Vàng |
|
||||
| `muscle` | Các nhóm cơ bắp vùng sau gối | `[0, 255, 255]` | 🟦 Lam sáng |
|
||||
| `femur` | Cấu trúc xương đùi sau | `[0, 255, 0]` | 🟩 Xanh lá |
|
||||
| `synovium` | Màng hoạt dịch mặt sau | `[255, 0, 255]` | 🟪 Tím |
|
||||
| `tendon` | Hệ thống gân cơ mặt sau | `[0, 0, 255]` | 🟦 Xanh dương |
|
||||
|
||||
---
|
||||
|
||||
### Phụ lục B: Thang điểm đánh giá mức độ nghiêm trọng của ổ viêm (Clinical Severity Score)
|
||||
|
||||
Hệ thống chấm điểm toán học tự động căn cứ trên trọng số diện tích và độ dày phân tách để đưa ra kết luận mức độ bệnh lý lâm sàng thông qua phương trình tuyến tính tổng hợp:
|
||||
|
||||
$$\text{combined\_score} = \text{effusion\_score} \times 0.6 + \text{synovium\_score} \times 0.4$$
|
||||
|
||||
Dựa trên kết quả giá trị của biến số $\text{combined\_score}$, hệ thống tự động phân cấp thành 4 ngưỡng trạng thái lâm sàng tương ứng:
|
||||
|
||||
* **Mức 0 - Rất nhẹ ($\text{score} < 3$):** Trạng thái ổ dịch khớp và cấu trúc màng hoạt dịch nằm hoàn toàn trong giới hạn sinh lý bình thường của cơ thể.
|
||||
* **Mức 1 - Nhẹ ($\text{score}$ từ $3$ đến $7.9$):** Xuất hiện hiện tượng tụ dịch khớp lớp mỏng, màng hoạt dịch có dấu hiệu tăng sinh nhẹ cấu trúc màng.
|
||||
* **Mức 2 - Trung bình ($\text{score}$ từ $8$ đến $15$):** Lượng dịch tụ khớp gối ở mức độ vừa phải, màng hoạt dịch bắt đầu phì đại và tăng sinh rõ nét.
|
||||
* **Mức 3 - Nặng ($\text{score} > 15$):** Lớp tụ dịch khớp gối dày kích thước lớn, màng hoạt dịch tăng sinh phì đại mạnh, lan rộng diện tích cấu trúc giải phẫu xung quanh.
|
||||
@@ -1,267 +0,0 @@
|
||||
Dưới đây là toàn bộ nội dung tài liệu về **Công nghệ siêu âm khớp gối (PILOT)** đã được làm sạch, đồng bộ hóa cấu trúc và làm giàu thông tin (enrichment). Các hình ảnh mất mát và sơ đồ quy trình đã được mã hóa chi tiết bằng ngôn ngữ **PlantUML** cùng với các mô tả dữ liệu cấu trúc trực quan nhằm phục vụ tối ưu cho việc huấn luyện, tích hợp hoặc phát triển hệ thống backend của kỹ sư AI/ML stack.
|
||||
|
||||
---
|
||||
|
||||
# CÔNG NGHỆ SIÊU ÂM KHỚP GỐI (PILOT)
|
||||
|
||||
**Tác giả:** Nguyễn Đăng Hà
|
||||
**Đơn vị phát triển:** VKIST (Viện Khoa học và Công nghệ Việt Nam - Hàn Quốc)
|
||||
|
||||
---
|
||||
|
||||
## 1. Giới thiệu chung & Mục tiêu nghiên cứu
|
||||
|
||||
### Giới thiệu chung
|
||||
|
||||
* Tràn dịch khớp gối là một trong những biểu hiện lâm sàng xuất hiện thường gặp của các bệnh lý liên quan đến cơ - xương - khớp.
|
||||
|
||||
|
||||
* Việc nhận diện chính xác và đánh giá chi tiết mức độ viêm khớp gối có ý nghĩa vô cùng quan trọng đối với quá trình điều trị cũng như phục hồi chức năng của bệnh nhân.
|
||||
|
||||
|
||||
* Hiện nay, phương pháp siêu âm được xem là giải pháp hiệu quả hàng đầu để đánh giá tình trạng này nhờ vào các đặc tính: an toàn, hoàn toàn không xâm lấn và tiết kiệm tối đa chi phí cho người bệnh.
|
||||
|
||||
|
||||
|
||||
### Mục tiêu nghiên cứu của dự án
|
||||
|
||||
1. Xây dựng một quy trình chuẩn hóa trong siêu âm chẩn đoán và thiết lập một cơ sở dữ liệu hình ảnh lớn (Large Dataset) chuyên biệt về tràn dịch khớp gối.
|
||||
|
||||
|
||||
2. Nghiên cứu và phát triển phần mềm ứng dụng Trí tuệ nhân tạo (AI) giúp hỗ trợ các bác sĩ lâm sàng chẩn đoán nhanh chóng, hiệu quả tình trạng tràn dịch khớp gối.
|
||||
|
||||
|
||||
3. Tiến hành thử nghiệm lâm sàng, kiểm thử và đánh giá độ chính xác của thuật toán AI trên các tập dữ liệu thực tế thu thập tại Bệnh viện E.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 2. Quy trình xử lý dữ liệu và Kiến trúc hệ thống (Pipeline AI)
|
||||
|
||||
Hệ thống xử lý ảnh siêu âm được thiết kế theo một chuỗi pipeline tuần tự từ ảnh thô (Raw Image) đầu vào cho đến khi kết xuất báo cáo định lượng. Dưới đây là sơ đồ kiến trúc luồng dữ liệu của hệ thống:
|
||||
|
||||
```plantuml
|
||||
@startuml
|
||||
skinparam handwritten false
|
||||
skinparam monochrome false
|
||||
skinparam packageStyle rect
|
||||
skinparam shadowing true
|
||||
|
||||
title SƠ ĐỒ PIPELINE XỬ LÝ ẢNH SIÊU ÂM KHỚP GỐI (AI BACKEND)
|
||||
|
||||
start
|
||||
|
||||
:Ảnh siêu âm đầu gối thô (Raw Image Input);
|
||||
note right: Tải lên từ client thông qua Web UI
|
||||
|
||||
partition "Khối Tiền Xử Lý (Preprocessing Block)" {
|
||||
:Tiền xử lý dữ liệu ảnh siêu âm;
|
||||
note right: Chuẩn hóa kích thước, lọc nhiễu, cân bằng độ tương phản
|
||||
}
|
||||
|
||||
partition "Khối Phân Loại (Classification Block)" {
|
||||
:Phân loại góc chụp (Scan View Classification);
|
||||
note right: Nhận diện mặt cắt: L SUPRAPAT LONG, L POST TRANS, vv.
|
||||
}
|
||||
|
||||
partition "Khối Nhận Diện & Phân Đoạn (Detection & Segmentation Block)" {
|
||||
:Nhận biết tình trạng viêm;
|
||||
if (Phát hiện thấy có viêm?) then (Có viêm)
|
||||
:Khoanh vùng viêm (Lesion Segmentation);
|
||||
note right: Sử dụng các mô hình: DeepLabV3, MedSAM, UltraSAM
|
||||
else (Không viêm)
|
||||
:Kết luận: Bình thường;
|
||||
detach
|
||||
endif
|
||||
}
|
||||
|
||||
partition "Khối Định Lượng & Phân Tích (Quantitative Analytics Block)" {
|
||||
:Phân tích mức độ viêm;
|
||||
note right: Tính toán diện tích vùng viêm,\ntỷ lệ phần trăm pixel viêm, quan hệ không gian với mô gân
|
||||
}
|
||||
|
||||
:Kết xuất kết quả (Output Analytics & Report);
|
||||
note right: Trả về file JSON kết quả, ảnh phân đoạn \nvà lưu trữ vào cơ sở dữ liệu hệ thống
|
||||
|
||||
stop
|
||||
@enduml
|
||||
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Các Mô hình Deep Learning áp dụng trong hệ thống
|
||||
|
||||
Mô hình AI được chia tách thành ba nhiệm vụ chính: Phân loại mặt cắt ảnh siêu âm, Phát hiện viêm (biến cố cố định) và Phân đoạn ngữ nghĩa (Semantic Segmentation) các vùng giải phẫu tổn thương.
|
||||
|
||||
### 3.1. Mô hình Phân loại Góc chụp (Scan View Classification - `angle_model`)
|
||||
|
||||
* **Nhiệm vụ:** Tự động nhận diện và phân loại cấu trúc tư thế đặt đầu dò siêu âm.
|
||||
* **Mô hình mặc định:** `convnext`
|
||||
* **Các kiến trúc hỗ trợ:**
|
||||
* **ConvNeXt Tiny:** Cấu hình gồm 4 lớp đầu ra (Mặc định).
|
||||
* **DenseNet-121:** Trích xuất đặc trưng sâu nhờ cơ chế kết nối dày đặc.
|
||||
* **ResNet-50:** Kiến trúc mạng phần dư chuẩn hóa giúp tối ưu hóa quá trình hội tụ.
|
||||
* **EfficientNet-B2:** Cân bằng tối ưu giữa hiệu năng và tài nguyên tính toán.
|
||||
* **Swin Transformer V2-S:** Mô hình dựa trên cơ chế Attention dịch chuyển cửa sổ, nâng cao khả năng học đặc trưng toàn cục.
|
||||
|
||||
* **Các mặt cắt mục tiêu chính:**
|
||||
* `L SUPRAPAT LONG` (Mặt cắt dọc trên xương bánh chè khớp gối trái).
|
||||
* `L POST TRANS` (Mặt cắt ngang phía sau khớp gối trái).
|
||||
|
||||
### 3.2. Mô hình Phát hiện Viêm (Inflammation Detection - `inflammation_model`)
|
||||
|
||||
* **Nhiệm vụ:** Tự động phát hiện và đánh giá tình trạng viêm (hiện cố định trong hệ thống).
|
||||
* **Mô hình mặc định:** `efficientnet_b0`
|
||||
|
||||
### 3.3. Mô hình Khoanh vùng & Phân đoạn tổn thương (Segmentation Models)
|
||||
|
||||
Hệ thống cho phép tùy chọn linh hoạt các kiến trúc mạng tiên tiến nhằm khoanh vùng chính xác các cấu trúc giải phẫu và ổ viêm theo từng nhóm góc chụp:
|
||||
|
||||
#### 3.3.1. Nhóm Phân đoạn SUP (Mặt cắt dọc trên xương bánh chè - `segment_model_sup`)
|
||||
|
||||
| Tên Mô Hình (Giá trị truyền vào) | Kiến trúc kĩ thuật & Vai trò trong hệ thống |
|
||||
| --- | --- |
|
||||
| **deeplabv3** | Kiến trúc DeepLabV3 ResNet-50 — 7 lớp (Mặc định). Phân đoạn phân cấp chuẩn giúp nhận diện ranh giới vùng mô mềm phức tạp. |
|
||||
| **unet_resnet101** | Kiến trúc UNet kết hợp với encoder ResNet-101 mạnh mẽ, tối ưu việc khôi phục chi tiết không gian ảnh. |
|
||||
| **efficientfeedback** | Kiến trúc EfficientFeedbackNetwork (tùy chỉnh), tăng cường cơ chế phản hồi đặc trưng để làm mịn biên tổn thương. |
|
||||
| **unet3plus** | Kiến trúc UNet3+ kết hợp cơ chế Attention (tùy chỉnh), tối ưu khả năng kết nối bỏ qua full-scale cho ảnh siêu âm. |
|
||||
|
||||
#### 3.3.2. Nhóm Phân đoạn POST (Mặt cắt ngang phía sau - `segment_model_post`)
|
||||
|
||||
| Tên Mô Hình (Giá trị truyền vào) | Kiến trúc kĩ thuật & Vai trò trong hệ thống |
|
||||
| --- | --- |
|
||||
| **deeplabv3_resnet101** | Kiến trúc DeepLabV3 ResNet-101 — 7 lớp (Mặc định). Chuyên biệt cho việc phân đoạn góc post-trans, tối ưu hóa độ chính xác biên tổn thương vùng khoeo sau gối. |
|
||||
|
||||
*(Lưu ý: Các mô hình MedSAM và UltraSAM không xuất hiện trong danh mục cấu hình kỹ thuật của phiên bản Pilot 2025 này, nên đã được thay thế bằng các mô hình thực tế hệ thống đang hỗ trợ bao gồm UNet và EfficientFeedback).*
|
||||
|
||||
### 3.3. Cơ chế phân tích định lượng mức độ viêm (Severity Analytics)
|
||||
|
||||
Sau khi mô hình phân đoạn (ví dụ: DeepLabV3) đưa ra mặt nạ phân đoạn (Segmentation Mask) , hệ thống thực hiện thuật toán đếm pixel để đưa ra báo cáo tự động:
|
||||
|
||||
* **Phân cấp độ viêm:** Hệ thống tự động phân loại thành 3 mức độ từ nhẹ đến nặng bao gồm: **Viêm mức 1**, **Viêm mức 2**, và **Viêm mức 3** (Nặng).
|
||||
|
||||
|
||||
* **Các chỉ số định lượng đầu ra (Output Metrics):**
|
||||
* **Tỷ lệ vùng viêm:** $\text{Tỷ lệ \%} = \frac{\text{Tổng số pixel vùng viêm}}{\text{Tổng số pixel của toàn bộ ảnh}} \times 100\%$.
|
||||
|
||||
|
||||
* **Phân tích quan hệ không gian (Spatial/Boundary Analysis):** Xác định vị trí tương quan của ổ viêm với các cấu trúc giải phẫu lân cận như: Vùng mô mỡ (`Fat`), Vùng gân (`Tendon`), Xương đùi (`Femur`).
|
||||
|
||||
|
||||
* *Ví dụ phân tích hệ thống:* "Vùng viêm nằm giữa vùng gân chiếm tỷ lệ $14.4\%$ vùng xung quanh".
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 4. Đặc tả chức năng của Phần mềm Web (Web UI/UX Specification)
|
||||
|
||||
Phần mềm được xây dựng dưới dạng ứng dụng Web tích hợp (Integrated Ultrasound Analytics Platform) sở hữu các chức năng cụ thể sau:
|
||||
|
||||
* **Tải lên ảnh đầu vào:** Giao diện hỗ trợ Kéo & thả ảnh siêu âm trực tiếp hoặc nút chọn file ảnh từ máy tính.
|
||||
|
||||
|
||||
* **Tiền xử lý tự động:** Tự động chuẩn hóa ảnh, khử nhiễu đầu vào.
|
||||
|
||||
|
||||
* **Nhận diện tự động:** Xác định góc chụp/mặt cắt ảnh tự động thông qua khối phân loại.
|
||||
|
||||
|
||||
* **Phân đoạn thông minh:** Tự động định vị, khoanh vùng tổn thương viêm bằng màu trực quan đè lên ảnh gốc (Overlay Mask).
|
||||
|
||||
|
||||
* **Định lượng đa chỉ số:** Trả về độ tin cậy của mô hình ($\%$ Confidence), phân loại mức độ nặng, tính toán diện tích pixel tổn thương.
|
||||
|
||||
|
||||
* **Tương tác lâm sàng:** Cho phép bác sĩ nhập ghi chú, nhận xét và các khuyến nghị điều trị trực tiếp trên giao diện.
|
||||
|
||||
|
||||
* **Lưu trữ & Kết xuất:** Lưu trữ kết quả phân tích vào DB và cho phép xuất báo cáo, tải ảnh phân đoạn về máy.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 5. Kết quả thử nghiệm dữ liệu thực tế tại Bệnh viện E
|
||||
|
||||
Hệ thống đã được đánh giá chéo giữa kết luận tự động của mô hình AI và chẩn đoán lâm sàng thực tế của các bác sĩ chuyên khoa tại Bệnh viện E mang lại kết quả có độ tương quan rất cao:
|
||||
|
||||
Bảng so sánh dữ liệu thử nghiệm thực tế 1
|
||||
|
||||
### Bảng 1: So sánh dữ liệu thử nghiệm thực tế 1
|
||||
|
||||
| Tiêu chí | Kết luận tự động từ AI | Kết luận lâm sàng của Bác sĩ |
|
||||
| :--- | :--- | :--- |
|
||||
| **Ca bệnh 1** | - Phân loại: **Có viêm**<br>- Độ tin cậy: **94.5%**<br>- Mức độ: **Nặng**<br>- Tỷ lệ diện tích viêm: **6.95%**<br>- Mô tả không gian: Vùng viêm lan rộng, nằm giữa vùng gân (chiếm 37.2% vùng xung quanh). | **Khớp gối phải:** Màng hoạt dịch dày, có dịch kích thước ~6.8mm, xuất hiện gai xương nhỏ khe đùi chày trong ngoài khớp.<br><br>**Khớp gối trái:** Màng hoạt dịch dày, có dịch ~11mm (dịch đồng nhất).<br><br>**KẾT LUẬN:** Viêm màng hoạt dịch, tràn dịch khớp gối hai bên kèm thoái hóa khớp gối hai bên. |
|
||||
| **Ca bệnh 2** | - Phân loại: **Có viêm**<br>- Độ tin cậy: **95.1%**<br>- Mức độ: **Nặng**<br>- Tỷ lệ diện tích viêm: **7.54%**<br>- Mô tả không gian: Vùng viêm lan rộng, nằm giữa vùng gân (chiếm 12.4% vùng xung quanh). | (Đồng nhất dữ liệu chẩn đoán lâm sàng tương tự ca bệnh 1 cho thấy độ tương quan cao về mặt định lượng của mô hình đối với các ca tràn dịch mức độ nặng). |
|
||||
|
||||
---
|
||||
|
||||
### Bảng 2: So sánh dữ liệu thử nghiệm thực tế 2
|
||||
|
||||
| Tiêu chí | Kết luận tự động từ AI | Kết luận lâm sàng của Bác sĩ |
|
||||
| :--- | :--- | :--- |
|
||||
| **Đánh giá Khớp Gối Phải** | - Trạng thái: **Có viêm**<br>- Độ tin cậy: **86.5%**<br>- Mức độ: **Nặng**<br>- Tỷ lệ vùng viêm: **6.77%**<br>- Vùng viêm nằm giữa vùng gân (24.6% vùng bao quanh). | - Màng hoạt dịch dày, có dịch kích thước **10mm** (tính chất dịch đồng nhất).<br>- Xuất hiện gai xương tại vị trí khe đùi chày ngoài khớp.<br>- Bề mặt sụn khớp đều, tuy nhiên cấu trúc âm không đồng nhất.<br>- Không có hình ảnh kén khoeo. |
|
||||
| **Đánh giá Khớp Gối Trái** | - Trạng thái: **Có viêm**<br>- Độ tin cậy: **91.7%**<br>- Mức độ: **Nặng**<br>- Tỷ lệ vùng viêm: **15.96%**<br>- Vùng viêm nằm giữa vùng gân (30.5% vùng bao quanh). | - Màng hoạt dịch dày, ghi nhận có ít dịch.<br>- Xuất hiện gai xương tại khe đùi chày trong ngoài khớp.<br>- Bề mặt sụn khớp đều, cấu trúc âm không đồng nhất.<br>- Không có hình ảnh kén khoeo. |
|
||||
| **KẾT LUẬN CHUNG** | *(AI hỗ trợ định lượng diện tích tổn thương chuyên sâu)* | **Hình ảnh:** Viêm màng hoạt dịch khớp gối hai bên, Tràn dịch khớp gối phải kèm theo Thoái hóa khớp gối hai bên. |
|
||||
|
||||
## 6. Hướng dẫn cài đặt và Triển khai mã nguồn Backend (ML Stack Deployment)
|
||||
|
||||
Dành cho coding agent vận hành hoặc deploy môi trường ứng dụng:
|
||||
|
||||
### 6.1. Yêu cầu hệ thống (Prerequisites)
|
||||
|
||||
* **Hệ điều hành hỗ trợ:** Ubuntu / Windows.
|
||||
|
||||
* **Trình quản lý gói:** Anaconda / Miniconda.
|
||||
|
||||
* **Hỗ trợ phần cứng:** CUDA Toolkit được cài đặt sẵn (Bắt buộc nếu cấu hình sử dụng GPU tăng tốc suy luận).
|
||||
|
||||
|
||||
|
||||
### 6.2. Các bước triển khai chi tiết
|
||||
|
||||
**Bước 1: Tải mã nguồn từ kho lưu trữ (Git Clone)**
|
||||
|
||||
```bash
|
||||
# Thực hiện clone thư mục chứa dự án từ Git Server hệ thống
|
||||
git clone https://vkist-hub.com/itvkist/vkist-ultrasound.git
|
||||
|
||||
# Di chuyển vào thư mục và cập nhật nếu đã tồn tại phiên bản trước đó
|
||||
cd vkist-ultrasound
|
||||
git pull
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
**Bước 2: Khởi tạo môi trường ảo và Cài đặt các thư viện phụ thuộc**
|
||||
|
||||
```bash
|
||||
# Tạo môi trường ảo conda mới chạy Python phiên bản ổn định 3.10
|
||||
conda create -n vkist-ultrasound python=3.10 -y
|
||||
|
||||
# Kích hoạt môi trường ảo vừa tạo
|
||||
conda activate vkist-ultrasound
|
||||
|
||||
# Tiến hành cài đặt tất cả các dependencies phụ thuộc nằm trong file requirements
|
||||
pip install -r requirements.txt
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
**Bước 3: Khởi chạy ứng dụng Web Server Backend**
|
||||
|
||||
```bash
|
||||
# Chạy file khởi tạo ứng dụng chính
|
||||
python app.py
|
||||
|
||||
```
|
||||
|
||||
|
||||
>
|
||||
> **Thông tin triển khai thành công:** Sau khi dòng lệnh thực thi hoàn tất, phần mềm cục bộ sẽ khởi chạy và lắng nghe kết nối tại địa chỉ URL mặc định: `http://localhost:8000`.
|
||||
>
|
||||
>
|
||||
@@ -1,256 +0,0 @@
|
||||
# VKIST ULTRASOUND - TÀI LIỆU HƯỚNG DẪN & GIỚI THIỆU DỰ ÁN
|
||||
|
||||
---
|
||||
|
||||
## 1. Giới thiệu chung
|
||||
|
||||
VKIST Ultrasound là hệ thống hỗ trợ chẩn đoán viêm khớp gối tiên tiến ứng dụng Trí tuệ nhân tạo (AI). Hệ thống giúp bác sĩ phân tích hình ảnh siêu âm một cách tự động, từ việc nhận diện chính xác góc chụp đến việc đo đạc các chỉ số bệnh lý phức tạp. Qua đó, giải pháp này giúp tối ưu hóa quy trình khám chữa bệnh, giảm thiểu sai sót chủ quan và nâng cao hiệu suất làm việc tại các cơ sở y tế.
|
||||
|
||||
---
|
||||
|
||||
## 2. Quy trình xử lý toàn diện (AI Pipeline Workflow)
|
||||
|
||||
Hệ thống vận hành theo một quy trình khép kín, tự động phân nhánh logic dựa trên tính chất hình ảnh đầu vào:
|
||||
|
||||
1. **Tiếp nhận & Tiền xử lý ảnh:** Tải ảnh siêu âm thô lên hệ thống. AI tự động áp dụng thuật toán tăng cường tương phản CLAHE để làm rõ nét các cấu trúc giải phẫu bị mờ hoặc nhiễu.
|
||||
|
||||
|
||||
2. **Phân loại góc chụp (Angle Classification):** Tự động xác định tư thế chụp/mặt cắt nhằm kích hoạt nhánh pipeline phân tích chuyên biệt.
|
||||
|
||||
|
||||
3. **Phát hiện Viêm (Inflammation Detection):** Đánh giá sơ bộ sự hiện diện của dịch khớp hoặc tình trạng tăng sinh màng hoạt dịch.
|
||||
|
||||
|
||||
4. **Phân vùng & Đo đạc (Segmentation & Measurement):** Tách biệt các lớp mô giải phẫu (xương, dịch, màng, gân...) và tự động xác định độ dày tổn thương tại các vùng trọng yếu.
|
||||
|
||||
|
||||
5. **Đánh giá mức độ nặng (Severity Scoring):** Tính toán chỉ số tổng hợp để phân cấp mức độ viêm khớp.
|
||||
|
||||
|
||||
6. **Quản lý hồ sơ & Báo cáo:** Lưu trữ dữ liệu chuẩn hóa vào hệ thống nội bộ và kết xuất phiếu kết quả khám bệnh dạng PDF.
|
||||
|
||||
|
||||
|
||||
Dưới đây là sơ đồ luồng logic của hệ thống được thiết kế bằng **PlantUML** giúp lập trình viên backend dễ dàng triển khai:
|
||||
|
||||
```plantuml
|
||||
@startuml
|
||||
skinparam handwritten false
|
||||
skinparam monochrome false
|
||||
skinparam packageStyle rect
|
||||
skinparam shadowing true
|
||||
|
||||
title SƠ ĐỒ ĐIỀU HƯỚNG LOGIC PIPELINE - VKIST ULTRASOUND AI
|
||||
|
||||
start
|
||||
|
||||
:Tiếp nhận ảnh siêu âm từ Web UI;
|
||||
:Áp dụng thuật toán tiền xử lý CLAHE (Tăng cường độ nét);
|
||||
|
||||
:Mô hình Phân loại góc chụp (Angle Model);
|
||||
note right: Tích hợp ConvNeXt / Swin / DenseNet
|
||||
|
||||
if (Góc chụp hợp lệ?) then (Không thuộc góc Sup_up_long / Post_trans)
|
||||
:Hiển thị nhãn góc chụp (ví dụ: Med-Lat Long);
|
||||
:Thông báo lỗi: Góc chụp đầu vào không hỗ trợ xử lý tiếp;
|
||||
stop
|
||||
else (Hợp lệ)
|
||||
:Mô hình Phát hiện viêm (EfficientNet-B0);
|
||||
|
||||
if (Trạng thái ổ khớp?) then (Không viêm)
|
||||
:Hiển thị trạng thái KHÔNG VIÊM trên giao diện;
|
||||
:Cho phép nhập thông tin và lưu trữ cơ bản;
|
||||
else (Có viêm)
|
||||
:Hiển thị trạng thái CÓ VIÊM;
|
||||
|
||||
if (Góc chụp phát hiện?) then (Post_trans)
|
||||
:Chạy phân đoạn vùng tổn thương (DeepLabV3-ResNet101);
|
||||
:Khoanh vùng, gắn nhãn màu Nang Baker (Baker's Cyst);
|
||||
else (Sup_up_long / Suprapat)
|
||||
:Chạy phân đoạn vùng tổn thương (DeepLabV3-ResNet50);
|
||||
partition "Thuật toán Đo đạc thông minh" {
|
||||
:Smart ROI (Tập trung 1/3 khu vực giữa);
|
||||
:Continuous Segment Search (Tìm đoạn dịch liên tục dài nhất);
|
||||
:Xác định độ dày ổ dịch/màng hoạt dịch (mm);
|
||||
}
|
||||
:Tính điểm toán học phân cấp Mức độ bệnh (Cấp 0 -> 3);
|
||||
endif
|
||||
fi
|
||||
endif
|
||||
|
||||
split
|
||||
:Nhập thông tin bệnh nhân & Chẩn đoán lâm sàng;
|
||||
:Lưu trữ cấu trúc thư mục nội bộ (patients/);
|
||||
split again
|
||||
:Xuất phiếu kết quả khám bệnh dạng PDF (report.pdf);
|
||||
end split
|
||||
|
||||
stop
|
||||
@enduml
|
||||
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Các Tính năng Chi tiết & Mô hình học máy
|
||||
|
||||
### 3.1. Hệ thống Mô hình AI Đa dạng (Multi-Model Architecture)
|
||||
|
||||
* **Khối Phân loại Góc chụp:** Tích hợp các kiến trúc mạng mạng nơ-ron tiên tiến (SOTA) bao gồm ConvNeXt-Tiny (đạt độ chính xác tuyệt đối 100% trên tập kiểm thử) , Swin Transformer, và DenseNet. Hệ thống mặc định sử dụng ConvNeXt làm lõi phân loại chính đầu tiên cho mọi ảnh.
|
||||
|
||||
|
||||
* **Khối Phân vùng (Semantic Segmentation):** Hỗ trợ linh hoạt các mô hình DeepLabV3+ (hoặc DeepLabV3-ResNet50 với độ chính xác 91.67%), UNet3+ Attention, và EfficientFeedback. Các cấu trúc này được tối ưu hóa sâu nhằm nhận diện chính xác đường biên ranh giới màng hoạt dịch phức tạp.
|
||||
|
||||
|
||||
|
||||
### 3.2. Tính năng Đo đạc Thông minh (Smart Measurement)
|
||||
|
||||
Nhánh xử lý góc `Sup_up_long` tích hợp hai thuật toán hình học độc quyền nhằm loại bỏ sai số đo đạc:
|
||||
|
||||
* **Smart ROI:** Hệ thống tự động khoanh vùng và tập trung phân tích vào khu vực 1/3 chính giữa của vùng nghi ngờ, nơi có giá trị và chỉ số chẩn đoán lâm sàng cao nhất.
|
||||
|
||||
|
||||
* **Continuous Segment Search:** Thuật toán tự động tìm kiếm đoạn mặt nạ (mask) liên tục dài nhất. Cơ chế này đảm bảo kết quả tính toán độ dày của màng và dịch khớp phản ánh đúng thực tế, khách quan và không bị ảnh hưởng bởi nhiễu ảnh cục bộ.
|
||||
|
||||
|
||||
|
||||
### 3.3. Phân cấp Mức độ Bệnh (Severity Classification)
|
||||
|
||||
Hệ thống tính toán chỉ số kết hợp giữa diện tích dịch khớp tụ và mức độ tăng sinh màng hoạt dịch để phân định thành 4 cấp độ bệnh lý rõ ràng:
|
||||
|
||||
| Cấp độ viêm | Phân loại lâm sàng | Đặc điểm cấu trúc hình ảnh siêu âm |
|
||||
| --- | --- | --- |
|
||||
| **Cấp 0** | Rất nhẹ | Lượng dịch khớp nằm trong giới hạn sinh lý bình thường.|
|
||||
| **Cấp 1** | Nhẹ | Lớp dịch khớp mỏng, xuất hiện tăng sinh màng hoạt dịch mức độ nhẹ.|
|
||||
| **Cấp 2** | Trung bình | Lượng dịch khớp và màng hoạt dịch tăng sinh phì đại rõ rệt.|
|
||||
| **Cấp 3** | Nặng | Ổ dịch khớp dày, màng hoạt dịch phát triển mạnh và xâm lấn sâu.|
|
||||
|
||||
---
|
||||
|
||||
## 4. Đặc tả Giao diện & Hướng dẫn Vận hành Hệ thống
|
||||
|
||||
Giao diện Web UI của ứng dụng được thiết kế phân cấp tối giản theo bố cục 3 luồng dọc: **Bảng trái (Cấu hình Mô hình)** $\rightarrow$ **Bảng giữa (Tải ảnh & Hiển thị)** $\rightarrow$ **Bảng phải (Nhập dữ liệu & Xem kết quả phân tích)**.
|
||||
|
||||
### Quy trình vận hành 4 bước dành cho Bác sĩ:
|
||||
|
||||
#### Bước 1: Cấu hình Mô hình AI (Left Panel)
|
||||
|
||||
* Trước khi tiến hành tải ảnh siêu âm, bác sĩ có thể linh hoạt tùy chọn thuật toán AI mong muốn cho từng block tác vụ tại bảng điều khiển bên trái. Hệ thống đã được cấu hình mặc định sẵn các mô hình tối ưu nhất.
|
||||
|
||||
|
||||
|
||||
#### Bước 2: Tải ảnh siêu âm lên hệ thống (Middle Panel)
|
||||
|
||||
* Bác sĩ thực hiện kéo thả tập tin ảnh trực tiếp vào vùng nhận diện hoặc nhấn nút `"Chọn ảnh"`.
|
||||
|
||||
|
||||
* Ngay sau khi tệp tin được nạp, hệ thống sẽ tự động kích hoạt toàn bộ Pipeline phân tích ngầm theo luồng xử lý tự động.
|
||||
|
||||
|
||||
|
||||
#### Bước 3: Đọc kết quả phân tích tự động từ AI (Right Panel & Modal Popup)
|
||||
|
||||
Hệ thống tự động hiển thị kết quả trực quan dựa trên nhãn mặt cắt nhận diện được:
|
||||
|
||||
*
|
||||
**Trường hợp góc chụp không hỗ trợ (ví dụ: `Med-Lat Long`):** Hệ thống lập tức xuất nhãn cảnh báo màu tím `"GÓC CHỤP: Med-Lat Long"` và dừng tiến trình xử lý tiếp theo.
|
||||
|
||||
|
||||
*
|
||||
**Trường hợp góc mặt sau `Post_trans`:** Hệ thống kiểm tra tình trạng viêm. Nếu có viêm, ảnh phân đoạn sẽ hiển thị màu overlay trực quan khoanh vùng chính xác ổ tổn thương **Nang Baker** (Baker's cyst) kèm bảng chú thích màu sắc mô giải phẫu tương ứng.
|
||||
|
||||
|
||||
*
|
||||
**Trường hợp góc mặt trước `Sup_up_long`:** Hệ thống tiến hành phân đoạn, hiển thị đường lưới đo đạc thông minh. Đồng thời tính toán chi tiết độ dày ổ dịch bằng đơn vị milimét ($mm$) và đưa ra chẩn đoán mức độ viêm chính xác.
|
||||
|
||||
|
||||
|
||||
#### Bước 4: Lưu trữ hồ sơ dữ liệu & Xuất bản phiếu khám lâm sàng
|
||||
|
||||
* Bác sĩ thực hiện điền đầy đủ thông tin vào form biểu mẫu bệnh nhân ở bảng bên phải (Trong đó hai trường dữ liệu **Mã bệnh nhân** và **Họ và tên** là bắt buộc). Ghi nhận thêm nhận xét lâm sàng cá nhân vào ô "Chẩn đoán của bác sĩ".
|
||||
|
||||
|
||||
* Nhấn nút `"Lưu dữ liệu"` để hệ thống đóng gói và lưu trữ nội bộ vào máy chủ.
|
||||
|
||||
|
||||
* Nhấn nút `"Xuất phiếu khám (PDF)"` để tải về máy mẫu phiếu in kết quả y khoa chính thức trả cho bệnh nhân.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 5. Đặc tả Kiến trúc Lưu trữ Dữ liệu Đầu ra
|
||||
|
||||
### 5.1. Cấu trúc cây thư mục lưu trữ hồ sơ bệnh nhân (Storage Structure)
|
||||
|
||||
Khi bác sĩ nhấn chọn chức năng lưu trữ dữ liệu , hệ thống tự động khởi tạo cấu trúc thư mục phân cấp động theo mốc thời gian thực tại phân vùng `patients/` nhằm tránh trùng lặp thông tin:
|
||||
|
||||
```text
|
||||
patients/
|
||||
└── <Mã_Bệnh_Nhân>_<Họ_Và_Tên_Không_Dấu>/
|
||||
└── <NamThangNgay>_<GioPhutGiay>/
|
||||
├── info.txt # Tệp tin cấu trúc lưu trữ siêu dữ liệu định lượng của ca bệnh
|
||||
├── original.png # File hình ảnh siêu âm gốc (hoặc ảnh đã tăng cường CLAHE)
|
||||
├── segmented.png # File ảnh overlay mặt nạ phân đoạn màu từ AI mô hình
|
||||
└── report.pdf # File phiếu kết quả chẩn đoán y khoa chính thức xuất cho bệnh nhân
|
||||
|
||||
```
|
||||
|
||||
*Ví dụ thực tế:* `patients\BN0001_Nguyen_Van_A\20260505_170011\`
|
||||
|
||||
### 5.2. Định dạng cấu trúc tệp dữ liệu `info.txt`
|
||||
|
||||
Tệp tin `info.txt` đóng vai trò lưu trữ toàn bộ thông tin tối giản, giúp hệ thống hoặc các agent xử lý số liệu có thể dễ dàng phân tích (parse) dữ liệu mà không cần đọc file PDF:
|
||||
|
||||
```text
|
||||
--- THÔNG TIN BỆNH NHÂN ---
|
||||
Mã bệnh nhân: BN0001
|
||||
Họ tên: Nguyễn Văn A
|
||||
Giới tính: Nam
|
||||
Tuổi: 88
|
||||
Ghi chú lâm sàng: Tràn dịch và có thể viêm
|
||||
|
||||
--- KẾT QUẢ PHÂN TÍCH AI ---
|
||||
Góc chụp: sup-up-long (99.93%)
|
||||
Viêm nhiễm: Có (94.44%)
|
||||
Độ dày màng: 6.53 mm (95 px)
|
||||
Vị trí x: 420
|
||||
Mức độ: Trung bình
|
||||
Mô tả: Dịch khớp trung bình (51px), màng hoạt dịch tăng sinh vừa
|
||||
|
||||
```
|
||||
|
||||
(Ghi chú: Nội dung trên được ánh xạ chính xác từ các chỉ số định lượng thu được qua mô hình phân đoạn hình học của hệ thống ).
|
||||
|
||||
### 5.3. Quy chuẩn nội dung phiếu kết quả y khoa `report.pdf`
|
||||
|
||||
Phiếu kết quả chẩn đoán hình ảnh dạng PDF được tạo tự động với cấu trúc bố cục chuẩn hóa y tế gồm 4 phần chính:
|
||||
|
||||
1.
|
||||
**Thông tin Cơ sở & Tiêu đề:** Biểu trưng nhận diện VKIST, tên "TRUNG TÂM CHẨN ĐOÁN HÌNH ẢNH VKIST", địa chỉ "Khu Công nghệ cao Hòa Lạc, Thạch Thất, Hà Nội" kèm tiêu đề lớn **"PHIẾU KẾT QUẢ SIÊU ÂM KHỚP GỐI"**.
|
||||
|
||||
|
||||
2. I. Thông tin bệnh nhân: Hiển thị chi tiết Họ tên, Giới tính, Mã BN, Tuổi của người bệnh.
|
||||
|
||||
|
||||
3. II. Hình ảnh siêu âm: Chèn song song hai khung hình trực quan bao gồm `Hình 1: Ảnh gốc / Tăng cường` và `Hình 2: Ảnh phân đoạn AI` (có kèm sơ đồ lưới đo đạc và chú thích màu).
|
||||
|
||||
|
||||
4. **III. Kết quả phân tích tự động (AI Metric):**
|
||||
* Góc chụp dự đoán đạt tỷ lệ tương ứng (Ví dụ: `sup-up-long` - Độ tin cậy: `99.93%`).
|
||||
|
||||
|
||||
* Tình trạng viêm ổ khớp (Ví dụ: `Có khả năng viêm` - Độ tin cậy: `94.44%`).
|
||||
|
||||
|
||||
* Chỉ số đo đạc vật lý: Độ dày dịch & màng hoạt dịch tính toán được đạt mức `6.53 mm`.
|
||||
|
||||
|
||||
* Đánh giá mức độ viêm tổng hợp: `Trung bình`.
|
||||
|
||||
|
||||
* Chi tiết mô tả định lượng: Dịch khớp trung bình ($51\text{ px}$), màng hoạt dịch tăng sinh vừa ($95\text{ px}$).
|
||||
|
||||
|
||||
|
||||
5. IV. Chẩn đoán và kết luận của Bác sĩ: Trích xuất nguyên vẹn nội dung ghi chú lâm sàng do bác sĩ trực tiếp nhập vào hệ thống (Ví dụ: `"Tràn dịch và có thể viêm"`).
|
||||
@@ -1,11 +0,0 @@
|
||||
# ML Model Architecture Report
|
||||
|
||||
| File | Architectures (code ranges) |
|
||||
|------|-----------------------------|
|
||||
| `ML/arch/unet3plus_att.py` | - **UNet3Plus_Attention** (87‑268)<br> - SelfAttention (7‑25)<br> - AttentionGate (28‑66)<br> - conv_block (69‑83) |
|
||||
| `ML/arch/efficientfeedback.py` | - **EfficientFeedbackNetwork** (171‑239)<br> - convblock (9‑14)<br> - DecoderBlock (17‑35)<br> - ASPP_module (37‑66)<br> - CAM_Module (68‑86)<br> - S_Module (93‑131)<br> - FeedbackSpatialAttention (134‑151)<br> - StageAttentionwCAM (153‑168) |
|
||||
| `ML/segment_anything/modeling/image_encoder.py` | - **ImageEncoderViT** (18‑119)<br> - PatchEmbed (389‑420)<br> - Block (122‑187)<br> - Attention (190‑254) |
|
||||
| `ML/segment_anything/modeling/prompt_encoder.py` | - **PromptEncoder** (17‑181)<br> - PositionEmbeddingRandom (183‑226) |
|
||||
| `ML/segment_anything/modeling/mask_decoder.py` | - **MaskDecoder** (17‑190)<br> - MLP (168‑190) |
|
||||
| `ML/segment_anything/modeling/transformer.py` | - **TwoWayTransformer** (17‑108)<br> - TwoWayAttentionBlock (110‑183)<br> - Attention (186‑243) |
|
||||
| `ML/segment_anything/modeling/sam.py` | - **Sam** (19‑181) *(no internal sub‑architectures; uses imported modules)* |
|
||||
@@ -1,46 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,46 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,52 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,51 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,41 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,39 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,42 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,40 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,31 +0,0 @@
|
||||
# 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
|
||||
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
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
|
||||
@@ -1,12 +0,0 @@
|
||||
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']
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,412 +0,0 @@
|
||||
#!/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()
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,215 +0,0 @@
|
||||
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"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,17 +0,0 @@
|
||||
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 ]
|
||||
}
|
||||
]
|
||||
@@ -1,346 +0,0 @@
|
||||
#!/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()
|
||||
@@ -1,29 +0,0 @@
|
||||
# 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
|
||||
|
||||
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
# 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
|
||||
@@ -1,29 +0,0 @@
|
||||
#!/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!"
|
||||
@@ -1,32 +0,0 @@
|
||||
#!/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!"
|
||||
@@ -1,216 +0,0 @@
|
||||
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
|
||||
@@ -1 +0,0 @@
|
||||
modal deploy PILOT_PROJECT/workspace/sprint_1_2/codebase/infra/implementation/triton_run/modal_triton.py
|
||||
@@ -1,41 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,43 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,50 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,50 +0,0 @@
|
||||
# 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)
|
||||
@@ -1,232 +0,0 @@
|
||||
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)))
|
||||
@@ -1,16 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,383 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,149 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,12 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,44 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,420 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,190 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,226 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,181 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,243 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,286 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,6 +0,0 @@
|
||||
# -*- 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.
|
||||
@@ -1,347 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,158 +0,0 @@
|
||||
# -*- 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
|
||||
@@ -1,111 +0,0 @@
|
||||
# -*- 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)
|
||||
@@ -1,269 +0,0 @@
|
||||
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