| import numpy as np | |
| import torch | |
| def generate_label(inputs, imsize=512): | |
| """Generate label maps from model outputs""" | |
| pred_batch = [] | |
| for input in inputs: | |
| input = input.unsqueeze(0) | |
| pred = np.squeeze(input.data.max(1)[1].cpu().numpy(), axis=0) | |
| pred_batch.append(pred) | |
| pred_batch = np.array(pred_batch) | |
| pred_batch = torch.from_numpy(pred_batch) | |
| label_batch = [] | |
| for p in pred_batch: | |
| p = p.view(1, imsize, imsize) | |
| label_batch.append(p.data.cpu()) | |
| label_batch = torch.cat(label_batch, 0) | |
| label_batch = label_batch.type(torch.LongTensor) | |
| return label_batch |