Spaces:
Runtime error
Runtime error
File size: 3,875 Bytes
7d0ed79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# ------------------------------------------------------------------------------------------
# Copyright (c) 2024 Baifeng Shi.
# All rights reserved.
#
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
# ------------------------------------------------------------------------------------------
import math
import torch
import torch.nn.functional as F
from einops import rearrange
from .utils import split_chessboard, merge_chessboard
def forward(model, input, scales=None, img_sizes=None, max_split_size=None, resize_output_to_idx=0, num_prefix_token=0,
output_shape='bnc'):
assert input.dim() == 4, "Input image must be in the shape of BxCxHxW."
assert input.shape[2] == input.shape[3], "Currently only square images are supported."
assert output_shape in ['bnc', 'bchw'], "Output shape should be either BxNxC (e.g., ViT) or BxCxHxW (e.g., ConvNet)."
assert output_shape == 'bnc' or num_prefix_token == 0, "For ConvNet there shouldn't be any prefix token."
b, c, input_size, _ = input.shape
# image size for each scale
assert scales is not None or img_sizes is not None, "Please assign either scales or img_sizes."
img_sizes = img_sizes or [int(input_size * scale) for scale in scales]
# prepare multiscale inputs
max_split_size = max_split_size or input_size # The maximum size of each split of image. Set as the input size by default
num_splits = [math.ceil(size / max_split_size) for size in img_sizes] # number of splits each scale
input_multiscale = []
for size, num_split in zip(img_sizes, num_splits):
x = F.interpolate(input.to(torch.float32), size=size, mode='bicubic').to(input.dtype)
x = split_chessboard(x, num_split=num_split)
input_multiscale.append(x)
# run feedforward on each scale
outs_multiscale = [model(x) for x in input_multiscale]
if num_prefix_token > 0:
outs_prefix_multiscale = [out[:, :num_prefix_token] for out in outs_multiscale]
outs_multiscale = [out[:, num_prefix_token:] for out in outs_multiscale]
if output_shape == 'bnc':
height = int(outs_multiscale[0].shape[1] ** 0.5)
if height**2 == outs_multiscale[0].shape[1]:
width = height
else:
width = int(outs_multiscale[0].shape[1]/height)
assert width*height == outs_multiscale[0].shape[1]
#print(height, width, outs_multiscale[0].shape[1])
# available by siglip
#outs_multiscale = [rearrange(out, 'b (h w) c -> b c h w', h=int(out.shape[1] ** 0.5), w=int(out.shape[1] ** 0.5))
# for out in outs_multiscale]
outs_multiscale = [rearrange(out, 'b (h w) c -> b c h w', h=height, w=width)
for out in outs_multiscale]
# merge outputs of different splits for each scale separately
outs_multiscale = [merge_chessboard(out, num_split=num_split) for num_split, out in zip(num_splits, outs_multiscale)]
# interpolate outputs from different scales and concat together
#output_size = outs_multiscale[resize_output_to_idx].shape[-2]
output_size = [height, width]
out = torch.cat([F.interpolate(outs_multiscale[i].to(torch.float32), size=output_size,
mode='area').to(outs_multiscale[i].dtype)
for i in range(len(outs_multiscale))], dim=1)
if output_shape == 'bnc':
out = rearrange(out, 'b c h w -> b (h w) c')
if num_prefix_token > 0:
# take the mean of prefix tokens from different splits for each scale
outs_prefix_multiscale = [torch.stack(out.split(b, dim=0), dim=0).mean(dim=0) for out in outs_prefix_multiscale]
out_prefix_multiscale = torch.cat(outs_prefix_multiscale, dim=-1)
out = torch.cat([out_prefix_multiscale, out], dim=1)
return out
|