| import json | |
| import os | |
| import random | |
| import re | |
| from pathlib import Path | |
| from PIL import Image | |
| from torch.utils.data import Dataset | |
| from mixofshow.data.pil_transform import PairCompose, build_transform | |
| class LoraDataset(Dataset): | |
| """ | |
| A dataset to prepare the instance and class images with the prompts for fine-tuning the model. | |
| It pre-processes the images and the tokenizes prompts. | |
| """ | |
| def __init__(self, opt): | |
| self.opt = opt | |
| self.instance_images_path = [] | |
| with open(opt['concept_list'], 'r') as f: | |
| concept_list = json.load(f) | |
| replace_mapping = opt.get('replace_mapping', {}) | |
| use_caption = opt.get('use_caption', False) | |
| use_mask = opt.get('use_mask', False) | |
| for concept in concept_list: | |
| instance_prompt = concept['instance_prompt'] | |
| caption_dir = concept.get('caption_dir') | |
| mask_dir = concept.get('mask_dir') | |
| instance_prompt = self.process_text(instance_prompt, replace_mapping) | |
| inst_img_path = [] | |
| for x in Path(concept['instance_data_dir']).iterdir(): | |
| if x.is_file() and x.name != '.DS_Store': | |
| basename = os.path.splitext(os.path.basename(x))[0] | |
| caption_path = os.path.join(caption_dir, f'{basename}.txt') if caption_dir is not None else None | |
| if use_caption and caption_path is not None and os.path.exists(caption_path): | |
| with open(caption_path, 'r') as fr: | |
| line = fr.readlines()[0] | |
| instance_prompt_image = self.process_text(line, replace_mapping) | |
| else: | |
| instance_prompt_image = instance_prompt | |
| if use_mask and mask_dir is not None: | |
| mask_path = os.path.join(mask_dir, f'{basename}.png') | |
| else: | |
| mask_path = None | |
| inst_img_path.append((x, instance_prompt_image, mask_path)) | |
| self.instance_images_path.extend(inst_img_path) | |
| random.shuffle(self.instance_images_path) | |
| self.num_instance_images = len(self.instance_images_path) | |
| self.instance_transform = PairCompose([ | |
| build_transform(transform_opt) | |
| for transform_opt in opt['instance_transform'] | |
| ]) | |
| def process_text(self, instance_prompt, replace_mapping): | |
| for k, v in replace_mapping.items(): | |
| instance_prompt = instance_prompt.replace(k, v) | |
| instance_prompt = instance_prompt.strip() | |
| instance_prompt = re.sub(' +', ' ', instance_prompt) | |
| return instance_prompt | |
| def __len__(self): | |
| return self.num_instance_images * self.opt['dataset_enlarge_ratio'] | |
| def __getitem__(self, index): | |
| example = {} | |
| instance_image, instance_prompt, instance_mask = self.instance_images_path[index % self.num_instance_images] | |
| instance_image = Image.open(instance_image).convert('RGB') | |
| extra_args = {'prompts': instance_prompt} | |
| if instance_mask is not None: | |
| instance_mask = Image.open(instance_mask).convert('L') | |
| extra_args.update({'mask': instance_mask}) | |
| instance_image, extra_args = self.instance_transform(instance_image, **extra_args) | |
| example['images'] = instance_image | |
| if 'mask' in extra_args: | |
| example['masks'] = extra_args['mask'] | |
| example['masks'] = example['masks'].unsqueeze(0) | |
| else: | |
| pass | |
| if 'img_mask' in extra_args: | |
| example['img_masks'] = extra_args['img_mask'] | |
| example['img_masks'] = example['img_masks'].unsqueeze(0) | |
| else: | |
| raise NotImplementedError | |
| example['prompts'] = extra_args['prompts'] | |
| return example | |