Spaces:
Build error
Build error
| 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() | |