Refactor EfficientSAM initialization and clean up imports

Author: jakmro

examples/models/efficient_sam/efficient_sam_core/__init__.py

@@ -4,17 +4,16 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the same directory.
 
-import math
-from typing import Any, List, Tuple, Type
+from typing import List, Tuple
 
 import torch
 import torch.nn.functional as F
 
-from torch import nn, Tensor
+from torch import nn
 
 from .efficient_sam_decoder import MaskDecoder, PromptEncoder
 from .efficient_sam_encoder import ImageEncoderViT
-from .two_way_transformer import TwoWayAttentionBlock, TwoWayTransformer
+from .two_way_transformer import TwoWayTransformer
 
 
 class EfficientSam(nn.Module):
@@ -27,8 +26,8 @@ def __init__(
         prompt_encoder: PromptEncoder,
         decoder_max_num_input_points: int,
         mask_decoder: MaskDecoder,
-        pixel_mean: List[float] = [0.485, 0.456, 0.406],
-        pixel_std: List[float] = [0.229, 0.224, 0.225],
+        pixel_mean: List[float] = None,
+        pixel_std: List[float] = None,
     ) -> None:
         """
         SAM predicts object masks from an image and input prompts.
@@ -47,6 +46,10 @@ def __init__(
         self.prompt_encoder = prompt_encoder
         self.decoder_max_num_input_points = decoder_max_num_input_points
         self.mask_decoder = mask_decoder
+        if pixel_mean is None:
+            pixel_mean = [0.485, 0.456, 0.406]
+        if pixel_std is None:
+            pixel_std = [0.229, 0.224, 0.225]
         self.register_buffer(
             "pixel_mean", torch.Tensor(pixel_mean).view(1, 3, 1, 1), False
         )

examples/models/efficient_sam/efficient_sam_core/efficient_sam.py

@@ -9,7 +9,6 @@
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from .mlp import MLPBlock
 

examples/models/efficient_sam/efficient_sam_core/efficient_sam_decoder.py

@@ -5,7 +5,7 @@
 # LICENSE file in the same directory.
 
 import math
-from typing import List, Optional, Tuple, Type
+from typing import List, Type
 
 import torch
 import torch.nn as nn
@@ -216,7 +216,7 @@ def __init__(
         num_positions = num_patches + 1
         self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, patch_embed_dim))
         self.blocks = nn.ModuleList()
-        for i in range(depth):
+        for _ in range(depth):
             vit_block = Block(patch_embed_dim, num_heads, mlp_ratio, True)
             self.blocks.append(vit_block)
         self.neck = nn.Sequential(

examples/models/efficient_sam/efficient_sam_core/efficient_sam_encoder.py

@@ -1,7 +1,9 @@
 import math
 from typing import Tuple, Type
+
 import torch
 from torch import nn, Tensor
+
 from .mlp import MLPBlock
 
 
@@ -84,7 +86,7 @@ def forward(
         keys = image_embedding
 
         # Apply transformer blocks and final layernorm
-        for idx, layer in enumerate(self.layers):
+        for _, layer in enumerate(self.layers):
             queries, keys = layer(
                 queries=queries,
                 keys=keys,

examples/models/efficient_sam/efficient_sam_core/two_way_transformer.py

@@ -8,10 +8,10 @@
 
 import torch
 
-from .efficient_sam_core.build_efficient_sam import build_efficient_sam_vitt
-
 from ..model_base import EagerModelBase
 
+from .efficient_sam_core.build_efficient_sam import build_efficient_sam_vitt
+
 
 class EfficientSAM(EagerModelBase):
     def __init__(self):