libs/data.py

import numpy
import numpy as np
import torch
import os
import random
import pandas as pd
import os.path as osp
import PIL.Image as Image
from torch.utils.data import Dataset
from pathlib import Path
from imagedream.ldm.util import add_random_background
from imagedream.camera_utils import get_camera_for_index
from libs.base_utils import do_resize_content, add_stroke

import torchvision.transforms as transforms


def to_rgb_image(maybe_rgba: Image.Image):
    if maybe_rgba.mode == "RGB":
        return maybe_rgba
    elif maybe_rgba.mode == "RGBA":
        rgba = maybe_rgba
        img = numpy.random.randint(
            127, 128, size=[rgba.size[1], rgba.size[0], 3], dtype=numpy.uint8
        )
        img = Image.fromarray(img, "RGB")
        img.paste(rgba, mask=rgba.getchannel("A"))
        return img
    else:
        raise ValueError("Unsupported image type.", maybe_rgba.mode)


def axis_rotate_xyz(img: Image.Image, rotate_axis="z", angle=90.0):
    img = img.convert("RGB")
    img = np.array(img) - 127
    img = img.astype(np.float32)
    # perform element-wise sin-cos rotation
    if rotate_axis == "z":
        img = np.stack(
            [
                img[..., 0] * np.cos(angle) - img[..., 1] * np.sin(angle),
                img[..., 0] * np.sin(angle) + img[..., 1] * np.cos(angle),
                img[..., 2],
            ],
            -1,
        )
    elif rotate_axis == "y":
        img = np.stack(
            [
                img[..., 0] * np.cos(angle) + img[..., 2] * np.sin(angle),
                img[..., 1],
                -img[..., 0] * np.sin(angle) + img[..., 2] * np.cos(angle),
            ],
            -1,
        )
    elif rotate_axis == "x":
        img = np.stack(
            [
                img[..., 0],
                img[..., 1] * np.cos(angle) - img[..., 2] * np.sin(angle),
                img[..., 1] * np.sin(angle) + img[..., 2] * np.cos(angle),
            ],
            -1,
        )

    return Image.fromarray(img.astype(np.uint8) + 127)


