Initialize a segmentor from config file. Args: config (str or :obj:`mmcv.Config`): Config file path or the config object. checkpoint (str, optional): Checkpoint path. If left as None, the model will not load any weights. device (str, optional) CPU
(config, checkpoint=None, device="cuda:0")
| 10 | |
| 11 | |
| 12 | def init_segmentor(config, checkpoint=None, device="cuda:0"): |
| 13 | """Initialize a segmentor from config file. |
| 14 | |
| 15 | Args: |
| 16 | config (str or :obj:`mmcv.Config`): Config file path or the config |
| 17 | object. |
| 18 | checkpoint (str, optional): Checkpoint path. If left as None, the model |
| 19 | will not load any weights. |
| 20 | device (str, optional) CPU/CUDA device option. Default 'cuda:0'. |
| 21 | Use 'cpu' for loading model on CPU. |
| 22 | Returns: |
| 23 | nn.Module: The constructed segmentor. |
| 24 | """ |
| 25 | if isinstance(config, str): |
| 26 | config = mmcv.Config.fromfile(config) |
| 27 | elif not isinstance(config, mmcv.Config): |
| 28 | raise TypeError("config must be a filename or Config object, but got {}".format(type(config))) |
| 29 | config.model.pretrained = None |
| 30 | config.model.train_cfg = None |
| 31 | model = build_segmentor(config.model, test_cfg=config.get("test_cfg")) |
| 32 | if checkpoint is not None: |
| 33 | checkpoint = load_checkpoint(model, checkpoint, map_location="cpu") |
| 34 | model.CLASSES = checkpoint["meta"]["CLASSES"] |
| 35 | model.PALETTE = checkpoint["meta"]["PALETTE"] |
| 36 | model.cfg = config # save the config in the model for convenience |
| 37 | model.to(device) |
| 38 | model.eval() |
| 39 | return model |
| 40 | |
| 41 | |
| 42 | class LoadImage: |
nothing calls this directly
no test coverage detected