| import random |
| from pathlib import Path |
| from typing import Dict, List, Optional, Tuple, Union |
|
|
| from PIL import Image |
| from torch import zeros_like |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
| import glob |
| from .preprocess_files import face_mask_google_mediapipe |
|
|
| OBJECT_TEMPLATE = [ |
| "a photo of a {}", |
| "a rendering of a {}", |
| "a cropped photo of the {}", |
| "the photo of a {}", |
| "a photo of a clean {}", |
| "a photo of a dirty {}", |
| "a dark photo of the {}", |
| "a photo of my {}", |
| "a photo of the cool {}", |
| "a close-up photo of a {}", |
| "a bright photo of the {}", |
| "a cropped photo of a {}", |
| "a photo of the {}", |
| "a good photo of the {}", |
| "a photo of one {}", |
| "a close-up photo of the {}", |
| "a rendition of the {}", |
| "a photo of the clean {}", |
| "a rendition of a {}", |
| "a photo of a nice {}", |
| "a good photo of a {}", |
| "a photo of the nice {}", |
| "a photo of the small {}", |
| "a photo of the weird {}", |
| "a photo of the large {}", |
| "a photo of a cool {}", |
| "a photo of a small {}", |
| ] |
|
|
| STYLE_TEMPLATE = [ |
| "a painting in the style of {}", |
| "a rendering in the style of {}", |
| "a cropped painting in the style of {}", |
| "the painting in the style of {}", |
| "a clean painting in the style of {}", |
| "a dirty painting in the style of {}", |
| "a dark painting in the style of {}", |
| "a picture in the style of {}", |
| "a cool painting in the style of {}", |
| "a close-up painting in the style of {}", |
| "a bright painting in the style of {}", |
| "a cropped painting in the style of {}", |
| "a good painting in the style of {}", |
| "a close-up painting in the style of {}", |
| "a rendition in the style of {}", |
| "a nice painting in the style of {}", |
| "a small painting in the style of {}", |
| "a weird painting in the style of {}", |
| "a large painting in the style of {}", |
| ] |
|
|
| NULL_TEMPLATE = ["{}"] |
|
|
| TEMPLATE_MAP = { |
| "object": OBJECT_TEMPLATE, |
| "style": STYLE_TEMPLATE, |
| "null": NULL_TEMPLATE, |
| } |
|
|
|
|
| def _randomset(lis): |
| ret = [] |
| for i in range(len(lis)): |
| if random.random() < 0.5: |
| ret.append(lis[i]) |
| return ret |
|
|
|
|
| def _shuffle(lis): |
|
|
| return random.sample(lis, len(lis)) |
|
|
|
|
| def _get_cutout_holes( |
| height, |
| width, |
| min_holes=8, |
| max_holes=32, |
| min_height=16, |
| max_height=128, |
| min_width=16, |
| max_width=128, |
| ): |
| holes = [] |
| for _n in range(random.randint(min_holes, max_holes)): |
| hole_height = random.randint(min_height, max_height) |
| hole_width = random.randint(min_width, max_width) |
| y1 = random.randint(0, height - hole_height) |
| x1 = random.randint(0, width - hole_width) |
| y2 = y1 + hole_height |
| x2 = x1 + hole_width |
| holes.append((x1, y1, x2, y2)) |
| return holes |
|
|
|
|
| def _generate_random_mask(image): |
| mask = zeros_like(image[:1]) |
| holes = _get_cutout_holes(mask.shape[1], mask.shape[2]) |
| for (x1, y1, x2, y2) in holes: |
| mask[:, y1:y2, x1:x2] = 1.0 |
| if random.uniform(0, 1) < 0.25: |
| mask.fill_(1.0) |
| masked_image = image * (mask < 0.5) |
| return mask, masked_image |
|
|
|
|
| class PivotalTuningDatasetCapation(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, |
| instance_data_root, |
| tokenizer, |
| token_map: Optional[dict] = None, |
| use_template: Optional[str] = None, |
| size=512, |
| h_flip=True, |
| color_jitter=False, |
| resize=True, |
| use_mask_captioned_data=False, |
| use_face_segmentation_condition=False, |
| train_inpainting=False, |
| blur_amount: int = 70, |
| ): |
| self.size = size |
| self.tokenizer = tokenizer |
| self.resize = resize |
| self.train_inpainting = train_inpainting |
|
|
| instance_data_root = Path(instance_data_root) |
| if not instance_data_root.exists(): |
| raise ValueError("Instance images root doesn't exists.") |
|
|
| self.instance_images_path = [] |
| self.mask_path = [] |
|
|
| assert not ( |
| use_mask_captioned_data and use_template |
| ), "Can't use both mask caption data and template." |
|
|
| |
| if use_mask_captioned_data: |
| src_imgs = glob.glob(str(instance_data_root) + "/*src.jpg") |
| for f in src_imgs: |
| idx = int(str(Path(f).stem).split(".")[0]) |
| mask_path = f"{instance_data_root}/{idx}.mask.png" |
|
|
| if Path(mask_path).exists(): |
| self.instance_images_path.append(f) |
| self.mask_path.append(mask_path) |
| else: |
| print(f"Mask not found for {f}") |
|
|
| self.captions = open(f"{instance_data_root}/caption.txt").readlines() |
|
|
| else: |
| possibily_src_images = ( |
| glob.glob(str(instance_data_root) + "/*.jpg") |
| + glob.glob(str(instance_data_root) + "/*.png") |
| + glob.glob(str(instance_data_root) + "/*.jpeg") |
| ) |
| possibily_src_images = ( |
| set(possibily_src_images) |
| - set(glob.glob(str(instance_data_root) + "/*mask.png")) |
| - set([str(instance_data_root) + "/caption.txt"]) |
| ) |
|
|
| self.instance_images_path = list(set(possibily_src_images)) |
| self.captions = [ |
| x.split("/")[-1].split(".")[0] for x in self.instance_images_path |
| ] |
|
|
| assert ( |
| len(self.instance_images_path) > 0 |
| ), "No images found in the instance data root." |
|
|
| self.instance_images_path = sorted(self.instance_images_path) |
|
|
| self.use_mask = use_face_segmentation_condition or use_mask_captioned_data |
| self.use_mask_captioned_data = use_mask_captioned_data |
|
|
| if use_face_segmentation_condition: |
|
|
| for idx in range(len(self.instance_images_path)): |
| targ = f"{instance_data_root}/{idx}.mask.png" |
| |
| if not Path(targ).exists(): |
| print(f"Mask not found for {targ}") |
|
|
| print( |
| "Warning : this will pre-process all the images in the instance data root." |
| ) |
|
|
| if len(self.mask_path) > 0: |
| print( |
| "Warning : masks already exists, but will be overwritten." |
| ) |
|
|
| masks = face_mask_google_mediapipe( |
| [ |
| Image.open(f).convert("RGB") |
| for f in self.instance_images_path |
| ] |
| ) |
| for idx, mask in enumerate(masks): |
| mask.save(f"{instance_data_root}/{idx}.mask.png") |
|
|
| break |
|
|
| for idx in range(len(self.instance_images_path)): |
| self.mask_path.append(f"{instance_data_root}/{idx}.mask.png") |
|
|
| self.num_instance_images = len(self.instance_images_path) |
| self.token_map = token_map |
|
|
| self.use_template = use_template |
| if use_template is not None: |
| self.templates = TEMPLATE_MAP[use_template] |
|
|
| self._length = self.num_instance_images |
|
|
| self.h_flip = h_flip |
| self.image_transforms = transforms.Compose( |
| [ |
| transforms.Resize( |
| size, interpolation=transforms.InterpolationMode.BILINEAR |
| ) |
| if resize |
| else transforms.Lambda(lambda x: x), |
| transforms.ColorJitter(0.1, 0.1) |
| if color_jitter |
| else transforms.Lambda(lambda x: x), |
| transforms.CenterCrop(size), |
| transforms.ToTensor(), |
| transforms.Normalize([0.5], [0.5]), |
| ] |
| ) |
|
|
| self.blur_amount = blur_amount |
|
|
| def __len__(self): |
| return self._length |
|
|
| def __getitem__(self, index): |
| example = {} |
| instance_image = Image.open( |
| self.instance_images_path[index % self.num_instance_images] |
| ) |
| if not instance_image.mode == "RGB": |
| instance_image = instance_image.convert("RGB") |
| example["instance_images"] = self.image_transforms(instance_image) |
|
|
| if self.train_inpainting: |
| ( |
| example["instance_masks"], |
| example["instance_masked_images"], |
| ) = _generate_random_mask(example["instance_images"]) |
|
|
| if self.use_template: |
| assert self.token_map is not None |
| input_tok = list(self.token_map.values())[0] |
|
|
| text = random.choice(self.templates).format(input_tok) |
| else: |
| text = self.captions[index % self.num_instance_images].strip() |
|
|
| if self.token_map is not None: |
| for token, value in self.token_map.items(): |
| text = text.replace(token, value) |
|
|
| print(text) |
|
|
| if self.use_mask: |
| example["mask"] = ( |
| self.image_transforms( |
| Image.open(self.mask_path[index % self.num_instance_images]) |
| ) |
| * 0.5 |
| + 1.0 |
| ) |
|
|
| if self.h_flip and random.random() > 0.5: |
| hflip = transforms.RandomHorizontalFlip(p=1) |
|
|
| example["instance_images"] = hflip(example["instance_images"]) |
| if self.use_mask: |
| example["mask"] = hflip(example["mask"]) |
|
|
| example["instance_prompt_ids"] = self.tokenizer( |
| text, |
| padding="do_not_pad", |
| truncation=True, |
| max_length=self.tokenizer.model_max_length, |
| ).input_ids |
|
|
| return example |
|
|
