model.py

"""
MAVNet (PyTorch rewrite) - model definitions.

This replaces the original Keras 1.x / tflearn implementation with a clean
PyTorch version, and fixes two conceptual issues found in the original paper:

1. Output formulation: the original code trained a single softmax over 7
   mutually-exclusive keystroke classes, while the paper described a 5-value
   multi-hot vector (movement + junction). Here we use two independent heads:
     - `movement_head`: softmax over {forward, yaw_left, yaw_right, halt}
     - `junction_head`: independent sigmoid (junction is a scene property,
        not a movement command, so it shouldn't share a softmax with it).

2. Architecture ambiguity: the original repo had three competing
   architectures (Inception-v3, AlexNet, LRCN) with no single canonical
   model. Here `MAVNetCNN` is the primary single-frame classifier (matching
   the paper's stated architecture), and `MAVNetLRCN` is an explicit,
   opt-in temporal variant (CNN backbone + LSTM) for later comparison -
   not silently swapped in like in the original train.py/test_model.py.
"""

import torch
import torch.nn as nn

NUM_MOVEMENT_CLASSES = 4  # forward, yaw_left, yaw_right, halt


class InceptionLiteBlock(nn.Module):
    """A pruned inception block: 1x1, 3x3 and pool branches only.

    The original paper argued that 5x5 convolutions were redundant given
    the dominance of 1x1 filters in Inception-v3, and removed them for
    speed. We keep that reasoning here (it's a defensible, testable claim,
    unlike the Radon-transform claim), but make it explicit and named.
    """

    def __init__(self, in_channels, c1x1, c3x3_reduce, c3x3, pool_proj):
        super().__init__()
        self.branch1x1 = nn.Sequential(
            nn.Conv2d(in_channels, c1x1, kernel_size=1),
            nn.BatchNorm2d(c1x1),
            nn.ReLU(inplace=True),
        )
        self.branch3x3 = nn.Sequential(
            nn.Conv2d(in_channels, c3x3_reduce, kernel_size=1),
            nn.BatchNorm2d(c3x3_reduce),
            nn.ReLU(inplace=True),
            nn.Conv2d(c3x3_reduce, c3x3, kernel_size=3, padding=1),
            nn.BatchNorm2d(c3x3),
            nn.ReLU(inplace=True),
        )
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1),
            nn.BatchNorm2d(pool_proj),
            nn.ReLU(inplace=True),
        )
        self.out_channels = c1x1 + c3x3 + pool_proj

    def forward(self, x):
        return torch.cat(
            [self.branch1x1(x), self.branch3x3(x), self.branch_pool(x)], dim=1
        )


class MAVNetBackbone(nn.Module):
    """Shared conv backbone used by both the single-frame and LRCN models."""

    def __init__(self, in_channels=1):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )
        self.block1 = InceptionLiteBlock(32, c1x1=16, c3x3_reduce=16, c3x3=32, pool_proj=16)
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.block2 = InceptionLiteBlock(
            self.block1.out_channels, c1x1=32, c3x3_reduce=32, c3x3=64, pool_proj=32
        )
        self.pool2 = nn.AdaptiveAvgPool2d(1)
        self.out_features = self.block2.out_channels

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.pool1(x)
        x = self.block2(x)
        x = self.pool2(x)
        return torch.flatten(x, 1)  # (batch, out_features)


class MAVNetCNN(nn.Module):
    """Primary single-frame MAVNet classifier (no temporal component)."""

    def __init__(self, in_channels=1, num_movement_classes=NUM_MOVEMENT_CLASSES):
        super().__init__()
        self.backbone = MAVNetBackbone(in_channels=in_channels)
        self.movement_head = nn.Linear(self.backbone.out_features, num_movement_classes)
        self.junction_head = nn.Linear(self.backbone.out_features, 1)

    def forward(self, x):
        """x: (batch, in_channels, H, W)"""
        features = self.backbone(x)
        movement_logits = self.movement_head(features)
        junction_logit = self.junction_head(features).squeeze(-1)
        return {"movement_logits": movement_logits, "junction_logit": junction_logit}


class MAVNetLRCN(nn.Module):
    """Optional temporal variant: shared CNN backbone + LSTM over timesteps.

    Kept separate and explicit (not silently swapped in) so any comparison
    against MAVNetCNN is an intentional, reported ablation.
    """

    def __init__(
        self,
        in_channels=1,
        num_movement_classes=NUM_MOVEMENT_CLASSES,
        lstm_hidden=128,
    ):
        super().__init__()
        self.backbone = MAVNetBackbone(in_channels=in_channels)
        self.lstm = nn.LSTM(
            input_size=self.backbone.out_features,
            hidden_size=lstm_hidden,
            batch_first=True,
        )
        self.movement_head = nn.Linear(lstm_hidden, num_movement_classes)
        self.junction_head = nn.Linear(lstm_hidden, 1)

    def forward(self, x):
        """x: (batch, timesteps, in_channels, H, W)"""
        batch, timesteps = x.shape[0], x.shape[1]
        x = x.view(batch * timesteps, *x.shape[2:])
        features = self.backbone(x)
        features = features.view(batch, timesteps, -1)
        lstm_out, _ = self.lstm(features)
        last = lstm_out[:, -1, :]  # prediction for the final timestep
        movement_logits = self.movement_head(last)
        junction_logit = self.junction_head(last).squeeze(-1)
        return {"movement_logits": movement_logits, "junction_logit": junction_logit}


if __name__ == "__main__":
    # Quick shape sanity check with dummy tensors (no real data required).
    cnn = MAVNetCNN(in_channels=1)
    dummy_frame = torch.randn(8, 1, 100, 100)
    out = cnn(dummy_frame)
    print("MAVNetCNN movement_logits:", out["movement_logits"].shape)
    print("MAVNetCNN junction_logit:", out["junction_logit"].shape)

    lrcn = MAVNetLRCN(in_channels=1)
    dummy_seq = torch.randn(8, 7, 1, 100, 100)
    out = lrcn(dummy_seq)
    print("MAVNetLRCN movement_logits:", out["movement_logits"].shape)
    print("MAVNetLRCN junction_logit:", out["junction_logit"].shape)