Load the specified network. Args: device (device): the torch device used model_path (str): path to saved model model_type (str): the type of the model to be loaded optimize (bool): optimize the model to half-integer on CUDA? height (int): inference encode
(device, model_path, model_type="dpt_large_384", optimize=True, height=None, square=False)
| 27 | |
| 28 | |
| 29 | def load_model(device, model_path, model_type="dpt_large_384", optimize=True, height=None, square=False): |
| 30 | """Load the specified network. |
| 31 | |
| 32 | Args: |
| 33 | device (device): the torch device used |
| 34 | model_path (str): path to saved model |
| 35 | model_type (str): the type of the model to be loaded |
| 36 | optimize (bool): optimize the model to half-integer on CUDA? |
| 37 | height (int): inference encoder image height |
| 38 | square (bool): resize to a square resolution? |
| 39 | |
| 40 | Returns: |
| 41 | The loaded network, the transform which prepares images as input to the network and the dimensions of the |
| 42 | network input |
| 43 | """ |
| 44 | if "openvino" in model_type: |
| 45 | from openvino.runtime import Core |
| 46 | |
| 47 | keep_aspect_ratio = not square |
| 48 | |
| 49 | if model_type == "dpt_beit_large_512": |
| 50 | model = DPTDepthModel( |
| 51 | path=model_path, |
| 52 | backbone="beitl16_512", |
| 53 | non_negative=True, |
| 54 | ) |
| 55 | net_w, net_h = 512, 512 |
| 56 | resize_mode = "minimal" |
| 57 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 58 | |
| 59 | elif model_type == "dpt_beit_large_384": |
| 60 | model = DPTDepthModel( |
| 61 | path=model_path, |
| 62 | backbone="beitl16_384", |
| 63 | non_negative=True, |
| 64 | ) |
| 65 | net_w, net_h = 384, 384 |
| 66 | resize_mode = "minimal" |
| 67 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 68 | |
| 69 | elif model_type == "dpt_beit_base_384": |
| 70 | model = DPTDepthModel( |
| 71 | path=model_path, |
| 72 | backbone="beitb16_384", |
| 73 | non_negative=True, |
| 74 | ) |
| 75 | net_w, net_h = 384, 384 |
| 76 | resize_mode = "minimal" |
| 77 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 78 | |
| 79 | elif model_type == "dpt_swin2_large_384": |
| 80 | model = DPTDepthModel( |
| 81 | path=model_path, |
| 82 | backbone="swin2l24_384", |
| 83 | non_negative=True, |
| 84 | ) |
| 85 | net_w, net_h = 384, 384 |
| 86 | keep_aspect_ratio = False |
nothing calls this directly
no test coverage detected