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
