Source code for main

#!/usr/local/bin/python
# python std
import argparse
import logging
import sys
import importlib

# installed packages
import cv2

# local imports
import pca
import aam
# import imm

from reconstruction import reconstruction

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s %(levelname)s %(name)s: %(message)s')
logger = logging.getLogger(__name__)


def add_parser_options():
    parser = argparse.ArgumentParser(description='IMMPoints tool')
    pca_group = parser.add_argument_group('show_pca')

    pca_group.add_argument(
        '--reconstruct', action='store_true',
        help='Reconstruct one face with a given pca model'
    )

    pca_group.add_argument(
        '--show_kivy', action='store_true',
        help='Reconstruct using kivy as a GUI'
    )

    pca_group.add_argument(
        '--generate_call_graph', action='store_true',
        help='Generate call graph from the reconstruction'
    )

    pca_group.add_argument(
        '--save_pca_shape', action='store_true',
        help='save the pca shape model'
    )

    pca_group.add_argument(
        '--save_pca_texture', action='store_true',
        help='save the pca texture model'
    )

    pca_group.add_argument(
        '--show_pca', action='store_true',
        help='Show and manipulate the saved PCA model'
    )

    pca_group.add_argument(
        '--files', nargs='+', help='files to process'
    )

    pca_group.add_argument(
        '--n_components', default=10, type=int,
        help='number of principle components to keep and are able to manipulate'
    )

    pca_group.add_argument(
        '--model_shape_file', type=str,
        help='pca model file that contains or is going to contain the pca shape model'
    )

    pca_group.add_argument(
        '--shape_type', type=str, choices=['imm'],
        help='type of shape, annotated dataset'
    )

    pca_group.add_argument(
        '--model_texture_file', type=str,
        help='pca model file that contains or is going to contain the pca texture model'
    )

    return parser


