基于U-Net的CT脊柱语义分割从CT脊柱语义分割数据集准备到GUI应用实现数据说明CT脊柱语义分割数据集608张原图jpg格式png标签背景和目标用0和1表示CT脊柱语义分割任务使用U-Net模型进行训练和推理。从安装依赖、准备数据集、训练模型、评估模型、推理代码到构建GUI应用程序进行详细说明并附上完整代码。—基于U-Net的CT脊柱语义分割全流程指南从数据准备到GUI应用实现**1. 安装依赖确保安装了必要的Python库pipinstalltorch torchvision matplotlib opencv-python-headless numpy scikit-image PyQt5 segmentation-models-pytorch albumentations—2. 准备数据集数据集描述原图数量608张。标签格式PNG像素值为背景0目标脊柱1假设你的数据集已组织如下spine_dataset/ ├── images/ │ ├── train/ │ │ ├── img1.jpg │ │ └── ... │ └── val/ │ ├── img401.jpg │ └── ... └── masks/ ├── train/ │ ├── mask1.png │ └── ... └── val/ ├── mask401.png └── ...3. 自定义数据集类我们将创建一个自定义的Dataset类来加载图像和对应的mask。importosfromPILimportImageimportnumpyasnpimporttorchfromtorch.utils.dataimportDatasetfromtorchvisionimporttransformsfromalbumentationsimportCompose,Resize,Normalize,HorizontalFlip,VerticalFlipclassSpineDataset(Dataset):def__init__(self,image_dir,mask_dir,transformNone):self.image_dirimage_dir self.mask_dirmask_dir self.transformtransform self.imagesos.listdir(image_dir)def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_pathos.path.join(self.image_dir,self.images[idx])mask_pathos.path.join(self.mask_dir,self.images[idx].replace(.jpg,.png))# 加载图像和maskimagenp.array(Image.open(img_path).convert(RGB))masknp.array(Image.open(mask_path),dtypenp.int64)# 确保mask是整数类型ifself.transform:augmentationsself.transform(imageimage,maskmask)imageaugmentations[image]maskaugmentations[mask]returnimage,mask4. 配置和训练 U-Net 模型数据增强与加载# 数据增强transformCompose([Resize(256,256),# 统一调整图像大小HorizontalFlip(p0.5),VerticalFlip(p0.5),Normalize(mean(0.485,0.456,0.406),std(0.229,0.224,0.225)),])# 加载训练集和验证集train_datasetSpineDataset(image_dir./path/to/train/images,mask_dir./path/to/train/masks,transformtransform)val_datasetSpineDataset(image_dir./path/to/val/images,mask_dir./path/to/val/masks,transformtransform)train_loadertorch.utils.data.DataLoader(train_dataset,batch_size8,shuffleTrue)val_loadertorch.utils.data.DataLoader(val_dataset,batch_size8,shuffleFalse)使用segmentation_models_pytorch加载 U-Netimporttorchimporttorch.nnasnnfromsegmentation_models_pytorchimportUnetfromsegmentation_models_pytorch.lossesimportDiceLossfromtorch.optimimportAdam# 定义模型modelUnet(encoder_nameresnet34,encoder_weightsimagenet,classes2,activationNone).cuda()# 损失函数和优化器criterionDiceLoss(modebinary)# 二分类任务optimizerAdam(model.parameters(),lr1e-4)# 训练函数deftrain_model(model,train_loader,val_loader,criterion,optimizer,num_epochs20):forepochinrange(num_epochs):model.train()running_loss0.0forimages,masksintrain_loader:imagesimages.cuda()masksmasks.long().cuda()outputsmodel(images)losscriterion(outputs,masks)optimizer.zero_grad()loss.backward()optimizer.step()running_lossloss.item()print(fEpoch [{epoch1}/{num_epochs}], Loss:{running_loss/len(train_loader):.4f})# 验证阶段model.eval()val_loss0.0withtorch.no_grad():forimages,masksinval_loader:imagesimages.cuda()masksmasks.long().cuda()outputsmodel(images)losscriterion(outputs,masks)val_lossloss.item()print(fValidation Loss:{val_loss/len(val_loader):.4f})运行训练代码train_model(model,train_loader,val_loader,criterion,optimizer,num_epochs20)5. 推理代码编写推理代码用于分割单张图片中的脊柱区域importcv2importnumpyasnpimporttorchfromtorchvisionimporttransformsdefsegment_image(image_path,model):transformtransforms.Compose([transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225]),])# 加载并预处理图像imageImage.open(image_path).convert(RGB)input_tensortransform(image).unsqueeze(0).cuda()# 模型预测model.eval()withtorch.no_grad():outputmodel(input_tensor)pred_masktorch.argmax(output,dim1).squeeze().cpu().numpy()# 将预测结果转换为彩色图像colored_masknp.zeros((pred_mask.shape[0],pred_mask.shape[1],3),dtypenp.uint8)colored_mask[pred_mask0][0,0,0]# 背景colored_mask[pred_mask1][255,0,0]# 脊柱returncolored_mask# 示例用法image_path./path/to/test/image.jpgcolored_masksegment_image(image_path,model)cv2.imshow(Segmented Image,colored_mask)cv2.waitKey(0)cv2.destroyAllWindows()6. GUI应用程序使用PyQt5构建一个简单的GUI应用程序用于选择图片并展示分割结果importsysfromPyQt5.QtWidgetsimportQApplication,QWidget,QVBoxLayout,QPushButton,QLabel,QFileDialogfromPyQt5.QtGuiimportQPixmapfromPyQt5.QtCoreimportQtimportcv2classApp(QWidget):def__init__(self,model):super().__init__()self.modelmodel self.initUI()definitUI(self):self.setWindowTitle(CT脊柱分割系统)layoutQVBoxLayout()self.image_labelQLabel(self)layout.addWidget(self.image_label)self.select_buttonQPushButton(选择图片,self)self.select_button.clicked.connect(self.select_image)layout.addWidget(self.select_button)self.setLayout(layout)self.show()defselect_image(self):optionsQFileDialog.Options()file_name,_QFileDialog.getOpenFileName(self,选择图片,,Image Files (*.jpg *.jpeg *.png),optionsoptions)iffile_name:self.segment_and_show(file_name)defsegment_and_show(self,image_path):segmented_imagesegment_image(image_path,self.model)height,width,channelsegmented_image.shape bytes_per_line3*width q_imgQImage(segmented_image.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()pixmapQPixmap.fromImage(q_img)self.image_label.setPixmap(pixmap)if__name____main__:appQApplication(sys.argv)exApp(model)# 将训练好的模型传入sys.exit(app.exec_())7. 运行确保所有路径和参数正确无误。训练模型后保存权重文件如best.pth。在推理代码和GUI中加载训练好的模型权重。启动GUI应用程序选择测试图片进行分割。