Spaces:
Paused
Paused
| from torch.utils.data import Dataset | |
| from PIL import Image | |
| import os | |
| import io | |
| import json | |
| import random | |
| import torch | |
| try: | |
| from aoss_client.client import Client | |
| except: | |
| try: | |
| from petrel_client.client import Client | |
| except: | |
| Client = None | |
| from glob import glob | |
| from xtuner.registry import BUILDER | |
| from src.datasets.utils import crop2square, resize_image_fix_pixels, resize_image_dynamic | |
| from einops import rearrange | |
| import numpy as np | |
| class CaptionDataset(Dataset): | |
| def __init__(self, | |
| data_path, | |
| image_folder=None, | |
| debug=False, | |
| image_processor=None, | |
| image_process='crop2square', | |
| ceph_folder=None, | |
| latents_ceph_folder=None, | |
| ceph_config=None, | |
| tokenizer=None, | |
| prompt_template=None, | |
| max_length=2048, | |
| min_image_size=80, | |
| image_size=256, | |
| image_length=256, | |
| unit_image_size=32, | |
| image_tokens_folder=None, | |
| image_latents_folder=None, | |
| cap_folder=None, | |
| cap_source='caption', | |
| tokenizer_kwargs=dict(add_special_tokens=True), | |
| unconditional=0.1 | |
| ): | |
| super().__init__() | |
| self.data_path = data_path | |
| self._load_data(data_path) | |
| self.image_folder = image_folder | |
| self.cap_folder = cap_folder | |
| self.cap_source = cap_source | |
| self.debug = debug | |
| if image_processor is not None: | |
| self.image_processor = BUILDER.build(image_processor) | |
| else: | |
| self.image_processor = None | |
| if tokenizer is not None: | |
| self.tokenizer = BUILDER.build(tokenizer) | |
| else: | |
| self.tokenizer = None | |
| self.prompt_template = prompt_template | |
| self.max_length = max_length | |
| self.image_process = image_process | |
| self.image_length = image_length | |
| self.image_tokens_folder = image_tokens_folder | |
| self.image_latents_folder = image_latents_folder | |
| self.min_image_size = min_image_size | |
| self.image_size = image_size | |
| self.unit_image_size = unit_image_size | |
| self.unconditional = unconditional | |
| self.tokenizer_kwargs = tokenizer_kwargs | |
| self.FILE_CLIENT = None | |
| self.ceph_folder = ceph_folder | |
| self.ceph_config = ceph_config | |
| self.latents_ceph_folder = latents_ceph_folder | |
| self.use_ceph = ((Client is not None) and (ceph_config is not None) and os.path.exists(ceph_config)) | |
| def _load_data(self, data_path: str): # image path and annotation path are saved in a json file | |
| if data_path.endswith('.json'): | |
| with open(data_path, 'r') as f: | |
| self.data_list = json.load(f) | |
| else: | |
| json_files = glob(f"{data_path}/*.json") | |
| data_list = [] | |
| for json_file in json_files: | |
| with open(json_file, 'r') as f: | |
| data_list += json.load(f) | |
| self.data_list = data_list | |
| print(f"Load {len(self.data_list)} data samples from {data_path}", flush=True) | |
| def __len__(self): | |
| return len(self.data_list) | |
| def _read_ceph(self, ceph_path): | |
| if self.FILE_CLIENT is None: | |
| self.FILE_CLIENT = Client(self.ceph_config) | |
| data_bytes = self.FILE_CLIENT.get(ceph_path) | |
| return io.BytesIO(data_bytes) | |
| def _read_image(self, image_file): | |
| if self.image_folder is None: | |
| assert self.use_ceph | |
| assert self.ceph_folder is not None | |
| image = Image.open( | |
| self._read_ceph( | |
| os.path.join(self.ceph_folder, image_file) | |
| ) | |
| ) | |
| else: | |
| image = Image.open( | |
| os.path.join(self.image_folder, image_file) | |
| ) | |
| assert image.width > self.min_image_size and image.height > self.min_image_size, f"Image: {image.size}" | |
| assert image.width / image.height > 0.1, f"Image: {image.size}" | |
| assert image.width / image.height < 10, f"Image: {image.size}" | |
| return image.convert('RGB') | |
| def _read_json(self, annotation_file): | |
| if self.cap_folder is None: | |
| assert self.use_ceph | |
| assert self.ceph_folder is not None | |
| annotation = json.load( | |
| self._read_ceph( | |
| os.path.join(self.ceph_folder, annotation_file) | |
| ) | |
| ) | |
| else: | |
| with open(os.path.join(self.cap_folder, annotation_file), 'r') as f: | |
| annotation = json.load(f) | |
| return annotation | |
| def _process_image(self, image): | |
| data = dict() | |
| if self.image_process == 'crop2square': | |
| image = crop2square(image) | |
| image = image.resize(size=(self.image_size, self.image_size)) | |
| elif self.image_process == 'dynamic': # dynamic and make sure the largest edge <= self.image_size | |
| image = resize_image_dynamic(x=image, image_size=self.image_size, unit_image_size=self.unit_image_size) | |
| elif self.image_process == 'fix_pixels': # fix pixels contain radio of image | |
| # import pdb; pdb.set_trace() | |
| image = resize_image_fix_pixels(x=image, image_size=self.image_size, unit_image_size=self.unit_image_size) | |
| elif self.image_process == 'resize2square': | |
| image = image.resize(size=(self.image_size, self.image_size)) | |
| else: | |
| raise NotImplementedError | |
| # assert image.width <= self.image_size | |
| # assert image.height <= self.image_size | |
| assert image.width % self.unit_image_size == 0 | |
| assert image.height % self.unit_image_size == 0 | |
| pixel_values = torch.from_numpy(np.array(image)).float() | |
| pixel_values = pixel_values / 255 | |
| pixel_values = 2 * pixel_values - 1 | |
| pixel_values = rearrange(pixel_values, 'h w c -> c h w') | |
| data.update(pixel_values=pixel_values) | |
| return data | |
| def _process_text(self, text): | |
| if self.tokenizer is None: | |
| return {} | |
| if random.uniform(0, 1) < self.unconditional: | |
| prompt = self.prompt_template['CFG'] | |
| else: | |
| prompt = self.prompt_template['GENERATION'].format(input=text.strip()) | |
| prompt = self.prompt_template['INSTRUCTION'].format(input=prompt) | |
| if self.prompt_template.get('IMG_START_TOKEN_FOR_GENERATION', True): | |
| prompt += self.prompt_template['IMG_START_TOKEN'] | |
| input_ids = self.tokenizer.encode(prompt, return_tensors='pt', **self.tokenizer_kwargs)[0] | |
| return dict(input_ids=input_ids[:self.max_length]) | |
| def _retry(self): | |
| return self.__getitem__(random.choice(range(self.__len__()))) | |
| def __getitem__(self, idx): | |
| if self.debug: | |
| idx = 0 | |
| try: | |
| data_sample = self.data_list[idx] | |
| if self.image_tokens_folder is not None: | |
| image_tokens = torch.load(os.path.join(self.image_tokens_folder, | |
| data_sample['image'] + '.pt')).long() | |
| data = dict(image_tokens=image_tokens) | |
| elif self.latents_ceph_folder is not None: | |
| image_latents = torch.load( | |
| self._read_ceph( | |
| os.path.join( | |
| self.latents_ceph_folder, data_sample['image'] + '.pt' | |
| ) | |
| ) | |
| ) | |
| data = dict(image_latents=image_latents) | |
| elif self.image_latents_folder is not None: | |
| image_latents = torch.load(os.path.join(self.image_latents_folder, | |
| data_sample['image'] + '.pt')) | |
| data = dict(image_latents=image_latents) | |
| else: | |
| image = self._read_image(data_sample['image']).convert('RGB') | |
| data = self._process_image(image) | |
| caption = self._read_json(data_sample['annotation'])[self.cap_source] | |
| # caption = self._read_json(data_sample['annotation']) | |
| # print(caption) | |
| data.update(self._process_text(caption)) | |
| data['pixel_init'] = image | |
| data.update(image_dir=self.image_folder, image_file=data_sample['image'], | |
| type='text2image',text=caption) | |
| return data | |
| except Exception as e: | |
| print(f"Error when reading {self.data_path}:{self.data_list[idx]}: {e}", flush=True) | |
| return self._retry() | |