This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user