Spaces:
No application file
No application file
File size: 3,830 Bytes
907462b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import pickle as pkl
from pathlib import Path
from scipy.io import loadmat
from scripts.annotations_conversion.common import parallel_map
ADE20K_STUFF_CLASSES = ['water', 'wall', 'snow', 'sky', 'sea', 'sand', 'road', 'route', 'river', 'path', 'mountain',
'mount', 'land', 'ground', 'soil', 'hill', 'grass', 'floor', 'flooring', 'field', 'earth',
'ground', 'fence', 'ceiling', 'wave', 'crosswalk', 'hay bale', 'bridge', 'span', 'building',
'edifice', 'cabinet', 'cushion', 'curtain', 'drape', 'drapery', 'mantle', 'pall', 'door',
'fencing', 'house', 'pole', 'seat', 'windowpane', 'window', 'tree', 'towel', 'table',
'stairs', 'steps', 'streetlight', 'street lamp', 'sofa', 'couch', 'lounge', 'skyscraper',
'signboard', 'sign', 'sidewalk', 'pavement', 'shrub', 'bush', 'rug', 'carpet']
def worker_annotations_loader(anno_pair, dataset_path):
image_id, folder = anno_pair
n_masks = len(list((dataset_path / folder).glob(f'{image_id}_*.png')))
# each image has several layers with instances,
# each layer has mask name and instance_to_class mapping
layers = [{
'mask_name': f'{image_id}_{suffix}.png',
'instance_to_class': {},
'object_instances': [],
'stuff_instances': []
} for suffix in ['seg'] + [f'parts_{i}' for i in range(1, n_masks)]]
# parse txt with instance to class mappings
with (dataset_path / folder / (image_id + "_atr.txt")).open('r') as f:
for line in f:
# instance_id layer_n is_occluded class_names class_name_raw attributes
line = line.strip().split('#')
inst_id, layer_n, class_names = int(line[0]), int(line[1]), line[3]
# there may be more than one class name for each instance
class_names = [name.strip() for name in class_names.split(',')]
# check if any of classes is stuff
if set(class_names) & set(ADE20K_STUFF_CLASSES):
layers[layer_n]['stuff_instances'].append(inst_id)
else:
layers[layer_n]['object_instances'].append(inst_id)
layers[layer_n]['instance_to_class'][inst_id] = class_names
return layers
def load_and_parse_annotations(dataset_path, dataset_split, n_jobs=1):
dataset_split_folder = 'training' if dataset_split == 'train' else 'validation'
orig_annotations = loadmat(dataset_path / 'index_ade20k.mat', squeeze_me=True, struct_as_record=True)
image_ids = [image_id.split('.')[0] for image_id in orig_annotations['index'].item()[0]
if dataset_split in image_id]
folders = [Path(folder).relative_to('ADE20K_2016_07_26') for folder in orig_annotations['index'].item()[1]
if dataset_split_folder in folder]
# list of dictionaries with filename and instance to class mapping
all_layers = parallel_map(list(zip(image_ids, folders)), worker_annotations_loader, n_jobs=n_jobs,
use_kwargs=False, const_args={
'dataset_path': dataset_path
})
return image_ids, folders, all_layers
def create_annotations(dataset_path, dataset_split='train', n_jobs=1):
anno_path = dataset_path / f'{dataset_split}-annotations-object-segmentation.pkl'
image_ids, folders, all_layers = load_and_parse_annotations(dataset_path, dataset_split, n_jobs=n_jobs)
# create dictionary with annotations
annotations = {}
for index, image_id in enumerate(image_ids):
annotations[image_id] = {
'folder': folders[index],
'layers': all_layers[index]
}
with anno_path.open('wb') as f:
pkl.dump(annotations, f)
return annotations
|