Doul0414's picture
Initial upload: HintsPrediction
343e05c verified
# 超声提示多标签分类模型
基于 **TransMIL + Query2Label** 混合架构的甲状腺超声图像多标签分类模型。
## 模型架构
- **Backbone**: ResNet-50 (预训练)
- **特征聚合**: TransMIL (Nystrom Attention)
- **多标签分类**: Query2Label (Transformer Decoder)
- **损失函数**: Asymmetric Loss (处理类别不平衡)
## 17类标签
| 序号 | 标签 | 序号 | 标签 |
|:---:|:---|:---:|:---|
| 1 | TI-RADS 1级 | 10 | 囊肿 |
| 2 | TI-RADS 2级 | 11 | 淋巴结 |
| 3 | TI-RADS 3级 | 12 | 胶质潴留 |
| 4 | TI-RADS 4a级 | 13 | 弥漫性病变 |
| 5 | TI-RADS 4b级 | 14 | 结节性甲状腺肿 |
| 6 | TI-RADS 4c级 | 15 | 桥本氏甲状腺炎 |
| 7 | TI-RADS 5级 | 16 | 反应性 |
| 8 | 钙化 | 17 | 转移性 |
| 9 | 甲亢 | | |
## 目录结构
```
HintsPrediction/
├── README.md # 本文件
├── requirements.txt # 依赖列表
├── config.yaml # 配置文件(需修改路径)
├── models/ # 模型代码
│ ├── __init__.py
│ ├── transmil_q2l.py # 主模型架构
│ ├── transformer.py # Transformer 组件
│ └── aslloss.py # 损失函数
├── checkpoints/
│ └── checkpoint_best.pth # 最佳模型权重
├── scripts/ # 懒人脚本
│ ├── train.sh
│ ├── evaluate.sh
│ └── infer_single.sh
├── train_hybrid.py # 训练代码
├── evaluate.py # 评估代码
├── thyroid_dataset.py # 数据集加载
└── infer_single_case.py # 单步推理代码
```
## 快速开始
### 1. 环境配置
```bash
# 安装依赖
pip install -r requirements.txt
```
### 2. 单步推理
```bash
用法:
# 指定多个图像文件
python infer_single_case.py --images /path/to/img1.png /path/to/img2.png --threshold 0.5
# 指定图像文件夹
python infer_single_case.py --image_dir /path/to/case_folder/ --threshold 0.5
# 或使用脚本
用法1: bash scripts/infer_single.sh /path/to/image1.png /path/to/image2.png ...
用法2: bash scripts/infer_single.sh --image_dir /path/to/case_folder/
```
### 3. 评估模型
```bash
# 先修改 config.yaml 中的数据路径
# 然后运行评估
python evaluate.py
# 或
bash scripts/evaluate.sh
```
### 4. 训练模型
```bash
# 先修改 config.yaml 中的数据路径
python train_hybrid.py --config config.yaml
# 或
bash scripts/train.sh
```
## 配置说明
使用前请修改 `config.yaml` 中的数据路径:
```yaml
data:
data_root: "/path/to/your/ReportData_ROI/"
annotation_csv: "/path/to/your/thyroid_multilabel_annotations.csv"
val_json: "/path/to/your/classification_val_set_single.json"
test_json: "/path/to/your/classification_test_set_single.json"
```
## 性能指标
在测试集上的性能(请参考 `checkpoints/evaluation_report.csv`
## 注意事项
1. 推理时需要 GPU(推荐),CPU 也可运行但较慢
2. 单病例可输入多张图像,模型会自动聚合特征
3. 默认阈值为 0.5,可根据需要调整