[docs]def import_dataset_module(shape_type): """ Includes the right implementation for the right dataset implementation for the given shape type, see --help for the available options. Args: shape_type(string): Name of the python file inside the `src/datasets` folder. """ return importlib.import_module('datasets.{}'.format(shape_type))
[docs]def save_pca_model_texture(args): """ save the U, s, Vt and mean of all the asf datafiles given by the asf files. It is saved in the following way: np.load(filename, np.assary([Vt, mean_values]) And accessed by: Vtm = np.load(args.model_file_texture) Vt = Vtm[0] mean_values = Vtm[1][0] """ assert args.files, '--files should be given' assert args.model_shape_file, '--model_texture_file needs to be provided to save the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) shape_model = pca.PcaModel(args.model_shape_file) mean_points = dataset_module.IMMPoints(points_list=shape_model.mean_values) textures = aam.build_texture_feature_vectors( args.files, dataset_module.get_imm_image_with_landmarks, # function mean_points, shape_model.triangles ) mean_texture = aam.get_mean(textures) _, s, Vt, n_components = pca.pca(textures, mean_texture) pca.save(Vt, s, n_components, mean_texture, shape_model.triangles, args.model_texture_file) logger.info('texture pca model saved in %s', args.model_texture_file)
[docs]def save_pca_model_shape(args): """ save the U, s, Vt and mean of all the asf datafiles given by the asf files. It is saved in the following way: np.load(filename, np.assary([Vt, mean_values]) And accessed by: Vtm = np.load(args.model_shape_file) Vt = Vtm[0] mean_values = Vtm[1][0] """ assert args.files, '--files should be given' assert args.model_shape_file, '--model_shape_file needs to be provided to save the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) points = aam.build_shape_feature_vectors( args.files, dataset_module.get_imm_points, flattened=True ) mean_values = aam.get_mean(points) _, s, Vt, n_components = pca.pca(points, mean_values) mean_xy = mean_values.reshape((-1, 2)) triangles = aam.get_triangles(mean_xy[:, 0], mean_xy[:, 1]) pca.save(Vt, s, n_components, mean_values, triangles, args.model_shape_file) logger.info('shape pca model saved in %s', args.model_shape_file + '_shape')
def reconstruct_with_model(args): assert args.files, '--files should be given to allow the image to be shown' assert args.model_shape_file, '--model_shape_file needs to be provided to get the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) # clear sys args. arguments are conflicting with parseargs # kivy will parse args upon import and will crash if it finds our # 'unsupported by kivy' arguments. sys.argv[1:] = [] from view.reconstruct import ReconstructApp Vt_shape, s, n_shape_components, mean_value_points, triangles = pca.load(args.model_shape_file) Vt_texture, s_texture, n_texture_components, mean_values_texture, _ = pca.load(args.model_texture_file) app = ReconstructApp() app.set_values( args=args, eigenv_shape=Vt_shape, eigenv_texture=Vt_texture, mean_value_points=mean_value_points, n_shape_components=n_shape_components, mean_values_texture=mean_values_texture, n_texture_components=n_texture_components, triangles=triangles ) app.run() def show_pca_model(args): assert args.model_shape_file, '--model_texture_file needs to be provided to save the pca model' assert args.model_texture_file, '--model_texture_file needs to be provided to save the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) from reconstruction.triangles import draw_shape, get_texture Vt_shape, s, n_shape_components, mean_value_points, triangles = pca.load(args.model_shape_file) Vt_texture, s_texture, n_texture_components, mean_values_texture, _ = pca.load(args.model_texture_file) imm_points = dataset_module.IMMPoints( filename='data/imm_face_db/40-1m.asf' ) input_image = imm_points.get_image() input_points = imm_points.get_points() h, w, c = input_image.shape input_points[:, 0] = input_points[:, 0] * w input_points[:, 1] = input_points[:, 1] * h mean_value_points = mean_value_points.reshape((58, 2)) mean_value_points[:, 0] = mean_value_points[:, 0] * w mean_value_points[:, 1] = mean_value_points[:, 1] * h while True: dst = get_texture(mean_value_points, mean_values_texture) cv2.imshow('input_image', input_image) cv2.imshow('image', dst) k = cv2.waitKey(0) & 0xFF if k == 27: break cv2.destroyAllWindows()
[docs]def generate_call_graph(args): """Performance debug function, will be (re)moved later. """ assert args.model_shape_file, '--model_texture_file needs to be provided to save the pca model' assert args.model_texture_file, '--model_texture_file needs to be provided to save the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) from pycallgraph import PyCallGraph from pycallgraph.output import GraphvizOutput graphviz = GraphvizOutput(output_file='filter_none.png') with PyCallGraph(output=graphviz): shape_model = pca.PcaModel(args.model_shape_file) texture_model = pca.PcaModel(args.model_texture_file) input_points = dataset_module.IMMPoints(filename='data/imm_face_db/40-3m.asf') input_image = input_points.get_image() mean_points = dataset_module.IMMPoints(points_list=shape_model.mean_values) mean_points.get_scaled_points(input_image.shape) reconstruction.reconstruct_texture( input_image, # src image input_image, # dst image texture_model, input_points, # shape points input mean_points, # shape points mean )
def show_reconstruction(args): assert args.model_shape_file, '--model_texture_file needs to be provided to save the pca model' assert args.model_texture_file, '--model_texture_file needs to be provided to save the pca model' assert args.shape_type, '--shape_type the type of dataset, see datasets module' dataset_module = import_dataset_module(args.shape_type) shape_model = pca.PcaModel(args.model_shape_file) texture_model = pca.PcaModel(args.model_texture_file) input_points = dataset_module.IMMPoints( filename='data/imm_face_db/40-3m.asf' ) input_image = input_points.get_image() mean_points = dataset_module.IMMPoints(points_list=shape_model.mean_values) mean_points.get_scaled_points(input_image.shape) while True: reconstruction.reconstruct_texture( input_image, # src image input_image, # dst image texture_model, input_points, # shape points input mean_points, # shape points mean ) dst = reconstruction.get_texture( mean_points, texture_model.mean_values ) cv2.imshow('original', input_points.get_image()) cv2.imshow('reconstructed', input_image) cv2.imshow('main face', dst) k = cv2.waitKey(0) & 0xFF if k == 27: break cv2.destroyAllWindows()
[docs]def main(): """main""" parser = add_parser_options() args = parser.parse_args() if args.show_pca: show_pca_model(args) elif args.save_pca_shape: save_pca_model_shape(args) elif args.save_pca_texture: save_pca_model_texture(args) elif args.reconstruct: #reconstruct_with_model(args) show_reconstruction(args) elif args.show_kivy: reconstruct_with_model(args) elif args.generate_call_graph: generate_call_graph(args)
if __name__ == '__main__': main()