Faster RCNN 对肺结节 目标检测
发布时间
阅读量：
阅读量
本项目介绍了基于Faster R-CNN模型对肺结节进行目标检测的应用。以下是核心内容：
项目结构：
- 数据集路径：data 包含训练图像和标签。
- 工具函数：包括从XML文件读取类别和边界框（getlabelfrom_xml）、图像归一化（normalization）以及计算iou值（miou）。
- 数据加载：使用自定义的数据加载器LungDetection处理DCM格式图像并进行预处理。
模型与训练：- 模型选择：使用Faster R-CNN-Mobilenet V3 backbone架构。
- 超参数：包括批量大小、学习率（0.00001）、设备设置等。
- 训练流程：采用交叉熵损失+iou损失作为总损失函数，并通过Adam优化器进行优化。
预测与结果展示：- 预测代码：加载预训练权重文件并进行推理。
- 结果展示：通过绘制真实边界框和预测边界框对比效果，并展示了在测试集上的表现。
评估指标：- 在验证集上获得平均iou为0.6524，在测试集中表现较好。
总结而言，该模型能够有效识别肺结节并提供较为准确的边界框标注。
目录
1. 介绍
2. utils 工具函数
2.1 从xml获取类别+边界框函数 get_label_from_xml
2.2 图像归一化函数 normalization
2.3 计算iou函数 miou
2.4 绘制边界框函数 draw_bounding_box
3. dataset
4. train 函数
4.1 函数代码
4.2 可视化
5. 预测
5.1 预测代码
5.2 结果展示
项目地址：Fast R-CNN算法用于 lung nodule的目标识别，请访问下载地址：[Fast R-CNN algorithm for lung nodule detection]( "Fast R-CNN algorithm for lung nodule detection")
1. 介绍
如图为数据集展示，其中被选中的部分就是病状
以下是文件的目录：
data 代表项目的数据集。
dataset 实现了数据加载功能。
train 展现了核心训练逻辑。
utils 包含了一系列辅助功能模块。
predict 用于生成预测结果。
log 存储了训练过程中的各种日志信息。
2. utils 工具函数
这里是需要导入的头文件
2.1 从xml获取类别+边界框函数 get_label_from_xml
get_label_from_xml 函数，传入的是xml文件地址，返回的是类名+边界框
这里有的name 是小写的，所以用upper 转换为大写字母
# 读取 xml 文件中的描述信息，返回 label：类名 ,边界框
def get_label_from_xml(xmlFile):
an_file = open(xmlFile, encoding='utf-8')
tree = ET.parse(an_file)
root = tree.getroot()
label = [] # 保存label
bounding_box = [] # 保存边界框
for objects in root.findall('object'): # label 在 object 的标签下
cancer_type = objects.find('name').text.upper() # 返回目标的名字,转为大写
cancer_type = class_to_id[str(cancer_type)]
xmin = objects.find('bndbox').find('xmin').text # 获取边界框
xmax = objects.find('bndbox').find('xmax').text
ymin = objects.find('bndbox').find('ymin').text
ymax = objects.find('bndbox').find('ymax').text
# 判断 bounding box 是否规范,错误的边界框不保存->continue
if int(xmin) >= 512 or int(xmax) >= 512 or int(ymin) >=512 or int(ymax) >=512: # 去除超出最大边界
continue
elif int(xmin) <=0 or int(xmax) <=0 or int(ymin) <=0 or int(ymax) <=0: # 去除超出最小边界
continue
elif int(xmin) == int(xmax) or int(ymin) ==int(ymax): # 去除成线的边框
continue
else:
label.append(cancer_type)
bounding_box.append([int(xmin),int(ymin),int(xmax),int(ymax)]) # 左上角 + 右下角 坐标
return label,bounding_box2.2 图像归一化函数 normalization
由于所处理的图像为DCM格式，在这种情况下单个像素的灰度范围可能达到数千甚至更高数值。因此，在常规归一化方法难以适用的情况下，必须实施一个量身定制化的归一化处理流程。
# 图像归一化
def normalization(x):
max_pixel = np.max(x)
min_pixel = np.min(x)
x = (x - min_pixel) / (max_pixel - min_pixel)
return x加了normalization 函数之后图像的灰度范围
没加的灰度范围
以及
2.3 计算iou函数 miou
代码
# 返回一个 batch 的 平均 iou
def miou(truth,prediction):
miou_list = [] # 存放 batch 的 iou 列表
for i in range(len(prediction)):
pred_boxes = prediction[i]['boxes'] # 预测的 bounding box
label_boxes = truth[i]['boxes'] # 真实的 bounding box
iou = torchvision.ops.box_iou(label_boxes,pred_boxes)
iou = torch.max(iou,dim = 1)[0].detach().cpu().numpy()
mean_iou = np.mean(iou) # 单个 sample的 iou
miou_list.append(mean_iou)
return np.array(miou_list).mean() # 一个 batch 的平均 iou2.4 绘制边界框函数 draw_bounding_box
由于原始的DCM数据为灰度图象，在此将其转换为彩色图像以便于后续显示彩色边界框。
# 绘制边界框
def draw_bounding_box(image,label,row = 2,col = 2):
for batch in range(len(image)):
img = np.transpose(image[batch].data.cpu().numpy(), (1, 2, 0)) # tensor -> numpy, 更改 channel
img = cv2.cvtColor(img,cv2.COLOR_GRAY2BGR) # 转为彩色图像，否则，彩色边框会显示不出
boxes = label[batch]['boxes'].data.cpu().numpy() # 取边界框
labels = label[batch]['labels'].data.cpu().numpy() # 取类别
for i in range(boxes.shape[0]): # 逐个打印边界框
x1,y1,x2,y2 = int(boxes[i][0]), int(boxes[i][1]), int(boxes[i][2]), int(boxes[i][3]) # 取点
classes = id_to_class[str(labels[i])] # 将类别显示成名称
cv2.rectangle(img,(x1,y1),(x2,y2),(255,0,0),2) # 绘制矩形边界框
cv2.putText(img, text=classes, org=(x1, y1+5), fontFace=cv2.FONT_HERSHEY_SIMPLEX, # 显示文本
fontScale=1.5, thickness=1, lineType=cv2.LINE_AA, color=(0, 0, 255))
# 绘制边框
plt.subplot(row,col,batch+1)
plt.imshow(img)
plt.show()3. dataset
代码如下
此处在处理过程中遇到的问题是原始的D_CM文件具有广泛的灰度范围。因此，在这种情况下，默认的ToTensor处理方式无法自动完成归一化。为了满足需求，在这种情况下必须自定义一个专门的归一化函数。值得注意的是，默认情况下（即使用ToTensor时），对于处于[0, 1]区间的数据不会再执行额外的归一化操作。因此，在这种情况下必须先将图像转换为浮点型数值类型
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
import torch
from torch.utils.data import Dataset
from PIL import Image
from utils import get_label_from_xml,normalization
import pydicom
import numpy as np
class LungDetection(Dataset):
def __init__(self,img_path,label_path,transform):
img_root = os.listdir(img_path)
label_root = os.listdir(label_path)
self.imgs = [os.path.join(img_path,img) for img in img_root] # 图像的路径
self.xmls = [os.path.join(label_path,label) for label in label_root] # xml 标签的路径
self.transform = transform
def __getitem__(self, index):
img_path = self.imgs[index] # 获取单个sample路径
label_xml = self.xmls[index]
img_open = pydicom.read_file(img_path) # 读取图像
img_array = img_open.pixel_array # dcm -> np
img_array = normalization(img_array)
img_array = np.array(img_array, dtype=np.float32)
img_pic = Image.fromarray(img_array) # np -> PIL
img_tensor = self.transform(img_pic) # PIL -> tensor
label,bbox = get_label_from_xml(label_xml) # 获取 label
bbox_tensor = torch.as_tensor(bbox,dtype=torch.float32) # 边界框类型是float
label_tensor = torch.as_tensor(label,dtype=torch.int64) # 分类的类名是整型
target = {} # target是一个字典，包含边界框boxes和检测的类名labels
target['boxes'] = bbox_tensor
target['labels'] = label_tensor
return img_tensor,target
def __len__(self):
return len(self.imgs)4. train 函数
训练的过程大同小异
4.1 函数代码
from torchvision.transforms import transforms
import torch
import torchvision.models
from dataset import LungDetection
from torch.utils.data import DataLoader
from utils import detection_collate, draw_bounding_box, miou
import torch.optim as optim
from tqdm import tqdm
import numpy as np
# 超参数设置
TRAIN_IMAGE_PATH = './data/train/image' # 训练集路径
TRAIN_LABEL_PATH = './data/train/label' # 训练集路径
TEST_IMAGE_PATH = './data/test/image' # 训练集路径
TEST_LABEL_PATH = './data/test/label' # 训练集路径
BATCH_SIZE, ROW, COLUMN = 4, 2, 2 # batch size + 可视化图像的行列数
LEARNING_RATE = 0.00001
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EPOCHS = 50
# 预处理
transformer = transforms.Compose([transforms.ToTensor()])
# 加载训练集
trainSet = LungDetection(img_path=TRAIN_IMAGE_PATH,label_path=TRAIN_LABEL_PATH,transform=transformer)
trainLoader = DataLoader(trainSet, BATCH_SIZE, shuffle=True, collate_fn=detection_collate)
# 加载验证集
testSet = LungDetection(img_path=TEST_IMAGE_PATH,label_path=TEST_LABEL_PATH,transform=transformer)
testLoader = DataLoader(testSet, batch_size=BATCH_SIZE, shuffle=False, collate_fn=detection_collate)
# 建立模型
model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weight=None, progress=True,num_classes=2)
# model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=False, progress=True, num_classes=2)
# model=torchvision.models.detection.ssd300_vgg16(pretrained=False, progress=True, num_classes=2)
model.to(DEVICE)
# 建立优化器
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
# 训练函数
def train_fn(model, optimizer, trainloader, testloader, device, epochs):
for epoch in range(epochs):
loss_epoch = [] # 一个epoch训练损失
iou_epoch = [] # 一个epoch训练 iou
for images, labels in tqdm(trainloader): # 读取训练集
model.train()
images = list(image.to(device) for image in images) # 将数据放到 cuda 上
labels = [{k: v.to(device) for k, v in t.items()} for t in labels]
loss_dict = model(images, labels) # 前向传播
losses = sum(loss for loss in loss_dict.values()) # 损失求和
optimizer.zero_grad() # 梯度清零
losses.backward() # 反向传播
optimizer.step() # 梯度下降
losses = losses.data.cpu().numpy()
loss_epoch.append(losses) # 保留损失
with torch.no_grad():
model.eval() # 测试模式
try:
pred = model(images)
batch_iou = miou(truth=labels, prediction=pred) # 计算一个batch的平均iou
iou_epoch.append(batch_iou)
except:
continue
test_loss_epoch = [] # 一个epoch 的 loss
test_iou_epoch = [] # 一个epoch 的 iou
with torch.no_grad(): # 验证的时候，不需要backward
for images, labels in tqdm(testloader):
model.train() # 训练模式
images = list(image.to(device) for image in images)
labels = [{k: v.to(device) for k, v in t.items()} for t in labels]
loss_dict = model(images, labels)
losses = sum(loss for loss in loss_dict.values())
losses = losses.data.cpu().numpy()
test_loss_epoch.append(losses) # 一个batch 的loss
model.eval() # 测试模式
try:
pred = model(images)
test_batch_iou = miou(truth=labels, prediction=pred) # 计算一个batch的平均 iou
test_iou_epoch.append(test_batch_iou)
except:
continue
train_iou = np.array(iou_epoch).mean() # 将不同 batch 的 iou 求平均
train_loss = np.array(loss_epoch).mean() # 将不同 batch 的 loss 求均值
test_iou = np.array(test_iou_epoch).mean()
test_loss = np.array(test_loss_epoch).mean()
# 保留权重文件
if test_iou > 0.65:
static_dict = model.state_dict()
torch.save(static_dict, './log/pretrained_epoch_{},train_iou_{},test_iou_{}.pth'.format(epoch + 1, train_iou, test_iou))
print('epoch:', epoch + 1)
print('epoch_train_loss:', train_loss)
print('epoch_train_iou:', train_iou)
print('epoch_test_loss:', test_loss)
print('epoch_test_iou:', test_iou)
if __name__ == '__main__':
#model.load_state_dict(torch.load('./log/epoch_35,train_iou_0.9122412204742432,test_iou_0.6948733925819397.pth'))
train_fn(model, optimizer, trainLoader, testLoader, DEVICE, EPOCHS)
print(' training over !!!! ')4.2 可视化
测试代码
#load a batch data + 可视化
img,label = next(iter(testLoader))
draw_bounding_box(img,label)显示结果：
5. 预测
5.1 预测代码
from torchvision.transforms import transforms
import torch
import torchvision.models
from dataset import LungDetection
from torch.utils.data import DataLoader
from utils import detection_collate, draw_bounding_box,normalization
import cv2
import numpy as np
import matplotlib.pyplot as plt
# 参数设定
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
TEST_IMAGE_PATH = './data/train/image' # 训练集路径
TEST_LABEL_PATH = './data/train/label' # 训练集路径
WIGHT = './log/pretrained_epoch_16,train_iou_0.978933572769165,test_iou_0.652420699596405.pth' # 权重文件
BATCH_SIZE = 4
ROW, COLUMN = 2, BATCH_SIZE # 结果展示2行，第一行为 truth，第二行为 predict
id_to_class = {'1':'A'}
# 预处理
transformer = transforms.Compose([transforms.ToTensor()])
# 加载数据
testSet = LungDetection(img_path=TEST_IMAGE_PATH,label_path=TEST_LABEL_PATH,transform=transformer)
testLoader = DataLoader(testSet, batch_size=BATCH_SIZE, shuffle=True, collate_fn=detection_collate)
# 取出一个batch 的数据
test_images,test_labels = next(iter(testLoader))
# 加载模型
model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weight=None, progress=True, num_classes=2)
model.load_state_dict(torch.load(WIGHT))
model.to(DEVICE)
# 网络预测
model.eval()
images = list(image.to(DEVICE) for image in test_images)
pred = model(images)
# 显示预测结果
for batch in range(BATCH_SIZE): # 遍历 batch 图像
img = np.transpose(images[batch].data.cpu().numpy(), (1, 2, 0)) # tensor -> numpy, 更改 channel
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # 转为彩色图像，否则，彩色边框会显示不出
boxes = pred[batch]['boxes'].data.cpu().numpy() # 取边界
labels = pred[batch]['labels'].data.cpu().numpy() # 取类别
scores = pred[batch]['scores'].data.cpu().numpy() # 显示置信度
for i in range(boxes.shape[0]): # 逐个打印边界框
if scores[i] > 0.5:
x1, y1, x2, y2 = int(boxes[i][0]), int(boxes[i][1]), int(boxes[i][2]), int(boxes[i][3]) # 取点
classes = id_to_class[str(labels[i])] # 取类别
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2) # 绘制矩形边界框
cv2.putText(img, text=classes, org=(x1, y1 + 10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, # 显示文本
fontScale=1, thickness=1, lineType=cv2.LINE_AA, color=(0, 0, 255))
# 绘制边框
plt.subplot(ROW, COLUMN, batch + 1 + COLUMN) # 将预测结果放到第二行
plt.imshow(img)
# 显示真实的结果
draw_bounding_box(images,test_labels,row=ROW,col=COLUMN)5.2 结果展示
这里呈现的是（训练数据集合）通过以下链接可获得：
第一行为truth，第二行为prediction
该测试集在测试集中的表现较为突出，在评估指标中呈现出良好的效果。
可能因为数据的原因吧，iou虽然不高，但是在testSet 的表现也比较好
全部评论 (0)
还没有任何评论哟~