| 12 | |
| 13 | |
| 14 | class RetrievalDatabase(nn.Module): |
| 15 | |
| 16 | def __init__(self, |
| 17 | num_retrieval=None, |
| 18 | use_motion=False, |
| 19 | use_text=False, |
| 20 | retrieval_file=None, |
| 21 | latent_dim=512, |
| 22 | output_dim=512, |
| 23 | num_layers=2, |
| 24 | max_seq_len=196, |
| 25 | num_heads=8, |
| 26 | ff_size=1024, |
| 27 | dropout=0): |
| 28 | super().__init__() |
| 29 | self.num_retrieval = num_retrieval |
| 30 | self.use_motion = use_motion |
| 31 | self.use_text = use_text |
| 32 | self.latent_dim = latent_dim |
| 33 | self.num_layers = num_layers |
| 34 | self.max_seq_len = max_seq_len |
| 35 | data = np.load(retrieval_file) |
| 36 | self.text_features = torch.Tensor(data['text_features']) |
| 37 | self.captions = data['captions'] |
| 38 | self.motions = data['motions'] |
| 39 | self.m_lengths = data['m_lengths'] |
| 40 | if 'text_seq_features' in data.keys(): |
| 41 | self.text_seq_features = data['text_seq_features'] |
| 42 | |
| 43 | self.latent_dim = latent_dim |
| 44 | self.output_dim = output_dim |
| 45 | if self.use_motion: |
| 46 | self.motion_proj = nn.Linear(self.motions.shape[-1], self.latent_dim) |
| 47 | self.motion_pos_embedding = nn.Parameter(torch.randn(max_seq_len, self.latent_dim)) |
| 48 | TransEncoderLayer = nn.TransformerEncoderLayer( |
| 49 | d_model=self.latent_dim, |
| 50 | nhead=num_heads, |
| 51 | dim_feedforward=ff_size, |
| 52 | dropout=dropout, |
| 53 | activation="gelu", |
| 54 | batch_first=True) |
| 55 | self.motion_encoder = nn.TransformerEncoder( |
| 56 | TransEncoderLayer, |
| 57 | num_layers=num_layers) |
| 58 | if self.use_text: |
| 59 | TransEncoderLayer = nn.TransformerEncoderLayer( |
| 60 | d_model=self.latent_dim, |
| 61 | nhead=num_heads, |
| 62 | dim_feedforward=ff_size, |
| 63 | dropout=dropout, |
| 64 | activation="gelu", |
| 65 | batch_first=True) |
| 66 | self.text_encoder = nn.TransformerEncoder( |
| 67 | TransEncoderLayer, |
| 68 | num_layers=num_layers) |
| 69 | self.results = {} |
| 70 | |
| 71 | def extract_text_feature(self, text, clip_model, device): |