Merge pull request #78 from shinning0821/main

The update of MobileSAM and other code optimizations.

Author: shinning0821

__pycache__/cfg.cpython-37.pyc

@@ -5,6 +5,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('-net', type=str, default='sam', help='net type')
     parser.add_argument('-baseline', type=str, default='unet', help='baseline net type')
+    parser.add_argument('-encoder', type=str, default='default', help='encoder type')
     parser.add_argument('-seg_net', type=str, default='transunet', help='net type')
     parser.add_argument('-mod', type=str, default='sam_adpt', help='mod type:seg,cls,val_ad')
     parser.add_argument('-exp_name', default='msa_test_isic', type=str, help='net type')

__pycache__/function.cpython-37.pyc

@@ -124,16 +124,20 @@ def train_sam(args, net: nn.Module, optimizer, train_loader,
             # imgs = imgs.to(dtype = mask_type,device = GPUdevice)
 
             '''Train'''
-            if args.net == 'sam' or args.net == 'efficient_sam':
+            if args.mod == 'sam_adpt':
                 for n, value in net.image_encoder.named_parameters(): 
                     if "Adapter" not in n:
                         value.requires_grad = False
                     else:
                         value.requires_grad = True
+            else:
+                for n, value in net.image_encoder.named_parameters(): 
+                    value.requires_grad = True
+                    
             imge= net.image_encoder(imgs)
 
             with torch.no_grad():
-                if args.net == 'sam':
+                if args.net == 'sam' or args.net == 'mobile_sam':
                     se, de = net.prompt_encoder(
                         points=pt,
                         boxes=None,
@@ -146,7 +150,7 @@ def train_sam(args, net: nn.Module, optimizer, train_loader,
                         labels=labels_torch,
                     )
 
-            if args.net == 'sam':
+            if args.net == 'sam' or args.net == 'mobile_sam':
                 pred, _ = net.mask_decoder(
                     image_embeddings=imge,
                     image_pe=net.prompt_encoder.get_dense_pe(), 
@@ -276,8 +280,7 @@ def validation_sam(args, val_loader, epoch, net: nn.Module, clean_dir=True):
                 '''test'''
                 with torch.no_grad():
                     imge= net.image_encoder(imgs)
-
-                    if args.net == 'sam':
+                    if args.net == 'sam' or args.net == 'mobile_sam':
                         se, de = net.prompt_encoder(
                             points=pt,
                             boxes=None,
@@ -290,7 +293,7 @@ def validation_sam(args, val_loader, epoch, net: nn.Module, clean_dir=True):
                             labels=labels_torch,
                         )
 
-                    if args.net == 'sam':
+                    if args.net == 'sam' or args.net == 'mobile_sam':
                         pred, _ = net.mask_decoder(
                             image_embeddings=imge,
                             image_pe=net.prompt_encoder.get_dense_pe(), 

__pycache__/utils.cpython-37.pyc

@@ -1 +1,2 @@
-from .adapter_block import AdapterBlock
+from .tinyvit.tiny_vit import TinyViT
+from .vit import AdapterBlock, Block