def __init__(self, directory, max_files=20000):
"""
SUN RGB-D Dataset reader
Note: First run find . | grep seg.mat > annotations.txt (in SUNRGBD folder)
@params directory: SUNRGBD directory listing with image/*.png, and seg.mat files
"""
self.directory_ = os.path.expanduser(directory)
with open(os.path.join(self.directory_, 'image.txt')) as f:
rgb_files = f.read().splitlines()
with open(os.path.join(self.directory_, 'depth.txt')) as f:
depth_files = f.read().splitlines()
assert(len(rgb_files) == len(depth_files))
self.rgb_files_ = [os.path.join(self.directory_, fn) for fn in fnmatch.filter(rgb_files,'*mit_*')][:max_files]
self.depth_files_ = [os.path.join(self.directory_, fn) for fn in fnmatch.filter(depth_files,'*mit_*')][:max_files]
self.label_files_ = [ os.path.join(
os.path.split(
os.path.split(fn)[0])[0], 'seg.mat') for fn in self.rgb_files_ ]
if not len(self.rgb_files_):
raise RuntimeError('{} :: Failed to load dataset'.format(self.__class__.__name__))
print('{} :: Loading {} image/depth/segmentation pairs'.format(self.__class__.__name__, len(self.rgb_files_)))
self.rgb_ = imap(lambda fn: self._pad_image(cv2.imread(fn, cv2.CV_LOAD_IMAGE_COLOR)), self.rgb_files_)
self.depth_ = imap(lambda fn: self._pad_image(cv2.imread(fn, -1)), self.depth_files_)
self.labels_ = imap(self._process_label, self.label_files_)
# self.target_hash_ = {item.encode('utf8'): idx+1
# for idx, item in enumerate(loadmat('data/sun3d/seg37list.mat', squeeze_me=True)['seg37list'])}
# self.target_unhash_ = {v:k for k,v in self.target_hash_.iteritems()}
# self.target_hash_ = SUNRGBDDataset.target_hash
# self.target_unhash_ = SUNRGBDDataset.target_unhash
# @property
# def target_unhash(self):
# return self.objects_.target_unhash
# @property
# def target_hash(self):
# return self.objects_.target_hash
评论列表
文章目录
