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 main():
use_cuda = torch.cuda.is_available()
path = os.path.expanduser('/home/yxk/Downloads/')
dataset = voc_loader.VOC2012ClassSeg(root=path,
split='val',
transform=True)
print("begin...")
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转换成图片
if i == 10:
break
import voc_loader
import loss
# import visualize
batch_size = 1
learning_rate = 1e-10
epoch_num = 30
best_test_loss = np.inf
pretrained = 'reload'
use_cuda = torch.cuda.is_available()
path = os.path.expanduser('/home/yxk/Downloads/')
n_class = 21
print('load data....')
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)