began defining visual communication model
Code Quality Pipeline / Test (push) Successful in 4m5s

This commit is contained in:
Brian Bjarke Jensen
2024-03-12 22:54:50 +01:00
parent 865a8ff2b7
commit b9f333c7d0
+29
View File
@@ -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