convert protobuf i/o to flat,row-like data
[face-privacy-filter.git] / face_privacy_filter / transform_region.py
1 #! python
2 # -*- coding: utf-8 -*-
3 # ===============LICENSE_START=======================================================
4 # Acumos Apache-2.0
5 # ===================================================================================
6 # Copyright (C) 2017-2018 AT&T Intellectual Property & Tech Mahindra. All rights reserved.
7 # ===================================================================================
8 # This Acumos software file is distributed by AT&T and Tech Mahindra
9 # under the Apache License, Version 2.0 (the "License");
10 # you may not use this file except in compliance with the License.
11 # You may obtain a copy of the License at
12 #
13 # http://www.apache.org/licenses/LICENSE-2.0
14 #
15 # This file is distributed on an "AS IS" BASIS,
16 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 # See the License for the specific language governing permissions and
18 # limitations under the License.
19 # ===============LICENSE_END=========================================================
20 """
21 Wrapper for region processing task; wrapped in classifier for pipieline terminus
22 """
23 import cv2
24 import pandas as pd
25 import numpy as np
26 from sklearn.base import BaseEstimator, ClassifierMixin
27 import base64
28
29 # NOTE: If this class were built in another model (e.g. another vendor, class, etc), we would need to
30 #       *exactly match* the i/o for the upstream (detection) and downstream (this processing)
31 # from face_privacy_filter.transform_detect import RegionTransform
32
33 from face_privacy_filter.transform_detect import FaceDetectTransform
34
35
36 class RegionTransform(BaseEstimator, ClassifierMixin):
37     '''
38     A sklearn classifier mixin that manpulates image content based on input
39     '''
40     CASCADE_DEFAULT_FILE = "data/haarcascade_frontalface_alt.xml.gz"
41
42     def __init__(self, transform_mode="pixelate"):
43         self.transform_mode = transform_mode    # specific image processing mode to utilize
44
45     def get_params(self, deep=False):
46         return {'transform_mode': self.transform_mode}
47
48     @staticmethod
49     def output_names_():
50         return [FaceDetectTransform.COL_IMAGE_MIME, FaceDetectTransform.COL_IMAGE_DATA]
51
52     @staticmethod
53     def generate_out_dict(bin_stream=b"", media=""):
54         return {FaceDetectTransform.COL_IMAGE_MIME: media, FaceDetectTransform.COL_IMAGE_DATA: bin_stream}
55
56     @staticmethod
57     def generate_in_df(idx=FaceDetectTransform.VAL_REGION_IMAGE_ID, x=0, y=0, w=0, h=0, image=0, bin_stream=b"", media=""):
58         return pd.DataFrame([], RegionTransform.generate_in_dict(idx=idx, x=x, y=y, h=h, w=w, image=image, bin_stream=bin_stream, media=media))
59
60     @staticmethod
61     def generate_in_dict(idx=FaceDetectTransform.VAL_REGION_IMAGE_ID, x=0, y=0, w=0, h=0, image=0, bin_stream=b"", media=""):
62         return FaceDetectTransform.generate_out_dict(idx=idx, x=x, y=y, h=h, w=w, image=image, bin_stream=bin_stream, media=media)
63
64     @property
65     def _type_in(self):
66         """Custom input type for this processing transformer"""
67         input_dict = RegionTransform.generate_in_dict()
68         return {k: type(input_dict[k]) for k in input_dict}, "DetectionFrame"
69
70     @property
71     def _type_out(self):
72         """Custom input type for this processing transformer"""
73         return {FaceDetectTransform.COL_IMAGE_MIME: str, FaceDetectTransform.COL_IMAGE_DATA: bytes}, "Image"
74
75     def score(self, X, y=None):
76         return 0
77
78     def fit(self, X, y=None):
79         return self
80
81     def predict(self, X, y=None):
82         """
83         Assumes a numpy array of [[mime_type, binary_string] ... ]
84            where mime_type is an image-specifying mime type and binary_string is the raw image bytes
85         """
86
87         # group by image index first
88         #   decode image at region -1
89         #   collect all remaining regions, operate with each on input image
90         #   generate output image, send to output
91
92         image_region_list = RegionTransform.transform_raw_sample(X)
93         listData = []
94         for image_data in image_region_list:
95             img = image_data['data']
96             for r in image_data['regions']:  # loop through regions
97                 x_max = min(r[0] + r[2], img.shape[1])
98                 y_max = min(r[1] + r[3], img.shape[0])
99                 if self.transform_mode == "pixelate":
100                     img[r[1]:y_max, r[0]:x_max] = \
101                         RegionTransform.pixelate_image(img[r[1]:y_max, r[0]:x_max])
102
103             # for now, we hard code to jpg output; TODO: add more encoding output (or try to match source?)
104             img_binary = cv2.imencode(".jpg", img)[1].tostring()
105             img_mime = 'image/jpeg'  # image_data['mime']
106
107             listData.append(RegionTransform.generate_out_dict(media=img_mime, bin_stream=img_binary))
108             # print("IMAGE {:} found {:} total rows".format(image_data['image'], len(image_data['regions'])))
109         return pd.DataFrame(listData, columns=RegionTransform.output_names_())
110
111     @staticmethod
112     def transform_raw_sample(raw_sample):
113         """Method to transform raw samples into dict of image and regions"""
114         raw_sample.sort_values([FaceDetectTransform.COL_IMAGE_IDX], ascending=True, inplace=True)
115         groupImage = raw_sample.groupby(FaceDetectTransform.COL_IMAGE_IDX)
116         return_set = []
117
118         for nameG, rowsG in groupImage:
119             local_image = {'image': -1, 'data': b"", 'regions': [], 'mime': ''}
120             image_row = rowsG[rowsG[FaceDetectTransform.COL_REGION_IDX] == FaceDetectTransform.VAL_REGION_IMAGE_ID]
121             if len(image_row) < 1:  # must have at least one image set
122                 print("Error: RegionTransform could not find a valid image reference for image set {:}".format(nameG))
123                 continue
124             if not len(image_row[FaceDetectTransform.COL_IMAGE_DATA]):  # must have valid image data
125                 print("Error: RegionTransform expected image data, but found empty binary string {:}".format(nameG))
126                 continue
127             image_byte = image_row[FaceDetectTransform.COL_IMAGE_DATA][0]
128             if type(image_byte) == str:
129                 image_byte = image_byte.encode()
130                 image_byte = bytearray(base64.b64decode(image_byte))
131             else:
132                 image_byte = bytearray(image_byte)
133             file_bytes = np.asarray(image_byte, dtype=np.uint8)
134             local_image['data'] = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
135             local_image['image'] = nameG
136             local_image['mime'] = image_row[FaceDetectTransform.COL_IMAGE_MIME]
137
138             # now proceed to loop around regions detected
139             for index, row in rowsG.iterrows():
140                 if row[FaceDetectTransform.COL_REGION_IDX] != FaceDetectTransform.VAL_REGION_IMAGE_ID:  # skip bad regions
141                     local_image['regions'].append([row[FaceDetectTransform.COL_FACE_X], row[FaceDetectTransform.COL_FACE_Y],
142                                                    row[FaceDetectTransform.COL_FACE_W], row[FaceDetectTransform.COL_FACE_H]])
143             return_set.append(local_image)
144         return return_set
145
146     ################################################################
147     # image processing routines (using opencv)
148
149     # http://www.jeffreythompson.org/blog/2012/02/18/pixelate-and-posterize-in-processing/
150     @staticmethod
151     def pixelate_image(img, blockSize=None):
152         if not img.shape[0] or not img.shape[1]:
153             return img
154         if blockSize is None:
155             blockSize = round(max(img.shape[0], img.shape[2]) / 8)
156         ratio = (img.shape[1] / img.shape[0]) if img.shape[0] < img.shape[1] else (img.shape[0] / img.shape[1])
157         blockHeight = round(blockSize * ratio)  # so that we cover all image
158         for x in range(0, img.shape[0], blockSize):
159             for y in range(0, img.shape[1], blockHeight):
160                 max_x = min(x + blockSize, img.shape[0])
161                 max_y = min(y + blockSize, img.shape[1])
162                 fill_color = img[x, y]  # img[x:max_x, y:max_y].mean()
163                 img[x:max_x, y:max_y] = fill_color
164         return img