import numpy as np import torch from torch.utils.data import Dataset from PIL import Image from typing import Tuple, Optional import cv2 import random import os import PIL from icecream import ic def add_margin(pil_img, color=0, size=256): width, height = pil_img.size result = Image.new(pil_img.mode, (size, size), color) result.paste(pil_img, ((size - width) // 2, (size - height) // 2)) return result def scale_and_place_object(image, scale_factor): assert np.shape(image)[-1]==4 # RGBA # Extract the alpha channel (transparency) and the object (RGB channels) alpha_channel = image[:, :, 3] # Find the bounding box coordinates of the object coords = cv2.findNonZero(alpha_channel) x, y, width, height = cv2.boundingRect(coords) # Calculate the scale factor for resizing original_height, original_width = image.shape[:2] if width > height: size = width original_size = original_width else: size = height original_size = original_height scale_factor = min(scale_factor, size / (original_size+0.0)) new_size = scale_factor * original_size scale_factor = new_size / size # Calculate the new size based on the scale factor new_width = int(width * scale_factor) new_height = int(height * scale_factor) center_x = original_width // 2 center_y = original_height // 2 paste_x = center_x - (new_width // 2) paste_y = center_y - (new_height // 2) # Resize the object (RGB channels) to the new size rescaled_object = cv2.resize(image[y:y+height, x:x+width], (new_width, new_height)) # Create a new RGBA image with the resized image new_image = np.zeros((original_height, original_width, 4), dtype=np.uint8) new_image[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = rescaled_object return new_image class SingleImageDataset(Dataset): def __init__(self, root_dir: str, num_views: int, img_wh: Tuple[int, int], bg_color: str, margin_size: int = 0, single_image: Optional[PIL.Image.Image] = None, num_validation_samples: Optional[int] = None, filepaths: Optional[list] = None, cond_type: Optional[str] = None, prompt_embeds_path: Optional[str] = None, gt_path: Optional[str] = None, crop_size: Optional[int] = 720, smpl_folder: Optional[str] = 'smpl_image_pymaf', ) -> None: """Create a dataset from a folder of images. If you pass in a root directory it will be searched for images ending in ext (ext can be a list) """ self.root_dir = root_dir self.num_views = num_views self.img_wh = img_wh self.margin_size = margin_size self.bg_color = bg_color self.cond_type = cond_type self.gt_path = gt_path self.crop_size = crop_size self.smpl_folder = smpl_folder if single_image is None: if filepaths is None: # Get a list of all files in the directory file_list = os.listdir(self.root_dir) else: file_list = filepaths # Filter the files that end with .png or .jpg self.file_list = [file for file in file_list if file.endswith(('.png', '.jpg', '.webp'))] else: self.file_list = [single_image] ic(len(self.file_list)) try: normal_prompt_embedding = torch.load(f'{prompt_embeds_path}/normal_embeds.pt') color_prompt_embedding = torch.load(f'{prompt_embeds_path}/clr_embeds.pt') if self.num_views == 7: self.normal_text_embeds = normal_prompt_embedding self.color_text_embeds = color_prompt_embedding elif self.num_views == 5: self.normal_text_embeds = torch.stack([normal_prompt_embedding[0], normal_prompt_embedding[2], normal_prompt_embedding[3], normal_prompt_embedding[4], normal_prompt_embedding[6]] , 0) self.color_text_embeds = torch.stack([color_prompt_embedding[0], color_prompt_embedding[2], color_prompt_embedding[3], color_prompt_embedding[4], color_prompt_embedding[6]] , 0) except: self.color_text_embeds = torch.load(f'{prompt_embeds_path}/embeds.pt') self.normal_text_embeds = None def __len__(self): return len(self.file_list) def get_face_info(self, file): file_name = file.split('.')[0] face_info = np.load(f'{self.root_dir}/{file_name}_face_info.npy', allow_pickle=True).item() return face_info def get_bg_color(self): if self.bg_color == 'white': bg_color = np.array([1., 1., 1.], dtype=np.float32) elif self.bg_color == 'black': bg_color = np.array([0., 0., 0.], dtype=np.float32) elif self.bg_color == 'gray': bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32) elif self.bg_color == 'random': bg_color = np.random.rand(3) elif isinstance(self.bg_color, float): bg_color = np.array([self.bg_color] * 3, dtype=np.float32) else: raise NotImplementedError return bg_color def load_smpl_images(self, smpl_path, bg_color, return_type='np'): if self.num_views - 1 == 4: _views = [0, 2, 4, 6] flip_views = [4, 6] elif self.num_views - 1 == 6: _views = [0, 1, 2, 4, 6, 7] flip_views = [4, 6] elif self.num_views - 1 == 8: _views = [0, 1, 2, 3, 4, 5, 6, 7] flip_views = [4, 5, 6, 7] imgs = [] alphas = [] for i in _views: smpl_image = Image.open(os.path.join(smpl_path, f'{i:03d}.png')) if i == 0: assert smpl_image.size[0] == self.img_wh[0] smpl_alpha = np.asarray(smpl_image)[...,3] coords = np.stack(np.nonzero(smpl_alpha), 1)[:, (1, 0)] min_x, min_y = np.min(coords, 0) max_x, max_y = np.max(coords, 0) crop_size = max(max_x - min_x, max_y - min_y) + self.margin_size # print(crop_size) smpl_image = np.asarray(smpl_image).astype(np.float32) / 255. # [0, 1] alpha = smpl_image[...,3:4] img = smpl_image[...,:3] * alpha + bg_color * (1 - alpha) if return_type == "np": pass elif return_type == "pt": img = torch.from_numpy(img) alpha = torch.from_numpy(alpha) else: raise NotImplementedError if i in flip_views: img = torch.flip(img, [1]) alpha = torch.flip(alpha, [1]) imgs.append(img) alphas.append(alpha) return imgs, crop_size, alphas def load_image(self, img_path, bg_color, crop_size, return_type='np', Imagefile=None): # pil always returns uint8 if Imagefile is None: image_input = Image.open(img_path) else: image_input = Imagefile image_size = self.img_wh[0] alpha_np = np.asarray(image_input)[:, :, 3] coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)] min_x, min_y = np.min(coords, 0) max_x, max_y = np.max(coords, 0) ref_img_ = image_input.crop((min_x, min_y, max_x, max_y)) h, w = ref_img_.height, ref_img_.width scale = crop_size / max(h, w) h_, w_ = int(scale * h), int(scale * w) ref_img_ = ref_img_.resize((w_, h_)) image_input = add_margin(ref_img_, size=image_size) # img = scale_and_place_object(img, self.scale_ratio) img = np.array(image_input) img = img.astype(np.float32) / 255. # [0, 1] assert img.shape[-1] == 4 # RGBA alpha = img[...,3:4] img = img[...,:3] * alpha + bg_color * (1 - alpha) if return_type == "np": pass elif return_type == "pt": img = torch.from_numpy(img) alpha = torch.from_numpy(alpha) else: raise NotImplementedError return img, alpha def load_face(self, img_path, bg_color, return_type='np', crop_size=-1): image_input = Image.open(img_path) image_size = self.img_wh[0] if crop_size > 0: # color image alpha_np = np.asarray(image_input)[:, :, 3] coords = np.stack(np.nonzero(alpha_np), 1)[:, (1, 0)] min_x, min_y = np.min(coords, 0) max_x, max_y = np.max(coords, 0) ref_img_ = image_input.crop((min_x, min_y, max_x, max_y)) h, w = ref_img_.height, ref_img_.width scale = crop_size / max(h, w) h_, w_ = int(scale * h), int(scale * w) ref_img_ = ref_img_.resize((w_, h_)) image_input = add_margin(ref_img_, size=image_size) image_input = image_input.crop((256, 0, 512, 256)).resize((self.img_wh[0], self.img_wh[1])) # img = scale_and_place_object(img, self.scale_ratio) img = np.array(image_input) img = img.astype(np.float32) / 255. # [0, 1] assert img.shape[-1] == 4 # RGBA alpha = img[...,3:4] img = img[...,:3] * alpha + bg_color * (1 - alpha) if return_type == "np": pass elif return_type == "pt": img = torch.from_numpy(img) alpha = torch.from_numpy(alpha) else: raise NotImplementedError return img def process_face(self, img_path, bbox, bg_color, normal_path=None, w2c=None, h=512, w=512): image = Image.open(img_path) bbox_w, bbox_h = bbox[2] - bbox[0], bbox[3] - bbox[1] if bbox_w > bbox_h: bbox[1] -= (bbox_w - bbox_h) // 2 bbox[3] += (bbox_w - bbox_h) // 2 else: bbox[0] -= (bbox_h - bbox_w) // 2 bbox[2] += (bbox_h - bbox_w) // 2 bbox[0:2] -= 20 bbox[2:4] += 20 image = image.crop(bbox) image = image.resize((w, h)) image = np.array(image) / 255. img, alpha = image[:, :, :3], image[:, :, 3:4] img = img * alpha + bg_color * (1 - alpha) padded_img = np.full((self.img_wh[0], self.img_wh[1], 3), bg_color, dtype=np.float32) dx = (self.img_wh[0] - w) // 2 dy = (self.img_wh[1] - h) // 2 padded_img[dy:dy+h, dx:dx+w] = img padded_img = torch.from_numpy(padded_img).permute(2,0,1) return padded_img def __getitem__(self, index): filename = self.file_list[index].split('.')[0] bg_color = self.get_bg_color() smpl_images, crop_size, smpl_alphas = self.load_smpl_images(f'{self.root_dir}/{self.smpl_folder}/{filename}', bg_color, return_type='pt') smpl_face = self.load_face(f'{self.root_dir}/{self.smpl_folder}/{filename}/000.png', bg_color, return_type='pt') image, _ = self.load_image(f'{self.root_dir}/{self.file_list[index]}', bg_color, crop_size, return_type='pt') # m face = self.load_face(f'{self.root_dir}/{self.file_list[index]}', bg_color, return_type='pt', crop_size=crop_size) # m img_tensors_in = [ image.permute(2, 0, 1) ] * (self.num_views-1) + [ face.permute(2, 0, 1)] smpl_tensors_in = [ tmp.permute(2, 0, 1) for tmp in smpl_images ] + [ smpl_face.permute(2, 0, 1) ] smpl_alphas = [ tmp.permute(2, 0, 1) for tmp in smpl_alphas ] # import pdb; pdb.set_trace() img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W) smpl_tensors_in = torch.stack(smpl_tensors_in, dim=0).float() # (Nv, 3, H, W) smpl_alphas = torch.stack(smpl_alphas, dim=0).float() # (Nv, 1, H, W) normal_prompt_embeddings = self.normal_text_embeds if hasattr(self, 'normal_text_embeds') else None color_prompt_embeddings = self.color_text_embeds if hasattr(self, 'color_text_embeds') else None if normal_prompt_embeddings is None: out = { 'imgs_in': img_tensors_in, 'smpl_imgs_in': smpl_tensors_in, 'smpl_alphas': smpl_alphas, 'color_prompt_embeddings': color_prompt_embeddings, 'filename': filename, } else: out = { 'imgs_in': img_tensors_in, 'smpl_imgs_in': smpl_tensors_in, 'smpl_alphas': smpl_alphas, 'normal_prompt_embeddings': normal_prompt_embeddings, 'color_prompt_embeddings': color_prompt_embeddings, 'filename': filename, } return out if __name__ == "__main__": # pass from torch.utils.data import DataLoader from torchvision.utils import make_grid from PIL import ImageDraw, ImageFont def draw_text(img, text, pos, color=(128, 128, 128)): draw = ImageDraw.Draw(img) # font = ImageFont.truetype(size= size) font = ImageFont.load_default() font = font.font_variant(size=10) draw.text(pos, text, color, font=font) return img random.seed(11) test_params = dict( root_dir='../../examples/CAPE', bg_color='white', img_wh=(768, 768), prompt_embeds_path='fixed_prompt_embeds_7view', num_views=7, # crop_size=740, margin_size=15, smpl_folder='gt_smpl_image', ) train_dataset = SingleImageDataset(**test_params) data_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0) for batch in data_loader: # batch = train_dataset.__getitem__(1) imgs = [] obj_name = batch['filename'][0] imgs_in = batch['imgs_in'][0] imgs_smpl_in = batch['smpl_imgs_in'][0] alphas_smpl = batch['smpl_alphas'][0] img0 = (imgs_in[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8) img1 = (imgs_smpl_in[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8) print(img0.shape, img1.shape) smpl_alpha = alphas_smpl[0].permute(1, 2, 0).repeat(1, 1, 3).numpy() img0[smpl_alpha > 0.5] = img1[smpl_alpha > 0.5] Image.fromarray(img0).save(f'../../debug/{obj_name}_rgb.png') # Image.fromarray(img1).save(f'../../debug/{obj_name}_smpl.png') exit() imgs_vis = torch.cat([imgs_in, imgs_smpl_in], 0) img_vis = make_grid(imgs_vis, nrow=4).permute(1, 2,0) img_vis = (img_vis.numpy() * 255).astype(np.uint8) img_vis = Image.fromarray(img_vis) img_vis = draw_text(img_vis, obj_name, (5, 1)) img_vis = torch.from_numpy(np.array(img_vis)).permute(2, 0, 1) / 255. imgs.append(img_vis) imgs = torch.stack(imgs, dim=0) img_grid = make_grid(imgs, nrow=4, padding=0) img_grid = img_grid.permute(1, 2, 0).numpy() img_grid = (img_grid * 255).astype(np.uint8) img_grid = Image.fromarray(img_grid) img_grid.save(f'../../debug/{obj_name}.png') print('finished.')