| 794 | |
| 795 | |
| 796 | class CaptionEmbTransform(AbstractTransform): |
| 797 | |
| 798 | def __init__(self, aligned_captions=True, no_aug=False): |
| 799 | self.aligned_captions = aligned_captions |
| 800 | self.no_aug = no_aug |
| 801 | |
| 802 | def load(self, path): |
| 803 | if path.endswith('.npz'): |
| 804 | sample = np.load(path) |
| 805 | sample = {'emb': sample['emb'], 'mask_valid': sample['mask_valid']} |
| 806 | else: |
| 807 | raise ValueError(f"Invalid file format for caption embedding: {path}") |
| 808 | return sample |
| 809 | |
| 810 | def preprocess(self, sample): |
| 811 | return sample |
| 812 | |
| 813 | def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, |
| 814 | rand_aug_idx: Optional[int], resample_mode: str = None): |
| 815 | |
| 816 | emb = val['emb'] |
| 817 | mask_valid = val['mask_valid'].astype(bool) |
| 818 | num_sequences = emb.shape[0] |
| 819 | |
| 820 | if num_sequences > 1: |
| 821 | if self.aligned_captions: |
| 822 | if rand_aug_idx is None: |
| 823 | emb, mask_valid = emb[0], mask_valid[0] |
| 824 | else: |
| 825 | emb, mask_valid = emb[rand_aug_idx], mask_valid[rand_aug_idx] |
| 826 | else: |
| 827 | if self.no_aug: |
| 828 | emb, mask_valid = emb[0], mask_valid[0] |
| 829 | else: |
| 830 | rand_idx = random.randint(0, num_sequences - 1) |
| 831 | emb, mask_valid = emb[rand_idx], mask_valid[rand_idx] |
| 832 | else: |
| 833 | emb, mask_valid = emb[0], mask_valid[0] |
| 834 | |
| 835 | emb = emb[mask_valid] # Keep only valid embeddings |
| 836 | |
| 837 | return emb |
| 838 | |
| 839 | def postprocess(self, sample): |
| 840 | return torch.tensor(sample) |
| 841 | |
| 842 | |
| 843 | class MetadataTransform(AbstractTransform): |