load from HF

hugging_face/app.py

@@ -416,9 +416,14 @@ sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[args.sam_model_type]
 model = MaskGenerator(sam_checkpoint, args)
 
 # initialize matanyone
-pretrain_model_url = "https://github.com/pq-yang/MatAnyone/releases/download/v1.0.0"
-ckpt_path = load_file_from_url(os.path.join(pretrain_model_url, 'matanyone.pth'), checkpoint_folder)
-matanyone_model = get_matanyone_model(ckpt_path, args.device)
+# load from ckpt
+# pretrain_model_url = "https://github.com/pq-yang/MatAnyone/releases/download/v1.0.0"
+# ckpt_path = load_file_from_url(os.path.join(pretrain_model_url, 'matanyone.pth'), checkpoint_folder)
+# matanyone_model = get_matanyone_model(ckpt_path, args.device)
+# load from Hugging Face
+from matanyone.model.matanyone import MatAnyone
+matanyone_model = MatAnyone.from_pretrained("PeiqingYang/MatAnyone")
+
 matanyone_model = matanyone_model.to(args.device).eval()
 matanyone_processor = InferenceCore(matanyone_model, cfg=matanyone_model.cfg)
 

matanyone/inference/inference_core.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Iterable, Dict
+from typing import List, Optional, Iterable
 import logging
 from omegaconf import DictConfig
 
@@ -302,7 +302,6 @@ class InferenceCore:
 
             mask, _ = pad_divide_by(mask, 16)
             if need_segment:
-                print("HERE!!!!!!!!!!!")
                 # merge predicted mask with the incomplete input mask
                 pred_prob_no_bg = pred_prob_with_bg[1:]
                 # use the mutual exclusivity of segmentation

matanyone/inference/memory_manager.py

@@ -2,12 +2,11 @@ import logging
 from omegaconf import DictConfig
 from typing import List, Dict
 import torch
-import cv2
 
 from matanyone.inference.object_manager import ObjectManager
 from matanyone.inference.kv_memory_store import KeyValueMemoryStore
 from matanyone.model.matanyone import MatAnyone
-from matanyone.model.utils.memory_utils import *
+from matanyone.model.utils.memory_utils import get_similarity, do_softmax
 
 log = logging.getLogger()
 
@@ -128,8 +127,6 @@ class MemoryManager:
         bs = pix_feat.shape[0]
         assert last_mask.shape[0] == bs
 
-        uncert_mask = uncert_output["mask"] if uncert_output is not None else None
-
         """
         Compute affinity and perform readout
         """
@@ -374,7 +371,6 @@ class MemoryManager:
             self.engaged = False
 
     def compress_features(self, bucket_id: int) -> None:
