Spaces:
Runtime error
Runtime error
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| import os | |
| cv2.setNumThreads(1) | |
| os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" | |
| class RandomResizedCropWithAutoCenteringAndZeroPadding (object): | |
| def __init__(self, output_size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), center_jitter=(0.1, 0.1), size_from_alpha_mask=True): | |
| assert isinstance(output_size, (int, tuple)) | |
| if isinstance(output_size, int): | |
| self.output_size = (output_size, output_size) | |
| else: | |
| assert len(output_size) == 2 | |
| self.output_size = output_size | |
| assert isinstance(scale, tuple) | |
| assert isinstance(ratio, tuple) | |
| if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): | |
| raise ValueError("Scale and ratio should be of kind (min, max)") | |
| self.size_from_alpha_mask = size_from_alpha_mask | |
| self.scale = scale | |
| self.ratio = ratio | |
| assert isinstance(center_jitter, tuple) | |
| self.center_jitter = center_jitter | |
| def __call__(self, sample): | |
| imidx, image = sample['imidx'], sample["image_np"] | |
| if "labels" in sample: | |
| label = sample["labels"] | |
| else: | |
| label = None | |
| im_h, im_w = image.shape[:2] | |
| if self.size_from_alpha_mask and image.shape[2] == 4: | |
| # compute bbox from alpha mask | |
| bbox_left, bbox_top, bbox_w, bbox_h = cv2.boundingRect( | |
| (image[:, :, 3] > 0).astype(np.uint8)) | |
| else: | |
| bbox_left, bbox_top = 0, 0 | |
| bbox_h, bbox_w = image.shape[:2] | |
| if bbox_h <= 1 and bbox_w <= 1: | |
| sample["bad"] = 0 | |
| else: | |
| # detect too small image here | |
| alpha_varea = np.sum((image[:, :, 3] > 0).astype(np.uint8)) | |
| image_area = image.shape[0]*image.shape[1] | |
| if alpha_varea/image_area < 0.001: | |
| sample["bad"] = alpha_varea | |
| # detect bad image | |
| if "bad" in sample: | |
| # baddata_dir = os.path.join(os.getcwd(), 'test_data', "baddata" + os.sep) | |
| # save_output(str(imidx)+".png",image,label,baddata_dir) | |
| bbox_h, bbox_w = image.shape[:2] | |
| sample["image_np"] = np.zeros( | |
| [self.output_size[0], self.output_size[1], image.shape[2]], dtype=image.dtype) | |
| if label is not None: | |
| sample["labels"] = np.zeros( | |
| [self.output_size[0], self.output_size[1], 4], dtype=label.dtype) | |
| return sample | |
| # compute default area by making sure output_size contains bbox_w * bbox_h | |
| jitter_h = np.random.uniform(-bbox_h * | |
| self.center_jitter[0], bbox_h*self.center_jitter[0]) | |
| jitter_w = np.random.uniform(-bbox_w * | |
| self.center_jitter[1], bbox_w*self.center_jitter[1]) | |
| # h/w | |
| target_aspect_ratio = np.exp( | |
| np.log(self.output_size[0]/self.output_size[1]) + | |
| np.random.uniform(np.log(self.ratio[0]), np.log(self.ratio[1])) | |
| ) | |
| source_aspect_ratio = bbox_h/bbox_w | |
| if target_aspect_ratio < source_aspect_ratio: | |
| # same w, target has larger h, use h to align | |
| target_height = bbox_h * \ | |
| np.random.uniform(self.scale[0], self.scale[1]) | |
| virtual_h = int( | |
| round(target_height)) | |
| virtual_w = int( | |
| round(target_height / target_aspect_ratio)) # h/w | |
| else: | |
| # same w, source has larger h, use w to align | |
| target_width = bbox_w * \ | |
| np.random.uniform(self.scale[0], self.scale[1]) | |
| virtual_h = int( | |
| round(target_width * target_aspect_ratio)) # h/w | |
| virtual_w = int( | |
| round(target_width)) | |
| # print("required aspect ratio:", target_aspect_ratio) | |
| virtual_top = int(round(bbox_top + jitter_h - (virtual_h-bbox_h)/2)) | |
| virutal_left = int(round(bbox_left + jitter_w - (virtual_w-bbox_w)/2)) | |
| if virtual_top < 0: | |
| top_padding = abs(virtual_top) | |
| crop_top = 0 | |
| else: | |
| top_padding = 0 | |
| crop_top = virtual_top | |
| if virutal_left < 0: | |
| left_padding = abs(virutal_left) | |
| crop_left = 0 | |
| else: | |
| left_padding = 0 | |
| crop_left = virutal_left | |
| if virtual_top+virtual_h > im_h: | |
| bottom_padding = abs(im_h-(virtual_top+virtual_h)) | |
| crop_bottom = im_h | |
| else: | |
| bottom_padding = 0 | |
| crop_bottom = virtual_top+virtual_h | |
| if virutal_left+virtual_w > im_w: | |
| right_padding = abs(im_w-(virutal_left+virtual_w)) | |
| crop_right = im_w | |
| else: | |
| right_padding = 0 | |
| crop_right = virutal_left+virtual_w | |
| # crop | |
| image = image[crop_top:crop_bottom, crop_left: crop_right] | |
| if label is not None: | |
| label = label[crop_top:crop_bottom, crop_left: crop_right] | |
| # pad | |
| if top_padding + bottom_padding + left_padding + right_padding > 0: | |
| padding = ((top_padding, bottom_padding), | |
| (left_padding, right_padding), (0, 0)) | |
| # print("padding", padding) | |
| image = np.pad(image, padding, mode='constant') | |
| if label is not None: | |
| label = np.pad(label, padding, mode='constant') | |
| if image.shape[0]/image.shape[1] - virtual_h/virtual_w > 0.001: | |
| print("virtual aspect ratio:", virtual_h/virtual_w) | |
| print("image aspect ratio:", image.shape[0]/image.shape[1]) | |
| assert (image.shape[0]/image.shape[1] - virtual_h/virtual_w < 0.001) | |
| sample["crop"] = np.array( | |
| [im_h, im_w, crop_top, crop_bottom, crop_left, crop_right, top_padding, bottom_padding, left_padding, right_padding, image.shape[0], image.shape[1]]) | |
| # resize | |
| if self.output_size[1] != image.shape[1] or self.output_size[0] != image.shape[0]: | |
| if self.output_size[1] > image.shape[1] and self.output_size[0] > image.shape[0]: | |
| # enlarging | |
| image = cv2.resize( | |
| image, (self.output_size[1], self.output_size[0]), interpolation=cv2.INTER_LINEAR) | |
| else: | |
| # shrinking | |
| image = cv2.resize( | |
| image, (self.output_size[1], self.output_size[0]), interpolation=cv2.INTER_AREA) | |
| if label is not None: | |
| label = cv2.resize(label, (self.output_size[1], self.output_size[0]), | |
| interpolation=cv2.INTER_NEAREST_EXACT) | |
| assert image.shape[0] == self.output_size[0] and image.shape[1] == self.output_size[1] | |
| sample['imidx'], sample["image_np"] = imidx, image | |
| if label is not None: | |
| assert label.shape[0] == self.output_size[0] and label.shape[1] == self.output_size[1] | |
| sample["labels"] = label | |
| return sample | |
| class FileDataset(Dataset): | |
| def __init__(self, image_names_list, fg_img_lbl_transform=None, shader_pose_use_gt_udp_test=True, shader_target_use_gt_rgb_debug=False): | |
| self.image_name_list = image_names_list | |
| self.fg_img_lbl_transform = fg_img_lbl_transform | |
| self.shader_pose_use_gt_udp_test = shader_pose_use_gt_udp_test | |
| self.shader_target_use_gt_rgb_debug = shader_target_use_gt_rgb_debug | |
| def __len__(self): | |
| return len(self.image_name_list) | |
| def get_gt_from_disk(self, idx, imname, read_label): | |
| if read_label: | |
| # read label | |
| with open(imname, mode="rb") as bio: | |
| if imname.find(".npz") > 0: | |
| label_np = np.load(bio, allow_pickle=True)[ | |
| 'i'].astype(np.float32, copy=False) | |
| else: | |
| label_np = cv2.cvtColor(cv2.imdecode(np.frombuffer(bio.read( | |
| ), np.uint8), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH | cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA) | |
| assert (4 == label_np.shape[2]) | |
| # fake image out of valid label | |
| image_np = (label_np*255).clip(0, 255).astype(np.uint8, copy=False) | |
| # assemble sample | |
| sample = {'imidx': np.array( | |
| [idx]), "image_np": image_np, "labels": label_np} | |
| else: | |
| # read image as unit8 | |
| with open(imname, mode="rb") as bio: | |
| image_np = cv2.cvtColor(cv2.imdecode(np.frombuffer( | |
| bio.read(), np.uint8), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA) | |
| # image_np = Image.open(bio) | |
| # image_np = np.array(image_np) | |
| assert (3 == len(image_np.shape)) | |
| if (image_np.shape[2] == 4): | |
| mask_np = image_np[:, :, 3:4] | |
| image_np = (image_np[:, :, :3] * | |
| (image_np[:, :, 3][:, :, np.newaxis]/255.0)).clip(0, 255).astype(np.uint8, copy=False) | |
| elif (image_np.shape[2] == 3): | |
| # generate a fake mask | |
| # Fool-proofing | |
| mask_np = np.ones( | |
| (image_np.shape[0], image_np.shape[1], 1), dtype=np.uint8)*255 | |
| print("WARN: transparent background is preferred for image ", imname) | |
| else: | |
| raise ValueError("weird shape of image ", imname, image_np) | |
| image_np = np.concatenate((image_np, mask_np), axis=2) | |
| sample = {'imidx': np.array( | |
| [idx]), "image_np": image_np} | |
| # apply fg_img_lbl_transform | |
| if self.fg_img_lbl_transform: | |
| sample = self.fg_img_lbl_transform(sample) | |
| if "labels" in sample: | |
| # return UDP as 4chn XYZV float tensor | |
| if "float" not in str(sample["labels"].dtype): | |
| sample["labels"] = sample["labels"].astype(np.float32) / np.iinfo(sample["labels"].dtype).max | |
| sample["labels"] = torch.from_numpy( | |
| sample["labels"].transpose((2, 0, 1))) | |
| assert (sample["labels"].dtype == torch.float32) | |
| if "image_np" in sample: | |
| # return image as 3chn RGB uint8 tensor and 1chn A uint8 tensor | |
| sample["mask"] = torch.from_numpy( | |
| sample["image_np"][:, :, 3:4].transpose((2, 0, 1))) | |
| assert (sample["mask"].dtype == torch.uint8) | |
| sample["image"] = torch.from_numpy( | |
| sample["image_np"][:, :, :3].transpose((2, 0, 1))) | |
| assert (sample["image"].dtype == torch.uint8) | |
| del sample["image_np"] | |
| return sample | |
| def __getitem__(self, idx): | |
| sample = { | |
| 'imidx': np.array([idx])} | |
| target = self.get_gt_from_disk( | |
| idx, imname=self.image_name_list[idx][0], read_label=self.shader_pose_use_gt_udp_test) | |
| if self.shader_target_use_gt_rgb_debug: | |
| sample["pose_images"] = torch.stack([target["image"]]) | |
| sample["pose_mask"] = target["mask"] | |
| elif self.shader_pose_use_gt_udp_test: | |
| sample["pose_label"] = target["labels"] | |
| sample["pose_mask"] = target["mask"] | |
| else: | |
| sample["pose_images"] = torch.stack([target["image"]]) | |
| if "crop" in target: | |
| sample["pose_crop"] = target["crop"] | |
| character_images = [] | |
| character_masks = [] | |
| for i in range(1, len(self.image_name_list[idx])): | |
| source = self.get_gt_from_disk( | |
| idx, self.image_name_list[idx][i], read_label=False) | |
| character_images.append(source["image"]) | |
| character_masks.append(source["mask"]) | |
| character_images = torch.stack(character_images) | |
| character_masks = torch.stack(character_masks) | |
| sample.update({ | |
| "character_images": character_images, | |
| "character_masks": character_masks | |
| }) | |
| # do not make fake labels in inference | |
| return sample | |