(model_type)
| 26 | |
| 27 | |
| 28 | def load_midas_transform(model_type): |
| 29 | # https://github.com/isl-org/MiDaS/blob/master/run.py |
| 30 | # load transform only |
| 31 | if model_type == "dpt_large": # DPT-Large |
| 32 | net_w, net_h = 384, 384 |
| 33 | resize_mode = "minimal" |
| 34 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 35 | |
| 36 | elif model_type == "dpt_hybrid": # DPT-Hybrid |
| 37 | net_w, net_h = 384, 384 |
| 38 | resize_mode = "minimal" |
| 39 | normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) |
| 40 | |
| 41 | elif model_type == "midas_v21": |
| 42 | net_w, net_h = 384, 384 |
| 43 | resize_mode = "upper_bound" |
| 44 | normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 45 | |
| 46 | elif model_type == "midas_v21_small": |
| 47 | net_w, net_h = 256, 256 |
| 48 | resize_mode = "upper_bound" |
| 49 | normalization = NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 50 | |
| 51 | else: |
| 52 | assert False, f"model_type '{model_type}' not implemented, use: --model_type large" |
| 53 | |
| 54 | transform = Compose( |
| 55 | [ |
| 56 | Resize( |
| 57 | net_w, |
| 58 | net_h, |
| 59 | resize_target=None, |
| 60 | keep_aspect_ratio=True, |
| 61 | ensure_multiple_of=32, |
| 62 | resize_method=resize_mode, |
| 63 | image_interpolation_method=cv2.INTER_CUBIC, |
| 64 | ), |
| 65 | normalization, |
| 66 | PrepareForNet(), |
| 67 | ] |
| 68 | ) |
| 69 | |
| 70 | return transform |
| 71 | |
| 72 | |
| 73 | def load_model(model_type): |
no test coverage detected