Spaces:
Running on Zero
Running on Zero
| import torch | |
| from .position import PositionEmbeddingSine | |
| def split_feature( | |
| feature, | |
| num_splits=2, | |
| channel_last=False, | |
| ): | |
| if channel_last: # [B, H, W, C] | |
| b, h, w, c = feature.size() | |
| assert h % num_splits == 0 and w % num_splits == 0 | |
| b_new = b * num_splits * num_splits | |
| h_new = h // num_splits | |
| w_new = w // num_splits | |
| feature = ( | |
| feature.view(b, num_splits, h // num_splits, num_splits, w // num_splits, c) | |
| .permute(0, 1, 3, 2, 4, 5) | |
| .reshape(b_new, h_new, w_new, c) | |
| ) # [B*K*K, H/K, W/K, C] | |
| else: # [B, C, H, W] | |
| b, c, h, w = feature.size() | |
| assert h % num_splits == 0 and w % num_splits == 0 | |
| b_new = b * num_splits * num_splits | |
| h_new = h // num_splits | |
| w_new = w // num_splits | |
| feature = ( | |
| feature.view(b, c, num_splits, h // num_splits, num_splits, w // num_splits) | |
| .permute(0, 2, 4, 1, 3, 5) | |
| .reshape(b_new, c, h_new, w_new) | |
| ) # [B*K*K, C, H/K, W/K] | |
| return feature | |
| def merge_splits( | |
| splits, | |
| num_splits=2, | |
| channel_last=False, | |
| ): | |
| if channel_last: # [B*K*K, H/K, W/K, C] | |
| b, h, w, c = splits.size() | |
| new_b = b // num_splits // num_splits | |
| splits = splits.view(new_b, num_splits, num_splits, h, w, c) | |
| return ( | |
| splits.permute(0, 1, 3, 2, 4, 5) | |
| .contiguous() | |
| .view(new_b, num_splits * h, num_splits * w, c) | |
| ) # [B, H, W, C] | |
| else: # [B*K*K, C, H/K, W/K] | |
| b, c, h, w = splits.size() | |
| new_b = b // num_splits // num_splits | |
| splits = splits.view(new_b, num_splits, num_splits, c, h, w) | |
| return ( | |
| splits.permute(0, 3, 1, 4, 2, 5) | |
| .contiguous() | |
| .view(new_b, c, num_splits * h, num_splits * w) | |
| ) # [B, C, H, W] | |
| def mv_feature_add_position(features, attn_splits, feature_channels): | |
| pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2) | |
| assert features.dim() == 4 # [B*V, C, H, W] | |
| if attn_splits > 1: # Add position inside split windows. | |
| features_splits = split_feature(features, num_splits=attn_splits) | |
| position = pos_enc(features_splits) | |
| features_splits = features_splits + position | |
| features = merge_splits(features_splits, num_splits=attn_splits) | |
| else: | |
| position = pos_enc(features) | |
| features = features + position | |
| return features | |