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')