Spaces:
Runtime error
Runtime error
File size: 2,230 Bytes
01bb3bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
import torch.nn as nn
from IndicPhotoOCR.detection.textbpn.network.backbone.resnet import deformable_resnet18,deformable_resnet50
import torch.utils.model_zoo as model_zoo
from IndicPhotoOCR.detection.textbpn.cfglib.config import config as cfg
model_urls = {
'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
class ResNet_DCN(nn.Module):
def __init__(self, name="deformable_resnet18", pretrain=False):
super().__init__()
if name == "deformable_resnet18":
self.base_net = deformable_resnet18(pretrained=False)
if pretrain:
print("load the {} weight from ./cache".format(name))
self.base_net.load_state_dict(
model_zoo.load_url(model_urls["resnet18"], model_dir="./cache",
map_location=torch.device(cfg.device)), strict=False)
elif name == "deformable_resnet50":
self.base_net = deformable_resnet50(pretrained=False)
if pretrain:
print("load the {} weight from ./cache".format(name))
self.base_net.load_state_dict(
model_zoo.load_url(model_urls["resnet50"], model_dir="./cache",
map_location=torch.device(cfg.device)), strict=False)
else:
print(" base model is not support !")
# print(base_net)
self.up2 = nn.ConvTranspose2d(64, 64, kernel_size=4, stride=2, padding=1)
def forward(self, x):
C1, C2, C3, C4, C5 = self.base_net(x)
# up2 --> 1/2
C1 = self.up2(C1)
return C1, C2, C3, C4, C5
if __name__ == '__main__':
import torch
input = torch.randn((4, 3, 512, 512))
net = ResNet_DCN()
C1, C2, C3, C4, C5 = net(input)
print(C1.size())
print(C2.size())
print(C3.size())
print(C4.size())
print(C5.size())
|