code/src/model.py

"""Word-level lip-reading baseline.

Architecture (Stafylakis & Tzimiropoulos, 2017 — the standard LRW baseline):

    input clip  (B, 1, T, 88, 88)   grayscale mouth ROI
        -> 3D-conv spatiotemporal stem            (B, 64, T, 22, 22)
        -> per-frame 2D ResNet-18 backbone        (B*T, 512)
        -> BiGRU over the T frame embeddings      (B, T, 2*hidden)
        -> temporal average pool + linear head    (B, num_classes)
"""
from __future__ import annotations

import torch
import torch.nn as nn
from torchvision.models import resnet18


class Conv3dStem(nn.Module):
    """Spatiotemporal front-end. Keeps the time dimension, downsamples space."""

    def __init__(self, in_ch: int = 1, out_ch: int = 64):
        super().__init__()
        self.conv = nn.Conv3d(in_ch, out_ch, kernel_size=(5, 7, 7),
                              stride=(1, 2, 2), padding=(2, 3, 3), bias=False)
        self.bn = nn.BatchNorm3d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        self.pool = nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2),
                                 padding=(0, 1, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # (B,1,T,H,W)
        return self.pool(self.relu(self.bn(self.conv(x))))  # (B,64,T,H/4,W/4)


class ResNet18Backbone(nn.Module):
    """torchvision ResNet-18 with its input stem removed.

    The Conv3dStem already produced 64-channel feature maps at 1/4 spatial
    resolution, which is exactly what ResNet's ``layer1`` expects, so we drop
    ResNet's own conv1/bn1/maxpool and feed straight into the residual stages.
    """

    out_dim = 512

    def __init__(self):
        super().__init__()
        net = resnet18(weights=None)
        self.layer1 = net.layer1
        self.layer2 = net.layer2
        self.layer3 = net.layer3
        self.layer4 = net.layer4
        self.avgpool = net.avgpool

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # (N,64,H',W')
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        return torch.flatten(x, 1)  # (N, 512)


class LipReadingModel(nn.Module):
    def __init__(self, num_classes: int, gru_hidden: int = 256,
                 gru_layers: int = 2, dropout: float = 0.4):
        super().__init__()
        self.stem = Conv3dStem()
        self.backbone = ResNet18Backbone()
        self.gru = nn.GRU(
            input_size=self.backbone.out_dim,
            hidden_size=gru_hidden,
            num_layers=gru_layers,
            batch_first=True,
            bidirectional=True,
            dropout=dropout if gru_layers > 1 else 0.0,
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(gru_hidden * 2, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # (B,1,T,H,W)
        b, _, t = x.size(0), x.size(1), x.size(2)
        x = self.stem(x)                              # (B,64,T,H',W')
        x = x.permute(0, 2, 1, 3, 4).contiguous()     # (B,T,64,H',W')
        x = x.view(b * t, x.size(2), x.size(3), x.size(4))
        x = self.backbone(x)                          # (B*T, 512)
        x = x.view(b, t, -1)                          # (B, T, 512)
        x, _ = self.gru(x)                            # (B, T, 2*hidden)
        x = x.mean(dim=1)                             # temporal average pool
        return self.classifier(x)                     # (B, num_classes)


def build_model(cfg) -> LipReadingModel:
    return LipReadingModel(
        num_classes=cfg.num_classes,
        gru_hidden=cfg.gru_hidden,
        gru_layers=cfg.gru_layers,
        dropout=cfg.dropout,
    )