基于U-Net的劃痕缺陷分割（課程設計）

論文：**U-Net:?Convolutional?Networks?for?Biomedical?Image?Segmentation**

https://arxiv.org/abs/1505.04597

包括U_Net/R2U_Net/AttU_Net/R2AttU_Net，基于Pytorch的代碼和數據集下載地址：下載地址

import argparse
import os
from solver import Solver
from data_loader import get_loader
from torch.backends import cudnn
import randomdef main(config):cudnn.benchmark = Trueif config.model_type not in ['U_Net','R2U_Net','AttU_Net','R2AttU_Net']:print('ERROR!! model_type should be selected in U_Net/R2U_Net/AttU_Net/R2AttU_Net')print('Your input for model_type was %s'%config.model_type)return# Create directories if not existif not os.path.exists(config.model_path):os.makedirs(config.model_path)if not os.path.exists(config.result_path):os.makedirs(config.result_path)config.result_path = os.path.join(config.result_path,config.model_type)if not os.path.exists(config.result_path):os.makedirs(config.result_path)#lr = random.random()*0.0005 + 0.0000005#augmentation_prob= random.random()*0.7#epoch = random.choice([100,150,200,250])#decay_ratio = random.random()*0.8#decay_epoch = int(epoch*decay_ratio)#config.augmentation_prob = config.augmentation_prob#config.num_epochs = config.epoch#config.lr = config.lr#config.num_epochs_decay = config.decay_epochprint(config)train_loader = get_loader(image_path=config.train_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='train',augmentation_prob=config.augmentation_prob)valid_loader = get_loader(image_path=config.valid_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='valid',augmentation_prob=0.)test_loader = get_loader(image_path=config.test_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='test',augmentation_prob=0.)solver = Solver(config, train_loader, valid_loader, test_loader)# Train and sample the imagesif config.mode == 'train':solver.train()elif config.mode == 'test':solver.test()if __name__ == '__main__':parser = argparse.ArgumentParser()# model hyper-parametersparser.add_argument('--image_size', type=int, default=224)parser.add_argument('--t', type=int, default=3, help='t for Recurrent step of R2U_Net or R2AttU_Net')# training hyper-parametersparser.add_argument('--img_ch', type=int, default=3)parser.add_argument('--output_ch', type=int, default=1)parser.add_argument('--num_epochs', type=int, default=100)parser.add_argument('--num_epochs_decay', type=int, default=70)parser.add_argument('--batch_size', type=int, default=1)parser.add_argument('--num_workers', type=int, default=8)parser.add_argument('--lr', type=float, default=0.0002)parser.add_argument('--beta1', type=float, default=0.5)        # momentum1 in Adamparser.add_argument('--beta2', type=float, default=0.999)      # momentum2 in Adam    parser.add_argument('--augmentation_prob', type=float, default=0.4)parser.add_argument('--log_step', type=int, default=2)parser.add_argument('--val_step', type=int, default=2)# miscparser.add_argument('--mode', type=str, default='train')parser.add_argument('--model_type', type=str, default='U_Net', help='U_Net/R2U_Net/AttU_Net/R2AttU_Net')parser.add_argument('--model_path', type=str, default='./models')parser.add_argument('--train_path', type=str, default='./dataset/train/')parser.add_argument('--valid_path', type=str, default='./dataset/valid/')parser.add_argument('--test_path', type=str, default='./dataset/test/')parser.add_argument('--result_path', type=str, default='./result/')parser.add_argument('--cuda_idx', type=int, default=0)config = parser.parse_args()main(config)

運行結果：

?

總結

以上是生活随笔為你收集整理的基于U-Net图像分割的划痕缺陷分割（课程设计）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。