def main():
use_cuda = torch.cuda.is_available()
path = os.path.expanduser('/home/yxk/data/')
dataset = voc_loader.VOC2012ClassSeg(root=path,
split='train',
transform=True)
vgg_model = models.VGGNet(requires_grad=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
fcn_model.load_state_dict(
torch.load('./pretrained_models/model120.pth', map_location='cpu'))
fcn_model.eval()
if use_cuda:
fcn_model.cuda()
criterion = CrossEntropyLoss2d()
for i in range(len(dataset)):
idx = random.randrange(0, len(dataset))
img, label = dataset[idx]
img_name = str(i)
img_src, _ = dataset.untransform(img, label) # whc
cv2.imwrite(path + 'image/%s_src.jpg' % img_name, img_src)
tools.labelTopng(label,
path + 'image/%s_label.png' % img_name) # 将label转换成图片
# a = tools.labelToimg(label)
#
# print(a)
if use_cuda:
img = img.cuda()
label = label.cuda()
img = Variable(img.unsqueeze(0), volatile=True)
label = Variable(label.unsqueeze(0), volatile=True)
# print("label: ", label.data)
out = fcn_model(img) # (1, 21, 320, 320)
loss = criterion(out, label)
# print(img_name, 'loss:', loss.data[0])
net_out = out.data.max(1)[1].squeeze_(0) # 320, 320
# print(out.data.max(1)[1].shape)
# print("out", net_out)
if use_cuda:
net_out = net_out.cpu()
tools.labelTopng(net_out,
path + 'image/%s_out.png' % img_name) # 将网络输出转换成图片
if i == 10:
break
def __init__(self):
parser = argparse.ArgumentParser(
description='Pytroch Cifar10 Training')
parser.add_argument('--lr',
default=0.01,
type=float,
help='learning rate\n default:0.01')
parser.add_argument('--bs',
default=256,
type=int,
help='batch size\n default:256')
parser.add_argument('--gpu',
default=True,
help='use gpu as accelerator')
parser.add_argument(
'--net',
default='vgg16',
choices=['vgg16', 'resnet18', 'sevgg16', 'seresnet18'],
help='provided network:vgg16, resnet18, sevgg16, seresnet18')
parser.add_argument('--verbose',
default=True,
help='print some useful info')
parser.add_argument('--e',
default=50,
type=int,
help='the num of training epoch\n default:50')
args = parser.parse_args()
# parse and preparing params
self.batch_size = args.bs
self.epoches = args.e
self.lr = args.lr
self.use_gpu = args.gpu
self.net = args.net
self.model = None
self.verbose = args.verbose
if self.net == 'vgg16':
self.model = models.VGGNet()
elif self.net == 'resnet18':
self.model = models.Resnet()
elif self.net == 'sevgg16':
self.model = models.SeVGGNet()
elif self.net == 'seresnet18':
self.model = models.SeResnet()
if self.use_gpu and torch.cuda.is_available():
self.model.to(self.gpu)
if self.verbose:
print(self.model)
self.loss_fn = nn.CrossEntropyLoss()
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=0.9)
def evaluate():
use_cuda = torch.cuda.is_available()
path = os.path.expanduser('/home/yxk/data/')
val_data = voc_loader.VOC2012ClassSeg(root=path,
split='val',
transform=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=1,
shuffle=False,
num_workers=5)
print('load model .....')
vgg_model = models.VGGNet(requires_grad=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
fcn_model.load_state_dict(torch.load('params.pth'))
if use_cuda:
fcn_model.cuda()
fcn_model.eval()
label_trues, label_preds = [], []
# for idx, (img, label) in enumerate(val_loader):
for idx in range(len(val_data)):
img, label = val_data[idx]
img = img.unsqueeze(0)
if use_cuda:
img = img.cuda()
img = Variable(img)
out = fcn_model(img) # 1, 21, 320, 320
pred = out.data.max(1)[1].squeeze_(1).squeeze_(0) # 320, 320
if use_cuda:
pred = pred.cpu()
label_trues.append(label.numpy())
label_preds.append(pred.numpy())
if idx % 30 == 0:
print('evaluate [%d/%d]' % (idx, len(val_loader)))
metrics = tools.accuracy_score(label_trues, label_preds)
metrics = np.array(metrics)
metrics *= 100
print('''\
Accuracy: {0}
Accuracy Class: {1}
Mean IU: {2}
FWAV Accuracy: {3}'''.format(*metrics))
def evaluate(modelPath):
labels = ["BG", "figure", "table", "text"]
use_cuda = torch.cuda.is_available()
val_data = data_loader.ClassSeg(root=data_path, split='test', transform=True, filePath='DSSE',chanelCat=in_channels_Nmuber)
val_loader = torch.utils.data.DataLoader(val_data,batch_size=1,shuffle=False,num_workers=5)
print('load model .....')
print("Using FCNs")
vgg_model = models.VGGNet(model='vgg_self', pretrained=False, in_channels=in_channels_Nmuber)
fcn_model = models.FCNs(pretrained_net=vgg_model, n_class=n_class, Attention=True)
fcn_model.load_state_dict(torch.load(modelPath))
if use_cuda:
fcn_model.cuda()
fcn_model.eval()
label_trues, label_preds = [], []
matrixs = np.zeros((n_class,n_class))
for idx, (img, label,_) in enumerate(val_loader):
img, label,Image_Path = val_data[idx]
img = img.unsqueeze(0)
if use_cuda:
img = img.cuda()
img = Variable(img)
out = fcn_model(img) # 1, 21, 320, 320
srcImage = mpimg.imread(Image_Path)
pred = out.data.max(1)[1].squeeze_(1).squeeze_(0) # 320, 320
if use_cuda:
pred = pred.cpu()
# 后处理
data = pred.numpy()
# CutX=int(data.shape[1]/32)
#
# for Cuti in range(data.shape[0]):
# for Cutj in range(0,data.shape[1],CutX):
# temp=data[Cuti,Cutj:Cutj+CutX-1]
# data[Cuti, Cutj:Cutj + CutX - 1]=stats.mode(temp)[0][0]
#
# dataT = data.T
# CutY = int(dataT.shape[1]/3)
#
# for Cuti in range(dataT.shape[0]):
# for Cutj in range(0,dataT.shape[1],CutY):
# temp=dataT[Cuti,Cutj:Cutj+CutX-1]
# data[Cutj:Cutj + CutX - 1,Cuti]=stats.mode(temp)[0][0]
# # -------------------------------------------------------------------------
#
# if len(srcImage.shape) == 3:
# image = cv2.cvtColor(srcImage, cv2.COLOR_BGR2GRAY) # 将图像转化为灰度图像
# else:
# image = srcImage
#
# h = int(max_height / 64) * 64
# w = int(image.shape[1] * (max_height / image.shape[0]) / 64) * 64
#
# image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LINEAR)
# # print(image.shape)
# sobelCombinedIMG = sobelCombined(image)
# data=np.multiply(data,sobelCombinedIMG)
#
# #-------------------------------------------------------------------------
label_trues.append(label.numpy())
label_preds.append(data)
if idx % 2 == 0:
print('evaluate [%d/%d]' % (idx, len(val_loader)))
label_matrix_T=label.numpy()
pre_matrix_T = data
label_matrix = label_matrix_T.flatten()
pre_matrix = pre_matrix_T.flatten()
matrix = metrics.confusion_matrix(label_matrix, pre_matrix)
if sum(sum(matrixs)) == 0:
for i in range(len(matrixs)):
for j in range(len(matrixs[0])):
matrixs[i][j] = matrix[i][j]
else:
# 迭代输出行
for i in range(len(matrix)):
# 迭代输出列
for j in range(len(matrix[0])):
# print(range(len(matrix[0])))
matrixs[i][j] = (matrixs[i][j] + matrix[i][j])/2
# Mymetrics = tools.accuracy_score(label_trues, label_preds,n_class)
# Mymetrics = np.array(Mymetrics)
# Mymetrics *= 100
# print('''\
# Accuracy: {0}
# Accuracy Class: {1}
# Mean IU: {2}
# FWAV Accuracy: {3}'''.format(*Mymetrics))
plot_confusion_matrix(matrixs, classes=labels, normalize=True, title='Normalized confusion matrix',cmap=plt.cm.Blues,yMax=3.5)
#
numberTotal = sum(sum(matrixs))
muberTrue = 0
PercisionList = []
RecallList = []
for i in range(len(matrixs)):
for j in range(len(matrixs[0])):
if i == j:
muberTrue = muberTrue + matrixs[i, j]
for i in range(len(matrixs)):
PercisionList.append(matrixs[i, i] / sum(matrixs[:, i]))
for i in range(len(matrixs)):
RecallList.append(matrixs[i, i] / sum(matrixs[i, :]))
Acurracy = muberTrue / numberTotal
Percision = sum(PercisionList) / len(PercisionList)
Recall = sum(RecallList) / len(RecallList)
F1 = (2 * Percision * Recall) / (Percision + Recall)
print(Acurracy)
print(Percision)
print(Recall)
print(F1)
split='val',
transform=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=batch_size,
shuffle=False,
num_workers=workNmber,
pin_memory=True)
# test_data = data_loader.ClassSeg(root=data_path, split='test', transform=True)
# test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=workNmber)
print('load model....')
if model_Type == 0:
print("Using FCN32s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN32s(pretrained_net=vgg_model, n_class=n_class)
test_model = models.FCN32s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 1:
print("Using FCN16s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN16s(pretrained_net=vgg_model, n_class=n_class)
test_model = models.FCN16s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 2:
print("Using FCN8s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
test_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 3:
print("Using FCN1s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
train_data = voc_loader.VOC2012ClassSeg(root=path, split='train', transform=True)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
num_workers=5)
val_data = voc_loader.VOC2012ClassSeg(root=path,
split='val',
transform=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=batch_size,
shuffle=False,
num_workers=5)
print('load model.....')
vgg_model = models.VGGNet(requires_grad=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
if use_cuda:
fcn_model.cuda()
print("begin loss")
criterion = loss.CrossEntropy2d()
# loss.CrossEntropy2d()
# create your optimizer
print("begin optimizer")
optimizer = torch.optim.SGD(fcn_model.parameters(), lr=0.01)
# optimizer = torch.optim.SGD(fcn_model.parameters(), lr=1e-3, momentum=0.9)
testset = data.dataset(root=args.data, train=False)
testloader = torch.utils.data.DataLoader(testset, batch_size=100,
shuffle=False, num_workers=2)
# define classes
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# define the model
print('==> building model',args.arch,'...')
if args.arch == 'nin':
model = models.NiN()
elif args.arch == 'AlexNet':
model = models.AlexNet()
elif args.arch == 'VGGNet':
model = models.VGGNet()
else:
raise Exception(args.arch+' is currently not supported')
# initialize the model
if not args.pretrained:
print('==> Initializing model parameters ...')
best_acc = 0
for m in model.modules():
if isinstance(m, nn.Conv2d):
m.weight.data.normal_(0, 0.05)
m.bias.data.zero_()
else:
print('==> Load pretrained model form', args.pretrained, '...')
pretrained_model = torch.load(args.pretrained)
best_acc = pretrained_model['best_acc']
def main():
use_cuda = torch.cuda.is_available()
path = os.path.expanduser(data_path)
dataset = data_loader.ClassSeg(root=data_path, split='val', transform=True)
if model_Type == 0:
print("Using FCN32s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN32s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 1:
print("Using FCN16s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN16s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 2:
print("Using FCN8s")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCN8s(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 3:
print("Using FCN1s")
vgg_model = models.VGGNet(model='vgg16', pretrained=False)
fcn_model = models.FCNss(pretrained_net=vgg_model,
n_class=n_class,
Time=False,
Space=False)
elif model_Type == 4:
print("Using FCNs")
vgg_model = models.VGGNet(model='vgg_self', pretrained=True)
fcn_model = models.FCNs(pretrained_net=vgg_model, n_class=n_class)
elif model_Type == 5:
print("Using FCNs")
vgg_model = models.VGGNet(model='vgg16', pretrained=True)
fcn_model = models.FCNss(pretrained_net=vgg_model, n_class=n_class)
# fcn_model.load_state_dict(torch.load('models/model50.pth'))
fcn_model.load_state_dict(torch.load('./models/temp.pth'))
fcn_model.eval()
if use_cuda:
fcn_model.cuda()
for i in range(len(dataset)):
idx = i
img, label, Image_Path = dataset[idx]
print("deal %s" % Image_Path[-14:])
labelImage = tools.labelToimg(label)
if use_cuda:
img = img.cuda()
img = Variable(img.unsqueeze(0))
out = fcn_model(img) # (1, 21, 320, 320)
net_out = out.data.max(1)[1].squeeze_(0) # 320, 320
if use_cuda:
net_out = net_out.cpu()
# outImage = tools.labelToimg(net_out) # 将网络输出转换成图片
# 后处理
data = net_out.numpy()
outImage = tools.labelToimg(torch.from_numpy(data)) # 将网络输出转换成图片
plt.imshow(labelImage)
plt.axis('off')
plt.savefig('./image/%s_P.jpg' % Image_Path[-14:-4],
bbox_inches='tight')