| 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 |
| |
