Upload 7 files

text_embedding_module/OCR/ocr_recog/RecModel.py

@@ -14,17 +14,17 @@ class RecModel(nn.Module):
     def __init__(self, config):
         super().__init__()
         assert "in_channels" in config, "in_channels must in model config"
-        backbone_type = config.backbone.pop("type")
+        backbone_type = config["backbone"].pop("type")
         assert backbone_type in backbone_dict, f"backbone.type must in {backbone_dict}"
-        self.backbone = backbone_dict[backbone_type](config.in_channels, **config.backbone)
+        self.backbone = backbone_dict[backbone_type](config['in_channels'], **config['backbone'])
 
-        neck_type = config.neck.pop("type")
+        neck_type = config['neck']("type")
         assert neck_type in neck_dict, f"neck.type must in {neck_dict}"
-        self.neck = neck_dict[neck_type](self.backbone.out_channels, **config.neck)
+        self.neck = neck_dict[neck_type](self.backbone.out_channels, **config['neck'])
 
-        head_type = config.head.pop("type")
+        head_type = config['head']("type")
         assert head_type in head_dict, f"head.type must in {head_dict}"
-        self.head = head_dict[head_type](self.neck.out_channels, **config.head)
+        self.head = head_dict[head_type](self.neck.out_channels, **config['head'])
 
         self.name = f"RecModel_{backbone_type}_{neck_type}_{head_type}"