| import numpy as np |
| import torch |
| from skimage import filters |
| from torchvision.transforms.functional import resize |
|
|
| from utils.saliency import decoder, resnet |
|
|
|
|
| def get_smap(image, path, filter_size=15): |
| """ |
| Compute the saliency map of the target image using EMLNet. |
| Reference: https://arxiv.org/abs/1805.01047 |
| Reference: https://github.com/SenJia/EML-NET-Saliency |
| """ |
| if image.shape[0] != 3: |
| raise ValueError("Saliency prediction only supports RGB images") |
| sod_res = (480, 640) |
| imagenet_model = resnet.resnet50(f"{path}/emlnet/res_imagenet.pth").cuda().eval() |
| places_model = resnet.resnet50(f"{path}/emlnet/res_places.pth").cuda().eval() |
| decoder_model = ( |
| decoder.build_decoder(f"{path}/emlnet/res_decoder.pth", sod_res, 5, 5) |
| .cuda() |
| .eval() |
| ) |
| image_sod = resize(image, sod_res).unsqueeze(0) |
| with torch.no_grad(): |
| imagenet_feat = imagenet_model(image_sod, decode=True) |
| places_feat = places_model(image_sod, decode=True) |
| smap = decoder_model([imagenet_feat, places_feat]) |
| smap = resize(smap.squeeze(0).detach().cpu(), image.shape[1:]).squeeze(0) |
|
|
| def post_process(smap): |
| smap = filters.gaussian(smap, filter_size) |
| smap -= smap.min() |
| smap /= smap.max() |
| return smap |
|
|
| return post_process(smap.numpy()).astype(np.float32) |
|
|
