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:
Laughing 2024-08-30 20:59:08 +08:00 committed by GitHub
parent 5d66140ce1
commit 7053169fd0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 70 additions and 7 deletions

View file

@ -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