|
from collections import Counter |
|
from pathlib import Path |
|
|
|
import numpy as np |
|
import torch |
|
from torch.utils.data import Dataset |
|
import torch.nn.functional as F |
|
|
|
from utils.file_utils import load_txt |
|
|
|
|
|
class CaptionDataset(Dataset): |
|
def __init__( |
|
self, |
|
name: str, |
|
dataset_dir: str, |
|
num_cams: int, |
|
num_feats: int, |
|
num_segments: int, |
|
sequential: bool, |
|
**kwargs, |
|
): |
|
super().__init__() |
|
self.modality = name |
|
self.name = name |
|
self.dataset_dir = Path(dataset_dir) |
|
|
|
for name, field in kwargs.items(): |
|
if isinstance(field, str): |
|
field = Path(field) |
|
if name == "feat_caption_dir": |
|
field = field / "seq" if sequential else field / "token" |
|
setattr(self, name, field) |
|
|
|
self.filenames = None |
|
|
|
self.clip_seq_dir = self.dataset_dir / "caption_clip" / "seq" |
|
self.num_cams = num_cams |
|
self.num_feats = num_feats |
|
self.num_segments = num_segments |
|
self.sequential = sequential |
|
|
|
def __len__(self): |
|
return len(self.filenames) |
|
|
|
def __getitem__(self, index): |
|
filename = self.filenames[index] |
|
|
|
|
|
if hasattr(self, "segment_dir"): |
|
raw_segments = torch.from_numpy( |
|
np.load((self.segment_dir / (filename + ".npy"))) |
|
) |
|
padded_raw_segments = F.pad( |
|
raw_segments, |
|
(0, self.num_cams - len(raw_segments)), |
|
value=self.num_segments, |
|
) |
|
if hasattr(self, "raw_caption_dir"): |
|
raw_caption = load_txt(self.raw_caption_dir / (filename + ".txt")) |
|
if hasattr(self, "feat_caption_dir"): |
|
feat_caption = torch.from_numpy( |
|
np.load((self.feat_caption_dir / (filename + ".npy"))) |
|
) |
|
if self.sequential: |
|
feat_caption = F.pad( |
|
feat_caption.to(torch.float32), |
|
(0, 0, 0, self.max_feat_length - feat_caption.shape[0]), |
|
) |
|
|
|
if self.modality == "caption": |
|
raw_data = {"caption": raw_caption, "segments": padded_raw_segments} |
|
feat_data = ( |
|
feat_caption.permute(1, 0) if feat_caption.dim() == 2 else feat_caption |
|
) |
|
elif self.modality == "segments": |
|
raw_data = {"segments": padded_raw_segments} |
|
|
|
feat_data = F.one_hot( |
|
padded_raw_segments, num_classes=self.num_segments + 1 |
|
).to(torch.float32) |
|
if self.sequential: |
|
feat_data = feat_data.permute(1, 0) |
|
else: |
|
feat_data = feat_data.reshape(-1) |
|
elif self.modality == "class": |
|
raw_data = {"segments": padded_raw_segments} |
|
most_frequent_segment = Counter(raw_segments).most_common(1)[0][0] |
|
feat_data = F.one_hot( |
|
torch.tensor(most_frequent_segment), num_classes=self.num_segments |
|
).to(torch.float32) |
|
else: |
|
raise ValueError(f"Modality {self.modality} not supported") |
|
|
|
clip_seq_caption = torch.from_numpy( |
|
np.load((self.clip_seq_dir / (filename + ".npy"))) |
|
) |
|
padding_mask = torch.ones((self.max_feat_length)) |
|
padding_mask[clip_seq_caption.shape[0] :] = 0 |
|
clip_seq_caption = F.pad( |
|
clip_seq_caption.to(torch.float32), |
|
(0, 0, 0, self.max_feat_length - clip_seq_caption.shape[0]), |
|
) |
|
raw_data["clip_seq_caption"] = clip_seq_caption |
|
raw_data["clip_seq_mask"] = padding_mask |
|
|
|
return filename, feat_data, raw_data |
|
|
