def __getitem__(self, index): elem = self.imgs_uniform[index] centroid = None if len(elem) == 4: img_path, mask_path, centroid, class_id = elem else: img_path, mask_path = elem img, mask = Image.open(img_path).convert('RGB'), Image.open(mask_path) img_name = os.path.splitext(os.path.basename(img_path))[0] mask = np.array(mask) mask_copy = mask.copy() for k, v in id_to_trainid.items(): mask_copy[mask == k] = v mask_trained = Image.fromarray(mask_copy.astype(np.uint8)) mask = mask_trained # Image Transformations if self.joint_transform_list is not None: for idx, xform in enumerate(self.joint_transform_list): if idx == 0 and centroid is not None: # HACK # We assume that the first transform is capable of taking # in a centroid img, mask = xform(img, mask, centroid) else: img, mask = xform(img, mask) # Debug if self.dump_images and centroid is not None: outdir = '../../dump_imgs_{}'.format(self.mode) os.makedirs(outdir, exist_ok=True) dump_img_name = trainid_to_name[class_id] + '_' + img_name out_img_fn = os.path.join(outdir, dump_img_name + '.png') out_msk_fn = os.path.join(outdir, dump_img_name + '_mask.png') mask_img = colorize_mask(np.array(mask)) img.save(out_img_fn) mask_img.save(out_msk_fn) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: mask = self.target_transform(mask) if self.edge_map: # _edgemap = np.array(mask_trained) _edgemap = mask[:-1, :, :] # c, h, w _edgemap = edge_utils.onehot_to_binary_edges( _edgemap, 2, num_classes) # h, w edgemap = torch.from_numpy(_edgemap).float() return img, mask, edgemap, img_name return img, mask, img_name
def __getitem__(self, index): item = self.files[index] img_name = item["name"] img_path = self.root + item['img'] label_path = self.root + item["label"] img = Image.open(img_path).convert('RGB') mask = np.array(Image.open(label_path)) mask = mask[:, :] mask_copy = self.convert_label(mask) if self.eval_mode: return self._eval_get_item(img, mask_copy, self.eval_scales, self.eval_flip), img_name mask = Image.fromarray(mask_copy.astype(np.uint8)) # Image Transformations if self.joint_transform is not None: img, mask = self.joint_transform(img, mask) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: mask = self.target_transform(mask) if self.mode == 'test': return img, mask, img_name, item['img'] _edgemap = mask.numpy() _edgemap = edge_utils.mask_to_onehot(_edgemap, num_classes) _edgemap = edge_utils.onehot_to_binary_edges(_edgemap, 2, num_classes) edgemap = torch.from_numpy(_edgemap).float() # Debug if self.dump_images: outdir = '../../dump_imgs_{}'.format(self.mode) os.makedirs(outdir, exist_ok=True) out_img_fn = os.path.join(outdir, img_name + '.png') out_msk_fn = os.path.join(outdir, img_name + '_mask.png') mask_img = colorize_mask(np.array(mask)) img.save(out_img_fn) mask_img.save(out_msk_fn) return img, mask, edgemap, img_name
def __getitem__(self, index): img_path, mask_path = self.imgs[index] img, mask = Image.open(img_path).convert('RGB'), Image.open(mask_path) img_name = os.path.splitext(os.path.basename(img_path))[0] mask = np.array(mask) mask_copy = mask.copy() for k, v in id_to_trainid.items(): mask_copy[mask == k] = v if self.eval_mode: return self._eval_get_item(img, mask_copy, self.eval_scales, self.eval_flip), img_name mask = Image.fromarray(mask_copy.astype(np.uint8)) # Image Transformations if self.joint_transform is not None: img, mask = self.joint_transform(img, mask) if self.transform is not None: img = self.transform(img) if self.target_transform is not None: mask = self.target_transform(mask) _edgemap = mask.numpy() _edgemap = edge_utils.mask_to_onehot(_edgemap, num_classes) _edgemap = edge_utils.onehot_to_binary_edges(_edgemap, 2, num_classes) edgemap = torch.from_numpy(_edgemap).float() # Debug if self.dump_images: outdir = '../../dump_imgs_{}'.format(self.mode) os.makedirs(outdir, exist_ok=True) out_img_fn = os.path.join(outdir, img_name + '.png') out_msk_fn = os.path.join(outdir, img_name + '_mask.png') mask_img = colorize_mask(np.array(mask)) img.save(out_img_fn) mask_img.save(out_msk_fn) return img, mask, edgemap, img_name
def __getitem__(self, index): elem = self.imgs_uniform[index] centroid = None if len(elem) == 4: img_path, mask_path, centroid, class_id = elem else: img_path, mask_path = elem if self.mode == 'test': img, mask = Image.open(img_path).convert('RGB'), None else: img, mask = Image.open(img_path).convert('RGB'), Image.open( mask_path) img_name = os.path.splitext(os.path.basename(img_path))[0] # kitti scale correction factor if self.mode == 'train' or self.mode == 'trainval': if self.scf: width, height = img.size img = img.resize((width * 2, height * 2), Image.BICUBIC) mask = mask.resize((width * 2, height * 2), Image.NEAREST) elif self.mode == 'val': width, height = 1242, 376 img = img.resize((width, height), Image.BICUBIC) mask = mask.resize((width, height), Image.NEAREST) elif self.mode == 'test': img_keepsize = img.copy() width, height = 1280, 384 img = img.resize((width, height), Image.BICUBIC) else: logging.info('Unknown mode {}'.format(mode)) sys.exit() if self.mode != 'test': mask = np.array(mask) mask_copy = mask.copy() for k, v in id_to_trainid.items(): mask_copy[mask == k] = v mask = Image.fromarray(mask_copy.astype(np.uint8)) # Image Transformations if self.joint_transform_list is not None: for idx, xform in enumerate(self.joint_transform_list): if idx == 0 and centroid is not None: # HACK # We assume that the first transform is capable of taking # in a centroid img, mask = xform(img, mask, centroid) else: img, mask = xform(img, mask) # Debug if self.dump_images and centroid is not None: outdir = './dump_imgs_{}'.format(self.mode) os.makedirs(outdir, exist_ok=True) dump_img_name = trainid_to_name[class_id] + '_' + img_name out_img_fn = os.path.join(outdir, dump_img_name + '.png') out_msk_fn = os.path.join(outdir, dump_img_name + '_mask.png') mask_img = colorize_mask(np.array(mask)) img.save(out_img_fn) mask_img.save(out_msk_fn) if self.transform is not None: img = self.transform(img) if self.mode == 'test': img_keepsize = self.transform(img_keepsize) mask = img_keepsize if self.target_transform is not None: if self.mode != 'test': mask = self.target_transform(mask) if self.edge_map: # _edgemap = np.array(mask_trained) # _edgemap = edge_utils.mask_to_onehot(_edgemap, num_classes) _edgemap = mask[:-1, :, :] _edgemap = edge_utils.onehot_to_binary_edges( _edgemap, 2, num_classes) edgemap = torch.from_numpy(_edgemap).float() return img, mask, edgemap, img_name return img, mask, img_name