# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import math import re, io import numpy as np import random, torch from PIL import Image import torchvision.transforms as T from collections import defaultdict from scepter.modules.data.dataset.registry import DATASETS from scepter.modules.data.dataset.base_dataset import BaseDataset from scepter.modules.transform.io import pillow_convert from scepter.modules.utils.directory import osp_path from scepter.modules.utils.file_system import FS from torchvision.transforms import InterpolationMode def load_image(prefix, img_path, cvt_type=None): if img_path is None or img_path == '': return None img_path = osp_path(prefix, img_path) with FS.get_object(img_path) as image_bytes: image = Image.open(io.BytesIO(image_bytes)) if cvt_type is not None: image = pillow_convert(image, cvt_type) return image def transform_image(image, std = 0.5, mean = 0.5): return (image.permute(2, 0, 1)/255. - mean)/std def transform_mask(mask): return mask.unsqueeze(0)/255. def ensure_src_align_target_h_mode(src_image, size, image_id, interpolation=InterpolationMode.BILINEAR): # padding mode H, W = size ret_image = [] for one_id in image_id: edit_image = src_image[one_id] _, eH, eW = edit_image.shape scale = H/eH tH, tW = H, int(eW * scale) ret_image.append(T.Resize((tH, tW), interpolation=interpolation, antialias=True)(edit_image)) return ret_image def ensure_src_align_target_padding_mode(src_image, size, image_id, size_h = [], interpolation=InterpolationMode.BILINEAR): # padding mode H, W = size ret_data = [] ret_h = [] for idx, one_id in enumerate(image_id): if len(size_h) < 1: rH = random.randint(int(H / 3), int(H)) else: rH = size_h[idx] ret_h.append(rH) edit_image = src_image[one_id] _, eH, eW = edit_image.shape scale = rH/eH tH, tW = rH, int(eW * scale) edit_image = T.Resize((tH, tW), interpolation=interpolation, antialias=True)(edit_image) # padding delta_w = 0 delta_h = H - tH padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2)) ret_data.append(T.Pad(padding, fill=0, padding_mode="constant")(edit_image).float()) return ret_data, ret_h def ensure_limit_sequence(image, max_seq_len = 4096, d = 16, interpolation=InterpolationMode.BILINEAR): # resize image for max_seq_len, while keep the aspect ratio H, W = image.shape[-2:] scale = min(1.0, math.sqrt(max_seq_len / ((H / d) * (W / d)))) rH = int(H * scale) // d * d # ensure divisible by self.d rW = int(W * scale) // d * d # print(f"{H} {W} -> {rH} {rW}") image = T.Resize((rH, rW), interpolation=interpolation, antialias=True)(image) return image @DATASETS.register_class() class ACEPlusDataset(BaseDataset): para_dict = { "DELIMITER": { "value": "#;#", "description": "The delimiter for records of data list." }, "FIELDS": { "value": ["data_type", "edit_image", "edit_mask", "ref_image", "target_image", "prompt"], "description": "The fields for every record." }, "PATH_PREFIX": { "value": "", "description": "The path prefix for every input image." }, "EDIT_TYPE_LIST": { "value": [], "description": "The edit type list to be trained for data list." }, "MAX_SEQ_LEN": { "value": 4096, "description": "The max sequence length for input image." }, "D": { "value": 16, "description": "Patch size for resized image." } } para_dict.update(BaseDataset.para_dict) def __init__(self, cfg, logger=None): super().__init__(cfg, logger=logger) delimiter = cfg.get("DELIMITER", "#;#") fields = cfg.get("FIELDS", []) prefix = cfg.get("PATH_PREFIX", "") edit_type_list = cfg.get("EDIT_TYPE_LIST", []) self.modify_mode = cfg.get("MODIFY_MODE", True) self.max_seq_len = cfg.get("MAX_SEQ_LEN", 4096) self.repaiting_scale = cfg.get("REPAINTING_SCALE", 0.5) self.d = cfg.get("D", 16) prompt_file = cfg.DATA_LIST self.items = self.read_data_list(delimiter, fields, prefix, edit_type_list, prompt_file) random.shuffle(self.items) use_num = int(cfg.get('USE_NUM', -1)) if use_num > 0: self.items = self.items[:use_num] def read_data_list(self, delimiter, fields, prefix, edit_type_list, prompt_file): with FS.get_object(prompt_file) as local_data: rows = local_data.decode('utf-8').strip().split('\n') items = list() dtype_level_num = {} for i, row in enumerate(rows): item = {"prefix": prefix} for key, val in zip(fields, row.split(delimiter)): item[key] = val edit_type = item["data_type"] if len(edit_type_list) > 0: for re_pattern in edit_type_list: if re.match(re_pattern, edit_type): items.append(item) if edit_type not in dtype_level_num: dtype_level_num[edit_type] = 0 dtype_level_num[edit_type] += 1 break else: items.append(item) if edit_type not in dtype_level_num: dtype_level_num[edit_type] = 0 dtype_level_num[edit_type] += 1 for edit_type in dtype_level_num: self.logger.info(f"{edit_type} has {dtype_level_num[edit_type]} samples.") return items def __len__(self): return len(self.items) def __getitem__(self, index): item = self._get(index) return self.pipeline(item) def _get(self, index): # normalize sample_id = index%len(self) index = self.items[index%len(self)] prefix = index.get("prefix", "") edit_image = index.get("edit_image", "") edit_mask = index.get("edit_mask", "") ref_image = index.get("ref_image", "") target_image = index.get("target_image", "") prompt = index.get("prompt", "") edit_image = load_image(prefix, edit_image, cvt_type="RGB") if edit_image != "" else None edit_mask = load_image(prefix, edit_mask, cvt_type="L") if edit_mask != "" else None ref_image = load_image(prefix, ref_image, cvt_type="RGB") if ref_image != "" else None target_image = load_image(prefix, target_image, cvt_type="RGB") if target_image != "" else None assert target_image is not None edit_id, ref_id, src_image_list, src_mask_list = [], [], [], [] # parse editing image if edit_image is None: edit_image = Image.new("RGB", target_image.size, (255, 255, 255)) edit_mask = Image.new("L", edit_image.size, 255) elif edit_mask is None: edit_mask = Image.new("L", edit_image.size, 255) src_image_list.append(edit_image) edit_id.append(0) src_mask_list.append(edit_mask) # parse reference image if ref_image is not None: src_image_list.append(ref_image) ref_id.append(1) src_mask_list.append(Image.new("L", ref_image.size, 0)) image = transform_image(torch.tensor(np.array(target_image).astype(np.float32))) if edit_mask is not None: image_mask = transform_mask(torch.tensor(np.array(edit_mask).astype(np.float32))) else: image_mask = Image.new("L", target_image.size, 255) image_mask = transform_mask(torch.tensor(np.array(image_mask).astype(np.float32))) src_image_list = [transform_image(torch.tensor(np.array(im).astype(np.float32))) for im in src_image_list] src_mask_list = [transform_mask(torch.tensor(np.array(im).astype(np.float32))) for im in src_mask_list] # decide the repainting scale for the editing task if len(ref_id) > 0: repainting_scale = 1.0 else: repainting_scale = self.repaiting_scale for e_i in edit_id: src_image_list[e_i] = src_image_list[e_i] * (1 - repainting_scale * src_mask_list[e_i]) size = image.shape[1:] ref_image_list, ret_h = ensure_src_align_target_padding_mode(src_image_list, size, image_id=ref_id, interpolation=InterpolationMode.NEAREST_EXACT) ref_mask_list, ret_h = ensure_src_align_target_padding_mode(src_mask_list, size, size_h=ret_h, image_id=ref_id, interpolation=InterpolationMode.NEAREST_EXACT) edit_image_list = ensure_src_align_target_h_mode(src_image_list, size, image_id=edit_id, interpolation=InterpolationMode.NEAREST_EXACT) edit_mask_list = ensure_src_align_target_h_mode(src_mask_list, size, image_id=edit_id, interpolation=InterpolationMode.NEAREST_EXACT) src_image_list = [torch.cat(ref_image_list + edit_image_list, dim=-1)] src_mask_list = [torch.cat(ref_mask_list + edit_mask_list, dim=-1)] image = torch.cat(ref_image_list + [image], dim=-1) image_mask = torch.cat(ref_mask_list + [image_mask], dim=-1) # limit max sequence length image = ensure_limit_sequence(image, max_seq_len = self.max_seq_len, d = self.d, interpolation=InterpolationMode.BILINEAR) image_mask = ensure_limit_sequence(image_mask, max_seq_len = self.max_seq_len, d = self.d, interpolation=InterpolationMode.NEAREST_EXACT) src_image_list = [ensure_limit_sequence(i, max_seq_len = self.max_seq_len, d = self.d, interpolation=InterpolationMode.BILINEAR) for i in src_image_list] src_mask_list = [ensure_limit_sequence(i, max_seq_len = self.max_seq_len, d = self.d, interpolation=InterpolationMode.NEAREST_EXACT) for i in src_mask_list] if self.modify_mode: # To be modified regions according to mask modify_image_list = [ii * im for ii, im in zip(src_image_list, src_mask_list)] # To be edited regions according to mask src_image_list = [ii * (1 - im) for ii, im in zip(src_image_list, src_mask_list)] else: src_image_list = src_image_list modify_image_list = src_image_list item = { "src_image_list": src_image_list, "src_mask_list": src_mask_list, "modify_image_list": modify_image_list, "image": image, "image_mask": image_mask, "edit_id": edit_id, "ref_id": ref_id, "prompt": prompt, "edit_key": index["edit_key"] if "edit_key" in index else "", "sample_id": sample_id } return item @staticmethod def collate_fn(batch): collect = defaultdict(list) for sample in batch: for k, v in sample.items(): collect[k].append(v) new_batch = dict() for k, v in collect.items(): if all([i is None for i in v]): new_batch[k] = None else: new_batch[k] = v return new_batch