the module for OCR. args: config (dict): the super parameters for module.
(self, config)
| 26 | |
| 27 | class BaseModel(nn.Layer): |
| 28 | def __init__(self, config): |
| 29 | """ |
| 30 | the module for OCR. |
| 31 | args: |
| 32 | config (dict): the super parameters for module. |
| 33 | """ |
| 34 | super(BaseModel, self).__init__() |
| 35 | in_channels = config.get("in_channels", 3) |
| 36 | model_type = config["model_type"] |
| 37 | # build transform, |
| 38 | # for rec, transform can be TPS,None |
| 39 | # for det and cls, transform should to be None, |
| 40 | # if you make model differently, you can use transform in det and cls |
| 41 | if "Transform" not in config or config["Transform"] is None: |
| 42 | self.use_transform = False |
| 43 | else: |
| 44 | self.use_transform = True |
| 45 | config["Transform"]["in_channels"] = in_channels |
| 46 | self.transform = build_transform(config["Transform"]) |
| 47 | in_channels = self.transform.out_channels |
| 48 | |
| 49 | # build backbone, backbone is need for del, rec and cls |
| 50 | if "Backbone" not in config or config["Backbone"] is None: |
| 51 | self.use_backbone = False |
| 52 | else: |
| 53 | self.use_backbone = True |
| 54 | config["Backbone"]["in_channels"] = in_channels |
| 55 | self.backbone = build_backbone(config["Backbone"], model_type) |
| 56 | in_channels = self.backbone.out_channels |
| 57 | |
| 58 | # build neck |
| 59 | # for rec, neck can be cnn,rnn or reshape(None) |
| 60 | # for det, neck can be FPN, BIFPN and so on. |
| 61 | # for cls, neck should be none |
| 62 | if "Neck" not in config or config["Neck"] is None: |
| 63 | self.use_neck = False |
| 64 | else: |
| 65 | self.use_neck = True |
| 66 | config["Neck"]["in_channels"] = in_channels |
| 67 | self.neck = build_neck(config["Neck"]) |
| 68 | in_channels = self.neck.out_channels |
| 69 | |
| 70 | # # build head, head is need for det, rec and cls |
| 71 | if "Head" not in config or config["Head"] is None: |
| 72 | self.use_head = False |
| 73 | else: |
| 74 | self.use_head = True |
| 75 | config["Head"]["in_channels"] = in_channels |
| 76 | self.head = build_head(config["Head"]) |
| 77 | |
| 78 | self.return_all_feats = config.get("return_all_feats", False) |
| 79 | |
| 80 | def forward(self, x, data=None): |
| 81 | y = dict() |
nothing calls this directly
no test coverage detected