-        HW = self.HW
 
         # perform memory consolidation
         prototype_key, prototype_value, prototype_shrinkage = self.consolidation(

matanyone/model/big_modules.py

@@ -8,14 +8,15 @@ g - usually denotes features that are not shared between objects
 The trailing number of a variable usually denotes the stride
 """
 
+from typing import Iterable
 from omegaconf import DictConfig
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from matanyone.model.group_modules import *
+from matanyone.model.group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d
 from matanyone.model.utils import resnet
-from matanyone.model.modules import *
+from matanyone.model.modules import SensoryDeepUpdater, SensoryUpdater_fullscale, DecoderFeatureProcessor, MaskUpsampleBlock
 
 class UncertPred(nn.Module):
     def __init__(self, model_cfg: DictConfig):
@@ -51,11 +52,14 @@ class PixelEncoder(nn.Module):
         super().__init__()
 
         self.is_resnet = 'resnet' in model_cfg.pixel_encoder.type
+        # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+        # else default to True
+        is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
         if self.is_resnet:
             if model_cfg.pixel_encoder.type == 'resnet18':
-                network = resnet.resnet18(pretrained=True)
+                network = resnet.resnet18(pretrained=is_pretrained_resnet)
             elif model_cfg.pixel_encoder.type == 'resnet50':
-                network = resnet.resnet50(pretrained=True)
+                network = resnet.resnet50(pretrained=is_pretrained_resnet)
             else:
                 raise NotImplementedError
             self.conv1 = network.conv1
@@ -127,10 +131,13 @@ class MaskEncoder(nn.Module):
         self.single_object = single_object
         extra_dim = 1 if single_object else 2
 
+        # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+        # else default to True
+        is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
         if model_cfg.mask_encoder.type == 'resnet18':
-            network = resnet.resnet18(pretrained=True, extra_dim=extra_dim)
+            network = resnet.resnet18(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
         elif model_cfg.mask_encoder.type == 'resnet50':
-            network = resnet.resnet50(pretrained=True, extra_dim=extra_dim)
+            network = resnet.resnet50(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
         else:
             raise NotImplementedError
         self.conv1 = network.conv1

matanyone/model/matanyone.py

@@ -1,21 +1,31 @@
-from typing import List, Dict
+from typing import List, Dict, Iterable
 import logging
 from omegaconf import DictConfig
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from huggingface_hub import PyTorchModelHubMixin
 
-from matanyone.model.modules import *
-from matanyone.model.big_modules import *
+from matanyone.model.big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder
 from matanyone.model.aux_modules import AuxComputer
-from matanyone.model.utils.memory_utils import *
+from matanyone.model.utils.memory_utils import get_affinity, readout
 from matanyone.model.transformer.object_transformer import QueryTransformer
 from matanyone.model.transformer.object_summarizer import ObjectSummarizer
 from matanyone.utils.tensor_utils import aggregate
 
 log = logging.getLogger()
-
-
-class MatAnyone(nn.Module):
+class MatAnyone(nn.Module,
+                PyTorchModelHubMixin,
+                library_name="matanyone",
+                repo_url="https://github.com/pq-yang/MatAnyone",
+                coders={
+                    DictConfig: (
+                        lambda x: OmegaConf.to_container(x),
+                        lambda data: OmegaConf.create(data),
+                    )
+                },
+        ):
 
     def __init__(self, cfg: DictConfig, *, single_object=False):
         super().__init__()
@@ -304,7 +314,7 @@ class MatAnyone(nn.Module):
             finetune a trained model with single object datasets.
             """
             if src_dict['mask_encoder.conv1.weight'].shape[1] == 5:
-                log.warning(f'Converting mask_encoder.conv1.weight from multiple objects to single object.'
+                log.warning('Converting mask_encoder.conv1.weight from multiple objects to single object.'
                             'This is not supposed to happen in standard training.')
                 src_dict['mask_encoder.conv1.weight'] = src_dict['mask_encoder.conv1.weight'][:, :-1]
                 src_dict['pixel_fuser.sensory_compress.weight'] = src_dict['pixel_fuser.sensory_compress.weight'][:, :-1]

matanyone/model/modules.py

@@ -1,8 +1,9 @@
 from typing import List, Iterable
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
-from matanyone.model.group_modules import *
+from matanyone.model.group_modules import MainToGroupDistributor, GroupResBlock, upsample_groups, GConv2d, downsample_groups
 
 
 class UpsampleBlock(nn.Module):
@@ -145,26 +146,4 @@ class ResBlock(nn.Module):
 
         g = self.downsample(g)
 
-        return out_g + g
-
-    def __init__(self, in_dim, reduction_dim, bins):
-        super(PPM, self).__init__()
-        self.features = []
-        for bin in bins:
-            self.features.append(nn.Sequential(
-                nn.AdaptiveAvgPool2d(bin),
-                nn.Conv2d(in_dim, reduction_dim, kernel_size=1, bias=False),
-                nn.PReLU()
-            ))
-        self.features = nn.ModuleList(self.features)
-        self.fuse = nn.Sequential(
-                nn.Conv2d(in_dim+reduction_dim*4, in_dim, kernel_size=3, padding=1, bias=False),
-                nn.PReLU())
-
-    def forward(self, x):
-        x_size = x.size()
-        out = [x]
-        for f in self.features:
-            out.append(F.interpolate(f(x), x_size[2:], mode='bilinear', align_corners=True))
-        out_feat = self.fuse(torch.cat(out, 1))
-        return out_feat 
+        return out_g + g

matanyone/model/transformer/object_summarizer.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Optional
+from typing import Optional
 from omegaconf import DictConfig
 
 import torch

matanyone/model/transformer/object_transformer.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 from matanyone.model.group_modules import GConv2d
 from matanyone.utils.tensor_utils import aggregate
 from matanyone.model.transformer.positional_encoding import PositionalEncoding
-from matanyone.model.transformer.transformer_layers import *
+from matanyone.model.transformer.transformer_layers import CrossAttention, SelfAttention, FFN, PixelFFN
 
 
 class QueryTransformerBlock(nn.Module):

matanyone/model/utils/resnet.py

@@ -15,7 +15,7 @@ def load_weights_add_extra_dim(target, source_state, extra_dim=1):
     new_dict = OrderedDict()
 
     for k1, v1 in target.state_dict().items():
-        if not 'num_batches_tracked' in k1:
+        if 'num_batches_tracked' not in k1:
             if k1 in source_state:
                 tar_v = source_state[k1]
 

matanyone/utils/get_default_model.py

@@ -6,9 +6,8 @@ from hydra import compose, initialize
 
 import torch
 from matanyone.model.matanyone import MatAnyone
-from matanyone.inference.utils.args_utils import get_dataset_cfg
 
-def get_matanyone_model(ckpt_path, device) -> MatAnyone:
+def get_matanyone_model(ckpt_path, device=None) -> MatAnyone:
     initialize(version_base='1.3.2', config_path="../config", job_name="eval_our_config")
     cfg = compose(config_name="eval_matanyone_config")
     
@@ -16,8 +15,13 @@ def get_matanyone_model(ckpt_path, device) -> MatAnyone:
         cfg['weights'] = ckpt_path
 
     # Load the network weights
-    matanyone = MatAnyone(cfg, single_object=True).to(device).eval()
-    model_weights = torch.load(cfg.weights, map_location=device)
+    if device is not None:
+        matanyone = MatAnyone(cfg, single_object=True).to(device).eval()
+        model_weights = torch.load(cfg.weights, map_location=device)
+    else:  # if device is not specified, `.cuda()` by default
+        matanyone = MatAnyone(cfg, single_object=True).cuda().eval()
+        model_weights = torch.load(cfg.weights)
+        
     matanyone.load_weights(model_weights)
 
     return matanyone