|
|
|
import json |
|
import os |
|
|
|
import torch |
|
from datasets import Dataset as HFDataset |
|
from datasets import DatasetDict |
|
from mmengine.config import Config, ConfigDict |
|
from PIL import Image |
|
from torch.utils.data import Dataset |
|
|
|
from xtuner.registry import BUILDER |
|
from .huggingface import process_hf_dataset |
|
from .utils import expand2square |
|
|
|
|
|
class LLaVADataset(Dataset): |
|
|
|
def __init__(self, |
|
data_path, |
|
image_folder, |
|
tokenizer, |
|
image_processor, |
|
max_dataset_length=None, |
|
dataset_map_fn=None, |
|
template_map_fn=None, |
|
max_length=2048, |
|
pad_image_to_square=False): |
|
super().__init__() |
|
|
|
json_data = json.load(open(data_path)) |
|
for idx in range(len(json_data)): |
|
if isinstance(json_data[idx]['id'], int): |
|
json_data[idx]['id'] = str(json_data[idx]['id']) |
|
json_data = DatasetDict({'train': HFDataset.from_list(json_data)}) |
|
self.text_data = process_hf_dataset( |
|
dataset=json_data, |
|
tokenizer=tokenizer, |
|
max_length=max_length, |
|
dataset_map_fn=dataset_map_fn, |
|
template_map_fn=template_map_fn, |
|
split='train', |
|
max_dataset_length=max_dataset_length, |
|
remove_unused_columns=False, |
|
pack_to_max_length=False, |
|
with_image_token=True) |
|
|
|
self.image_folder = image_folder |
|
if isinstance(image_processor, dict) or isinstance( |
|
image_processor, Config) or isinstance(image_processor, |
|
ConfigDict): |
|
self.image_processor = BUILDER.build(image_processor) |
|
else: |
|
self.image_processor = image_processor |
|
self.pad_image_to_square = pad_image_to_square |
|
|
|
@property |
|
def modality_length(self): |
|
length_list = [] |
|
for data_dict in self.text_data: |
|
cur_len = len(data_dict['input_ids']) |
|
if data_dict.get('image', None) is None: |
|
cur_len = -cur_len |
|
length_list.append(cur_len) |
|
return length_list |
|
|
|
def __len__(self): |
|
return len(self.text_data) |
|
|
|
def __getitem__(self, index): |
|
data_dict = self.text_data[index] |
|
if data_dict.get('image', None) is not None: |
|
image_file = data_dict['image'] |
|
image = Image.open(os.path.join(self.image_folder, |
|
image_file)).convert('RGB') |
|
if self.pad_image_to_square: |
|
image = expand2square( |
|
image, |
|
tuple( |
|
int(x * 255) for x in self.image_processor.image_mean)) |
|
image = self.image_processor.preprocess( |
|
image, return_tensors='pt')['pixel_values'][0] |
|
data_dict['pixel_values'] = image |
|
else: |
|
size = self.image_processor.size |
|
data_dict['pixel_values'] = torch.zeros(3, size['height'], |
|
size['width']) |
|
return data_dict |
|
|
