| 140 | |
| 141 | |
| 142 | class DPTDepthModel(DPT): |
| 143 | def __init__(self, path=None, non_negative=True, **kwargs): |
| 144 | features = kwargs["features"] if "features" in kwargs else 256 |
| 145 | head_features_1 = kwargs["head_features_1"] if "head_features_1" in kwargs else features |
| 146 | head_features_2 = kwargs["head_features_2"] if "head_features_2" in kwargs else 32 |
| 147 | kwargs.pop("head_features_1", None) |
| 148 | kwargs.pop("head_features_2", None) |
| 149 | |
| 150 | head = nn.Sequential( |
| 151 | nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1), |
| 152 | Interpolate(scale_factor=2, mode="bilinear", align_corners=True), |
| 153 | nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1), |
| 154 | nn.ReLU(True), |
| 155 | nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0), |
| 156 | nn.ReLU(True) if non_negative else nn.Identity(), |
| 157 | nn.Identity(), |
| 158 | ) |
| 159 | |
| 160 | super().__init__(head, **kwargs) |
| 161 | |
| 162 | if path is not None: |
| 163 | self.load(path) |
| 164 | |
| 165 | def forward(self, x): |
| 166 | return super().forward(x).squeeze(dim=1) |
no outgoing calls
no test coverage detected