Spaces:
Running
Running
File size: 1,202 Bytes
8320ccc 63f3cf2 8320ccc 63f3cf2 8320ccc |
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 |
import sys
from pathlib import Path
import torchvision.transforms as tvf
from .. import logger
from ..utils.base_model import BaseModel
tp_path = Path(__file__).parent / "../../third_party"
sys.path.append(str(tp_path))
from pram.nets.sfd2 import load_sfd2
class SFD2(BaseModel):
default_conf = {
"max_keypoints": 4096,
"model_name": "sfd2_20230511_210205_resnet4x.79.pth",
"conf_th": 0.001,
}
required_inputs = ["image"]
def _init(self, conf):
self.conf = {**self.default_conf, **conf}
self.norm_rgb = tvf.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
)
model_path = tp_path / "pram" / "weights" / self.conf["model_name"]
self.net = load_sfd2(weight_path=model_path).eval()
logger.info("Load SFD2 model done.")
def _forward(self, data):
pred = self.net.extract_local_global(
data={"image": self.norm_rgb(data["image"])}, config=self.conf
)
out = {
"keypoints": pred["keypoints"][0][None],
"scores": pred["scores"][0][None],
"descriptors": pred["descriptors"][0][None],
}
return out
|