30 lines
726 B
Python
30 lines
726 B
Python
import torch
|
|
from torch import nn
|
|
from torchvision import models
|
|
|
|
|
|
class KPDetector(nn.Module):
|
|
"""
|
|
Predict K*5 keypoints.
|
|
"""
|
|
|
|
def __init__(self, num_tps, **kwargs):
|
|
super(KPDetector, self).__init__()
|
|
self.num_tps = num_tps
|
|
|
|
self.fg_encoder = models.resnet18(pretrained=False)
|
|
num_features = self.fg_encoder.fc.in_features
|
|
self.fg_encoder.fc = nn.Linear(num_features, num_tps * 5 * 2)
|
|
|
|
def forward(self, image):
|
|
fg_kp = self.fg_encoder(image)
|
|
(
|
|
bs,
|
|
_,
|
|
) = fg_kp.shape
|
|
fg_kp = torch.sigmoid(fg_kp)
|
|
fg_kp = fg_kp * 2 - 1
|
|
out = {"fg_kp": fg_kp.view(bs, self.num_tps * 5, -1)}
|
|
|
|
return out
|