- fix i/o to match more consistent dataframe
[face-privacy-filter.git] / face_privacy_filter / filter_image.py
1 #! python
2 # -*- coding: utf-8 -*-
3 """
4 Wrapper for image emotion classification task
5 """
6
7 import os.path
8 import sys
9
10 import numpy as np
11 import pandas as pd
12
13 from face_privacy_filter.transform_detect import FaceDetectTransform
14 from face_privacy_filter.transform_region import RegionTransform
15 from face_privacy_filter._version import MODEL_NAME
16
17
18 def model_create_pipeline(transformer):
19     from acumos.session import Requirements
20     from acumos.modeling import Model, List, create_namedtuple
21     import sklearn
22     import cv2
23     from os import path
24
25     # derive the input type from the transformer
26     type_list, type_name = transformer._type_in  # it looked like this {'test': int, 'tag': str}
27     input_type = [(k, List[type_list[k]]) for k in type_list]
28     type_in = create_namedtuple(type_name, input_type)
29
30     # derive the output type from the transformer
31     type_list, type_name = transformer._type_out
32     output_type = [(k, List[type_list[k]]) for k in type_list]
33     type_out = create_namedtuple(type_name, output_type)
34
35     def predict_class(value: type_in) -> type_out:
36         '''Returns an array of float predictions'''
37         df = pd.DataFrame(np.column_stack(value), columns=value._fields)
38         tags_df = transformer.predict(df)
39         tags_list = type_out(*(col for col in tags_df.values.T))  # flatten to tag set
40         return tags_list
41
42     # compute path of this package to add it as a dependency
43     package_path = path.dirname(path.realpath(__file__))
44     return Model(transform=predict_class), Requirements(packages=[package_path], reqs=[pd, np, sklearn],
45                                                         req_map={cv2: 'opencv-python'})
46
47
48 def main(config={}):
49     import argparse
50     parser = argparse.ArgumentParser()
51     parser.add_argument('-p', '--predict_path', type=str, default='', help="save detections from model (model must be provided via 'dump_model')")
52     parser.add_argument('-i', '--input', type=str, default='', help='absolute path to input data (image or csv, only during prediction / dump)')
53     parser.add_argument('-c', '--csv_input', dest='csv_input', action='store_true', default=False, help='input as CSV format not an image')
54     parser.add_argument('-s', '--suppress_image', dest='suppress_image', action='store_true', default=False, help='do not create an extra row for a returned image')
55     parser.add_argument('-f', '--function', type=str, default='detect', help='which type of model to generate', choices=['detect', 'pixelate'])
56     parser.add_argument('-a', '--push_address', help='server address to push the model (e.g. http://localhost:8887/v2/models)', default='')
57     parser.add_argument('-d', '--dump_model', help='dump model to a pickle directory for local running', default='')
58     config.update(vars(parser.parse_args()))     # pargs, unparsed = parser.parse_known_args()
59
60     if not config['predict_path']:
61         print("Attempting to create new model for dump or push...")
62
63         # refactor the raw samples from upstream image classifier
64         if config['function'] == "detect":
65             transform = FaceDetectTransform(include_image=not config['suppress_image'])
66         elif config['function'] == "pixelate":
67             transform = RegionTransform()
68         else:
69             print("Error: Functional mode '{:}' unknown, aborting create".format(config['function']))
70         inputDf = transform.generate_in_df()
71         pipeline, reqs = model_create_pipeline(transform)
72
73         # formulate the pipeline to be used
74         model_name = MODEL_NAME + "_" + config['function']
75         if 'push_address' in config and config['push_address']:
76             from acumos.session import AcumosSession
77             print("Pushing new model to '{:}'...".format(config['push_address']))
78             session = AcumosSession(push_api=config['push_address'], auth_api=config['auth_address'])
79             session.push(pipeline, model_name, reqs)  # creates ./my-iris.zip
80
81         if 'dump_model' in config and config['dump_model']:
82             from acumos.session import AcumosSession
83             from os import makedirs
84             if not os.path.exists(config['dump_model']):
85                 makedirs(config['dump_model'])
86             print("Dumping new model to '{:}'...".format(config['dump_model']))
87             session = AcumosSession()
88             session.dump(pipeline, model_name, config['dump_model'], reqs)  # creates ./my-iris.zip
89
90     else:
91         if not config['dump_model'] or not os.path.exists(config['dump_model']):
92             print("Attempting to predict from a dumped model, but model not found.".format(config['dump_model']))
93             sys.exit(-1)
94         if not os.path.exists(config['input']):
95             print("Predictino requested but target input '{:}' was not found, please check input arguments.".format(config['input']))
96             sys.exit(-1)
97
98         print("Attempting predict/transform on input sample...")
99         from acumos.wrapped import load_model
100         model = load_model(config['dump_model'])
101         if not config['csv_input']:
102             inputDf = FaceDetectTransform.generate_in_df(config['input'])
103         else:
104             inputDf = pd.read_csv(config['input'], converters={FaceDetectTransform.COL_IMAGE_DATA: FaceDetectTransform.read_byte_arrays})
105
106         type_in = model.transform._input_type
107         transform_in = type_in(*tuple(col for col in inputDf.values.T))
108         transform_out = model.transform.from_wrapped(transform_in).as_wrapped()
109         dfPred = pd.DataFrame(np.column_stack(transform_out), columns=transform_out._fields)
110
111         if not config['csv_input']:
112             dfPred = FaceDetectTransform.suppress_image(dfPred)
113         print("ALMOST DONE")
114         print(dfPred)
115
116         if config['predict_path']:
117             print("Writing prediction to file '{:}'...".format(config['predict_path']))
118             if not config['csv_input']:
119                 dfPred.to_csv(config['predict_path'], sep=",", index=False)
120             else:
121                 FaceDetectTransform.generate_out_image(dfPred, config['predict_path'])
122
123         if dfPred is not None:
124             print("Predictions:\n{:}".format(dfPred))
125
126
127 if __name__ == '__main__':
128     main()