| 753 | |
| 754 | |
| 755 | class CaptionTransform(AbstractTransform): |
| 756 | |
| 757 | def __init__(self, aligned_captions=True, no_aug=False): |
| 758 | self.aligned_captions = aligned_captions |
| 759 | self.no_aug = no_aug |
| 760 | |
| 761 | def load(self, path): |
| 762 | # Caption can either be stored as .txt or .json.gz (in which case it's a list of dicts) |
| 763 | if path.endswith('.txt'): |
| 764 | sample = Path(path).read_text() |
| 765 | elif path.endswith('.json'): |
| 766 | with open(path, 'r') as f: |
| 767 | sample = json.load(f) |
| 768 | elif path.endswith('.json.gz'): |
| 769 | with gzip.open(path, 'rb') as f: |
| 770 | sample = json.load(f) |
| 771 | return sample |
| 772 | |
| 773 | def preprocess(self, sample): |
| 774 | return sample |
| 775 | |
| 776 | def image_augment(self, val, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, |
| 777 | rand_aug_idx: Optional[int], resample_mode: str = None): |
| 778 | |
| 779 | if isinstance(val, list) or isinstance(val, tuple): |
| 780 | if self.aligned_captions: |
| 781 | val = val[0] if rand_aug_idx is None else val[rand_aug_idx] |
| 782 | else: |
| 783 | val = random.choice(val) if not self.no_aug else val[0] |
| 784 | |
| 785 | if isinstance(val, dict): |
| 786 | # If each caption is saved as a dict, extract the string |
| 787 | val = val["caption"] |
| 788 | assert isinstance(val, str) |
| 789 | |
| 790 | return val |
| 791 | |
| 792 | def postprocess(self, sample): |
| 793 | return sample |
| 794 | |
| 795 | |
| 796 | class CaptionEmbTransform(AbstractTransform): |
no outgoing calls
no test coverage detected