Upload RecModel.py
Browse files
text_embedding_module/OCR/ocr_recog/RecModel.py
CHANGED
|
@@ -18,11 +18,11 @@ class RecModel(nn.Module):
|
|
| 18 |
assert backbone_type in backbone_dict, f"backbone.type must in {backbone_dict}"
|
| 19 |
self.backbone = backbone_dict[backbone_type](config['in_channels'], **config['backbone'])
|
| 20 |
|
| 21 |
-
neck_type = config['neck']("type")
|
| 22 |
assert neck_type in neck_dict, f"neck.type must in {neck_dict}"
|
| 23 |
self.neck = neck_dict[neck_type](self.backbone.out_channels, **config['neck'])
|
| 24 |
|
| 25 |
-
head_type = config['head']("type")
|
| 26 |
assert head_type in head_dict, f"head.type must in {head_dict}"
|
| 27 |
self.head = head_dict[head_type](self.neck.out_channels, **config['head'])
|
| 28 |
|
|
|
|
| 18 |
assert backbone_type in backbone_dict, f"backbone.type must in {backbone_dict}"
|
| 19 |
self.backbone = backbone_dict[backbone_type](config['in_channels'], **config['backbone'])
|
| 20 |
|
| 21 |
+
neck_type = config['neck'].pop("type")
|
| 22 |
assert neck_type in neck_dict, f"neck.type must in {neck_dict}"
|
| 23 |
self.neck = neck_dict[neck_type](self.backbone.out_channels, **config['neck'])
|
| 24 |
|
| 25 |
+
head_type = config['head'].pop("type")
|
| 26 |
assert head_type in head_dict, f"head.type must in {head_dict}"
|
| 27 |
self.head = head_dict[head_type](self.neck.out_channels, **config['head'])
|
| 28 |
|