Skip to content

network

xvr.model.network

PoseRegressor 

PoseRegressor(
    model_name,
    parameterization,
    convention=None,
    pretrained=False,
    height=256,
    unit_conversion_factor=1000.0,
    **kwargs
)

A PoseRegressor is comprised of a pretrained backbone model that extracts features from an input X-ray and two linear layers that decode these features into rotational and translational camera pose parameters, respectively.

Source code in src/xvr/model/network.py 
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
def __init__(
    self,
    model_name,
    parameterization,
    convention=None,
    pretrained=False,
    height=256,
    unit_conversion_factor=1000.0,
    **kwargs,
):
    super().__init__()

    self.parameterization = parameterization
    self.convention = convention
    n_angular_components = N_ANGULAR_COMPONENTS[parameterization]

    # Get the size of the output from the backbone
    self.backbone = timm.create_model(
        model_name,
        pretrained,
        num_classes=0,
        in_chans=1,
        **kwargs,
    )
    output = self.backbone(torch.randn(1, 1, height, height)).shape[-1]
    self.xyz_regression = torch.nn.Linear(output, 3)
    self.rot_regression = torch.nn.Linear(output, n_angular_components)

    # E.g., if 1000.0, converts output from meters to millimeters
    self.unit_conversion_factor = unit_conversion_factor

load_model 

load_model(ckptpath, meta=False)

Load a pretrained pose regression model

Source code in src/xvr/model/network.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def load_model(ckptpath, meta=False):
    """Load a pretrained pose regression model"""
    ckpt = torch.load(ckptpath, weights_only=False)
    config = ckpt["config"]

    model_state_dict = ckpt["model_state_dict"]
    model = PoseRegressor(
        model_name=config["model_name"],
        parameterization=config["parameterization"],
        convention=config["convention"],
        norm_layer=config["norm_layer"],
        height=config["height"],
        unit_conversion_factor=config.get("unit_conversion_factor", 1.0),
    ).cuda()
    model.load_state_dict(model_state_dict)
    model.eval()

    if meta:
        return model, config, ckpt["date"]
    else:
        return model, config