Method __iter__

altogether/infer.py:80–134

(self)

        return f"{self.root_dir}/{shard_id % 100}/{shard_id}.tar"

    def __iter__(self):
        shard_id = self.shard_id

        # v2 374px is face blurred. iterate over all image-text pairs until needs legal mitigation;
        tarball_path = self._get_tarball_path(shard_id)

        with tarfile.open(tarball_path) as tar:        
            img_uuid, json_uuid = None, None
            members = tar.getmembers()
            # metaclip_v1 can be iterative but the paper uses mmap for random access.
            for member in members:
                # read jpeg first and json next 
                if member.name.endswith(".jpeg"):
                    img_uuid = member.name[:-len(".jpeg")]
                    if img_uuid.startswith("./"):
                        img_uuid = img_uuid[len("./"):]
                    with tar.extractfile(member) as f:
                        img = f.read()

                elif member.name.endswith(".json"):
                    json_uuid = member.name[:-len(".json")]
                    if json_uuid.startswith("./"):
                        json_uuid = json_uuid[len("./"):]
                    with tar.extractfile(member) as f:
                        text_json = json.load(f)
                else:
                    print(f"unknown file ext {member.name}")
                    continue

                if img_uuid is None or json_uuid is None:
                    continue

                assert img_uuid == json_uuid
    
                with Image.open(BytesIO(img)) as img:
                    image = img.convert("RGB")
                    # assert "face_bbox" in text_json
                    # image = fairblurbbox(image, text_json["face_bbox"])
                    image = self.transform(image)
                
                pad_token_id = self.args.clipcap_args["pad_token_id"]
                prefix_length = self.args.clipcap_args["prefix_length"]

                alt_text = "; ".join(list(set([txt_tuple[1] for txt_tuple in text_json["texts"]])))

                prompt = alt_text

                prompt_input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)[:self.args.rewrite_prompt]
                padding_len = self.args.rewrite_prompt - len(prompt_input_ids)
                prompt_input_ids = [pad_token_id] * padding_len + prompt_input_ids
                tokens = prompt_input_ids
                tokens = torch.tensor(tokens, dtype=torch.long)
                
                yield image, tokens, img_uuid
                img_uuid, json_uuid = None, None


def get_dataloader(args, batch_size, data_path, shard_id, transform, tokenize):