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 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 setup_model(mode, pretrained_path=None, pretrained_n_input_ch=2,
pretrained_n_output_ch=21, mat_scale=4):
# Create empty model
model = models.create(mode, mat_scale=mat_scale)
# Copy from pretrained model
if pretrained_path is not None:
if mode == 'seg' or mode == 'seg+' or mode == 'seg_tri' or \
mode == 'mat':
# FCN8s
logger.info('Load pretrained FCN8s model (%s)', pretrained_path)
pretrained = models.FCN8s(n_input_ch=pretrained_n_input_ch,
n_output_ch=pretrained_n_output_ch)
chainer.serializers.load_npz(pretrained_path, pretrained)
model.init_from_fcn8s(pretrained)
else:
logger.error('Unknown mode')
return model
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)
fcn_model = models.FCNss(pretrained_net=vgg_model,
n_class=n_class,
Time=False,
Space=False)
test_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)
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')
