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