Spaces:
Sleeping
Sleeping
File size: 2,162 Bytes
9d665dd | 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 | """
PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
Official implementation of the paper:
"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
Licensed under a modified MIT license
"""
from .prima import PRIMA
def load_prima(checkpoint_path):
from pathlib import Path
from ..configs import get_config
model_cfg = str(Path(checkpoint_path).parent.parent / '.hydra/config.yaml')
model_cfg = get_config(model_cfg)
# Override some config values, to crop bbox correctly
if (model_cfg.MODEL.BACKBONE.TYPE == 'vit') and ('BBOX_SHAPE' not in model_cfg.MODEL):
model_cfg.defrost()
assert model_cfg.MODEL.IMAGE_SIZE == 256, f"MODEL.IMAGE_SIZE ({model_cfg.MODEL.IMAGE_SIZE}) should be 256 for ViT backbone"
model_cfg.MODEL.BBOX_SHAPE = [192, 256]
model_cfg.freeze()
if (model_cfg.MODEL.BACKBONE.TYPE == 'dinov3') and ('BBOX_SHAPE' not in model_cfg.MODEL):
model_cfg.defrost()
assert model_cfg.MODEL.IMAGE_SIZE == 256, f"MODEL.IMAGE_SIZE ({model_cfg.MODEL.IMAGE_SIZE}) should be 256 for dino backbone"
model_cfg.MODEL.BBOX_SHAPE = [256, 256]
model_cfg.freeze()
if (model_cfg.MODEL.BACKBONE.TYPE == 'dinov2') and ('BBOX_SHAPE' not in model_cfg.MODEL):
model_cfg.defrost()
assert model_cfg.MODEL.IMAGE_SIZE == 252, f"MODEL.IMAGE_SIZE ({model_cfg.MODEL.IMAGE_SIZE}) should be 252 for dino backbone"
model_cfg.MODEL.BBOX_SHAPE = [252, 252]
model_cfg.freeze()
# Update config to be compatible with demo
if ('PRETRAINED_WEIGHTS' in model_cfg.MODEL.BACKBONE):
model_cfg.defrost()
model_cfg.MODEL.BACKBONE.pop('PRETRAINED_WEIGHTS')
model_cfg.freeze()
# Offscreen training renderer is not needed for demo/inference startup and
# can fail on some local OpenGL backends.
model = PRIMA.load_from_checkpoint(
checkpoint_path,
strict=False,
cfg=model_cfg,
map_location='cpu',
init_renderer=False,
)
return model, model_cfg
|