modules/ace_plus_dataset.py

# -*- 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