Spaces:
Running on Zero
Running on Zero
File size: 1,895 Bytes
2407511 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """Changer feature-exchange layers.
Adapted from Open-CD's Apache-2.0 licensed
``opencd/models/utils/interaction_layer.py`` at the pinned commit.
"""
from __future__ import annotations
import torch
from torch import nn
class TwoIdentity(nn.Module):
def forward(self, first: torch.Tensor, second: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
return first, second
class SpatialExchange(nn.Module):
def __init__(self, p: float = 0.5) -> None:
super().__init__()
if not 0 < p <= 1:
raise ValueError("p must be in (0, 1]")
self.period = int(1 / p)
def forward(self, first: torch.Tensor, second: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
mask = torch.arange(first.shape[-1], device=first.device) % self.period == 0
output_first = torch.zeros_like(first)
output_second = torch.zeros_like(second)
output_first[..., ~mask] = first[..., ~mask]
output_second[..., ~mask] = second[..., ~mask]
output_first[..., mask] = second[..., mask]
output_second[..., mask] = first[..., mask]
return output_first, output_second
class ChannelExchange(nn.Module):
def __init__(self, p: float = 0.5) -> None:
super().__init__()
if not 0 < p <= 1:
raise ValueError("p must be in (0, 1]")
self.period = int(1 / p)
def forward(self, first: torch.Tensor, second: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
mask = torch.arange(first.shape[1], device=first.device) % self.period == 0
output_first = torch.zeros_like(first)
output_second = torch.zeros_like(second)
output_first[:, ~mask] = first[:, ~mask]
output_second[:, ~mask] = second[:, ~mask]
output_first[:, mask] = second[:, mask]
output_second[:, mask] = first[:, mask]
return output_first, output_second
|