Spaces:
Running
Running
File size: 1,652 Bytes
499e141 |
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 |
import sys
from pathlib import Path
import torch
from .. import logger
from ..utils.base_model import BaseModel
example_path = Path(__file__).parent / "../../third_party/example"
sys.path.append(str(example_path))
# import some modules here
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Example(BaseModel):
# change to your default configs
default_conf = {
"name": "example",
"keypoint_threshold": 0.1,
"max_keypoints": 2000,
"model_name": "model.pth",
}
required_inputs = ["image"]
def _init(self, conf):
# set checkpoints paths if needed
model_path = example_path / "checkpoints" / f'{conf["model_name"]}'
if not model_path.exists():
logger.info(f"No model found at {model_path}")
# init model
self.net = callable
# self.net = ExampleNet(is_test=True)
state_dict = torch.load(model_path, map_location="cpu")
self.net.load_state_dict(state_dict["model_state"])
logger.info("Load example model done.")
def _forward(self, data):
# data: dict, keys: 'image'
# image color mode: RGB
# image value range in [0, 1]
image = data["image"]
# B: batch size, N: number of keypoints
# keypoints shape: B x N x 2, type: torch tensor
# scores shape: B x N, type: torch tensor
# descriptors shape: B x 128 x N, type: torch tensor
keypoints, scores, descriptors = self.net(image)
return {
"keypoints": keypoints,
"scores": scores,
"descriptors": descriptors,
}
|