Upload cnn_bigru/models/multimodal_encoders.py with huggingface_hub

cnn_bigru/models/multimodal_encoders.py

@@ -0,0 +1,141 @@
+"""multimodal_encoders.py — Encoders para imagem e áudio.
+
+Para o projeto CNN-BiGRU multimodal, precisamos de encoders que produzam
+sequências de features compatíveis com o núcleo cooperativo CNN-BiGRU.
+
+    - ImageEncoder: CNN 2D -> flatten espacial -> projeção linear -> [B, T_img, C]
+    - AudioEncoder: CNN 1D sobre espectrograma -> [B, T_aud, C]
+
+A fusão multimodal acontece trocando o stream B do núcleo cooperativo pela
+saída do encoder visual/áudio, ou adicionando-os como streams extras numa
+versão estendida. Aqui usamos uma fusão por concatenação + projeção.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class ImageEncoder(nn.Module):
+    """CNN 2D simples que codifica imagem em sequência de features.
+
+    Input:  [B, C, H, W]
+    Output: [B, T_img, out_dim]   onde T_img = (H/4) * (W/4)
+
+    Usa 2 blocos Conv2d+ReLU+MaxPool para reduzir dimensão espacial por 4.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        out_dim: int = 64,
+        hidden_dim: int = 32,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim * 2, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2)
+        self.proj = nn.Linear(hidden_dim * 2, out_dim)
+        self.out_dim = out_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, C, H, W]
+        x = F.relu(self.conv1(x))
+        x = self.pool(x)            # [B, C', H/2, W/2]
+        x = F.relu(self.conv2(x))
+        x = self.pool(x)            # [B, C'', H/4, W/4]
+        B, C, H, W = x.shape
+        x = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)  # [B, T, C]
+        x = self.proj(x)            # [B, T, out_dim]
+        return x
+
+
+class AudioEncoder(nn.Module):
+    """Encoder de áudio: CNN 1D sobre espectrograma (T freq bins como canais).
+
+    Input:  [B, 1, T_time, F_freq]
+    Output: [B, T_time, out_dim]
+
+    Trata a dimensão de frequência como canais e processa a dimensão temporal
+    com Conv1d, produzindo uma sequência compatível com BiGRU.
+    """
+
+    def __init__(
+        self,
+        n_freq: int = 40,
+        out_dim: int = 64,
+        hidden_dim: int = 32,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv1d(n_freq, hidden_dim, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(hidden_dim, hidden_dim * 2, kernel_size=3, padding=1)
+        self.proj = nn.Linear(hidden_dim * 2, out_dim)
+        self.out_dim = out_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, 1, T, F] -> queremos [B, F, T] para Conv1d
+        B, _, T, F_ = x.shape
+        x = x.view(B, F_, T)  # [B, F, T]
+        x = F.relu(self.conv1(x))     # [B, hidden, T]
+        x = F.relu(self.conv2(x))     # [B, hidden*2, T]
+        x = x.permute(0, 2, 1)        # [B, T, hidden*2]
+        x = self.proj(x)              # [B, T, out_dim]
+        return x
+
+
+class MultimodalFusion(nn.Module):
+    """Fusão de features de texto, imagem e áudio.
+
+    Recebe três tensores [B, D_text], [B, D_img], [B, D_aud] e produz
+    um único tensor [B, out_dim] usando concatenação + projeção +
+    uma porta de fusão (gated) que aprende a importância de cada modalidade.
+    """
+
+    def __init__(
+        self,
+        d_text: int,
+        d_img: int,
+        d_aud: int,
+        out_dim: int = 128,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        d_concat = d_text + d_img + d_aud
+        self.proj = nn.Linear(d_concat, out_dim)
+        # Porta de fusão: aprende pesos por modalidade
+        self.gate_text = nn.Linear(d_text, 1)
+        self.gate_img = nn.Linear(d_img, 1)
+        self.gate_aud = nn.Linear(d_aud, 1)
+        self.dropout = nn.Dropout(dropout)
+        self.out_dim = out_dim
+
+    def forward(
+        self,
+        text_feat: torch.Tensor,
+        img_feat: torch.Tensor,
+        aud_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        # Gates (Sigmoid -> [0, 1])
+        g_t = torch.sigmoid(self.gate_text(text_feat))  # [B, 1]
+        g_i = torch.sigmoid(self.gate_img(img_feat))
+        g_a = torch.sigmoid(self.gate_aud(aud_feat))
+
+        # Aplica gates como máscara suave
+        text_g = text_feat * g_t
+        img_g = img_feat * g_i
+        aud_g = aud_feat * g_a
+
+        # Concatena e projeta
+        concat = torch.cat((text_g, img_g, aud_g), dim=-1)
+        out = self.proj(concat)
+        out = self.dropout(out)
+        return out
+
+
+__all__ = ["ImageEncoder", "AudioEncoder", "MultimodalFusion"]