class DataHQCRelative(Dataset):
    """
    - base_dir
        - uid1
            - 000.png
            - 001.png
            - ...
        - uid2
    - xyz_base
        - uid1
            - xyz_new_000.png
            - xyz_new_001.png
            - ...
    accepte caption data(in csv format)
    """

    def __init__(
        self,
        base_dir,
        caption_csv,
        ref_indexs=[0],
        ref_position=-1,
        xyz_base=None,
        camera_views=[3, 6, 9, 12, 15],  # camera views are relative views, not abs
        split="train",
        image_size=256,
        random_background=False,
        resize_rate=1,
        num_frames=5,
        repeat=100,
        outer_file=None,
        debug=False,
        eval_size=100,
    ):
        print(__class__)
        OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
        OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)
        df = pd.read_csv(caption_csv, sep=",", names=["id", "caption"])
        id_to_caption = {}
        for i in range(len(df.index)):
            item = df.iloc[i]
            id_to_caption[item["id"]] = item["caption"]

        # outer file is txt file, containing each ident per line, specific idents not included in the train process
        outer_set = (
            set(open(outer_file, "r").read().strip().split("\n"))
            if outer_file is not None
            else set()
        )
        xyz_set = set(os.listdir(xyz_base)) if xyz_base is not None else set()
        common_keys = set(id_to_caption.keys()) & set(os.listdir(base_dir))
        common_keys = common_keys & xyz_set if xyz_base is not None else common_keys
        common_keys = common_keys - outer_set
        self.common_keys = common_keys
        self.id_to_caption = id_to_caption
        final_dict = {key: id_to_caption[key] for key in common_keys}
        self.image_size = image_size
        self.base_dir = Path(base_dir)
        self.xyz_base = xyz_base
        self.repeat = repeat
        self.num_frames = num_frames
        self.camera_views = camera_views[:num_frames]
        self.split = split
        self.ref_indexs = ref_indexs
        self.ref_position = ref_position
        self.resize_rate = resize_rate
        self.random_background = random_background
        self.debug = debug
        assert split in ["train", "eval"]

        clip_size = 224
        self.transfrom_clip = transforms.Compose(
            [
                transforms.Resize(
                    (clip_size, clip_size),
                    interpolation=Image.BICUBIC,
                    antialias="warn",
                ),
                transforms.ToTensor(),
                transforms.Normalize(mean=OPENAI_DATASET_MEAN, std=OPENAI_DATASET_STD),
            ]
        )

        self.transfrom_vae = transforms.Compose(
            [
                transforms.Resize((image_size, image_size)),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
            ]
        )
        # 对于第i个视角作为参考时，左边，下面，背面，右边，上面，的图片名称index
        import torchvision.transforms.functional as TF
        from functools import partial as PA

        self.index_mapping = [
            #    正  左 下 背  右 上
            [0, 1, 2, 3, 4, 5],  # 0
            [1, 3, 2, 4, 0, 5],  # 1
            [2, 1, 3, 5, 4, 0],  # 2
            [3, 4, 2, 0, 1, 5],  # 3
            [4, 0, 2, 1, 3, 5],  # 4
            [5, 1, 0, 2, 4, 3],  # 5
        ]
        TT = {
            "r90": PA(TF.rotate, angle=-90.0),  # 顺时针90
            "r180": PA(TF.rotate, angle=-180.0),  # 顺时针180
            "r270": PA(TF.rotate, angle=-270.0),  # 顺时针270
            "s90": PA(TF.rotate, angle=90.0),  # 逆时针90
            "s180": PA(TF.rotate, angle=180.0),  # 逆时针180
            "s270": PA(TF.rotate, angle=270.0),  # 逆时针270
        }

        self.transfroms_mapping = [
            # 正 左 下 背  右 上
            [None, None, None, None, None, None],  # 0
            [None, None, TT["r90"], None, None, TT["s90"]],  # 1
            [None, TT["s90"], TT["s180"], TT["r180"], TT["r90"], None],  # 2
            [None, None, TT["r180"], None, None, TT["s180"]],  # 3
            [None, None, TT["s90"], None, None, TT["r90"]],  # 4
            [None, TT["r90"], None, TT["r180"], TT["s90"], TT["s180"]],  # 5
        ]

        XT = {  # xyz transforms
            "zRrota90": PA(axis_rotate_xyz, rotate_axis="z", angle=np.pi / 2),
            "zSrota90": PA(axis_rotate_xyz, rotate_axis="z", angle=-np.pi / 2),
            "zrota180": PA(axis_rotate_xyz, rotate_axis="z", angle=-np.pi),
            "xRrota90": PA(axis_rotate_xyz, rotate_axis="x", angle=np.pi / 2),
            "xSrota90": PA(axis_rotate_xyz, rotate_axis="x", angle=-np.pi / 2),
        }
        
        self.xyz_transforms_mapping = [
            # 正 左 下 背  右 上
            [None,] * 6, # 0
            [XT["zRrota90"],] * 6, # 1
            [XT["xRrota90"],] * 6, # 2
            [XT["zrota180"],] * 6, # 3
            [XT["zSrota90"],] * 6, # 4
            [XT["xSrota90"],] * 6, # 5
        ]
        
        total_items = [
            {
                "path": os.path.join(base_dir, k),
                "xyz_path": os.path.join(xyz_base, k) if xyz_base is not None else None,
                "caption": v,
            }
            for k, v in final_dict.items()
        ]
        total_items.sort(key=lambda x: x["path"])

        if len(total_items) > eval_size:
            if split == "train":
                self.items = total_items[eval_size:]
            else:
                self.items = total_items[:eval_size]
        else:
            self.items = total_items

        print("============= length of dataset %d =============" % len(self.items))

    def __len__(self):
        return len(self.items) * self.repeat

    def __getitem__(self, index):
        """
        choose index for target 6 images
        select one of them as input image
        target_images_vae: batch of `num_frame` images of one object from different views, processed by vae_processor
        ref_ip: ref image in piexl space
        ref_ip_img:
        camera views decide the logical camera pose of images:
                000 is front , ev: 0, azimuth: 0
                001 is left  , ev: 0, azimuth: -90
                002 is down  , ev: -90, azimuth: 0
                003 is back  , ev: 0, azimuth: 180
                004 is right , ev: 0, azimuth: 90
                005 is top   , ev: 90, azimuth: 0
        ref_index decides which image choose to be input image

        for example when camera views = [1,   2,   3,   4,  5,  0], ref_position=5
        then dataset return the instance images in order as [left, down, back, right, top, front]
        in which view[ref_position] = view[5] = 0, so the refrence image is the front image

        as all the faces can be rotated to the front face, so any image can be placed to ref_position as ref image(need some tramsforms)
        to do a better control of which image can be placed to ref_position, we can set ref_indexs.
        ref_indexs set [0] default, that means only 000 named images will be placed to ref_position.
        on the situation of ref_indexs=[0, 1, 3, 4], only 000, 001, 003, 004 named images will be placed to ref_position.
        """
        index_mapping = self.index_mapping
        transfroms_mapping = self.transfroms_mapping
        index = index % len(self.items)

        target_dir = self.items[index]["path"]
        target_xyz_dir = self.items[index]["xyz_path"]
        caption = self.items[index]["caption"]

        bg_color = np.random.rand() * 255
        target_images = []
        target_xyz_images = []
        raw_images = []
        raw_xyz_images = []
        alpha_masks = []
        ref_index = random.choice(self.ref_indexs)
        cur_index_mapping = index_mapping[ref_index]
        cur_transfroms_mapping = transfroms_mapping[ref_index]
        cur_xyz_transfroms_mapping = self.xyz_transforms_mapping[ref_index]
        for relative_view in self.camera_views:
            image_index = cur_index_mapping[relative_view]
            trans = cur_transfroms_mapping[relative_view]
            trans_xyz = cur_xyz_transfroms_mapping[relative_view]
            # open
            img = Image.open(
                os.path.join(target_dir, f"{image_index:03d}.png")
            ).convert("RGBA")
            if trans is not None:
                img = trans(img)
            img = do_resize_content(img, self.resize_rate)
            alpha_mask = img.getchannel("A")
            alpha_masks.append(alpha_mask)
            if self.random_background:
                img = add_random_background(img, bg_color)
            img = img.convert("RGB")
            target_images.append(self.transfrom_vae(img))

            raw_images.append(img)

            if self.xyz_base is not None:
                img_xyz = Image.open(
                    os.path.join(target_xyz_dir, f"xyz_new_{image_index:03d}.png")
                ).convert("RGBA")
                img_xyz = trans_xyz(img_xyz) if trans_xyz is not None else img_xyz
                img_xyz = trans(img_xyz) if trans is not None else img_xyz
                img_xyz = do_resize_content(img_xyz, self.resize_rate)
                img_xyz.putalpha(alpha_mask)
                if self.random_background:
                    img_xyz = add_random_background(img_xyz, bg_color)
                img_xyz = img_xyz.convert("RGB")
                target_xyz_images.append(self.transfrom_vae(img_xyz))
                if self.debug:
                    raw_xyz_images.append(img_xyz)

        cameras = [get_camera_for_index(i).squeeze() for i in self.camera_views]
        if self.ref_position is not None:
            cameras[self.ref_position] = torch.zeros_like(
                cameras[self.ref_position]
            )  # set ref camera to zero

        cameras = torch.stack(cameras)

        input_img = Image.open(
            os.path.join(target_dir, f"{ref_index:03d}.png")
        ).convert("RGBA")
        input_img = do_resize_content(input_img, self.resize_rate)
        if self.random_background:
            input_img = add_random_background(input_img, bg_color)
        input_img = input_img.convert("RGB")

        clip_cond = self.transfrom_clip(input_img)
        vae_cond = self.transfrom_vae(input_img)

        vae_target = torch.stack(target_images, dim=0)
        if self.xyz_base is not None:
            xyz_vae_target = torch.stack(target_xyz_images, dim=0)
        else:
            xyz_vae_target = []

        if self.debug:
            print(f"debug!!,{bg_color}")
            return {
                "target_images": raw_images,
                "target_images_xyz": raw_xyz_images,
                "input_img": input_img,
                "cameras": cameras,
                "caption": caption,
                "item": self.items[index],
                "alpha_masks": alpha_masks,
            }

        if self.split == "train":
            return {
                "target_images_vae": vae_target,
                "target_images_xyz_vae": xyz_vae_target,
                "clip_cond": clip_cond,
                "vae_cond": vae_cond,
                "cameras": cameras,
                "caption": caption,
            }
        else:  # eval
            path = os.path.join(target_dir, f"{ref_index:03d}.png")
            return dict(
                path=path,
                target_dir=target_dir,
                cond_raw_images=raw_images,
                cond=input_img,
                ref_index=ref_index,
                ident=f"{index}-{Path(target_dir).stem}",
            )


