Spaces:
Running
Running
| import webdataset as wds | |
| import glob | |
| import os | |
| from tqdm import tqdm | |
| import orjson as json | |
| import itertools | |
| from PIL import Image | |
| import numpy as np | |
| from typing import List | |
| class Generator(): | |
| def __init__(self, dataset_name): | |
| self.dataset_name = dataset_name | |
| self.is_end = False | |
| class CC3MGenerator(Generator): | |
| def __init__(self, root: str, dataset_name="cc3m"): | |
| super().__init__(dataset_name=dataset_name) | |
| self.tars = glob.glob(os.path.join(root, "cc3m_*", "*.tar")) | |
| def __len__(self): | |
| return 3000000 | |
| def __iter__(self): | |
| for tar in self.tars: | |
| dataset = wds.WebDataset(tar).decode("pilrgb").to_tuple("jpg;png;jpeg", "txt") | |
| for data in dataset: | |
| yield [self.dataset_name] + list(data) | |
| self.is_end = True | |
| class CC12MGenerator(CC3MGenerator): | |
| def __init__(self, root: str): | |
| super().__init__(root, "cc12m") | |
| self.tars = glob.glob(os.path.join(root, "*.tar")) | |
| def __len__(self): | |
| return 12000000 | |
| class COCOGenerator(Generator): | |
| def __init__(self, anno: str, image_dir): | |
| super().__init__(dataset_name="coco") | |
| data = json.loads(open(anno).read()) | |
| self.annotations = data["annotations"] | |
| self.image_id_to_filename = {} | |
| for image in data["images"]: | |
| file_name = image["file_name"] | |
| image_id = image["id"] | |
| self.image_id_to_filename[image_id] = os.path.join(image_dir, file_name) | |
| def __len__(self): | |
| return len(self.annotations) | |
| def __iter__(self): | |
| for anno in self.annotations: | |
| image_id = anno["image_id"] | |
| caption = anno["caption"] | |
| try: | |
| image = Image.open(self.image_id_to_filename[image_id]) | |
| except: | |
| continue | |
| yield [self.dataset_name, image, caption] | |
| self.is_end = True | |
| class KarpathyCOCOGenerator(Generator): | |
| def __init__(self, anno="/gpfs/u/home/LMCG/LMCGljnn/scratch/code/multimodal/tools/coco_karpathy_train.json", image_dir="/gpfs/u/home/LMCG/LMCGljnn/scratch/.cache/lavis/coco/images"): | |
| super().__init__(dataset_name="coco") | |
| data = json.loads(open(anno).read()) | |
| self.annotations = data | |
| self.image_id_to_filename = {} | |
| for d in data: | |
| self.image_id_to_filename[d["image_id"]] = os.path.join(image_dir, d["image"]) | |
| def __len__(self): | |
| return len(self.annotations) | |
| def __iter__(self): | |
| for anno in self.annotations: | |
| image_id = anno["image_id"] | |
| caption = anno["caption"] | |
| try: | |
| image = Image.open(self.image_id_to_filename[image_id]) | |
| except: | |
| print(self.image_id_to_filename[image_id]) | |
| yield [self.dataset_name, image, caption] | |
| self.is_end = True | |
| class VisualGenomeGenerator(Generator): | |
| def __init__(self, root: str): | |
| super().__init__(dataset_name="vg") | |
| data = json.loads(open(os.path.join(root, "region_descriptions.json")).read()) | |
| image_data = json.loads(open(os.path.join(root, "image_data.json")).read()) | |
| self.image_id_to_filename = {} | |
| self.image_id_to_wh = {} | |
| for image in image_data: | |
| image_id = image["image_id"] | |
| subfolder, filename = image['url'].split("/")[-2:] | |
| self.image_id_to_filename[image_id] = os.path.join(root, subfolder, filename) | |
| self.image_id_to_wh[image_id] = (image["width"], image["height"]) | |
| self.regions = [] | |
| total = 0 | |
| total_image = 0 | |
| used_image = 0 | |
| for xx in data: | |
| total_image += 1 | |
| flag = False | |
| for region in xx['regions']: | |
| total += 1 | |
| region_w = int(region["width"]) | |
| region_h = int(region["height"]) | |
| image_w = self.image_id_to_wh[region["image_id"]][0] | |
| image_h = self.image_id_to_wh[region["image_id"]][1] | |
| if region_w * region_h < (image_w * image_h) * 0.2: | |
| continue | |
| self.regions.append(region) | |
| flag = True | |
| if flag: | |
| used_image += 1 | |
| print("valid region", len(self.regions), total, len(self.regions) / total) | |
| print("valid image", used_image, total_image, used_image / total_image) | |
| def __len__(self): | |
| return len(self.regions) | |
| def __iter__(self): | |
| for region in self.regions: | |
| image_id = region["image_id"] | |
| phrase = region["phrase"] | |
| try: | |
| image = Image.open(self.image_id_to_filename[image_id]) | |
| except: | |
| continue | |
| yield [self.dataset_name, image, phrase] | |
| self.is_end = True | |
| class ShuffleGenerator(): | |
| def __init__(self, generators: List[Generator], p: List[int]): | |
| self.generators = generators | |
| self.p = list(np.array(p) / sum(p)) | |
| self.ids = list(range(len(self.generators))) | |
| print("rebalance", self.ids, self.p) | |
| def __len__(self): | |
| return sum([len(g) for g in self.generators]) | |
| def __iter__(self): | |
| while True: | |
| if len(self.ids) == 0: | |
| break | |
| id = np.random.choice(self.ids, p=self.p) | |
| gen = self.generators[id] | |
| if gen.is_end: | |
| print(gen.dataset_name, "is all done") | |
| del self.ids[id] | |
| del self.p[id] | |
| self.p = list(np.array(self.p) / sum(p)) | |
| print("rebalance", self.ids, self.p) | |
| else: | |
| return iter(gen) | |
| if __name__ == "__main__": | |
| OUT_DIR = "/gpfs/u/home/LMCG/LMCGljnn/scratch-shared/junyan/raw/karpathy_coco_wds_full" | |
| os.makedirs(OUT_DIR, exist_ok=True) | |
| # cc3m_generator = CC3MGenerator("/gpfs/u/home/LMCG/LMCGljnn/scratch-shared/junyan/raw/cc3m") | |
| # cc12m_generator = CC12MGenerator("/gpfs/u/home/LMCG/LMCGljnn/scratch-shared/junyan/raw/cc12m/tars") | |
| coco_generator = KarpathyCOCOGenerator() | |
| # visual_genome_generator = VisualGenomeGenerator("/gpfs/u/home/LMCG/LMCGljnn/scratch/datasets/raw/vg") | |
| # generators = [cc3m_generator, cc12m_generator, coco_generator, visual_genome_generator] | |
| # p = [len(generator) for generator in generators] | |
| # dataset = ShuffleGenerator(generators, p) | |
| with wds.ShardWriter(os.path.join(OUT_DIR, "%05d.tar"), maxcount=8500) as sink: | |
| sink.verbose = False | |
| for i, data in enumerate(tqdm(coco_generator)): | |
| dataset_name, image, caption = data | |
| sink.write({"__key__": f"{dataset_name}_{i}", "jpg": image, "txt": caption}) | |