feat: complete Sprint 5 training infrastructure

configs/default.yaml

@@ -6,38 +6,56 @@ paths:
   raw_data: data/raw
   processed_data: data/processed
   sequences: data/intermediate/sequences
+
+  outputs: outputs
   checkpoints: outputs/checkpoints
   logs: outputs/logs
   metrics: outputs/metrics
   predictions: outputs/predictions
 
 dataset:
   name: FFPP_CelebDF
+
   image_size: 224
+
   sequence_length: 30
+
   fps: 5
+
+dataloader:
   batch_size: 8
-  num_workers: 4
+
+  num_workers: 0
+
+  pin_memory: true
+
+  drop_last: false
 
 training:
   epochs: 30
+
   learning_rate: 0.0001
+
   weight_decay: 0.0001
+
   optimizer: AdamW
+
   scheduler: CosineAnnealingLR
 
 model:
-  backbone: vit_b_16
+  name: vit_temporal_pooling
+
   pretrained: true
+
   freeze_backbone: true
-  temporal_pooling: mean
+
   num_classes: 2
 
-dataloader:
-  batch_size: 8
-  num_workers: 4
-  pin_memory: true
-  drop_last: false
+  pooling: mean
+
+  hidden_dim: 256
+
+  dropout: 0.5
 
 evaluation:
   metrics:

src/models/__init__.py

@@ -0,0 +1 @@
+from .model_factory import ModelFactory

src/models/base_model.py

@@ -0,0 +1,35 @@
+"""
+DeepVision AI
+
+Base Model Interface
+"""
+
+from abc import ABC, abstractmethod
+
+import torch.nn as nn
+
+
+class BaseModel(nn.Module, ABC):
+    """
+    Base class for all DeepVision AI models.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    @abstractmethod
+    def forward(self, x):
+        """
+        Forward pass.
+
+        Args:
+            x (torch.Tensor)
+
+        Returns:
+            torch.Tensor
+        """
+        pass
+
+    @property
+    def name(self):
+        return self.__class__.__name__

src/models/efficientnet_bilstm.py

@@ -0,0 +1,18 @@
+"""
+Placeholder.
+
+Will be integrated in Sprint 4.2
+"""
+
+import torch.nn as nn
+
+
+class EfficientNetBiLSTM(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        raise NotImplementedError(
+            "EfficientNet-BiLSTM integration begins in Sprint 4.2"
+        )

src/models/model_factory.py

@@ -0,0 +1,37 @@
+"""
+DeepVision AI
+
+Model Factory
+"""
+
+from src.models.vit_temporal_pooling import ViTTemporalPooling
+from src.models.efficientnet_bilstm import EfficientNetBiLSTM
+
+
+class ModelFactory:
+
+    @staticmethod
+    def create(config):
+
+        model_name = config["model"]["name"].lower()
+
+
+        if model_name == "vit_temporal_pooling":
+
+            return ViTTemporalPooling(
+                image_size=config["dataset"]["image_size"],
+                num_classes=config["model"]["num_classes"],
+                pretrained=config["model"]["pretrained"],
+                freeze_backbone=config["model"]["freeze_backbone"],
+                pooling=config["model"]["pooling"],
+                hidden_dim=config["model"]["hidden_dim"],
+                dropout=config["model"]["dropout"],
+        )
+
+        elif model_name == "efficientnet_bilstm":
+
+            return EfficientNetBiLSTM()
+
+        raise ValueError(
+            f"Unsupported model: {model_name}"
+        )

src/models/vit_temporal_pooling.py

@@ -0,0 +1,95 @@
+"""
+DeepVision AI
+
+Vision Transformer with Temporal Mean Pooling
+"""
+
+import torch
+import torch.nn as nn
+from torchvision.models import vit_b_16, ViT_B_16_Weights
+
+from src.models.base_model import BaseModel
+
+
+class ViTTemporalPooling(BaseModel):
+    """
+    Vision Transformer + Temporal Pooling
+    """
+
+    def __init__(
+        self,
+        image_size=224,
+        num_classes=2,
+        pretrained=True,
+        freeze_backbone=True,
+        pooling="mean",
+        hidden_dim=256,
+        dropout=0.5,
+    ):
+        super().__init__()
+
+        self.image_size = image_size
+        self.pooling = pooling
+
+        # Load ViT backbone
+        if pretrained:
+            weights = ViT_B_16_Weights.IMAGENET1K_V1
+        else:
+            weights = None
+
+        self.backbone = vit_b_16(weights=weights)
+
+        # Remove classification head
+        self.backbone.heads = nn.Identity()
+
+        # Feature dimension of ViT-B/16
+        self.feature_dim = 768
+
+        # Freeze backbone if requested
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+        self.classifier = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_classes),
+        )
+
+    def temporal_pool(self, features):
+        """
+        features:
+            (B,T,768)
+        """
+
+        if self.pooling == "mean":
+            return features.mean(dim=1)
+
+        elif self.pooling == "max":
+            return features.max(dim=1)[0]
+
+        else:
+            raise ValueError(
+                f"Unsupported pooling: {self.pooling}"
+            )
+
+    def forward(self, x):
+        """
+        x:
+            (B,T,C,H,W)
+        """
+
+        B, T, C, H, W = x.shape
+
+        x = x.reshape(B * T, C, H, W)
+
+        features = self.backbone(x)
+
+        features = features.reshape(B, T, self.feature_dim)
+
+        video_features = self.temporal_pool(features)
+
+        logits = self.classifier(video_features)
+
+        return logits