diff --git a/model/src/models/visual_communication.py b/model/src/models/visual_communication.py new file mode 100644 index 0000000..e11a81f --- /dev/null +++ b/model/src/models/visual_communication.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import torch.nn as nn + +from .resnet18_head import ResNet18Head +from .visual_syntax import VisualSyntaxTail +from core.data_models import VisualSyntaxData + + +class VisualCommunicationModel(nn.Module): + def __init__(self): + super().__init__() + # store other models + self.resnet_head = ResNet18Head() + self.visual_syntax_tail = VisualSyntaxTail() + + def forward(self, x): + # generate visual representation + vis_rep = self.resnet_head(x) + # prepare result map + results = {} + # predict visual syntax + visual_syntax_pred = self.visual_syntax_tail(vis_rep).cpu() + # visual_syntax_pred = float() + results['visual_syntax'] = VisualSyntaxData.from_list( + visual_syntax_pred, + ) + + return results