mfrashad's picture
Init code
97069e1
import os, errno, numpy, torch, csv, re, shutil, os, zipfile
from collections import OrderedDict
from torchvision.datasets.folder import default_loader
from torchvision import transforms
from scipy import ndimage
from urllib.request import urlopen
class BrodenDataset(torch.utils.data.Dataset):
'''
A multicategory segmentation data set.
Returns three streams:
(1) The image (3, h, w).
(2) The multicategory segmentation (labelcount, h, w).
(3) A bincount of pixels in the segmentation (labelcount).
Net dissect also assumes that the dataset object has three properties
with human-readable labels:
ds.labels = ['red', 'black', 'car', 'tree', 'grid', ...]
ds.categories = ['color', 'part', 'object', 'texture']
ds.label_category = [0, 0, 2, 2, 3, ...] # The category for each label
'''
def __init__(self, directory='dataset/broden', resolution=384,
split='train', categories=None,
transform=None, transform_segment=None,
download=False, size=None, include_bincount=True,
broden_version=1, max_segment_depth=6):
assert resolution in [224, 227, 384]
if download:
ensure_broden_downloaded(directory, resolution, broden_version)
self.directory = directory
self.resolution = resolution
self.resdir = os.path.join(directory, 'broden%d_%d' %
(broden_version, resolution))
self.loader = default_loader
self.transform = transform
self.transform_segment = transform_segment
self.include_bincount = include_bincount
# The maximum number of multilabel layers that coexist at an image.
self.max_segment_depth = max_segment_depth
with open(os.path.join(self.resdir, 'category.csv'),
encoding='utf-8') as f:
self.category_info = OrderedDict()
for row in csv.DictReader(f):
self.category_info[row['name']] = row
if categories is not None:
# Filter out unused categories
categories = set([c for c in categories if c in self.category_info])
for cat in list(self.category_info.keys()):
if cat not in categories:
del self.category_info[cat]
categories = list(self.category_info.keys())
self.categories = categories
# Filter out unneeded images.
with open(os.path.join(self.resdir, 'index.csv'),
encoding='utf-8') as f:
all_images = [decode_index_dict(r) for r in csv.DictReader(f)]
self.image = [row for row in all_images
if index_has_any_data(row, categories) and row['split'] == split]
if size is not None:
self.image = self.image[:size]
with open(os.path.join(self.resdir, 'label.csv'),
encoding='utf-8') as f:
self.label_info = build_dense_label_array([
decode_label_dict(r) for r in csv.DictReader(f)])
self.labels = [l['name'] for l in self.label_info]
# Build dense remapping arrays for labels, so that you can
# get dense ranges of labels for each category.
self.category_map = {}
self.category_unmap = {}
self.category_label = {}
for cat in self.categories:
with open(os.path.join(self.resdir, 'c_%s.csv' % cat),
encoding='utf-8') as f:
c_data = [decode_label_dict(r) for r in csv.DictReader(f)]
self.category_unmap[cat], self.category_map[cat] = (
build_numpy_category_map(c_data))
self.category_label[cat] = build_dense_label_array(
c_data, key='code')
self.num_labels = len(self.labels)
# Primary categories for each label is the category in which it
# appears with the maximum coverage.
self.label_category = numpy.zeros(self.num_labels, dtype=int)
for i in range(self.num_labels):
maxcoverage, self.label_category[i] = max(
(self.category_label[cat][self.category_map[cat][i]]['coverage']
if i < len(self.category_map[cat])
and self.category_map[cat][i] else 0, ic)
for ic, cat in enumerate(categories))
def __len__(self):
return len(self.image)
def __getitem__(self, idx):
record = self.image[idx]
# example record: {
# 'image': 'opensurfaces/25605.jpg', 'split': 'train',
# 'ih': 384, 'iw': 384, 'sh': 192, 'sw': 192,
# 'color': ['opensurfaces/25605_color.png'],
# 'object': [], 'part': [],
# 'material': ['opensurfaces/25605_material.png'],
# 'scene': [], 'texture': []}
image = self.loader(os.path.join(self.resdir, 'images',
record['image']))
segment = numpy.zeros(shape=(self.max_segment_depth,
record['sh'], record['sw']), dtype=int)
if self.include_bincount:
bincount = numpy.zeros(shape=(self.num_labels,), dtype=int)
depth = 0
for cat in self.categories:
for layer in record[cat]:
if isinstance(layer, int):
segment[depth,:,:] = layer
if self.include_bincount:
bincount[layer] += segment.shape[1] * segment.shape[2]
else:
png = numpy.asarray(self.loader(os.path.join(
self.resdir, 'images', layer)))
segment[depth,:,:] = png[:,:,0] + png[:,:,1] * 256
if self.include_bincount:
bincount += numpy.bincount(segment[depth,:,:].flatten(),
minlength=self.num_labels)
depth += 1
if self.transform:
image = self.transform(image)
if self.transform_segment:
segment = self.transform_segment(segment)
if self.include_bincount:
bincount[0] = 0
return (image, segment, bincount)
else:
return (image, segment)
def build_dense_label_array(label_data, key='number', allow_none=False):
'''
Input: set of rows with 'number' fields (or another field name key).
Output: array such that a[number] = the row with the given number.
'''
result = [None] * (max([d[key] for d in label_data]) + 1)
for d in label_data:
result[d[key]] = d
# Fill in none
if not allow_none:
example = label_data[0]
def make_empty(k):
return dict((c, k if c is key else type(v)())
for c, v in example.items())
for i, d in enumerate(result):
if d is None:
result[i] = dict(make_empty(i))
return result
def build_numpy_category_map(map_data, key1='code', key2='number'):
'''
Input: set of rows with 'number' fields (or another field name key).
Output: array such that a[number] = the row with the given number.
'''
results = list(numpy.zeros((max([d[key] for d in map_data]) + 1),
dtype=numpy.int16) for key in (key1, key2))
for d in map_data:
results[0][d[key1]] = d[key2]
results[1][d[key2]] = d[key1]
return results
def index_has_any_data(row, categories):
for c in categories:
for data in row[c]:
if data: return True
return False
def decode_label_dict(row):
result = {}
for key, val in row.items():
if key == 'category':
result[key] = dict((c, int(n))
for c, n in [re.match('^([^(]*)\(([^)]*)\)$', f).groups()
for f in val.split(';')])
elif key == 'name':
result[key] = val
elif key == 'syns':
result[key] = val.split(';')
elif re.match('^\d+$', val):
result[key] = int(val)
elif re.match('^\d+\.\d*$', val):
result[key] = float(val)
else:
result[key] = val
return result
def decode_index_dict(row):
result = {}
for key, val in row.items():
if key in ['image', 'split']:
result[key] = val
elif key in ['sw', 'sh', 'iw', 'ih']:
result[key] = int(val)
else:
item = [s for s in val.split(';') if s]
for i, v in enumerate(item):
if re.match('^\d+$', v):
item[i] = int(v)
result[key] = item
return result
class ScaleSegmentation:
'''
Utility for scaling segmentations, using nearest-neighbor zooming.
'''
def __init__(self, target_height, target_width):
self.target_height = target_height
self.target_width = target_width
def __call__(self, seg):
ratio = (1, self.target_height / float(seg.shape[1]),
self.target_width / float(seg.shape[2]))
return ndimage.zoom(seg, ratio, order=0)
def scatter_batch(seg, num_labels, omit_zero=True, dtype=torch.uint8):
'''
Utility for scattering semgentations into a one-hot representation.
'''
result = torch.zeros(*((seg.shape[0], num_labels,) + seg.shape[2:]),
dtype=dtype, device=seg.device)
result.scatter_(1, seg, 1)
if omit_zero:
result[:,0] = 0
return result
def ensure_broden_downloaded(directory, resolution, broden_version=1):
assert resolution in [224, 227, 384]
baseurl = 'http://netdissect.csail.mit.edu/data/'
dirname = 'broden%d_%d' % (broden_version, resolution)
if os.path.isfile(os.path.join(directory, dirname, 'index.csv')):
return # Already downloaded
zipfilename = 'broden1_%d.zip' % resolution
download_dir = os.path.join(directory, 'download')
os.makedirs(download_dir, exist_ok=True)
full_zipfilename = os.path.join(download_dir, zipfilename)
if not os.path.exists(full_zipfilename):
url = '%s/%s' % (baseurl, zipfilename)
print('Downloading %s' % url)
data = urlopen(url)
with open(full_zipfilename, 'wb') as f:
f.write(data.read())
print('Unzipping %s' % zipfilename)
with zipfile.ZipFile(full_zipfilename, 'r') as zip_ref:
zip_ref.extractall(directory)
assert os.path.isfile(os.path.join(directory, dirname, 'index.csv'))
def test_broden_dataset():
'''
Testing code.
'''
bds = BrodenDataset('dataset/broden', resolution=384,
transform=transforms.Compose([
transforms.Resize(224),
transforms.ToTensor()]),
transform_segment=transforms.Compose([
ScaleSegmentation(224, 224)
]),
include_bincount=True)
loader = torch.utils.data.DataLoader(bds, batch_size=100, num_workers=24)
for i in range(1,20):
print(bds.label[i]['name'],
list(bds.category.keys())[bds.primary_category[i]])
for i, (im, seg, bc) in enumerate(loader):
print(i, im.shape, seg.shape, seg.max(), bc.shape)
if __name__ == '__main__':
test_broden_dataset()