(model_type)
| 75 | |
| 76 | |
| 77 | def load_model(model_type): |
| 78 | # https://github.com/isl-org/MiDaS/blob/master/run.py |
| 79 | # load network |
| 80 | model_path = ISL_PATHS[model_type] |
| 81 | if model_type == "dpt_large": # DPT-Large |
| 82 | model = DPTDepthModel( |
| 83 | path=model_path, |
| 84 | backbone="vitl16_384", |
| 85 | non_negative=True, |
| 86 | ) |
| 87 | net_w, net_h = 384, 384 |
| 88 | resize_mode = "minimal" |
| 89 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 90 | |
| 91 | elif model_type == "dpt_hybrid": # DPT-Hybrid |
| 92 | if not os.path.exists(model_path): |
| 93 | from basicsr.utils.download_util import load_file_from_url |
| 94 | load_file_from_url(remote_model_path, model_dir=annotator_ckpts_path) |
| 95 | |
| 96 | model = DPTDepthModel( |
| 97 | path=model_path, |
| 98 | backbone="vitb_rn50_384", |
| 99 | non_negative=True, |
| 100 | ) |
| 101 | net_w, net_h = 384, 384 |
| 102 | resize_mode = "minimal" |
| 103 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 104 | |
| 105 | elif model_type == "midas_v21": |
| 106 | model = MidasNet(model_path, non_negative=True) |
| 107 | net_w, net_h = 384, 384 |
| 108 | resize_mode = "upper_bound" |
| 109 | normalization = NormalizeImage( |
| 110 | mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] |
| 111 | ) |
| 112 | |
| 113 | elif model_type == "midas_v21_small": |
| 114 | model = MidasNet_small(model_path, features=64, backbone="efficientnet_lite3", exportable=True, |
| 115 | non_negative=True, blocks={'expand': True}) |
| 116 | net_w, net_h = 256, 256 |
| 117 | resize_mode = "upper_bound" |
| 118 | normalization = NormalizeImage( |
| 119 | mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] |
| 120 | ) |
| 121 | |
| 122 | else: |
| 123 | print(f"model_type '{model_type}' not implemented, use: --model_type large") |
| 124 | assert False |
| 125 | |
| 126 | transform = Compose( |
| 127 | [ |
| 128 | Resize( |
| 129 | net_w, |
| 130 | net_h, |
| 131 | resize_target=None, |
| 132 | keep_aspect_ratio=True, |
| 133 | ensure_multiple_of=32, |
| 134 | resize_method=resize_mode, |
no test coverage detected