Upload hubconf.py

hubconf.py

@@ -0,0 +1,37 @@
+dependencies = ["efficientnet_pytorch", "pretrainedmodels",
+                "timm", "torch", "torchvision"]
+import torch
+from utils.utils import Params
+from backbone import HybridNetsBackbone
+from pathlib import Path
+import os
+
+
+def hybridnets(pretrained=True, compound_coef=3, device='cpu'):
+    """Creates a HybridNets model
+
+    Arguments:
+        pretrained (bool): load pretrained weights into the model
+        compound_coef (int): compound coefficient of efficientnet backbone
+        device (str): 'cuda:0' or 'cpu'
+
+    Returns:
+        HybridNets model
+    """
+    params = Params(os.path.join(Path(__file__).resolve().parent, "projects/bdd100k.yml"))
+    model = HybridNetsBackbone(num_classes=len(params.obj_list), compound_coef=compound_coef,
+                               ratios=eval(params.anchors_ratios), scales=eval(params.anchors_scales),
+                               seg_classes=len(params.seg_list))
+    if pretrained and compound_coef == 3:
+        weight_url = 'https://github.com/datvuthanh/HybridNets/releases/download/v1.0/hybridnets.pth'
+        model.load_state_dict(torch.hub.load_state_dict_from_url(weight_url, map_location=device))
+    model = model.to(device)
+    return model
+
+
+if __name__ == "__main__":
+    model = hybridnets(device='cpu')
+    img = torch.rand(1, 3, 384, 640)
+
+    result = model(img)
+    print(result)