model.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from util.conf import Config
from util.register import register_model


@register_model('hmmfn')
class SpatialFusionModel(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.base = config.train.base
        self.config = config
        self.ndvi_encoder = SpatialEncoder(in_channels=1, base=self.base) 
        self.weather_encoder = SpatialEncoder(in_channels=2, base=self.base) 

        feat_dim = self.base * 4
        self.H_e3, self.W_e3 = self.config.data.height // 4, self.config.data.width // 4
        H_map, W_map = self.config.hyper.map_h, self.config.hyper.map_w  # 映射到较小的空间尺寸

        self.ndvi_temporal_encoder = TemporalTransformerEncoder(
            feat_dim, self.config.hyper.num_heads, self.config.hyper.num_layers, self.config.hyper.dropout
        )
        self.weather_temporal_encoder = TemporalTransformerEncoder(
            feat_dim, self.config.hyper.num_heads, self.config.hyper.num_layers, self.config.hyper.dropout)

        self.temporal_feat_mapper = TemporalFeatureMapper(
            feat_dim=feat_dim, base=self.base, H_out=H_map, W_out=W_map
        )
        self.fuse_conv = nn.Conv2d(self.base*8, self.base*4, kernel_size=1)
        self.decoder = UNetDecoder(self.config)

        self.fuse_weight_e1 = nn.Parameter(torch.tensor(0.5))
        self.fuse_weight_e2 = nn.Parameter(torch.tensor(0.5))

        self.st_w = nn.Parameter(torch.tensor(0.5))
 
    def forward(self, ndvi_seq, weather_seq):
        B, T, C_ndvi, H, W = ndvi_seq.shape
        _, _, C_weather, _, _ = weather_seq.shape

        ndvi_seq_reshaped = ndvi_seq.view(B * T, C_ndvi, H, W)
        weather_seq_reshaped = weather_seq.view(B * T, C_weather, H, W)

        ndvi_e1, ndvi_e2, ndvi_e3 = self.ndvi_encoder(ndvi_seq_reshaped)
        weather_e1, weather_e2, weather_e3 = self.weather_encoder(weather_seq_reshaped)

        ndvi_e3_pooled = F.adaptive_avg_pool2d(ndvi_e3, (1, 1)).view(B, T, -1)
        weather_e3_pooled = F.adaptive_avg_pool2d(weather_e3, (1, 1)).view(B, T, -1)

        ndvi_temporal_feat = self.ndvi_temporal_encoder(ndvi_e3_pooled)
        weather_temporal_feat = self.weather_temporal_encoder(weather_e3_pooled)

        ndvi_map = self.temporal_feat_mapper(ndvi_temporal_feat)
        weather_map = self.temporal_feat_mapper(weather_temporal_feat)

        fused = torch.cat([ndvi_map, weather_map], dim=1)
        fused = self.fuse_conv(fused)

        fused = F.interpolate(fused, size=(self.H_e3, self.W_e3), mode='bilinear', align_corners=False)
        
        ndvi_e1 = ndvi_e1.view(B, T, ndvi_e1.size(1), ndvi_e1.size(2), ndvi_e1.size(3))
        ndvi_e1_mean = ndvi_e1.mean(dim=1)
        ndvi_e1 = self.st_w * ndvi_e1_mean + (1 - self.st_w) * ndvi_e1[:, -1]

        ndvi_e2 = ndvi_e2.view(B, T, ndvi_e2.size(1), ndvi_e2.size(2), ndvi_e2.size(3))
        ndvi_e2_mean = ndvi_e2.mean(dim=1)
        ndvi_e2 = self.st_w * ndvi_e2_mean + (1 - self.st_w) * ndvi_e2[:, -1]

        weather_e1 = weather_e1.view(B, T, weather_e1.size(1), weather_e1.size(2), weather_e1.size(3))
        weather_e1_mean = weather_e1.mean(dim=1)
        weather_e1 = self.st_w * weather_e1_mean + (1 - self.st_w) * weather_e1[:, -1]

        weather_e2 = weather_e2.view(B, T, weather_e2.size(1), weather_e2.size(2), weather_e2.size(3))
        weather_e2_mean = weather_e2.mean(dim=1)
        weather_e2 = self.st_w * weather_e2_mean + (1 - self.st_w) * weather_e2[:, -1]


        w1 = torch.sigmoid(self.fuse_weight_e1)
        w2 = torch.sigmoid(self.fuse_weight_e2)

        fused_e1 = w1 * ndvi_e1 + (1-w1) * weather_e1
        fused_e2 = w2 * ndvi_e2 + (1-w2) * weather_e2

        out = self.decoder(fused_e1, fused_e2, fused)
        out = torch.tanh(out)
        out = out.clamp(-1, 1)

        return out


class SpatialEncoder(nn.Module):
    def __init__(self, in_channels, base):
        super().__init__()
        self.enc1 = nn.Sequential(nn.Conv2d(in_channels, base, 3, padding=1), nn.ReLU(), nn.BatchNorm2d(base))
        self.enc2 = nn.Sequential(nn.Conv2d(base, base*2, 3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(base*2))
        self.enc3 = nn.Sequential(nn.Conv2d(base*2, base*4, 3, stride=2, padding=1), nn.ReLU(), nn.BatchNorm2d(base*4))

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(e1)
        e3 = self.enc3(e2)
        return e1, e2, e3


class TemporalTransformerEncoder(nn.Module):
    def __init__(self, feat_dim, nhead, num_layers, dropout):
        super().__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model=feat_dim, nhead=nhead, dropout=dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.pos_encoder = PositionalEncoding(feat_dim)

    def forward(self, x):
        x = self.pos_encoder(x)
        x = x.permute(1,0,2)
        output = self.transformer_encoder(x)
        output = output.mean(dim=0)
        return output


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]  # type: ignore
        return x


class TemporalFeatureMapper(nn.Module):
    def __init__(self, feat_dim, base, H_out, W_out):
        super().__init__()
        self.H_out = H_out
        self.W_out = W_out
        self.base = base
 
        hidden_dim = feat_dim * 4
        self.mlp = nn.Sequential(
            nn.Linear(feat_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, base * 4 * H_out * W_out),
            nn.ReLU()
        )
 
    def forward(self, x):
        B = x.size(0)
        x = self.mlp(x)
        x = x.view(B, self.base * 4, self.H_out, self.W_out)
        return x


class UNetDecoder(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.base = config.train.base
        self.size = (config.data.height, config.data.width)
        self.up1 = nn.ConvTranspose2d(self.base*4, self.base*2, 2, stride=2)
        self.conv1 = nn.Sequential(nn.Conv2d(self.base*4, self.base*2, 3, padding=1), nn.ReLU())
        self.up2 = nn.ConvTranspose2d(self.base*2, self.base, 2, stride=2)
        self.conv2 = nn.Sequential(nn.Conv2d(self.base*2, self.base, 3, padding=1), nn.ReLU())
        self.out = nn.Conv2d(self.base, 1, 1)

    def forward(self, e1, e2, e3):
        d1 = self.up1(e3)
        # 尺寸对齐
        if d1.size(2) != e2.size(2) or d1.size(3) != e2.size(3):
            d1 = F.interpolate(d1, size=(e2.size(2), e2.size(3)), mode='bilinear', align_corners=False)
        d1 = self.conv1(torch.cat([d1, e2], dim=1))
 
        d2 = self.up2(d1)
        if d2.size(2) != e1.size(2) or d2.size(3) != e1.size(3):
            d2 = F.interpolate(d2, size=(e1.size(2), e1.size(3)), mode='bilinear', align_corners=False)
        d2 = self.conv2(torch.cat([d2, e1], dim=1))
        
        out = self.out(d2)
        return F.interpolate(out, size=self.size, mode='bilinear', align_corners=False)