#model = WNet()
if is_cuda:
device = torch.device("cuda:1")
else:
device = torch.device("cpu")
model.to(device)
#model.cuda()
model.eval()
#model_downscale = False
mode = 'test'
optimizer = torch.optim.Adam(model.parameters(), lr=config.init_lr)
#optimizer
with open(config.model_tested, 'rb') as f:
para = torch.load(f, "cpu")
#para = torch.load(f,"cuda:0")
model.load_state_dict(para['state_dict'])
for step, [x] in enumerate(dataloader):
print('Step' + str(step + 1))
#print(x.shape)
x = x.to(device)
pred, pad_pred = model(x, mode, config.ModelDownscale)
print(pred.shape)
seg = pred.argmax(dim=1)
#seg = pred
print(seg.shape)
#seg = np.reshape(seg, (seg.shape[1],seg.shape[2],seg.shape[3]))
#x = np.reshape(x, (x.shape[1],x.shape[2],x.shape[3], x.shape[4]))
#print(x.max())
print(x.shape)
if result_upsample:
import torchvision
import pdb
from PIL import Image
config = Config()
if __name__ == '__main__':
dataset = DataLoader(config.datapath, "test")
dataloader = dataset.torch_loader()
model = WNet()
model.cuda(config.cuda_dev)
optimizer = torch.optim.SGD(model.parameters(), lr=config.init_lr)
#optimizer
with open(config.model_tested, 'rb') as f:
para = torch.load(f, "cuda:0")
pdb.set_trace()
model.load_state_dict(para['state_dict'], False)
for step, [x] in enumerate(dataloader):
print('Step' + str(step + 1))
#NCuts Loss
x = x.cuda(config.cuda_dev)
pred, rec_image = model(x)
seg = (pred.argmax(dim=1).to(torch.float) / 3 *
255).cpu().detach().numpy()
rec_image = rec_image.cpu().detach().numpy() * 255
x = x.cpu().detach().numpy() * 255
x = np.transpose(x.astype(np.uint8), (0, 2, 3, 1))
rec_image = np.transpose(rec_image.astype(np.uint8), (0, 2, 3, 1))
seg = seg.astype(np.uint8)
#pdb.set_trace()
for i in range(seg.shape[0]):
#reconstr = torch.nn.MSELoss().cuda(config.cuda_dev)
if config.useSSIMLoss:
import pytorch_ssim
reconstr = pytorch_ssim.SSIM()
else:
reconstr = torch.nn.MSELoss()
Ncuts = NCutsLoss()
mode = 'train'
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=config.lr_decay_iter,
gamma=config.lr_decay)
if checkname is not None:
with open(checkname, 'rb') as f:
checkpoint_data = torch.load(f, "cpu")
model.load_state_dict(checkpoint_data['state_dict'])
if "pretrain" not in checkname:
print('a non-pretrain checkpoint is now being loaded')
start_epoch = int(checkpoint_data['epoch'])
optimizer.load_state_dict(checkpoint_data['optimizer'])
scheduler.load_state_dict(checkpoint_data['scheduler'])
else:
print('a pretrain checkpoint is now being loaded')
start_epoch = 0
print('checkpoint loaded')
else:
start_epoch = 0
for epoch in range(start_epoch, config.max_iter):
print("Epoch: " + str(epoch + 1))
Ave_Ncuts = 0.0