ultralytics 8.2.84 new SAM flexible imgsz inference (#15882)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
5d66140ce1
commit
7053169fd0
6 changed files with 70 additions and 7 deletions
|
|
@ -982,10 +982,31 @@ class TinyViT(nn.Module):
|
|||
layer = self.layers[i]
|
||||
x = layer(x)
|
||||
batch, _, channel = x.shape
|
||||
x = x.view(batch, 64, 64, channel)
|
||||
x = x.view(batch, self.patches_resolution[0] // 4, self.patches_resolution[1] // 4, channel)
|
||||
x = x.permute(0, 3, 1, 2)
|
||||
return self.neck(x)
|
||||
|
||||
def forward(self, x):
|
||||
"""Performs the forward pass through the TinyViT model, extracting features from the input image."""
|
||||
return self.forward_features(x)
|
||||
|
||||
def set_imgsz(self, imgsz=[1024, 1024]):
|
||||
"""
|
||||
Set image size to make model compatible with different image sizes.
|
||||
|
||||
Args:
|
||||
imgsz (Tuple[int, int]): The size of the input image.
|
||||
"""
|
||||
imgsz = [s // 4 for s in imgsz]
|
||||
self.patches_resolution = imgsz
|
||||
for i, layer in enumerate(self.layers):
|
||||
input_resolution = (
|
||||
imgsz[0] // (2 ** (i - 1 if i == 3 else i)),
|
||||
imgsz[1] // (2 ** (i - 1 if i == 3 else i)),
|
||||
)
|
||||
layer.input_resolution = input_resolution
|
||||
if layer.downsample is not None:
|
||||
layer.downsample.input_resolution = input_resolution
|
||||
if isinstance(layer, BasicLayer):
|
||||
for b in layer.blocks:
|
||||
b.input_resolution = input_resolution
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue