Upload libs/data.py with huggingface_hub

libs/data.py

@@ -0,0 +1,580 @@
+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")
+        )