| | import torch |
| | from location_encoder import get_neural_network, get_positional_encoding, LocationEncoder |
| |
|
| |
|
| | def get_satclip_loc_encoder(ckpt_path, device): |
| | ckpt = torch.load(ckpt_path,map_location=device) |
| | hp = ckpt['hyper_parameters'] |
| |
|
| | posenc = get_positional_encoding( |
| | hp['le_type'], |
| | hp['legendre_polys'], |
| | hp['harmonics_calculation'], |
| | hp['min_radius'], |
| | hp['max_radius'], |
| | hp['frequency_num'] |
| | ) |
| | |
| | nnet = get_neural_network( |
| | hp['pe_type'], |
| | posenc.embedding_dim, |
| | hp['embed_dim'], |
| | hp['capacity'], |
| | hp['num_hidden_layers'] |
| | ) |
| |
|
| | |
| | state_dict = ckpt['state_dict'] |
| | state_dict = {k[k.index('nnet'):]:state_dict[k] |
| | for k in state_dict.keys() if 'nnet' in k} |
| | |
| | loc_encoder = LocationEncoder(posenc, nnet).double() |
| | loc_encoder.load_state_dict(state_dict) |
| | loc_encoder.eval() |
| |
|
| | return loc_encoder |
| | |
