Source code for wbia_cnn.models.aoi

# -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function
import functools
import six
import numpy as np
import utool as ut
from wbia_cnn import ingest_data
from Lasagne.lasagne import layers, nonlinearities
from theano import tensor as T  # NOQA
from wbia_cnn.models import abstract_models

print, rrr, profile = ut.inject2(__name__)


[docs]def augment_parallel(X, y, w): return augment_wrapper( [X], None if y is None else [y], None if w is None else [w], )
[docs]def augment_wrapper(Xb, yb=None, wb=None): Xb_ = [] yb_ = [] wb_ = [] for index in range(len(Xb)): X = Xb[index] y = yb[index] for values in y: bbox = values[:4] Xb_.append(np.append(X, bbox)) yb_.append(values[4]) wb_.append(1.0) Xb_ = np.array(Xb_, dtype=Xb.dtype) yb_ = np.array(yb_, dtype=yb.dtype) wb_ = np.array(wb_, dtype=np.float32) return Xb_, yb_, wb_
[docs]@six.add_metaclass(ut.ReloadingMetaclass) class AoIModel(abstract_models.AbstractVectorVectorModel): def __init__( model, autoinit=False, batch_size=128, data_shape=(64, 64, 3), name='aoi', **kwargs ): super(AoIModel, model).__init__( batch_size=batch_size, data_shape=data_shape, name=name, **kwargs )
[docs] def augment(model, Xb, yb=None, wb=None, parallel=False): if not parallel: return augment_wrapper(Xb, yb, wb) # Run in parallel if yb is None: yb = [None] * len(Xb) if wb is None: wb = [None] * len(Xb) arg_iter = list(zip(Xb, yb, wb)) result_list = ut.util_parallel.generate2( augment_parallel, arg_iter, ordered=True, verbose=False ) result_list = list(result_list) X = [result[0][0] for result in result_list] X = np.array(X) if yb is None: y = None else: y = [result[1] for result in result_list] y = np.hstack(y) if wb is None: w = None else: w = [result[2] for result in result_list] w = np.hstack(w) return X, y, w
[docs] def get_aoi_def(model, verbose=ut.VERBOSE, **kwargs): # _CaffeNet = abstract_models.PretrainedNetwork('caffenet') _P = functools.partial hidden_initkw = { 'nonlinearity': nonlinearities.LeakyRectify(leakiness=(1.0 / 10.0)), } (_, input_size) = model.input_shape network_layers_def = [ _P( layers.InputLayer, shape=( None, input_size, ), ), _P(layers.DenseLayer, num_units=1024, name='F0', **hidden_initkw), _P(layers.FeaturePoolLayer, pool_size=2, name='FP0'), _P(layers.DropoutLayer, p=0.5, name='D0'), _P(layers.DenseLayer, num_units=512, name='F1', **hidden_initkw), _P(layers.FeaturePoolLayer, pool_size=2, name='FP1'), _P(layers.DropoutLayer, p=0.5, name='D1'), _P(layers.DenseLayer, num_units=256, name='F2', **hidden_initkw), _P(layers.FeaturePoolLayer, pool_size=2, name='FP2'), _P(layers.DropoutLayer, p=0.5, name='D2'), _P(layers.DenseLayer, num_units=128, name='F3', **hidden_initkw), _P(layers.FeaturePoolLayer, pool_size=2, name='FP3'), _P(layers.DropoutLayer, p=0.5, name='D3'), _P( layers.DenseLayer, num_units=model.output_dims, name='F4', nonlinearity=nonlinearities.sigmoid, ), ] return network_layers_def
[docs] def init_arch(model, verbose=True, **kwargs): r"""""" (_, input_size) = model.input_shape if verbose: print('[model] Initialize aoi model architecture') print('[model] * batch_size = %r' % (model.batch_size,)) print('[model] * input_width = %r' % (input_size,)) print('[model] * output_dims = %r' % (model.output_dims,)) network_layers_def = model.get_aoi_def(verbose=verbose, **kwargs) # connect and record layers from wbia_cnn import custom_layers network_layers = custom_layers.evaluate_layer_list( network_layers_def, verbose=verbose ) # model.network_layers = network_layers output_layer = network_layers[-1] model.output_layer = output_layer return output_layer
[docs]def train_aoi(output_path, data_fpath, labels_fpath): r""" CommandLine: python -m wbia_cnn.train --test-train_aoi Example: >>> # DISABLE_DOCTEST >>> from wbia_cnn.train import * # NOQA >>> result = train_aoi() >>> print(result) """ era_size = 256 batch_size = 16 max_epochs = era_size * 16 hyperparams = ut.argparse_dict( { 'era_size': era_size, 'learning_rate': 0.01, 'rate_schedule': 0.75, 'momentum': 0.9, 'weight_decay': 0.0001, 'augment_on': True, 'augment_on_validate': True, 'whiten_on': False, 'max_epochs': max_epochs, 'stopping_patience': max_epochs, 'class_weight': None, } ) ut.colorprint('[netrun] Ensuring Dataset', 'yellow') dataset = ingest_data.get_numpy_dataset2('aoi', data_fpath, labels_fpath, output_path) X_train, y_train = dataset.subset('train') X_valid, y_valid = dataset.subset('valid') print('dataset.training_dpath = %r' % (dataset.training_dpath,)) input_shape = ( batch_size, dataset.data_shape[0] + 4, ) ut.colorprint('[netrun] Architecture Specification', 'yellow') model = AoIModel( input_shape=input_shape, training_dpath=dataset.training_dpath, **hyperparams ) ut.colorprint('[netrun] Initialize architecture', 'yellow') model.output_dims = 1 model.input_shape = ( None, dataset.data_shape[0] + 4, ) model.batch_size = batch_size model.init_arch() ut.colorprint('[netrun] * Initializing new weights', 'lightgray') if model.has_saved_state(): model.load_model_state() ut.colorprint('[netrun] Training Requested', 'yellow') # parse training arguments config = ut.argparse_dict( dict( era_size=era_size, max_epochs=max_epochs, show_confusion=False, ) ) model.monitor_config.update(**config) print('\n[netrun] Model Info') model.print_layer_info() ut.colorprint('[netrun] Begin training', 'yellow') model.fit(X_train, y_train, X_valid=X_valid, y_valid=y_valid) model_path = model.save_model_state() return model_path
if __name__ == '__main__': """ CommandLine: python -m wbia_cnn.models.aoi python -m wbia_cnn.models.aoi --allexamples python -m wbia_cnn.models.aoi --allexamples --noface --nosrc """ import multiprocessing multiprocessing.freeze_support() # for win32 import utool as ut # NOQA ut.doctest_funcs()