mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
@@ -25,8 +25,8 @@ class Sam(nn.Module):
|
||||
image_encoder: Union[ImageEncoderViT, TinyViT],
|
||||
prompt_encoder: PromptEncoder,
|
||||
mask_decoder: MaskDecoder,
|
||||
pixel_mean: List[float] | None = None,
|
||||
pixel_std: List[float] | None = None,
|
||||
pixel_mean: List[float] = None,
|
||||
pixel_std: List[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
SAM predicts object masks from an image and input prompts.
|
||||
|
||||
@@ -79,7 +79,7 @@ class TwoWayTransformer(nn.Module):
|
||||
torch.Tensor: the processed image_embedding
|
||||
"""
|
||||
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
|
||||
_bs, _c, _h, _w = image_embedding.shape
|
||||
bs, c, h, w = image_embedding.shape
|
||||
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
|
||||
image_pe = image_pe.flatten(2).permute(0, 2, 1)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user