load a pretrained pointersect model from bolt. Args: filename: filename of the model checkpoint device: device to load the model Returns: pointersect model
(
filename: str,
device: torch.device = torch.device('cpu'),
)
| 40 | |
| 41 | |
| 42 | def load_pointersect( |
| 43 | filename: str, |
| 44 | device: torch.device = torch.device('cpu'), |
| 45 | ) -> T.Tuple[SimplePointersect, T.Dict[str, T.Any]]: |
| 46 | """ |
| 47 | load a pretrained pointersect model from bolt. |
| 48 | |
| 49 | Args: |
| 50 | filename: |
| 51 | filename of the model checkpoint |
| 52 | device: |
| 53 | device to load the model |
| 54 | |
| 55 | Returns: |
| 56 | pointersect model |
| 57 | """ |
| 58 | |
| 59 | model_dict, checkpoint = cds_model_utils.load_model( |
| 60 | filename=filename, |
| 61 | model_names='model', |
| 62 | model_classes=SimplePointersect, |
| 63 | model_params_names='model_info', |
| 64 | device=device, |
| 65 | ) |
| 66 | # assume load only one model |
| 67 | assert len(model_dict) == 1 |
| 68 | model_name = 'model' |
| 69 | if isinstance(model_name, (list, tuple)): |
| 70 | model_name = model_name[0] |
| 71 | model = model_dict[model_name] |
| 72 | |
| 73 | model.eval() |
| 74 | model.to(device=device) |
| 75 | |
| 76 | model_info = dict( |
| 77 | filename=filename, |
| 78 | ) |
| 79 | return model, model_info |
| 80 | |
| 81 | |
| 82 | def save_imgs( |