develop - hack_2026
This commit is contained in:
parent
c234193eb4
commit
1b82992a6d
420
README.md
420
README.md
|
|
@ -1,144 +1,154 @@
|
|||
# 🦴 Bone Quality Assessment
|
||||
# 🦴 DXA Quality Assessment
|
||||
|
||||
[](https://www.python.org/)
|
||||
[](https://fastapi.tiangolo.com/)
|
||||
[](https://pytorch.org/)
|
||||
[](https://www.docker.com/)
|
||||
[](https://www.python.org/)
|
||||
[](https://fastapi.tiangolo.com/)
|
||||
[](https://pytorch.org/)
|
||||
|
||||
## 📋 Описание
|
||||
## Описание
|
||||
|
||||
Сервис искусственного интеллекта для автоматизированной оценки качества денситометрических изображений и их разметки. Система получает на вход рентгеновское денситометрическое исследование в формате DICOM, и оценивает качество выполнения исследования по стандартным критериям, а также корректность разметки анатомических структур на изображениях.
|
||||
Сервис для автоматизированной оценки качества денситометрических исследований (DXA). Система анализирует DICOM-изображения костной денситометрии и определяет качество исследования по следующим критериям:
|
||||
|
||||
- **Артефакты** — движение, размытость, металлические объекты, имплантаты
|
||||
- **Позиционирование** — правильное расположение анатомической области в кадре
|
||||
- **Полнота изображения** — видимость всех анатомических структур (позвонки L1-L4, бедро)
|
||||
- **Ротация** — корректный угол поворота (для исследования бедра)
|
||||
- **ROI-валидация** — правильность расположения области интереса
|
||||
|
||||
### Основные возможности
|
||||
|
||||
- 🖼️ **Анализ изображений** — загрузка и обработка медицинских изображений
|
||||
- 🧠 **Сегментация объектов** — выделение анатомических структур (позвонки, кости)
|
||||
- 📊 **Оценка качества** — проверка по 3 критериям:
|
||||
- Артефакты (движение, шум, размытость)
|
||||
- Позиционирование (правильное расположение объекта)
|
||||
- Контрастность (качество изображения)
|
||||
- 🔍 **Детекция нарушений** — определение типа нарушения для некачественных исследований
|
||||
- 🌐 **Web-интерфейс** — удобная загрузка и визуализация результатов
|
||||
- 📡 **REST API** — интеграция с внешними системами
|
||||
- 📈 **Визуализация маски** — отображение сегментации на изображении
|
||||
- 🔬 **Анализ DICOM** — загрузка и обработка медицинских изображений
|
||||
- 🧠 **Классификация** — бинарная оценка качества (OK / Violation)
|
||||
- 🔍 **Детекция нарушений** — определение типа нарушения:
|
||||
- `correct` — качество соответствует норме
|
||||
- `artifact_motion` — артефакт движения
|
||||
- `artifact_other` — прочие артефакты
|
||||
- `position_error` — ошибка позиционирования
|
||||
- `rotation` — нарушение ротации
|
||||
- `incomplete_view` — неполный вид
|
||||
- `roi_error` — ошибка ROI
|
||||
- `labeling_error` — ошибка разметки
|
||||
- 🌐 **REST API** — интеграция с внешними системами
|
||||
- 📊 **Веб-интерфейс** — загрузка и визуализация результатов
|
||||
- 📈 **Экспорт** — выгрузка результатов в XLSX
|
||||
|
||||
## 🏗️ Архитектура
|
||||
---
|
||||
|
||||
|
||||
## Архитектура
|
||||
|
||||
## 🎯 Доступные режимы
|
||||
|
||||
### DXA Режим (хакатон)
|
||||
Анализ денситометрических исследований:
|
||||
```bash
|
||||
# Обучение
|
||||
python src/dxa/train.py --epochs 10
|
||||
|
||||
# Инференс
|
||||
python src/dxa/inference.py --input-path dataset_hack/Для\ теста --output-path results.xlsx
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ FastAPI Server │
|
||||
│ (port 8000) │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ /api/v1/analyze → Basic quality prediction │
|
||||
│ /api/v1/analyze/detailed → Full report with metrics │
|
||||
│ /api/v1/analyze/sr → DICOM SR (Structured Report) │
|
||||
│ /api/v1/batch → Batch processing │
|
||||
│ /api/v1/export → XLSX export │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌──────────────┐ ┌─────────────────┐ │
|
||||
│ │ ResNet18 │───▶│ Quality Model │ │
|
||||
│ │ (pretrained) │ │ (binary class) │ │
|
||||
│ └──────────────┘ └────────┬────────┘ │
|
||||
│ │ │
|
||||
│ ▼ │
|
||||
│ ┌─────────────────────┐ │
|
||||
│ │ Detailed Assessment │ │
|
||||
│ │ - Motion detection │ │
|
||||
│ │ - Artifact detection│ │
|
||||
│ │ - ROI validation │ │
|
||||
│ │ - View completeness│ │
|
||||
│ └─────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Быстрый старт
|
||||
## Быстрый старт
|
||||
|
||||
### Требования
|
||||
|
||||
- Python 3.10+
|
||||
- PyTorch 2.1+
|
||||
- Docker (опционально)
|
||||
- PyTorch 2.0+
|
||||
- 4GB+ RAM
|
||||
- (опционально) GPU CUDA/MPS для ускорения
|
||||
|
||||
### Локальная установка
|
||||
### Установка
|
||||
|
||||
```bash
|
||||
# 1. Клонирование репозитория
|
||||
# Клонирование
|
||||
git clone https://github.com/yourusername/bone-quality-assessment.git
|
||||
cd bone-quality-assessment
|
||||
|
||||
# 2. Создание виртуального окружения
|
||||
# Создание виртуального окружения
|
||||
python -m venv venv
|
||||
source venv/bin/activate # Linux/Mac
|
||||
# или
|
||||
venv\Scripts\activate # Windows
|
||||
# venv\Scripts\activate # Windows
|
||||
|
||||
# 3. Установка зависимостей
|
||||
# Установка зависимостей
|
||||
pip install -r requirements.txt
|
||||
|
||||
# 4. Загрузка обученной модели (опционально)
|
||||
# Поместите модель в папку models/dxa_model.pth
|
||||
# Загрузка модели (опционально)
|
||||
# Поместите файл модели в models/dxa_model.pth
|
||||
```
|
||||
|
||||
# 5. Запуск сервера
|
||||
### Запуск сервера
|
||||
|
||||
```bash
|
||||
# Локальный запуск
|
||||
python -m uvicorn src.main:app --host 0.0.0.0 --port 8000
|
||||
|
||||
# Или через run.py
|
||||
python run.py
|
||||
```
|
||||
|
||||
Docker
|
||||
После запуска:
|
||||
- Web-интерфейс: http://localhost:8000
|
||||
- Swagger UI: http://localhost:8000/docs
|
||||
- ReDoc: http://localhost:8000/redoc
|
||||
|
||||
### Docker
|
||||
|
||||
```bash
|
||||
# 1. Сборка образа
|
||||
docker build -t bone-quality-api .
|
||||
# Сборка
|
||||
docker build -t dxa-quality-api .
|
||||
|
||||
# 2. Запуск контейнера
|
||||
docker run -p 8000:8000 bone-quality-api
|
||||
|
||||
# 3. Или используя docker-compose
|
||||
docker-compose up -d
|
||||
# Запуск
|
||||
docker run -p 8000:8000 dxa-quality-api
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
📡 API Endpoints
|
||||
## API Endpoints
|
||||
|
||||
| Метод | Эндпоинт | Описание |
|
||||
|--------|-------------------|--------------------------|
|
||||
| GET | / | Главная страница
|
||||
| GET | /docs | Swagger UI документация
|
||||
| GET | /redoc | ReDoc документация
|
||||
| GET | /api/v1/health | Проверка статуса сервиса
|
||||
| POST | /api/v1/analyze | Базовый анализ изображения
|
||||
| POST | /api/v1/analyze/detailed | Детальный анализ с метриками
|
||||
| POST | /api/v1/analyze/sr | Анализ с DICOM SR отчётом
|
||||
| POST | /api/v1/batch | Пакетный анализ файлов
|
||||
| POST | /api/v1/export | Анализ и экспорт в XLSX
|
||||
| Метод | Эндпоинт | Описание |
|
||||
|-------|----------|----------|
|
||||
| GET | `/` | Главная страница (веб-интерфейс) |
|
||||
| GET | `/api/v1/health` | Проверка статуса сервиса |
|
||||
| POST | `/api/v1/analyze` | Базовый анализ изображения |
|
||||
| POST | `/api/v1/analyze/detailed` | Детальный анализ с метриками |
|
||||
| POST | `/api/v1/analyze/sr` | DICOM SR отчёт |
|
||||
| POST | `/api/v1/batch` | Пакетный анализ |
|
||||
| POST | `/api/v1/export` | Анализ и экспорт в XLSX |
|
||||
|
||||
### Пример использования
|
||||
|
||||
Пример запроса
|
||||
```bash
|
||||
# Анализ файла
|
||||
curl -X POST "http://localhost:8000/api/v1/analyze" \
|
||||
-H "accept: application/json" \
|
||||
-H "Content-Type: multipart/form-data" \
|
||||
-F "file=@/path/to/image.jpg"
|
||||
-F "file=@/path/to/image.dcm"
|
||||
|
||||
# Детальный анализ
|
||||
curl -X POST "http://localhost:8000/api/v1/analyze/detailed" \
|
||||
-H "accept: application/json" \
|
||||
-H "Content-Type: multipart/form-data" \
|
||||
-F "file=@/path/to/image.dcm"
|
||||
```
|
||||
|
||||
Пример ответа
|
||||
|
||||
```json
|
||||
{
|
||||
"overall_quality": "GOOD",
|
||||
"severity": "LOW",
|
||||
"issues": [],
|
||||
"confidence": 0.9,
|
||||
"metrics": {
|
||||
"artifact": {
|
||||
"artifact": false,
|
||||
"num_objects": 1,
|
||||
"edge_energy": 0.234,
|
||||
"object_size": 0.123
|
||||
},
|
||||
"position": {
|
||||
"position": [0.45, 0.52],
|
||||
"valid": true,
|
||||
"deviation": 0.032
|
||||
},
|
||||
"contrast": {
|
||||
"valid": true,
|
||||
"contrast": 0.456
|
||||
}
|
||||
},
|
||||
"mask": "base64_encoded_mask_image"
|
||||
}
|
||||
```
|
||||
|
||||
### Детальный анализ (/api/v1/analyze/detailed)
|
||||
|
||||
Возвращает расширенную информацию с детализацией нарушений:
|
||||
### Ответ детального анализа
|
||||
|
||||
```json
|
||||
{
|
||||
|
|
@ -156,7 +166,7 @@ curl -X POST "http://localhost:8000/api/v1/analyze" \
|
|||
"metrics": {
|
||||
"motion": {
|
||||
"motion_detected": true,
|
||||
"blur_laplacian": 0.001,
|
||||
"blur_laplacian": 0.0008,
|
||||
"severity": "HIGH"
|
||||
},
|
||||
"artifacts": {
|
||||
|
|
@ -167,147 +177,153 @@ curl -X POST "http://localhost:8000/api/v1/analyze" \
|
|||
"valid": true
|
||||
}
|
||||
},
|
||||
"spine_completeness": {
|
||||
"valid": true,
|
||||
"num_vertebrae": 4
|
||||
},
|
||||
"overall_quality": "POOR",
|
||||
"severity": "HIGH"
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
📊 Интерфейс
|
||||
Веб-интерфейс доступен по адресу http://localhost:8000/:
|
||||
## Структура проекта
|
||||
|
||||
- 📤 Drag-and-drop загрузка изображений
|
||||
- 🔍 Автоматический анализ
|
||||
- 🎯 Визуализация маски сегментации
|
||||
- 📈 Детальные метрики качества
|
||||
- 🏷️ Подробный отчет о нарушениях
|
||||
- 👆 Клик по строке для просмотра деталей
|
||||
|
||||
|
||||
🛠️ Технологии
|
||||
|
||||
|Компонент | Технология
|
||||
|-|-|
|
||||
|Бэкенд | Python 3.10, FastAPI, Uvicorn
|
||||
|ML | PyTorch, NumPy, SciPy
|
||||
|Обработка изображений | PIL, OpenCV
|
||||
|Визуализация | HTML5, CSS3, Canvas API
|
||||
|Контейнеризация | Docker, Docker Compose
|
||||
|Документация | Swagger UI, ReDoc
|
||||
|
||||
|
||||
📁 Структура проекта
|
||||
```text
|
||||
```
|
||||
bone_2026/
|
||||
├── src/
|
||||
│ ├── dxa/ # DXA Quality модуль
|
||||
│ │ ├── model.py # ResNet18 классификатор
|
||||
│ │ ├── dataset.py # Загрузчик данных
|
||||
│ │ ├── train.py # Обучение
|
||||
│ │ └── inference.py # Инференс
|
||||
│ ├── api/ # REST API
|
||||
│ ├── quality/ # Оценка качества
|
||||
│ └── main.py # FastAPI приложение
|
||||
├── models/
|
||||
│ └── dxa_model.pth # Обученная модель
|
||||
├── dataset_hack/ # DICOM датасет
|
||||
│ ├── main.py # FastAPI приложение
|
||||
│ ├── run.py # Запуск сервера
|
||||
│ ├── dxa/ # DXA модуль
|
||||
│ │ ├── model.py # ResNet18 классификатор
|
||||
│ │ ├── dataset.py # Загрузчик данных
|
||||
│ │ ├── train.py # Обучение модели
|
||||
│ │ └── inference.py # Инференс и batch-обработка
|
||||
│ ├── quality/ # Оценка качества
|
||||
│ │ ├── quality_scorer.py # Базовый скорer
|
||||
│ │ └── detailed_assessment.py # Детальный анализ
|
||||
│ ├── api/ # REST API
|
||||
│ │ ├── endpoints.py # Дополнительные эндпоинты
|
||||
│ │ └── static/ # Веб-интерфейс
|
||||
│ └── utils/ # Утилиты
|
||||
├── models/ # Обученные модели
|
||||
│ └── dxa_model.pth # Модель классификатора
|
||||
├── dataset_hack/ # Датасет для обучения/тестирования
|
||||
├── docs/ # Документация
|
||||
├── Dockerfile
|
||||
├── requirements.txt
|
||||
├── run.sh
|
||||
└── README.md
|
||||
```
|
||||
|
||||
🧪 Тестирование
|
||||
```bash
|
||||
# Запуск тестов (если есть)
|
||||
pytest tests/
|
||||
```
|
||||
---
|
||||
|
||||
# Проверка API
|
||||
## Обучение модели
|
||||
|
||||
### Подготовка данных
|
||||
|
||||
1. Разместите DICOM-файлы в `dataset_hack/НД_для_обучения/Исследования/`
|
||||
2. Подготовьте Excel-файл разметки `dataset_hack/НД_для_обучения/разметка.xlsx`
|
||||
|
||||
Столбцы разметки:
|
||||
- `study_uid` — ID исследования
|
||||
- `позвоночник_укладка`, `позвоночник_ось`, `позвоночник_артефакты` — критерии для позвоночника
|
||||
- `бедро_позиция_лев`, `бедро_roi_лев` — критерии для левого бедра
|
||||
- `бедро_позиция_прав`, `бедро_roi_прав` — критерии для правого бедра
|
||||
- `итог_позвоночник`, `итог_бедро_лев`, `итог_бедро_прав` — итоговая оценка (0/1)
|
||||
|
||||
### Запуск обучения
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/api/v1/health
|
||||
python src/dxa/train.py \
|
||||
--epochs 20 \
|
||||
--batch-size 8 \
|
||||
--backbone resnet18 \
|
||||
--output-dir models
|
||||
```
|
||||
|
||||
## 📊 Метрики качества
|
||||
### Аргументы
|
||||
|
||||
Артефакты
|
||||
| Параметр | По умолчанию | Описание |
|
||||
|----------|-------------|----------|
|
||||
| `--data-root` | `dataset_hack` | Путь к директории с данными |
|
||||
| `--annotation-path` | `dataset_hack/НД_для_обучения/разметка.xlsx` | Путь к файлу разметки |
|
||||
| `--epochs` | 20 | Количество эпох |
|
||||
| `--batch-size` | 8 | Размер батча |
|
||||
| `--backbone` | `resnet18` | Архитектура (resnet18/resnet34/efficientnet_b0) |
|
||||
| `--input-size` | 224 | Размер входного изображения |
|
||||
| `--output-dir` | `models` | Директория для сохранения модели |
|
||||
|
||||
*Описание:* Обнаружение шума, размытости, фрагментации
|
||||
---
|
||||
|
||||
Проверка:
|
||||
## Инференс
|
||||
|
||||
- Энергия границ (edge_energy)
|
||||
- Количество объектов (num_objects)
|
||||
- Размер объекта (object_size)
|
||||
### Одиночный файл
|
||||
|
||||
### Позиционирование
|
||||
|
||||
*Описание:* Проверка правильности расположения объекта
|
||||
|
||||
Проверка:
|
||||
|
||||
- Центр масс объекта
|
||||
- Отклонение от центра изображения
|
||||
- Минимальный размер объекта
|
||||
|
||||
### Контраст
|
||||
|
||||
*Описание:* Оценка качества изображения
|
||||
|
||||
Проверка:
|
||||
|
||||
- Контраст внутри объекта
|
||||
- Отношение средних (объект/фон)
|
||||
- Интенсивность пикселей
|
||||
|
||||
## 🚀 Деплой
|
||||
На сервер
|
||||
```bash
|
||||
# Копирование на сервер
|
||||
scp -r ./bone-quality-assessment user@server:/var/www/
|
||||
python src/dxa/inference.py \
|
||||
--input-path path/to/image.dcm \
|
||||
--output-path result.xlsx
|
||||
```
|
||||
|
||||
# Запуск в фоновом режиме
|
||||
nohup python run.py > logs/out.log 2>&1 &
|
||||
### Директория
|
||||
|
||||
Использование с Nginx
|
||||
```nginx
|
||||
location /api/ {
|
||||
proxy_pass http://localhost:8000;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
}
|
||||
```bash
|
||||
python src/dxa/inference.py \
|
||||
--input-path dataset_hack/Для\ теста \
|
||||
--output-path results.xlsx \
|
||||
--model-path models/dxa_model.pth
|
||||
```
|
||||
|
||||
## 🤝 Вклад в проект
|
||||
- Fork репозитория
|
||||
- Создайте ветку для вашей фичи (git checkout -b feature/amazing-feature)
|
||||
- Commit изменений (git commit -m 'Add amazing feature')
|
||||
- Push в ветку (git push origin feature/amazing-feature)
|
||||
- Откройте Pull Request
|
||||
### Выходной формат (XLSX/CSV)
|
||||
|
||||
| Колонка | Описание |
|
||||
|---------|----------|
|
||||
| `path_to_study` | Путь к директории исследования |
|
||||
| `study_uid` | StudyInstanceUID |
|
||||
| `image_uid` | SOPInstanceUID |
|
||||
| `anatomical_region` | Анатомическая область (spine/hip_left/hip_right) |
|
||||
| `quality_class` | Класс качества (0 — OK, 1 — Violation) |
|
||||
| `violation_type` | Тип нарушения |
|
||||
| `processing_status` | Статус обработки |
|
||||
| `time_of_processing` | Время обработки (сек) |
|
||||
|
||||
**📄 Лицензия**
|
||||
---
|
||||
|
||||
## Метрики качества
|
||||
|
||||
### Детекция движения
|
||||
|
||||
- **Laplacian variance** — дисперсия лапласиана (меньше = сильнее размытие)
|
||||
- **FFT high-frequency ratio** — отношение высокочастотной энергии (меньше = размытие)
|
||||
- **Edge duplication** — проверка "призрачных" контуров
|
||||
|
||||
### Детекция артефактов
|
||||
|
||||
- **Metal detection** — яркие области (>99.5 перцентиль)
|
||||
- **Implant detection** — линейные структуры (морфологические операции)
|
||||
- **Cement detection** — локальные яркие пятна в ROI
|
||||
- **Calcification** — малые яркие области вне ROI
|
||||
|
||||
### Полнота изображения
|
||||
|
||||
- **Spine**: подсчет позвонков (ожидается 3-4), проверка межпозвоночных промежутков
|
||||
- **Hip**: проверка видимости шейки бедра, большого/малого вертелов
|
||||
|
||||
### Валидация ROI
|
||||
|
||||
- Проверка отступа от краев (>10 пикселей)
|
||||
- Проверка размера ROI (>30% высоты, >20% ширины изображения)
|
||||
|
||||
---
|
||||
|
||||
## Лицензия
|
||||
|
||||
MIT License
|
||||
|
||||
👥 Команда
|
||||
Грачев Денис — Разработка - GitHub
|
||||
## Команда
|
||||
|
||||
### 🙏 Благодарности
|
||||
- **Грачев Денис** — Разработка
|
||||
- **Грачев Татьяна** — Капитан
|
||||
|
||||
***Oxford-IIIT Pet Dataset*** для обучения модели
|
||||
---
|
||||
|
||||
***Сообществу PyTorch и FastAPI***
|
||||
|
||||
📞 Контакты
|
||||
- 📧 Email: your.email@example.com
|
||||
- 🐦 Telegram: @oxydencher
|
||||
- 🐙 GitHub: gdg6
|
||||
|
||||
<div align="center"> <sub>Built with ❤️ for the Bone Quality Assessment Hackathon</sub> </div>
|
||||
<div align="center">
|
||||
<sub>Built for Bone Quality Assessment Hackathon 2026</sub>
|
||||
</div>
|
||||
|
|
|
|||
Binary file not shown.
|
Before Width: | Height: | Size: 538 KiB |
|
|
@ -62,6 +62,7 @@ timm==1.0.29
|
|||
torch==2.8.0
|
||||
torchvision==0.23.0
|
||||
TotalSegmentator==2.18.0
|
||||
requests==2.32.3
|
||||
tqdm==4.70.0
|
||||
typer==0.23.2
|
||||
typing-inspection==0.4.2
|
||||
|
|
|
|||
|
|
@ -7,16 +7,12 @@ from typing import Dict, Any
|
|||
from src.classifiers.bone_condition_classifier import BoneConditionClassifier
|
||||
from src.core.config import PipelineConfig
|
||||
from src.core.registry import ModelRegistry
|
||||
from src.classifiers.pet_breed_classifier import PetBreedClassifier
|
||||
from src.quality.universal_scorer import UniversalQualityScorer
|
||||
from src.segmentators.totalsegmentator_wrapper import TotalSegmentatorWrapper
|
||||
from src.segmentators.unet_segmentator import UNetSegmentator
|
||||
|
||||
# Регистрируем модели
|
||||
ModelRegistry.register_segmentator("unet", UNetSegmentator)
|
||||
ModelRegistry.register_segmentator("totalsegmentator", TotalSegmentatorWrapper)
|
||||
|
||||
ModelRegistry.register_classifier("pet_breed", PetBreedClassifier)
|
||||
ModelRegistry.register_classifier("bone_condition", BoneConditionClassifier)
|
||||
|
||||
|
||||
|
|
@ -43,7 +39,11 @@ class Orchestrator:
|
|||
# Сегментатор
|
||||
segmentator_class = ModelRegistry.get_segmentator(self.config.segmentator)
|
||||
if segmentator_class:
|
||||
self.segmentator = segmentator_class(device=self.device)
|
||||
# Передаем device для TotalSegmentator
|
||||
if self.config.segmentator == "totalsegmentator":
|
||||
self.segmentator = segmentator_class(device=self.device)
|
||||
else:
|
||||
self.segmentator = segmentator_class(device=self.device)
|
||||
else:
|
||||
raise ValueError(f"Сегментатор {self.config.segmentator} не найден")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,16 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
DXA Quality Inference - Batch processing with Excel output
|
||||
|
||||
Этот модуль выполняет инференс модели классификации качества DXA исследований.
|
||||
Основные функции:
|
||||
- Загрузка и предобработка DICOM изображений
|
||||
- Определение анатомической области (позвоночник/бедро)
|
||||
- Бинарная классификация качества (OK/Violation)
|
||||
- Пакетная обработка с экспортом в Excel
|
||||
|
||||
Использование:
|
||||
python src/dxa/inference.py --input-path <path> --output-path <output.xlsx>
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -22,7 +32,15 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
|
|||
|
||||
|
||||
def get_device():
|
||||
"""Get best available device"""
|
||||
"""
|
||||
Определение доступного устройства для вычислений.
|
||||
|
||||
Порядок приоритета: MPS (Apple Silicon) -> CUDA (NVIDIA GPU) -> CPU.
|
||||
Это нужно для максимальной производительности на доступном железе.
|
||||
|
||||
Returns:
|
||||
str: Устройство ('mps', 'cuda' или 'cpu')
|
||||
"""
|
||||
if torch.backends.mps.is_available():
|
||||
return 'mps'
|
||||
elif torch.cuda.is_available():
|
||||
|
|
@ -32,7 +50,21 @@ def get_device():
|
|||
|
||||
|
||||
def load_model(model_path: str, backbone: str = 'resnet18', device: str = 'cpu'):
|
||||
"""Load trained model"""
|
||||
"""
|
||||
Загрузка обученной модели классификатора качества DXA.
|
||||
|
||||
Модель использует предобученный ResNet18 в качестве backbone и
|
||||
добавляет классификационную голову для бинарной классификации
|
||||
(качество OK vs Violation).
|
||||
|
||||
Args:
|
||||
model_path: Путь к файлу модели (.pth)
|
||||
backbone: Архитектура backbone (resnet18/resnet34/efficientnet_b0)
|
||||
device: Устройство для загрузки модели
|
||||
|
||||
Returns:
|
||||
DXAQualityModel: Обертка модели с методами predict и load
|
||||
"""
|
||||
from src.dxa.model import create_model
|
||||
|
||||
model = create_model(backbone=backbone, pretrained=False, device=device)
|
||||
|
|
@ -42,7 +74,24 @@ def load_model(model_path: str, backbone: str = 'resnet18', device: str = 'cpu')
|
|||
|
||||
|
||||
def load_dicom_image(dcm_path: str, input_size: int = 224) -> torch.Tensor:
|
||||
"""Load and preprocess DICOM image"""
|
||||
"""
|
||||
Загрузка и предобработка DICOM изображения для модели.
|
||||
|
||||
Этапы предобработки:
|
||||
1. Чтение DICOM и извлечение pixel_array
|
||||
2. Нормализация интенсивности в диапазон [0, 1]
|
||||
3. Преобразование в 3 канала (дублирование для RGB)
|
||||
4. Изменение размера до input_size x input_size
|
||||
5. Нормализация для ImageNet (деление на 255)
|
||||
6. Преобразование в PyTorch тензор
|
||||
|
||||
Args:
|
||||
dcm_path: Путь к DICOM файлу
|
||||
input_size: Целевой размер изображения (по умолчанию 224 для ResNet)
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Тензор изображения формы (1, 3, 224, 224)
|
||||
"""
|
||||
ds = pydicom.dcmread(dcm_path)
|
||||
img = ds.pixel_array.astype(np.float32)
|
||||
|
||||
|
|
@ -95,14 +144,20 @@ def extract_anatomical_region_from_filename(dcm_path: str) -> str:
|
|||
|
||||
def determine_region_from_image(img: np.ndarray) -> str:
|
||||
"""
|
||||
Determine anatomical region from DICOM image content.
|
||||
Определение анатомической области (позвоночник/бедро) по содержимому изображения.
|
||||
|
||||
Uses multiple features:
|
||||
- Bright region aspect ratio (main discriminator for spine vs hip)
|
||||
- Image symmetry
|
||||
- Horizontal position for left/right hip distinction
|
||||
Алгоритм использует анализ формы яркой области на изображении:
|
||||
- Позвоночник: яркая область более квадратная (aspect ratio ~1.2)
|
||||
- Бедро: яркая область вытянута вертикально (aspect ratio > 1.5)
|
||||
|
||||
Returns: 'spine', 'hip_left', 'hip_right', or 'hip' (unknown side)
|
||||
Дополнительно для определения левого/правого бедра:
|
||||
- Сравнение яркости левой и правой половин изображения
|
||||
|
||||
Args:
|
||||
img: Нормализованное изображение (np.array)
|
||||
|
||||
Returns:
|
||||
str: 'spine', 'hip_left', 'hip_right' или 'hip' (неопределенная сторона)
|
||||
"""
|
||||
h, w = img.shape
|
||||
|
||||
|
|
@ -215,8 +270,23 @@ def determine_anatomical_region(dcm_path: str) -> str:
|
|||
|
||||
|
||||
def process_dicom_files(args):
|
||||
"""Run inference on DICOM files"""
|
||||
"""
|
||||
Основная функция пакетного инференса DXA изображений.
|
||||
|
||||
Этапы обработки:
|
||||
1. Определение устройства (MPS/CUDA/CPU)
|
||||
2. Загрузка обученной модели
|
||||
3. Поиск DICOM файлов в указанной директории
|
||||
4. Инференс для каждого файла:
|
||||
- Предобработка изображения
|
||||
- Предикт модели (бинарная классификация)
|
||||
- Определение анатомической области
|
||||
- Сохранение метаданных DICOM
|
||||
5. Экспорт результатов в Excel/CSV
|
||||
|
||||
Args:
|
||||
args: Аргументы командной строки (input_path, output_path, model_path и т.д.)
|
||||
"""
|
||||
# Setup
|
||||
device = get_device()
|
||||
print(f"Using device: {device}")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,20 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Training script for DXA Quality Classifier
|
||||
=========================================
|
||||
|
||||
Скрипт обучения модели классификации качества DXA изображений.
|
||||
|
||||
Основные этапы:
|
||||
1. Загрузка данных из Excel разметки
|
||||
2. Создание DataLoader с разбиением на train/val
|
||||
3. Инициализация модели (ResNet18 с предобучением на ImageNet)
|
||||
4. Обучение с использованием CrossEntropyLoss
|
||||
5. Валидация после каждой эпохи
|
||||
6. Сохранение лучшей модели по F1-score
|
||||
|
||||
Использование:
|
||||
python src/dxa/train.py --epochs 20 --batch-size 8
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -31,7 +45,22 @@ def get_device():
|
|||
|
||||
|
||||
def compute_metrics(preds, labels):
|
||||
"""Compute classification metrics"""
|
||||
"""
|
||||
Вычисление метрик классификации.
|
||||
|
||||
Рассчитывает:
|
||||
- Accuracy: доля правильных предсказаний
|
||||
- Precision: точность (доля TP среди предсказанных positive)
|
||||
- Recall: полнота (доля TP среди реальных positive)
|
||||
- F1: гармоническое среднее precision и recall
|
||||
|
||||
Args:
|
||||
preds: Предсказания модели (numpy array)
|
||||
labels: Истинные метки (numpy array)
|
||||
|
||||
Returns:
|
||||
Dict с метриками
|
||||
"""
|
||||
preds = np.array(preds)
|
||||
labels = np.array(labels)
|
||||
|
||||
|
|
@ -62,7 +91,24 @@ def compute_metrics(preds, labels):
|
|||
|
||||
|
||||
def train(args):
|
||||
"""Main training loop"""
|
||||
"""
|
||||
Основной цикл обучения модели.
|
||||
|
||||
Этапы:
|
||||
1. Определение устройства (MPS/CUDA/CPU)
|
||||
2. Создание директории для сохранения модели
|
||||
3. Загрузка данных (DataLoader)
|
||||
4. Создание модели
|
||||
5. Цикл обучения по эпохам:
|
||||
- Обучение на train set
|
||||
- Валидация на val set
|
||||
- Расчет метрик (accuracy, precision, recall, F1)
|
||||
- Сохранение лучшей модели по F1
|
||||
6. Сохранение финальной модели
|
||||
|
||||
Args:
|
||||
args: Аргументы командной строки
|
||||
"""
|
||||
|
||||
# Setup
|
||||
device = get_device()
|
||||
|
|
|
|||
86
src/main.py
86
src/main.py
|
|
@ -1,5 +1,24 @@
|
|||
"""
|
||||
Main FastAPI application for DXA Quality Assessment
|
||||
=================================================
|
||||
|
||||
FastAPI сервер для оценки качества DXA исследований.
|
||||
|
||||
Основные функции:
|
||||
- REST API для анализа DICOM изображений
|
||||
- Веб-интерфейс для загрузки и визуализации
|
||||
- Поддержка одиночного и пакетного анализа
|
||||
- Экспорт результатов в XLSX
|
||||
- Генерация DICOM SR (Structured Report)
|
||||
|
||||
Эндпоинты:
|
||||
- GET / - Главная страница (веб-интерфейс)
|
||||
- GET /api/v1/health - Проверка статуса
|
||||
- POST /api/v1/analyze - Базовый анализ
|
||||
- POST /api/v1/analyze/detailed - Детальный анализ
|
||||
- POST /api/v1/analyze/sr - DICOM SR отчет
|
||||
- POST /api/v1/batch - Пакетный анализ
|
||||
- POST /api/v1/export - Анализ и экспорт в XLSX
|
||||
"""
|
||||
from fastapi import FastAPI, File, UploadFile, APIRouter, Query
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
|
@ -43,7 +62,15 @@ device = None
|
|||
|
||||
|
||||
def load_model():
|
||||
"""Load DXA model"""
|
||||
"""
|
||||
Загрузка модели классификатора качества DXA.
|
||||
|
||||
Модель загружается глобально при первом запросе и сохраняется в памяти.
|
||||
Это позволяет избежать повторной загрузки при каждом запросе.
|
||||
|
||||
Returns:
|
||||
DXAQualityModel: Обученная модель или None при ошибке
|
||||
"""
|
||||
global dxa_model, device
|
||||
|
||||
if dxa_model is None:
|
||||
|
|
@ -67,7 +94,25 @@ def load_model():
|
|||
|
||||
|
||||
def preprocess_dicom(dcm_bytes: bytes, input_size: int = 224):
|
||||
"""Preprocess DICOM for model input"""
|
||||
"""
|
||||
Предобработка DICOM изображения для модели.
|
||||
|
||||
Этапы:
|
||||
1. Сохранение байтов во временный файл
|
||||
2. Чтение DICOM (pydicom)
|
||||
3. Нормализация значений пикселей
|
||||
4. Создание 3-канального изображения
|
||||
5. Изменение размера до input_size x input_size
|
||||
6. Нормализация для PyTorch (деление на 255)
|
||||
7. Преобразование в тензор
|
||||
|
||||
Args:
|
||||
dcm_bytes: Байты DICOM файла
|
||||
input_size: Целевой размер (по умолчанию 224)
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, pydicom.Dataset]: Тензор изображения и метаданные DICOM
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix='.dcm', delete=False) as f:
|
||||
|
|
@ -192,14 +237,18 @@ async def analyze_dicom_detailed(
|
|||
include_visualization: bool = Query(False, description="Include base64 mask visualization")
|
||||
):
|
||||
"""
|
||||
Detailed analysis of DICOM file with comprehensive quality assessment.
|
||||
Детальный анализ DICOM файла с комплексной оценкой качества.
|
||||
|
||||
Это основной эндпоинт для получения полной информации о качестве исследования.
|
||||
В отличие от базового /analyze, возвращает:
|
||||
- violation_type: конкретный тип нарушения
|
||||
- reason: человекочитаемое объяснение
|
||||
- confidence_per_class: вероятности для каждого класса
|
||||
- view_quality: полный/частичный вид
|
||||
- metrics: детальные метрики по каждому критерию
|
||||
|
||||
Returns:
|
||||
- violation_type: Specific type of violation (correct, position_error, artifact_motion, etc.)
|
||||
- reason: Human-readable explanation of why the image is non-compliant
|
||||
- confidence_per_class: Probability for each class
|
||||
- view_quality: Full or partial view assessment
|
||||
- Detailed metrics for motion, artifacts, position, ROI
|
||||
JSON с детальным отчетом о качестве исследования
|
||||
"""
|
||||
try:
|
||||
# Load model
|
||||
|
|
@ -484,7 +533,26 @@ async def analyze_dicom_sr(file: UploadFile = File(...)):
|
|||
|
||||
|
||||
def generate_dicom_sr_text(study_uid: str, image_uid: str, quality_report: Dict) -> str:
|
||||
"""Generate DICOM SR text representation"""
|
||||
"""
|
||||
Генерация текстового представления DICOM SR (Structured Report).
|
||||
|
||||
DICOM SR - это стандартизированный формат отчета в DICOM.
|
||||
Здесь генерируется текстовое представление для совместимости.
|
||||
|
||||
Структура отчета:
|
||||
- Study/Series Instance UID
|
||||
- Procedure Report (регион, качество, уверенность)
|
||||
- Findings (тип нарушения, причина, метрики)
|
||||
- Completion Flag
|
||||
|
||||
Args:
|
||||
study_uid: StudyInstanceUID из DICOM
|
||||
image_uid: SOPInstanceUID из DICOM
|
||||
quality_report: Результат generate_quality_report()
|
||||
|
||||
Returns:
|
||||
str: Текстовое представление SR отчета
|
||||
"""
|
||||
|
||||
# Map violation types to DICOM codes (simplified)
|
||||
violation_code_map = {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,28 @@
|
|||
"""
|
||||
Detailed Quality Assessment Module for DXA images
|
||||
Implements checks from condition.txt and condition_doctor.txt
|
||||
================================================
|
||||
|
||||
Модуль детальной оценки качества DXA изображений.
|
||||
Реализует проверки по критериям из condition.txt и condition_doctor.txt.
|
||||
|
||||
Основные функции:
|
||||
- detect_motion_blur(): Детекция артефактов движения
|
||||
- detect_artifacts(): Детекция металлических объектов, имплантатов, цемента
|
||||
- check_spine_completeness(): Проверка видимости позвонков L1-L4
|
||||
- check_hip_completeness(): Проверка видимости бедра
|
||||
- check_hip_rotation(): Проверка угла ротации бедра
|
||||
- check_roi_boundaries(): Валидация границ ROI
|
||||
- generate_quality_report(): Генерация итогового отчета
|
||||
|
||||
Типы нарушений (violation_types):
|
||||
- correct: Качество соответствует норме
|
||||
- artifact_motion: Артефакт движения
|
||||
- artifact_other: Прочие артефакты (металл, имплантаты)
|
||||
- position_error: Ошибка позиционирования
|
||||
- rotation: Нарушение ротации
|
||||
- incomplete_view: Неполный вид
|
||||
- roi_error: Ошибка ROI
|
||||
- labeling_error: Ошибка разметки
|
||||
"""
|
||||
import numpy as np
|
||||
from scipy import ndimage
|
||||
|
|
@ -82,12 +104,29 @@ def calculate_blur_fft(image: np.ndarray) -> float:
|
|||
|
||||
def detect_motion_blur(image: np.ndarray, segmentation: Optional[np.ndarray] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Detect motion artifacts in the image.
|
||||
Детекция артефактов движения на DXA изображении.
|
||||
|
||||
Checks:
|
||||
- Blur using Laplacian variance
|
||||
- Blur using FFT high-frequency energy
|
||||
- Edge duplication (ghost edges)
|
||||
Артефакты движения проявляются в виде:
|
||||
1. Размытия (blur) - снижение четкости границ
|
||||
2. Дублирования краев (edge duplication) - "призрачные" контуры
|
||||
|
||||
Методы детекции:
|
||||
- Laplacian variance: дисперсия лапласиана изображения
|
||||
(низкое значение = сильное размытие)
|
||||
- FFT high-frequency energy: отношение высокочастотной энергии
|
||||
(низкое значение = размытие)
|
||||
- Edge duplication: проверка повторяющихся границ
|
||||
|
||||
Args:
|
||||
image: Исходное изображение (нормализованное, 0-1)
|
||||
segmentation: Бинарная маска сегментации (опционально)
|
||||
|
||||
Returns:
|
||||
Dict с ключами:
|
||||
- motion_detected: bool - обнаружено ли движение
|
||||
- blur_laplacian: float - дисперсия лапласиана
|
||||
- blur_fft: float - FFT high-frequency ratio
|
||||
- severity: str - степень (NONE/MEDIUM/HIGH)
|
||||
"""
|
||||
if len(image.shape) == 3:
|
||||
gray = np.mean(image, axis=2)
|
||||
|
|
@ -136,11 +175,25 @@ def detect_motion_blur(image: np.ndarray, segmentation: Optional[np.ndarray] = N
|
|||
|
||||
def detect_artifacts(image: np.ndarray, segmentation: Optional[np.ndarray] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Detect various artifacts in the image:
|
||||
- Metal objects (high intensity regions)
|
||||
- Implants, screws, plates
|
||||
- Cement
|
||||
- Calcifications
|
||||
Детекция артефактов на DXA изображении.
|
||||
|
||||
Типы артефактов:
|
||||
- Metal (металл): очень яркие области (>99.5 перцентиль), небольшого размера
|
||||
- Implants (имплантаты): регулярные геометрические структуры (пластины, винты)
|
||||
- Cement (цемент): локальные яркие пятна в области ROI
|
||||
- Calcifications (кальцификаты): малые яркие области вне ROI
|
||||
|
||||
Args:
|
||||
image: Исходное изображение (нормализованное)
|
||||
segmentation: Бинарная маска сегментации (опционально)
|
||||
|
||||
Returns:
|
||||
Dict с ключами:
|
||||
- metal_detected: bool
|
||||
- implant_detected: bool
|
||||
- cement_detected: bool
|
||||
- calcification_detected: bool
|
||||
- any_detected: bool - любой артефакт обнаружен
|
||||
"""
|
||||
if len(image.shape) == 3:
|
||||
gray = np.mean(image, axis=2)
|
||||
|
|
@ -212,12 +265,28 @@ def detect_artifacts(image: np.ndarray, segmentation: Optional[np.ndarray] = Non
|
|||
|
||||
def check_spine_completeness(segmentation: np.ndarray) -> Dict[str, Any]:
|
||||
"""
|
||||
Check if spine vertebrae are fully visible (L1-L4).
|
||||
Проверка полноты изображения позвоночника.
|
||||
|
||||
Verifies:
|
||||
- Number of vertebrae (should be 3-4 for lumbar spine)
|
||||
- Vertebrae are not cropped
|
||||
- Intervertebral spaces visible
|
||||
Для качественного исследования поясничного отдела позвоночника должны быть видны:
|
||||
- 3-4 позвонка (L1-L4)
|
||||
- Межпозвоночные промежутки
|
||||
- Позвонки не должны быть обрезаны краями изображения
|
||||
|
||||
Алгоритм:
|
||||
1. Выделяет связные компоненты (отдельные позвонки)
|
||||
2. Проверяет количество (3-4)
|
||||
3. Проверяет каждый позвонок на обрезанность
|
||||
4. Проверяет выравнивание позвонков
|
||||
|
||||
Args:
|
||||
segmentation: Бинарная маска сегментации позвоночника
|
||||
|
||||
Returns:
|
||||
Dict с ключами:
|
||||
- valid: bool - изображение полное
|
||||
- num_vertebrae: int - количество видимых позвонков
|
||||
- vertebrae: list - информация о каждом позвонке
|
||||
- issues: list - список проблем
|
||||
"""
|
||||
if segmentation is None or segmentation.sum() == 0:
|
||||
return {
|
||||
|
|
@ -303,10 +372,27 @@ def check_spine_labels(segmentation: np.ndarray, image: Optional[np.ndarray] = N
|
|||
|
||||
def check_hip_completeness(segmentation: np.ndarray) -> Dict[str, Any]:
|
||||
"""
|
||||
Check if hip is fully visible:
|
||||
- Femur neck visible
|
||||
- Femur head visible
|
||||
- Greater and lesser trochanters visible
|
||||
Проверка полноты изображения бедра.
|
||||
|
||||
Для качественного исследования бедра должны быть видны:
|
||||
- Шейка бедра (femur neck)
|
||||
- Головка бедра (femur head)
|
||||
- Большой и малый вертелы
|
||||
|
||||
Алгоритм:
|
||||
1. Проверяет, что сегментация не пустая
|
||||
2. Проверяет обрезанность (касание краев изображения)
|
||||
3. Проверяет соотношение сторон (для бедра должно быть >1.5)
|
||||
|
||||
Args:
|
||||
segmentation: Бинарная маска сегментации бедра
|
||||
|
||||
Returns:
|
||||
Dict с ключами:
|
||||
- valid: bool - изображение полное
|
||||
- is_cropped: bool - обрезано ли изображение
|
||||
- aspect_ratio: float - соотношение высоты к ширине
|
||||
- issues: list - список проблем
|
||||
"""
|
||||
if segmentation is None or segmentation.sum() == 0:
|
||||
return {
|
||||
|
|
@ -349,12 +435,27 @@ def check_hip_completeness(segmentation: np.ndarray) -> Dict[str, Any]:
|
|||
|
||||
def check_hip_rotation(segmentation: np.ndarray, image: np.ndarray) -> Dict[str, Any]:
|
||||
"""
|
||||
Check hip rotation from the image.
|
||||
Проверка ротации бедра на изображении.
|
||||
|
||||
Proper positioning:
|
||||
- Femur should be slightly internally rotated (15-20 degrees)
|
||||
- Lesser trochanter should be slightly visible medially
|
||||
- Greater trochanter should not be prominent
|
||||
Правильное позиционирование бедра:
|
||||
- Легкая внутренняя ротация (15-20 градусов)
|
||||
- Малый вертел слегка виден медиально
|
||||
- Большой вертел не должен быть выражен
|
||||
|
||||
Алгоритм:
|
||||
1. Находит главную ость сегментированного бедра
|
||||
2. Вычисляет угол отклонения от вертикали
|
||||
3. Проверяет, что угол в пределах допустимого диапазона (±30°)
|
||||
|
||||
Args:
|
||||
segmentation: Бинарная маска сегментации бедра
|
||||
image: Исходное изображение
|
||||
|
||||
Returns:
|
||||
Dict с ключами:
|
||||
- valid: bool - ротация корректна
|
||||
- rotation_angle: float - угол ротации в градусах
|
||||
- issues: list - список проблем
|
||||
"""
|
||||
if segmentation is None or segmentation.sum() == 0:
|
||||
return {
|
||||
|
|
@ -634,7 +735,32 @@ def generate_quality_report(
|
|||
model_confidence: float = 0.5
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate comprehensive quality assessment report.
|
||||
Генерация комплексного отчета о качестве исследования.
|
||||
|
||||
Этапы формирования отчета:
|
||||
1. Детекция движения (Laplacian, FFT)
|
||||
2. Детекция артефактов (металл, имплантаты, цемент)
|
||||
3. Проверка полноты (spine/hip)
|
||||
4. Проверка ротации (только для hip)
|
||||
5. Валидация ROI
|
||||
6. Определение типа нарушения
|
||||
7. Формирование итогового вердикта
|
||||
|
||||
Args:
|
||||
image: Исходное изображение (нормализованное)
|
||||
segmentation: Бинарная маска сегментации
|
||||
region: Анатомическая область (spine/hip_left/hip_right)
|
||||
model_prediction: Предсказание модели (0=OK, 1=Violation)
|
||||
model_confidence: Уверенность модели (0-1)
|
||||
|
||||
Returns:
|
||||
Dict с полным отчетом о качестве, включая:
|
||||
- quality_class: 0 или 1
|
||||
- violation_type: тип нарушения
|
||||
- reason: человекочитаемое описание
|
||||
- metrics: детальные метрики по каждому критерию
|
||||
- overall_quality: GOOD/WARNING/POOR
|
||||
- severity: NONE/LOW/MEDIUM/HIGH
|
||||
"""
|
||||
# Calculate all quality metrics
|
||||
quality_metrics = {}
|
||||
|
|
|
|||
|
|
@ -1,35 +1,147 @@
|
|||
from src.segmentators.base import BaseSegmentator
|
||||
import numpy as np
|
||||
import tempfile
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
|
||||
class TotalSegmentatorWrapper(BaseSegmentator):
|
||||
"""Обертка для TotalSegmentator (для костей)"""
|
||||
"""Обертка для TotalSegmentator (костная сегментация)"""
|
||||
|
||||
TASK_MAP = {
|
||||
"spine": "vertebrae_body",
|
||||
"hip": "appendicular_bones",
|
||||
"hip_left": "appendicular_bones",
|
||||
"hip_right": "appendicular_bones",
|
||||
"vertebrae": "vertebrae_body",
|
||||
"appendicular": "appendicular_bones",
|
||||
"total": "total",
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device: str = "cpu",
|
||||
task: str = "vertebrae_body",
|
||||
anatomy_region: str = None
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
device: устройство для инференса ('cpu', 'gpu', 'mps')
|
||||
task: задача сегментации ('vertebrae_body', 'appendicular_bones', 'total')
|
||||
anatomy_region: анатомическая область ('spine', 'hip', 'hip_left', 'hip_right')
|
||||
"""
|
||||
# Определяем задачу на основе анатомической области
|
||||
if anatomy_region and anatomy_region in self.TASK_MAP:
|
||||
task = self.TASK_MAP[anatomy_region]
|
||||
|
||||
self.device = device
|
||||
self.task = task
|
||||
self.loaded = False
|
||||
self.model = None
|
||||
self.totalsegmentator_func = None
|
||||
|
||||
# Lazy load - откладываем до первого использования
|
||||
# self._load_model()
|
||||
|
||||
def _load_model(self):
|
||||
"""Ленивая загрузка модели при первом вызове"""
|
||||
if self.loaded:
|
||||
return
|
||||
|
||||
def __init__(self):
|
||||
try:
|
||||
from totalsegmentator import TotalSegmentator
|
||||
self.model = TotalSegmentator()
|
||||
import nibabel as nib
|
||||
from totalsegmentator.python_api import totalsegmentator
|
||||
|
||||
self.totalsegmentator_func = totalsegmentator
|
||||
self.nibabel = nib
|
||||
self.loaded = True
|
||||
print("✅ TotalSegmentator загружен")
|
||||
except ImportError:
|
||||
print("⚠️ TotalSegmentator не установлен")
|
||||
print(f"✅ TotalSegmentator API загружен (task={self.task})")
|
||||
except ImportError as e:
|
||||
self.loaded = False
|
||||
print(f"⚠️ TotalSegmentator: {e}")
|
||||
except Exception as e:
|
||||
self.loaded = False
|
||||
print(f"⚠️ TotalSegmentator: ошибка инициализации - {e}")
|
||||
|
||||
def segment(self, image):
|
||||
"""
|
||||
Сегментация 2D изображения (numpy array или PIL Image)
|
||||
|
||||
Args:
|
||||
image: numpy array [H, W] или PIL Image
|
||||
|
||||
Returns:
|
||||
Бинарная маска [H, W] где 1 = кость
|
||||
"""
|
||||
# Ленивая загрузка модели при первом вызове
|
||||
if not self.loaded:
|
||||
return np.zeros((512, 512), dtype=np.int64)
|
||||
self._load_model()
|
||||
|
||||
# TotalSegmentator ожидает DICOM или NIfTI
|
||||
# Здесь адаптация под входной формат
|
||||
# Реальная реализация зависит от TotalSegmentator API
|
||||
if not self.loaded:
|
||||
h, w = 512, 512
|
||||
if hasattr(image, 'shape'):
|
||||
if len(image.shape) >= 2:
|
||||
h, w = image.shape[:2]
|
||||
return np.zeros((h, w), dtype=np.int64)
|
||||
|
||||
# Пока возвращаем заглушку
|
||||
return np.ones((512, 512), dtype=np.int64)
|
||||
# Конвертация в numpy если PIL
|
||||
if hasattr(image, 'convert'):
|
||||
image = np.array(image)
|
||||
|
||||
# Проверка что это 2D изображение
|
||||
if len(image.shape) == 3:
|
||||
image = image[:, :, 0] # Берем первый канал
|
||||
|
||||
# Сохраняем во временный файл как NIfTI
|
||||
import nibabel as nib
|
||||
|
||||
# Создаем 3D объем (1 срез)
|
||||
image_3d = np.expand_dims(image.astype(np.float32), axis=0)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
input_path = Path(tmpdir) / "input.nii.gz"
|
||||
output_path = Path(tmpdir) / "output"
|
||||
output_path.mkdir()
|
||||
|
||||
# Сохраняем как NIfTI
|
||||
nifti_img = nib.Nifti1Image(image_3d, np.eye(4))
|
||||
nib.save(nifti_img, str(input_path))
|
||||
|
||||
# Запускаем TotalSegmentator
|
||||
try:
|
||||
result = self.totalsegmentator_func(
|
||||
input=str(input_path),
|
||||
output=str(output_path),
|
||||
task=self.task,
|
||||
device=self.device,
|
||||
quiet=True,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Загружаем результат
|
||||
result_files = list(output_path.glob("*.nii.gz"))
|
||||
if result_files:
|
||||
result_nifti = nib.load(str(result_files[0]))
|
||||
mask_3d = result_nifti.get_fdata()
|
||||
# Берем первый (и единственный) срез
|
||||
if mask_3d.ndim == 3:
|
||||
mask_2d = mask_3d[0, :, :]
|
||||
else:
|
||||
mask_2d = mask_3d
|
||||
return (mask_2d > 0).astype(np.int64)
|
||||
|
||||
except Exception as e:
|
||||
pass # Молча возвращаем пустую маску
|
||||
|
||||
return np.zeros(image.shape, dtype=np.int64)
|
||||
|
||||
def get_info(self):
|
||||
return {
|
||||
"name": "TotalSegmentator",
|
||||
"type": "segmentation",
|
||||
"target": "bones",
|
||||
"task": self.task,
|
||||
"device": self.device,
|
||||
"loaded": self.loaded
|
||||
}
|
||||
Loading…
Reference in New Issue