class DataRelativeStroke(DataHQCRelative):
    """a temp dataset for add sync base using fov data as ref image"""

    def __init__(
        self,
        base_dir,
        caption_csv,
        ref_indexs=[0],
        ref_position=-1,
        xyz_base=None,
        camera_views=[3, 6, 9, 12, 15],  # camera views are relative views, not abs
        split="train",
        image_size=256,
        random_background=False,
        resize_rate=1,
        num_frames=5,
        repeat=100,
        outer_file=None,
        debug=False,
        eval_size=100,
        stroke_p=0.3,
        resize_range=None,
    ):
        print(__class__)
        super().__init__(
            base_dir,
            caption_csv,
            ref_indexs=ref_indexs,
            ref_position=ref_position,
            xyz_base=xyz_base,
            camera_views=camera_views,
            split=split,
            image_size=image_size,
            random_background=random_background,
            resize_rate=resize_rate,
            num_frames=num_frames,
            repeat=repeat,
            outer_file=outer_file,
            debug=debug,
            eval_size=eval_size,
        )
        self.stroke_p = stroke_p
        assert (
            resize_range is None or len(resize_range) == 2
        ), "resize_range should be a tuple of 2 elements"
        self.resize_range = resize_range

    def __len__(self):
        return len(self.items) * self.repeat

    def __getitem__(self, index):
        index_mapping = self.index_mapping
        transfroms_mapping = self.transfroms_mapping
        index = index % len(self.items)

        target_dir = self.items[index]["path"]
        target_xyz_dir = self.items[index]["xyz_path"]
        caption = self.items[index]["caption"]

        bg_color = np.random.rand() * 255
        target_images = []
        target_xyz_images = []
        raw_images = []
        raw_xyz_images = []
        alpha_masks = []
        ref_index = random.choice(self.ref_indexs)
        cur_index_mapping = index_mapping[ref_index]
        cur_transfroms_mapping = transfroms_mapping[ref_index]
        cur_xyz_transfroms_mapping = self.xyz_transforms_mapping[ref_index]
        cur_resize_rate = (
            random.uniform(*self.resize_range) * self.resize_rate
            if self.resize_range is not None
            else self.resize_rate
        )
        for relative_view in self.camera_views:
            image_index = cur_index_mapping[relative_view]
            trans = cur_transfroms_mapping[relative_view]
            trans_xyz = cur_xyz_transfroms_mapping[relative_view]
            # open
            img = Image.open(
                os.path.join(target_dir, f"{image_index:03d}.png")
            ).convert("RGBA")
            if trans is not None:
                img = trans(img)
            img = do_resize_content(img, cur_resize_rate)
            alpha_mask = img.getchannel("A")
            alpha_masks.append(alpha_mask)
            if self.random_background:
                img = add_random_background(img, bg_color)

            img = img.convert("RGB")
            target_images.append(self.transfrom_vae(img))
            raw_images.append(img)

            if self.xyz_base is not None:
                img_xyz = Image.open(
                    os.path.join(target_xyz_dir, f"xyz_new_{image_index:03d}.png")
                ).convert("RGBA")
                img_xyz = trans_xyz(img_xyz) if trans_xyz is not None else img_xyz
                img_xyz = trans(img_xyz) if trans is not None else img_xyz
                img_xyz = do_resize_content(img_xyz, cur_resize_rate)
                img_xyz.putalpha(alpha_mask)
                if self.random_background:
                    img_xyz = add_random_background(img_xyz, bg_color)
                img_xyz = img_xyz.convert("RGB")
                target_xyz_images.append(self.transfrom_vae(img_xyz))
                if self.debug:
                    raw_xyz_images.append(img_xyz)

        cameras = [get_camera_for_index(i).squeeze() for i in self.camera_views]
        if self.ref_position is not None:
            cameras[self.ref_position] = torch.zeros_like(
                cameras[self.ref_position]
            )  # set ref camera to zero

        cameras = torch.stack(cameras)

        input_img = Image.open(
            os.path.join(target_dir, f"{ref_index:03d}.png")
        ).convert("RGBA")
        input_img = do_resize_content(input_img, cur_resize_rate)
        if random.random() < self.stroke_p:
            ## random rgb color
            color = (
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
            )
            radius = random.randint(1, 3)
            input_img = add_stroke(input_img, color=color, stroke_radius=radius)
        if self.random_background:
            input_img = add_random_background(input_img, bg_color)
        input_img = input_img.convert("RGB")

        clip_cond = self.transfrom_clip(input_img)
        vae_cond = self.transfrom_vae(input_img)

        vae_target = torch.stack(target_images, dim=0)
        if self.xyz_base is not None:
            xyz_vae_target = torch.stack(target_xyz_images, dim=0)
        else:
            xyz_vae_target = []

        if self.debug:
            print(f"debug!!,{bg_color}")
            return {
                "target_images": raw_images,
                "target_images_xyz": raw_xyz_images,
                "input_img": input_img,
                "cameras": cameras,
                "caption": caption,
                "item": self.items[index],
                "alpha_masks": alpha_masks,
                "cur_resize_rate": cur_resize_rate,
            }

        if self.split == "train":
            return {
                "target_images_vae": vae_target,
                "target_images_xyz_vae": xyz_vae_target,
                "clip_cond": clip_cond,
                "vae_cond": vae_cond,
                "cameras": cameras,
                "caption": caption,
            }
        else:  # eval
            path = os.path.join(target_dir, f"{ref_index:03d}.png")
            return dict(
                path=path,
                target_dir=target_dir,
                cond_raw_images=raw_images,
                cond=input_img,
                ref_index=ref_index,
                ident=f"{index}-{Path(target_dir).stem}",
            )


class InTheWildImages(Dataset):
    """
    a data set for in the wild images,
    receive base floders, image path ls, path files as input
    """

    def __init__(self, base_dirs=[], image_paths=[], path_files=[]):
        print(__class__)
        self.base_dirs = base_dirs
        self.image_paths = image_paths
        self.path_files = path_files
        self.init_item()

    def init_item(self):
        items = []
        for d in self.base_dirs:
            items += [osp.join(d, f) for f in os.listdir(d)]
        items = items + self.image_paths

        for file in self.path_files:
            with open(file, "r") as f:
                items += [line.strip() for line in f.readlines()]
        items.sort()
        self.items = items

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        item = self.items[index]
        img = Image.open(item)
        background = Image.new("RGBA", img.size, (0, 0, 0, 0))
        cond = Image.alpha_composite(background, img)
        return dict(
            path=item, ident=f"{index}-{Path(item).stem}", cond=cond.convert("RGB")
        )