Doul0414/HintsPredictionModel

超声提示多标签分类模型

基于 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. 环境配置

# 安装依赖
pip install -r requirements.txt

2. 单步推理

用法:
    # 指定多个图像文件
    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. 评估模型

# 先修改 config.yaml 中的数据路径
# 然后运行评估
python evaluate.py
# 或
bash scripts/evaluate.sh

4. 训练模型

# 先修改 config.yaml 中的数据路径
python train_hybrid.py --config config.yaml
# 或
bash scripts/train.sh

配置说明

使用前请修改 config.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，可根据需要调整