Spaces:
Running
Running
File size: 1,979 Bytes
95f8bbc aa34300 95f8bbc fb96f4f 95f8bbc aa34300 95f8bbc fb96f4f 95f8bbc |
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 61 62 63 64 65 66 67 68 |
import sys
import torch
import torch._utils
import torch.nn as nn
import torch.utils.data
import torch.utils.data.distributed
from SPPE.src.models.FastPose import createModel
from SPPE.src.utils.img import flip, shuffleLR
try:
torch._utils._rebuild_tensor_v2
except AttributeError:
def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
tensor.requires_grad = requires_grad
tensor._backward_hooks = backward_hooks
return tensor
torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2
class InferenNet(nn.Module):
def __init__(self, kernel_size, dataset):
super(InferenNet, self).__init__()
model = createModel()
print('Loading pose model from {}'.format('joints_detectors/Alphapose/models/sppe/duc_se.pth'))
sys.stdout.flush()
model.load_state_dict(torch.load('joints_detectors/Alphapose/models/sppe/duc_se.pth', map_location=torch.device('cpu')))
model.eval()
self.pyranet = model
self.dataset = dataset
def forward(self, x):
out = self.pyranet(x)
out = out.narrow(1, 0, 17)
flip_out = self.pyranet(flip(x))
flip_out = flip_out.narrow(1, 0, 17)
flip_out = flip(shuffleLR(
flip_out, self.dataset))
out = (flip_out + out) / 2
return out
class InferenNet_fast(nn.Module):
def __init__(self, kernel_size, dataset):
super(InferenNet_fast, self).__init__()
model = createModel()
print('Loading pose model from {}'.format('models/sppe/duc_se.pth'))
model.load_state_dict(torch.load('models/sppe/duc_se.pth', map_location=torch.device('cpu')))
model.eval()
self.pyranet = model
self.dataset = dataset
def forward(self, x):
out = self.pyranet(x)
out = out.narrow(1, 0, 17)
return out
|