| """Transformers-compatible wrapper for harpertoken/move. |
| |
| The model classifies KTH actions from 16 frames with max-pooling over time, |
| not a transformer. This subclasses PreTrainedModel so the weights serialize |
| honestly as config.json plus model.safetensors and load through AutoModel. |
| Importing this module registers the architecture. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel |
|
|
|
|
| class MoveConfig(PretrainedConfig): |
| model_type = "move_temporal" |
|
|
| def __init__(self, num_labels=6, dropout=0.25, **kwargs): |
| super().__init__(**kwargs) |
| self.num_labels = num_labels |
| self.dropout = dropout |
|
|
|
|
| LABELS = ["walking", "jogging", "running", "boxing", "handwaving", "handclapping"] |
|
|
|
|
| class MoveTemporal(PreTrainedModel): |
| config_class = MoveConfig |
| base_model_prefix = "move" |
|
|
| |
| |
| |
| all_tied_weights_keys = {} |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.cnn = nn.ModuleDict( |
| { |
| "conv1": nn.Conv2d(1, 32, 3, padding=1), |
| "conv2": nn.Conv2d(32, 64, 3, padding=1), |
| "fc1": nn.Linear(64 * 8 * 8, 128), |
| } |
| ) |
| self.pool = nn.MaxPool2d(2, 2) |
| self.adapt = nn.AdaptiveAvgPool2d((8, 8)) |
| self.head = nn.Linear(128, config.num_labels) |
| self.drop = nn.Dropout(config.dropout) |
|
|
| def forward(self, clip): |
| b, t, _, h, w = clip.shape |
| x = clip.view(b * t, 1, h, w) |
| x = self.pool(torch.relu(self.cnn["conv1"](x))) |
| x = self.pool(torch.relu(self.cnn["conv2"](x))) |
| x = self.adapt(x) |
| feats = torch.relu(self.cnn["fc1"](x.view(b * t, -1))).view(b, t, -1) |
| return self.head(self.drop(feats.amax(1))) |
|
|
| def predict_label(self, clip): |
| """clip: 16x64x64 float array in [0, 1]. Returns the action.""" |
| self.eval() |
| x = torch.as_tensor(clip, dtype=torch.float32) |
| if x.dim() == 3: |
| x = x.unsqueeze(0).unsqueeze(0) |
| elif x.dim() == 4: |
| x = x.unsqueeze(0) |
| with torch.no_grad(): |
| return LABELS[int(self.forward(x).argmax(1)[0])] |
|
|
|
|
| def register(): |
| AutoConfig.register(MoveConfig.model_type, MoveConfig, exist_ok=True) |
| AutoModel.register(MoveConfig, MoveTemporal, exist_ok=True) |
|
|
|
|
| register() |
|
|
