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") )