| 86 | |
| 87 | |
| 88 | class DPTDepthModel(DPT): |
| 89 | def __init__(self, path=None, non_negative=True, **kwargs): |
| 90 | features = kwargs["features"] if "features" in kwargs else 256 |
| 91 | |
| 92 | head = nn.Sequential( |
| 93 | nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1), |
| 94 | Interpolate(scale_factor=2, mode="bilinear", align_corners=True), |
| 95 | nn.Conv2d(features // 2, 32, kernel_size=3, stride=1, padding=1), |
| 96 | nn.ReLU(True), |
| 97 | nn.Conv2d(32, 1, kernel_size=1, stride=1, padding=0), |
| 98 | nn.ReLU(True) if non_negative else nn.Identity(), |
| 99 | nn.Identity(), |
| 100 | ) |
| 101 | |
| 102 | super().__init__(head, **kwargs) |
| 103 | |
| 104 | if path is not None: |
| 105 | self.load(path) |
| 106 | |
| 107 | def forward(self, x): |
| 108 | return super().forward(x).squeeze(dim=1) |
| 109 | |