Update all files for SegEarth-OV

OV/upsamplers.py

@@ -0,0 +1,251 @@
+"""
+SimFeatUp upsamplers for dense feature restoration.
+From SegEarth-OV/OV-2 simfeatup_dev. Used by CLIP-based variants (OV, OV-2).
+"""
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from featup.adaptive_conv_cuda.adaptive_conv import AdaptiveConv
+except Exception:
+    AdaptiveConv = None
+
+
+def adaptive_conv_py_simple(input, filters):
+    """Pure PyTorch fallback when featup CUDA is unavailable."""
+    b, c, h1, w1 = input.shape
+    b, h2, w2, f1, f2 = filters.shape
+    assert f1 == f2
+    t_filters = filters.reshape(b, h2, w2, f1 * f2)
+    patches = torch.nn.Unfold(f1)(input).view((b, c, f1 * f2, h2, w2))
+    return torch.einsum("bhwf,bcfhw->bchw", t_filters, patches)
+
+
+def _meshgrid(device, diameter):
+    dist_range = torch.linspace(-1, 1, diameter, device=device)
+    x, y = torch.meshgrid(dist_range, dist_range, indexing="ij")
+    return torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0)
+
+
+class Bilinear(torch.nn.Module):
+    def forward(self, source, guidance):
+        _, _, h, w = guidance.shape
+        return F.interpolate(source, (h, w), mode="bilinear")
+
+
+class LayeredResizeConv(torch.nn.Module):
+    def __init__(self, dim, kernel_size=1, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.conv1 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same")
+        self.conv2 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same")
+        self.conv3 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same")
+        self.conv4 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same")
+
+    def apply_conv(self, source, guidance, conv, activation):
+        big_source = F.interpolate(source, scale_factor=2, mode="bilinear")
+        _, _, h, w = big_source.shape
+        small_guidance = F.interpolate(guidance, (h, w), mode="bilinear")
+        output = activation(conv(torch.cat([big_source, small_guidance], dim=1)))
+        return big_source + output
+
+    def forward(self, source, guidance):
+        source_2 = self.apply_conv(source, guidance, self.conv1, F.relu)
+        source_4 = self.apply_conv(source_2, guidance, self.conv2, F.relu)
+        source_8 = self.apply_conv(source_4, guidance, self.conv3, F.relu)
+        source_16 = self.apply_conv(source_8, guidance, self.conv4, lambda x: x)
+        return source_16
+
+
+class SimpleImplicitFeaturizer(torch.nn.Module):
+    def __init__(self, n_freqs=20):
+        super().__init__()
+        self.n_freqs = n_freqs
+        self.dim_multiplier = 2
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        dtype = x.dtype
+        grid_h = torch.linspace(-1, 1, h, device=x.device, dtype=dtype)
+        grid_w = torch.linspace(-1, 1, w, device=x.device, dtype=dtype)
+        feats = torch.stack(torch.meshgrid(grid_h, grid_w, indexing="ij")).unsqueeze(0)
+        feats = feats.broadcast_to((b, feats.shape[1], h, w))
+        freqs = torch.exp(torch.linspace(-2, 10, self.n_freqs, device=x.device)).to(dtype).reshape(
+            1, self.n_freqs, 1, 1, 1
+        )
+        feats = (feats.unsqueeze(1) * freqs).reshape(b, self.n_freqs * self.dim_multiplier, h, w)
+        return torch.cat([torch.sin(feats), torch.cos(feats), x], dim=1)
+
+
+class IFA(torch.nn.Module):
+    def __init__(self, feat_dim, num_scales=20):
+        super().__init__()
+        self.feat_dim = feat_dim
+        self.sin_feats = SimpleImplicitFeaturizer()
+        self.mlp = nn.Sequential(
+            nn.Conv2d(feat_dim + (num_scales * 4) + 2, feat_dim, 1),
+            nn.BatchNorm2d(feat_dim),
+            nn.LeakyReLU(),
+            nn.Conv2d(feat_dim, feat_dim, 1),
+        )
+
+    def _upsample_2x(self, source):
+        b, c, h, w = source.shape
+        dtype = source.dtype
+        up_source = F.interpolate(source, (h * 2, w * 2), mode="nearest")
+        lr_cord = torch.linspace(0, h, steps=h, device=source.device, dtype=dtype)
+        hr_cord = torch.linspace(0, h, steps=2 * h, device=source.device, dtype=dtype)
+        lr_coords = torch.stack(torch.meshgrid(lr_cord, lr_cord, indexing="ij")).unsqueeze(0)
+        hr_coords = torch.stack(torch.meshgrid(hr_cord, hr_cord, indexing="ij")).unsqueeze(0)
+        up_lr_coords = F.interpolate(lr_coords, (h * 2, w * 2), mode="nearest")
+        coord_diff = up_lr_coords - hr_coords
+        coord_diff_feats = self.sin_feats(coord_diff).to(dtype)
+        bcast_coord_feats = coord_diff_feats.broadcast_to((b, coord_diff_feats.shape[1], h * 2, w * 2))
+        return self.mlp(torch.cat([up_source, bcast_coord_feats], dim=1))
+
+    def forward(self, source, guidance):
+        _, _, gh, gw = guidance.shape
+        x = source
+        while x.shape[2] < gh or x.shape[3] < gw:
+            x = self._upsample_2x(x)
+        if x.shape[2] != gh or x.shape[3] != gw:
+            x = F.interpolate(x, (gh, gw), mode="bilinear")
+        return x
+
+
+class JBULearnedRange(torch.nn.Module):
+    def __init__(self, guidance_dim, feat_dim, key_dim, scale=2, radius=3):
+        super().__init__()
+        self.scale = scale
+        self.radius = radius
+        self.diameter = self.radius * 2 + 1
+        self.guidance_dim = guidance_dim
+        self.key_dim = key_dim
+        self.feat_dim = feat_dim
+        self.range_temp = nn.Parameter(torch.tensor(0.0))
+        self.range_proj = nn.Sequential(
+            nn.Conv2d(guidance_dim, key_dim, 1, 1),
+            nn.GELU(),
+            nn.Dropout2d(0.1),
+            nn.Conv2d(key_dim, key_dim, 1, 1),
+        )
+        self.fixup_proj = nn.Sequential(
+            nn.Conv2d(guidance_dim + self.diameter ** 2, self.diameter ** 2, 1, 1),
+            nn.GELU(),
+            nn.Dropout2d(0.1),
+            nn.Conv2d(self.diameter ** 2, self.diameter ** 2, 1, 1),
+        )
+        self.sigma_spatial = nn.Parameter(torch.tensor(1.0))
+
+    def get_range_kernel(self, x):
+        GB, GC, GH, GW = x.shape
+        proj_x = self.range_proj(x)
+        proj_x_padded = F.pad(proj_x, pad=[self.radius] * 4, mode="reflect")
+        queries = (
+            torch.nn.Unfold(self.diameter)(proj_x_padded)
+            .reshape((GB, self.key_dim, self.diameter * self.diameter, GH, GW))
+            .permute(0, 1, 3, 4, 2)
+        )
+        pos_temp = self.range_temp.exp().clamp_min(1e-4).clamp_max(1e4)
+        return F.softmax(pos_temp * torch.einsum("bchwp,bchw->bphw", queries, proj_x), dim=1)
+
+    def get_spatial_kernel(self, device):
+        patch = _meshgrid(device, self.diameter)
+        return torch.exp(-patch.square().sum(0) / (2 * self.sigma_spatial ** 2)).reshape(
+            1, self.diameter * self.diameter, 1, 1
+        )
+
+    def forward(self, source, guidance):
+        GB, GC, GH, GW = guidance.shape
+        SB, SC, SH, SQ = source.shape
+        assert SB == GB
+        dtype = source.dtype
+        guidance = guidance.to(dtype)
+        spatial_kernel = self.get_spatial_kernel(source.device).to(dtype)
+        range_kernel = self.get_range_kernel(guidance).to(dtype)
+        combined_kernel = (range_kernel * spatial_kernel).to(dtype)
+        combined_kernel /= combined_kernel.sum(1, keepdim=True).clamp(1e-7)
+        combined_kernel += 0.1 * self.fixup_proj(torch.cat([combined_kernel, guidance], dim=1))
+        combined_kernel = combined_kernel.permute(0, 2, 3, 1).reshape(
+            GB, GH, GW, self.diameter, self.diameter
+        )
+        hr_source = F.interpolate(source, size=(GH, GW), mode="bicubic", align_corners=False)
+        hr_source_padded = F.pad(hr_source, pad=[self.radius] * 4, mode="reflect")
+        combined_kernel = combined_kernel.to(hr_source_padded.dtype)
+        if AdaptiveConv is not None:
+            result = AdaptiveConv.apply(hr_source_padded, combined_kernel)
+        else:
+            result = adaptive_conv_py_simple(hr_source_padded, combined_kernel)
+        return result
+
+
+class JBUStack(torch.nn.Module):
+    def __init__(self, feat_dim, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.up1 = JBULearnedRange(3, feat_dim, 32, radius=3)
+        self.up2 = JBULearnedRange(3, feat_dim, 32, radius=3)
+        self.up3 = JBULearnedRange(3, feat_dim, 32, radius=3)
+        self.up4 = JBULearnedRange(3, feat_dim, 32, radius=3)
+        self.fixup_proj = nn.Sequential(
+            nn.Dropout2d(0.2),
+            nn.Conv2d(feat_dim, feat_dim, kernel_size=1),
+        )
+
+    def upsample(self, source, guidance, up):
+        _, _, h, w = source.shape
+        small_guidance = F.adaptive_avg_pool2d(guidance, (h * 2, w * 2))
+        return up(source, small_guidance)
+
+    def forward(self, source, guidance):
+        source_2 = self.upsample(source, guidance, self.up1)
+        source_4 = self.upsample(source_2, guidance, self.up2)
+        source_8 = self.upsample(source_4, guidance, self.up3)
+        source_16 = self.upsample(source_8, guidance, self.up4)
+        return self.fixup_proj(source_16) * 0.1 + source_16
+
+
+class JBUOne(torch.nn.Module):
+    def __init__(self, feat_dim, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.up = JBULearnedRange(3, feat_dim, 32, radius=5)
+        self.fixup_proj = nn.Sequential(
+            nn.Dropout2d(0.2),
+            nn.Conv2d(feat_dim, feat_dim, kernel_size=1),
+        )
+
+    def upsample(self, source, guidance, up):
+        _, _, h, w = source.shape
+        small_guidance = F.adaptive_avg_pool2d(guidance, (h * 2, w * 2))
+        return up(source, small_guidance)
+
+    def forward(self, source, guidance):
+        source_2 = self.upsample(source, guidance, self.up)
+        source_4 = self.upsample(source_2, guidance, self.up)
+        source_8 = self.upsample(source_4, guidance, self.up)
+        source_16 = self.upsample(source_8, guidance, self.up)
+        return self.fixup_proj(source_16) * 0.1 + source_16
+
+
+FEATUP_CHECKPOINTS = {
+    "jbu_one": "simfeatup/xclip_jbu_one_million_aid.ckpt",
+    "jbu_stack": "simfeatup/clip_jbu_stack_cocostuff.ckpt",
+    "jbu_stack_maskclip": "simfeatup/maskclip_jbu_stack_cocostuff.ckpt",
+}
+
+
+def get_upsampler(name: str, feat_dim: int):
+    if name == "bilinear":
+        return Bilinear()
+    elif name == "jbu_one":
+        return JBUOne(feat_dim)
+    elif name == "jbu_stack":
+        return JBUStack(feat_dim)
+    elif name == "resize_conv":
+        return LayeredResizeConv(feat_dim, 1)
+    elif name == "ifa":
+        return IFA(feat_dim)
+    else:
+        raise ValueError(f"Unknown upsampler: {name}. Use: bilinear, jbu_one, jbu_stack, resize_conv, ifa")