import torch import torch.nn as nn import torch.nn.functional as F from lib.exceptions import EmptyTensorError from lib.utils import interpolate_dense_features, upscale_positions def process_multiscale(image, model, scales=[.5, 1, 2]): b, _, h_init, w_init = image.size() device = image.device assert(b == 1) all_keypoints = torch.zeros([3, 0]) all_descriptors = torch.zeros([ model.dense_feature_extraction.num_channels, 0 ]) all_scores = torch.zeros(0) previous_dense_features = None banned = None for idx, scale in enumerate(scales): current_image = F.interpolate( image, scale_factor=scale, mode='bilinear', align_corners=True ) _, _, h_level, w_level = current_image.size() dense_features = model.dense_feature_extraction(current_image) del current_image _, _, h, w = dense_features.size() # Sum the feature maps. if previous_dense_features is not None: dense_features += F.interpolate( previous_dense_features, size=[h, w], mode='bilinear', align_corners=True ) del previous_dense_features # Recover detections. detections = model.detection(dense_features) if banned is not None: banned = F.interpolate(banned.float(), size=[h, w]).bool() detections = torch.min(detections, ~banned) banned = torch.max( torch.max(detections, dim=1)[0].unsqueeze(1), banned ) else: banned = torch.max(detections, dim=1)[0].unsqueeze(1) fmap_pos = torch.nonzero(detections[0].cpu()).t() del detections # Recover displacements. displacements = model.localization(dense_features)[0].cpu() displacements_i = displacements[ 0, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :] ] displacements_j = displacements[ 1, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :] ] del displacements mask = torch.min( torch.abs(displacements_i) < 0.5, torch.abs(displacements_j) < 0.5 ) fmap_pos = fmap_pos[:, mask] valid_displacements = torch.stack([ displacements_i[mask], displacements_j[mask] ], dim=0) del mask, displacements_i, displacements_j fmap_keypoints = fmap_pos[1 :, :].float() + valid_displacements del valid_displacements try: raw_descriptors, _, ids = interpolate_dense_features( fmap_keypoints.to(device), dense_features[0] ) except EmptyTensorError: continue fmap_pos = fmap_pos.to(device) fmap_keypoints = fmap_keypoints.to(device) fmap_pos = fmap_pos[:, ids] fmap_keypoints = fmap_keypoints[:, ids] del ids keypoints = upscale_positions(fmap_keypoints, scaling_steps=2) del fmap_keypoints descriptors = F.normalize(raw_descriptors, dim=0).cpu() del raw_descriptors keypoints[0, :] *= h_init / h_level keypoints[1, :] *= w_init / w_level fmap_pos = fmap_pos.cpu() keypoints = keypoints.cpu() keypoints = torch.cat([ keypoints, torch.ones([1, keypoints.size(1)]) * 1 / scale, ], dim=0) scores = dense_features[ 0, fmap_pos[0, :], fmap_pos[1, :], fmap_pos[2, :] ].cpu() / (idx + 1) del fmap_pos all_keypoints = torch.cat([all_keypoints, keypoints], dim=1) all_descriptors = torch.cat([all_descriptors, descriptors], dim=1) all_scores = torch.cat([all_scores, scores], dim=0) del keypoints, descriptors previous_dense_features = dense_features del dense_features del previous_dense_features, banned keypoints = all_keypoints.t().detach().numpy() del all_keypoints scores = all_scores.detach().numpy() del all_scores descriptors = all_descriptors.t().detach().numpy() del all_descriptors return keypoints, scores, descriptors