def points_single_category(subset, category='airplane', dims='x_y_z', # combinations of 'x', 'y', 'z', 'nx', 'ny', 'nz' and 'one' read_cache=True, write_cache=True, cache_dir='', shuffle=False, root=SHAPENET3D_DATA_DIR): if not category.startswith('0'): category = SN_CATEGORIES[SN_CATEGORY_NAMES.index(category)] object_label = SN_CATEGORIES.index(category) if not cache_dir: cache_dir = os.path.join(root, 'cache') if read_cache or write_cache: if not os.path.exists(cache_dir): os.makedirs(cache_dir) #os.makedirs(cache_dir, exist_ok=True) cache_path = os.path.join(cache_dir, '{}_{}.cache'.format(subset, category)) if read_cache and os.path.exists(cache_path): with open(cache_path, mode='rb') as f: feat_list, label_list, hash_list = pickle.load(f) else: feat_list, label_list, hash_list = [], [], [] if not feat_list: data_dir = os.path.join(root, subset, category) hash_list = sorted([s[:-4] for s in filter(lambda s: s.endswith('.ply'), os.listdir(data_dir))]) for s in hash_list: data = np.loadtxt(os.path.join(data_dir, s + '.ply'), skiprows=14) feat_list.append(data[:, :6]) label_list.append(data[:, -1]) if write_cache: with open(cache_path, mode='wb') as f: pickle.dump((feat_list, label_list, hash_list), f) # adding 'one' as an additional feature feat_dict = dict(zip('x_y_z_nx_ny_nz_one'.split('_'), range(7))) feat_idxs = [feat_dict[f] for f in dims.split('_')] if 'one' in dims.split('_'): feat_list = [np.concatenate((d, np.ones((len(d), 1))), axis=1) for d in feat_list] feat_list = [d[:, feat_idxs] for d in feat_list] # shuffle if shuffle: idx = np.random.permutation(len(hash_list)) feat_list, label_list, hash_list = feat_list[idx], label_list[idx], hash_list[idx] return feat_list, [object_label] * len(feat_list), label_list, hash_list
def category_mask(category): """ Get a mask vector for part categories corresponding to a specific shape category :param category: (string) shape category name or id :return: [50] numpy array """ if type(category) == str: if category.startswith('0'): category = SN_CATEGORIES.index(category) else: category = SN_CATEGORY_NAMES.index(category) mask = np.zeros(sum(SN_NUM_PART_CATEGORIES)) mask[sum(SN_NUM_PART_CATEGORIES[:category]):sum(SN_NUM_PART_CATEGORIES[:category + 1])] = 1 return mask
def partseg_seq( arch_str='64_128_256_256', batchnorm=True, skip_str=( ), # tuple of strings like '4_1_ga' - relu4 <- relu1 w/ options 'ga' bilateral_nbr=1, conv_weight_filler='xavier', bltr_weight_filler='gauss_0.001', dataset='shapenet', dataset_params=None, category='airplane', sample_size=3000, batch_size=32, feat_dims_str='x_y_z', lattice_dims_str=None, deploy=False, create_prototxt=True, save_path=None): n = caffe.NetSpec() arch_str = [(v[0], int(v[1:])) if v[0] in {'b', 'c'} else ('c', int(v)) for v in arch_str.split('_')] num_bltr_layers = sum(v[0] == 'b' for v in arch_str) if num_bltr_layers > 0: if type(lattice_dims_str) == str: lattice_dims_str = (lattice_dims_str, ) * num_bltr_layers elif len(lattice_dims_str) == 1: lattice_dims_str = lattice_dims_str * num_bltr_layers else: assert len( lattice_dims_str ) == num_bltr_layers, '{} lattices should be provided'.format( num_bltr_layers) feat_dims = parse_channel_scale(feat_dims_str, channel_str=True)[0] lattice_dims = [ parse_channel_scale(s, channel_str=True)[0] for s in lattice_dims_str ] input_dims_w_dup = feat_dims + reduce(lambda x, y: x + y, lattice_dims) input_dims = reduce(lambda x, y: x if y in x else x + [y], input_dims_w_dup, []) feat_dims_str = map_channel_scale(feat_dims_str, input_dims) lattice_dims_str = [ map_channel_scale(s, input_dims) for s in lattice_dims_str ] input_dims_str = '_'.join(input_dims) else: feat_dims = parse_channel_scale(feat_dims_str, channel_str=True)[0] input_dims = feat_dims feat_dims_str = map_channel_scale(feat_dims_str, input_dims) input_dims_str = '_'.join(input_dims) # dataset specific settings: nclass, datalayer_train, datalayer_test if dataset == 'shapenet': from splatnet.configs import SN_CATEGORIES, SN_CATEGORY_NAMES, SN_NUM_PART_CATEGORIES if not category.startswith('0'): category = SN_CATEGORIES[SN_CATEGORY_NAMES.index(category)] nclass = SN_NUM_PART_CATEGORIES[SN_CATEGORIES.index(category)] # default dataset params dataset_params_new = {} if not dataset_params else dataset_params dataset_params = dict(subset_train='train', subset_test='val') dataset_params.update(dataset_params_new) dataset_params['feat_dims'] = input_dims_str dataset_params['sample_size'] = sample_size dataset_params['batch_size'] = batch_size dataset_params['category'] = category # dataset params type casting for v in {'jitter_xyz', 'jitter_rotation', 'jitter_stretch'}: if v in dataset_params: dataset_params[v] = float(dataset_params[v]) for v in {'sample_size', 'batch_size'}: if v in dataset_params: dataset_params[v] = int(dataset_params[v]) # training time dataset params dataset_params_train = dataset_params.copy() dataset_params_train['subset'] = dataset_params['subset_train'] del dataset_params_train['subset_train'], dataset_params_train[ 'subset_test'] # testing time dataset params: turn off all data augmentations dataset_params_test = dataset_params.copy() dataset_params_test['subset'] = dataset_params['subset_test'] dataset_params_test['jitter_xyz'] = 0.0 dataset_params_test['jitter_stretch'] = 0.0 dataset_params_test['jitter_rotation'] = 0.0 del dataset_params_test['subset_train'], dataset_params_test[ 'subset_test'] # data layers datalayer_train = L.Python(name='data', include=dict(phase=caffe.TRAIN), ntop=2, python_param=dict( module='dataset_shapenet', layer='InputShapenet', param_str=repr(dataset_params_train))) datalayer_test = L.Python(name='data', include=dict(phase=caffe.TEST), ntop=0, top=['data', 'label'], python_param=dict( module='dataset_shapenet', layer='InputShapenet', param_str=repr(dataset_params_test))) else: raise ValueError('Dataset {} unknown'.format(dataset)) # Input/Data layer if deploy: n.data = L.Input(shape=dict(dim=[1, len(input_dims), 1, sample_size])) else: n.data, n.label = datalayer_train n.test_data = datalayer_test n.data_feat = L.Python(n.data, python_param=dict(module='custom_layers', layer='PickAndScale', param_str=feat_dims_str)) top_prev = n.data_feat if conv_weight_filler in {'xavier', 'msra'}: conv_weight_filler = dict(type=conv_weight_filler) elif conv_weight_filler.startswith('gauss_'): conv_weight_filler = dict(type='gaussian', std=float(conv_weight_filler.split('_')[1])) else: conv_weight_filler = eval(conv_weight_filler) assert bltr_weight_filler.startswith('gauss_') bltr_weight_filler = dict(type='gaussian', std=float(bltr_weight_filler.split('_')[1])) # multiple 1x1 conv-(bn)-relu blocks, optionally with a single global pooling somewhere among them idx = 1 bltr_idx = 0 lattices = dict() last_in_block = dict() for (layer_type, n_out) in arch_str: if layer_type == 'c': n['conv' + str(idx)] = L.Convolution( top_prev, convolution_param=dict(num_output=n_out, kernel_size=1, stride=1, pad=0, weight_filler=conv_weight_filler, bias_filler=dict(type='constant', value=0)), param=[dict(lr_mult=1), dict(lr_mult=0.1)]) elif layer_type == 'b': lattice_dims_str_curr = lattice_dims_str[bltr_idx] if lattice_dims_str_curr in lattices: top_data_lattice, top_lattice = lattices[lattice_dims_str_curr] n['conv' + str(idx)] = L.Permutohedral( top_prev, top_data_lattice, top_data_lattice, top_lattice, permutohedral_param=dict(num_output=n_out, group=1, neighborhood_size=bilateral_nbr, bias_term=True, norm_type=P.Permutohedral.AFTER, offset_type=P.Permutohedral.NONE, filter_filler=bltr_weight_filler, bias_filler=dict(type='constant', value=0)), param=[{ 'lr_mult': 1, 'decay_mult': 1 }, { 'lr_mult': 2, 'decay_mult': 0 }]) else: top_data_lattice = L.Python( n.data, python_param=dict(module='custom_layers', layer='PickAndScale', param_str=lattice_dims_str_curr)) n['data_lattice' + str(len(lattices))] = top_data_lattice if lattice_dims_str.count(lattice_dims_str_curr) > 1: n['conv' + str(idx)], top_lattice = L.Permutohedral( top_prev, top_data_lattice, top_data_lattice, ntop=2, permutohedral_param=dict( num_output=n_out, group=1, neighborhood_size=bilateral_nbr, bias_term=True, norm_type=P.Permutohedral.AFTER, offset_type=P.Permutohedral.NONE, filter_filler=bltr_weight_filler, bias_filler=dict(type='constant', value=0)), param=[{ 'lr_mult': 1, 'decay_mult': 1 }, { 'lr_mult': 2, 'decay_mult': 0 }]) n['lattice' + str(len(lattices))] = top_lattice else: n['conv' + str(idx)] = L.Permutohedral( top_prev, top_data_lattice, top_data_lattice, permutohedral_param=dict( num_output=n_out, group=1, neighborhood_size=bilateral_nbr, bias_term=True, norm_type=P.Permutohedral.AFTER, offset_type=P.Permutohedral.NONE, filter_filler=bltr_weight_filler, bias_filler=dict(type='constant', value=0)), param=[{ 'lr_mult': 1, 'decay_mult': 1 }, { 'lr_mult': 2, 'decay_mult': 0 }]) top_lattice = None lattices[lattice_dims_str_curr] = (top_data_lattice, top_lattice) bltr_idx += 1 top_prev = n['conv' + str(idx)] if batchnorm: n['bn' + str(idx)] = L.BatchNorm(top_prev) top_prev = n['bn' + str(idx)] n['relu' + str(idx)] = L.ReLU(top_prev, in_place=True) top_prev = n['relu' + str(idx)] # skip connection & global pooling if skip_str is None: skip_str = () skip_tos = [v.split('_')[0] for v in skip_str] if str(idx) in skip_tos: skip_idxs = list( filter(lambda i: skip_tos[i] == str(idx), range(len(skip_tos)))) skip_params = [skip_str[i].split('_') for i in skip_idxs] if len(skip_params[0]) == 2: assert all(len(v) == 2 for v in skip_params) else: assert all(v[2] == skip_params[0][2] for v in skip_params) if len(skip_params[0]) > 2 and 'g' in skip_params[0][ 2]: # global pooling on current layer n['gpool' + str(idx)] = L.Python(top_prev, python_param=dict( module='custom_layers', layer='GlobalPooling')) top_prev = n['gpool' + str(idx)] if len(skip_params[0]) > 2 and 'a' in skip_params[0][ 2]: # addition instead of concatenation n['add' + str(idx)] = L.Eltwise( top_prev, *[last_in_block[int(v[1])] for v in skip_params], eltwise_param=dict(operation=P.Eltwise.SUM)) top_prev = n['add' + str(idx)] else: n['concat' + str(idx)] = L.Concat( top_prev, *[last_in_block[int(v[1])] for v in skip_params]) top_prev = n['concat' + str(idx)] last_in_block[idx] = top_prev idx += 1 # classification & loss n['conv' + str(idx)] = L.Convolution( top_prev, convolution_param=dict(num_output=nclass, kernel_size=1, stride=1, pad=0, weight_filler=conv_weight_filler, bias_filler=dict(type='constant', value=0)), param=[dict(lr_mult=1), dict(lr_mult=0.1)]) top_prev = n['conv' + str(idx)] if deploy: n.prob = L.Softmax(top_prev) else: n.loss = L.SoftmaxWithLoss(top_prev, n.label) n.accuracy = L.Accuracy(top_prev, n.label) net = n.to_proto() if create_prototxt: net = get_prototxt(net, save_path) return net