Spaces:
Sleeping
Sleeping
| from res.impl.HRNetV2 import HRNetV2 | |
| import torch | |
| class Config: | |
| pass | |
| class HRNetV2Wrapper: | |
| def __init__(self): | |
| config = Config() | |
| config.data_len = 5000 | |
| config.kernel_size = 5 | |
| config.dilation = 1 | |
| config.num_stages = 3 | |
| config.num_blocks = 6 | |
| config.num_modules = [1, 1, 1, 4, 3] | |
| config.use_bottleneck = [1, 0, 0, 0, 0] | |
| config.stage1_channels = 128 | |
| config.num_channels_init = 48 | |
| config.interpolate_mode = "linear" | |
| config.output_size = 3 | |
| self.model = HRNetV2(config) | |
| weights = torch.load("./res/models/hrnetv2/weights.pth") | |
| self.model.load_state_dict(weights) | |
| self.model = self.model.to("cpu").eval() | |