Source code for datasets.imm

"""
.. module:: datasets
   :platform: Unix, Windows
   :synopsis: Contains imm dataset abstraction layer


"""

from matplotlib.tri import Triangulation

import cv2
import numpy as np
import argparse
import os

import aam


[docs]class IMMPoints(aam.AAMPoints): """Accepts IMM datapoint file which can be shown or used""" def __init__(self, filename=None, points_list=None): """ Args: filename: optional .asf file with the imm format points: optional list of x,y points """ assert filename is not None or points_list is not None, 'filename or \ a ndarray of points list should be given' self.filename = filename if filename: points_list = self.import_file(filename) aam.AAMPoints.__init__( self, normalized_flattened_points_list=points_list.flatten(), actual_shape=(58, 2) )
[docs] def get_points(self): """ Get the flattened list of points Returns: ndarray. flattened array of points, see AAMPoints for more information. """ return self.normalized_flattened_points_list
[docs] def get_image(self): """ Get the image corresponding to the filename If filename == image_1.asf, then we read image_1.jpg from disk and return this to the user. Returns: ndarray image """ return cv2.imread(self.image_file)
[docs] def import_file(self, filename): """ Import an .asf filename. Load the points into a list of points and store the relative path to image file. Returns: ndarray(float). Numpy array of landmark locations as stated in the .asf files. """ points_list = [] with open(filename, 'r') as f: lines = f.readlines() data = lines[16:74] dir_name = os.path.dirname(filename) self.image_file = "{}/{}".format(dir_name, lines[-1].strip()) for d in data: points_list.append(d.split()[2:4]) return np.asarray(points_list, dtype='f')
def draw_triangles(self, image, points, multiply=True): if multiply: h, w, c = image.shape points[:, 0] = points[:, 0] * w points[:, 1] = points[:, 1] * h point_indices = list(range(0, 58)) triangles = Triangulation(points[:, 0], points[:, 1]) for t, tri in enumerate(triangles.triangles): p1, p2, p3 = points[tri] cv2.line(image, tuple(p1), tuple(p2), (255, 0, 100), 1) cv2.line(image, tuple(p2), tuple(p3), (255, 0, 100), 1) cv2.line(image, tuple(p3), tuple(p1), (255, 0, 100), 1) for i, p in enumerate(points): point_index = int(point_indices[i]) cv2.putText(image, str(point_index), tuple((p[0], p[1])), cv2.FONT_HERSHEY_SIMPLEX, .5, (100, 0, 255)) cv2.circle(image, tuple(p), 3, color=(0, 255, 100)) def show_on_image(self, image, window_name='image', multiply=True): self.draw_triangles(image, self.points_list, multiply=multiply)
[docs] def show(self, window_name='image'): """show the image and datapoints on the image""" assert(len(self.points_list) > 0) assert(len(self.filename) > 0) image = self.get_image() self.draw_triangles(image, self.points_list)
[docs]def get_imm_points(files): """ This function does something. Args: files (array): Array of .asf full or relative path to .asf files. Returns: ndarray. Array of landmarks. """ points = [] for f in files: imm = IMMPoints(filename=f) points.append(imm.get_points()) return np.asarray(points)
[docs]def get_imm_image_with_landmarks(filename): """ Get Points with image and landmarks/points Args: filename(fullpath): .asf file Returns: image, points """ imm = IMMPoints(filename=filename) return imm.get_image(), imm.get_points()
def add_parser_options(): parser = argparse.ArgumentParser(description='IMMPoints tool') # asf files parser.add_argument( 'asf', type=str, nargs='+', help='asf files to process' ) return parser if __name__ == '__main__': parser = add_parser_options() args = parser.parse_args() for f in args.asf: imm = IMMPoints(f) imm.show()