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