Weird matt images #10
Description
Hi,
Thanks for sharing the code. I have created the inference script to get the mask. But the masks are no where related to the inputs. And may I know if the output mask shape is 224*224. Attached the output and below is the inference code.
CODE
opt = get_option()
torch.manual_seed(opt.seed)
module = importlib.import_module("model.{}".format(opt.model.lower()))
logger = LogWritter(opt)
dev = torch.device("cuda:{}".format(opt.GPU_ID) if torch.cuda.is_available() else "cpu")
net = module.Net(opt)
net = net.to(dev)
msg = "# params:{}\n".format(
sum(map(lambda x: x.numel(), net.parameters())))
aug = Augment(opt)
if opt.save_result:
save_root = os.path.join(opt.save_root, opt.save_msg)
os.makedirs(save_root, exist_ok=True)
path = './ckpt/ckpt_rs1/best_DUTS-TE.pt'
print('loading model from: {}'.format(path))
state_dict = torch.load(path, map_location=lambda storage, loc: storage)
net.load_state_dict(state_dict)
IMG = cv2.imread('./images/IMG_3042.jpg')
IMG = aug.resize(IMG)
IMG = np.ascontiguousarray(IMG.transpose((2, 0, 1)))
IMG = aug.norm(aug.to_tensor(IMG) / 255.)
IMG = np.expand_dims(IMG, axis=0)
b, c, h, w = IMG.shape
IMG = torch.from_numpy(IMG).to(dev)
pred = net(IMG)
pred_sal = F.pixel_shuffle(pred[-1], 4)
pred_sal = F.interpolate(pred_sal, (h,w), mode='bilinear', align_corners=False)
pred_sal = torch.sigmoid(pred_sal).squeeze()
pred_sal = (pred_sal * 255.).detach().cpu().numpy().astype('uint8')
matt_img = pred[0].repeat(1,256,1,1)
matt_img = F.pixel_shuffle(matt_img, 16)
matt_img = F.interpolate(matt_img, (h,w), mode='bilinear', align_corners=False)
matt_img = torch.sigmoid(matt_img)
matt_img = (matt_img*255.).squeeze().detach().cpu().numpy().astype('uint8')
NAME = 'IMG_3042'
save_path_msk = os.path.join(save_root, "{}_msk.png".format(NAME))
save_path_matt = os.path.join(save_root, "{}_matt.png".format(NAME))
io.imsave(save_path_msk, mask)
io.imsave(save_path_matt, matt_img)