def _process(self, task):
self.logger.info('processing task {}: {}, {}, {}...'.format(
task.task_id, task.func_group, task.func_name, task.func_params))
process_ok = False
if not self._engine:
self._engine_init()
if self._engine:
image = load_task_obj_data(app_config.SAVE_FILE_NAME['image_save'],
task.m_id, app_config, logger)
face_data = load_task_obj_data(
app_config.SAVE_FILE_NAME['face_data'], task.m_id, app_config,
logger)
# 根据 task 内容选择相应引擎处理函数，对图片进行处理
image_output = image_process_wrap(image,
face_data,
self._func_process_image,
self._engine,
task.func_class,
task.func_group,
task.func_name,
task.func_params,
worker=self)
image.close()
if image_output:
save_task_obj_data(image_output,
app_config.SAVE_FILE_NAME['image_temp'],
task.m_id, app_config, logger)
process_ok = True
self.logger.info('task {} process done.'.format(task.task_id))
# 对于 demo 图片，定期更新缓存结果
if 'img_hash' in face_data:
img_hash = face_data['img_hash']
if img_hash in app_config.IMG_HASH_APP_DEMO: # APP 前端所用的四个 demo 图片之一
if task.func_params and len(task.func_params) > 0:
func_args = str(task.func_params[0])
else:
func_args = None
cache_file_name = make_cache_file_name(
task.func_class, task.func_group, task.func_name,
func_args)
save_cache_demo_image(image_output, img_hash,
app_config.DEMO_CACHE_DIR,
cache_file_name)
else:
self.logger.info(
'task {} process failed. Unsupported picture?'.format(
task.task_id))
return process_ok
def run_worker(self):
"""定时循环获取新任务，对任务图片进行预处理，并将处理完的任务节点进行迁移"""
# 首先执行人脸检测和特征识别辅助引擎初始化
self.logger.info("initializing face detector engine...")
self.face_detector = FaceDetector()
self.logger.info("initializing face attribute engine...")
self.face_attribute = face_attribute_adapter.init_model()
func_list = self._mm_client.shared_obj_func_list('task_controller')
self.logger.info('task controller functions: {}'.format(func_list))
func_list = self._mm_client.shared_obj_func_list('engine_controller')
self.logger.info('engine controller functions: {}'.format(func_list))
while True:
time.sleep(0.2) # 间隔定时处理，避免空耗CPU
result = self._mm_client.shared_obj_func_call(
'task_controller', 'new_task_get')
if result and result['return'] == 'success' and result['data']:
task = result['data']
self.logger.info('got a new task: {}'.format(task.task_id))
task_stat = TaskStat.SUCCESS # 预处理成功后的状态
task_msg = None
# 1、读取新提交的任务图片
image = load_task_obj_data(
app_config.SAVE_FILE_NAME['image_base'], task.m_id,
app_config, logger)
# 2、对图片进行归一化处理，并保存到文件
image_output = image_process_wrap(image) # 不指定引擎函数即只做归一化
for ff in (app_config.SAVE_FILE_NAME['image_norm'],
app_config.SAVE_FILE_NAME['image_temp'],
app_config.SAVE_FILE_NAME['image_save']):
save_task_obj_data(image_output, ff, task.m_id, app_config,
logger)
image.close()
# 3、调用引擎获取人脸地图及人脸特征，保存到文件
#（大部分 AI 模型只支持正方形的面部图片，所以需要进行人脸识别，并截取）
# 不管是否有人脸图均返回成功，对于无人脸图片引擎处理时再判断报错
face_data = {}
image_output_np = np.asarray(image_output)
# 获取人脸位置及图像（如果人脸部分尺寸小于模型尺寸，则返回指定的最小尺寸为模型尺寸）
faces_images = self.face_detector.get_faces_images(
image_output_np, app_config.FACE_MIN_SIZE)
if len(faces_images) == 0:
task_msg = "no face found in this picture."
self.logger.info(task_msg)
else:
self.logger.info("got %d faces_images" % len(faces_images))
# 识别第一张人脸的特征，并保存 #FIXME: 或者取中间的人脸？
for (x, y, w, h, face_image) in faces_images:
attr_all = face_attribute_adapter.get_face_attrs(
self.face_attribute, face_image)
self.logger.info(
'face_attribute_adapter get_face_attrs: %s' %
attr_all)
# 从 celeba 全部 40 个属性中选取所需的部分
name_all = "5_o_Clock_Shadow Arched_Eyebrows Attractive Bags_Under_Eyes Bald Bangs Big_Lips Big_Nose Black_Hair Blond_Hair Blurry Brown_Hair Bushy_Eyebrows Chubby Double_Chin Eyeglasses Goatee Gray_Hair Heavy_Makeup High_Cheekbones Male Mouth_Slightly_Open Mustache Narrow_Eyes No_Beard Oval_Face Pale_Skin Pointy_Nose Receding_Hairline Rosy_Cheeks Sideburns Smiling Straight_Hair Wavy_Hair Wearing_Earrings Wearing_Hat Wearing_Lipstick Wearing_Necklace Wearing_Necktie Young".split(
)
# default attrs: Bald Bangs Black_Hair Blond_Hair Brown_Hair Bushy_Eyebrows Eyeglasses Male Mouth_Slightly_Open Mustache No_Beard Pale_Skin Young
attr = [
attr_all[name_all.index(n)]
for n in app_config.ATTRS_ENABLED
]
attr_dict = {}
for i in range(len(attr)):
if attr[i] == 1:
attr_dict[app_config.ATTRS_ENABLED[i]] = True
else:
attr_dict[app_config.ATTRS_ENABLED[i]] = False
self.logger.info(
'selected attr: {}, attr_dict: {}'.format(
attr, attr_dict))
face_data['x'] = x
face_data['y'] = y
face_data['w'] = w
face_data['h'] = h
face_data['attr'] = attr # 人脸属性（0 或 1 的列表）
face_data['attr_dict'] = attr_dict # 人脸属性（名称、键值字典）
face_data['face_image'] = face_image # 截取的人脸图片
#print('--> task.m_id:', task.m_id) # 20200107-07-45-32f6e7262b0e3c9b99c59f710d652b2a
face_data['img_hash'] = task.m_id.split('-')[-1]
#cv2.imwrite(join(app_config.SAVE_FILE_NAME['face_data'], '.jpg'), face_image)
task_msg = face_data['attr_dict']
break # 只处理第一张人脸
# 保存人脸图片及位置等信息到 pickle 文件
save_task_obj_data(face_data,
app_config.SAVE_FILE_NAME['face_data'],
task.m_id, app_config, logger)
# 4、更新任务节点状态
result = self._mm_client.shared_obj_func_call(
'task_controller', 'task_update',
(task.task_id, TaskQueue.NEW, {
'task_stat': task_stat,
'msg': task_msg
}))
self.logger.info('task_update {} result: {}'.format(
task.task_id, result))
# 5、迁移到”处理任务队列“（任务预处理引擎已在 new_task_get 中做迁移了）
result = self._mm_client.shared_obj_func_call(
'task_controller', 'task_move', (
task.task_id,
TaskQueue.NEW,
TaskQueue.PROCESS,
))
self.logger.info('task_move {} result: {}'.format(
task.task_id, result))
# 对于 demo 图片，定期更新缓存结果
if task_stat != TaskStat.FAIL and 'img_hash' in face_data:
img_hash = face_data['img_hash']
if img_hash in app_config.IMG_HASH_APP_DEMO: # APP 前端所用的四个 demo 图片之一
cache_file_name = make_cache_file_name()
save_cache_demo_image(image_output, img_hash,
app_config.DEMO_CACHE_DIR,
cache_file_name)
# 上述预处理文件，也需要拷贝到缓存目录下
dir_from = os.path.join(app_config.UPLOAD_FOLDER,
task.m_id.replace('-', '/'))
img_hash = task.m_id.split('-')[-1]
for ff in (app_config.SAVE_FILE_NAME['image_norm'],
app_config.SAVE_FILE_NAME['image_temp'],
app_config.SAVE_FILE_NAME['image_save'],
app_config.SAVE_FILE_NAME['face_data']):
file_from = os.path.join(dir_from, ff)
file_to = os.path.join(app_config.DEMO_CACHE_DIR,
img_hash, ff)
shutil.copyfile(file_from, file_to)
else:
#print('no new task found.')
pass
def task_apply_image_save():
"""保存任务处理后图片
"""
data = {}
if not current_app.config['ENABLE_APPLY_IMAGE_SAVE']:
data['msg'] = 'Function disabled. Bypass task_apply_image_save().'
current_app.logger.info(data['msg'])
return resp_result(data)
m_id = request.args.get('m_id')
if m_id:
try:
demo_cahce_file = None
img_hash = m_id.split('-')[-1]
if img_hash in current_app.config['IMG_HASH_APP_DEMO']: # 前端所用四个 demo 图片之一
current_app.logger.info('task_apply_image_save() got request args: {}'.format(repr(request.args)))
f_group, f_class, f_name, f_param0 = __get_func_args_from_request_args(request)
demo_cahce_file = os.path.join(current_app.config['DEMO_CACHE_DIR'], img_hash, make_cache_file_name(f_class, f_group, f_name, f_param0))
save_task_image_from_selected(m_id, current_app.config, current_app.logger, demo_cahce_file)
except Exception as e:
current_app.logger.error('task_apply_image_save failed: {}'.format(e))
data['m_id'] = m_id
data['msg'] = str(e)
return resp_error(data)
return resp_result(data)
def task_add():
"""根据客户端请求生成图片处理任务"""
data = {}
if request.args.get("tk") == "harekrishna10800": # FIXME: 改为正规的令牌校验机制
if request.method == "POST":
m_id, params, msg, is_new_image = get_request_data(request)
current_app.logger.info('task_add() got post request. params: {}, msg: {}, '
'm_id: {}, is_new_image: {}'
.format(params, msg, m_id, is_new_image))
if not params or not m_id:
return resp_error(msg)
# 根据 m_id 和 params 生成新任务，提交后端进行处理
task = {}
task.update(params)
task['m_id'] = m_id
#print('--> repr(task):', repr(task))
# 新提交图片 task 示例：{'user_id': 'zzz123', 'func_group': None, 'm_id': '20200108-10-35-12d2b6c4a62be56c494324809441be41', 'func_class': None, 'func_params': None, 'func_name': None}
# 图片任务 task 示例 {'func_params': ['avatar', ['face', 'l_ear', 'r_ear', 'nose', 'u_lip', 'l_lip']], 'func_class': 'fc_editor', 'func_name': 'chg_color', 'func_group': 'impression', 'user_id': 'zzz123'}
data['m_id'] = m_id
# 如果是 demo 图片，且请求的处理已经有未过期缓存，则直接返回一个特殊且固定的 task_id
img_hash = m_id.split('-')[-1]
if img_hash in current_app.config['IMG_HASH_APP_DEMO']: # 前端所用四个 demo 图片之一
if task['func_params'] and type(task['func_params']) is list \
and len(task['func_params']) > 0:
func_args = str(task['func_params'][0])
else:
func_args = None
cache_file = os.path.join(current_app.config['DEMO_CACHE_DIR'], img_hash, make_cache_file_name(task['func_class'], task['func_group'],
task['func_name'], func_args))
current_app.logger.info('task_add() check cache file {}'.format(cache_file))
if os.path.isfile(cache_file) and time.time() - os.path.getctime(cache_file) <= current_app.config['DEMO_CACHE_UPDATE_TERM']:
data['task_id'] = 'THE_CACHE_TASK_ID' # 其他 task_id 为浮点数
current_app.logger.info('task_add() found valid cache, bypass new_task_add.')
# 从缓存目录拷贝预处理文件到任务目录下，以便后续没有做过缓存的可以正常处理
dir_to = os.path.join(current_app.config['UPLOAD_FOLDER'],
m_id.replace('-', '/'))
for ff in (current_app.config['SAVE_FILE_NAME']['image_norm'],
current_app.config['SAVE_FILE_NAME']['image_temp'],
current_app.config['SAVE_FILE_NAME']['image_save'],
current_app.config['SAVE_FILE_NAME']['face_data']):
file_from = os.path.join(current_app.config['DEMO_CACHE_DIR'],
img_hash, ff)
file_to = os.path.join(dir_to, ff)
shutil.copyfile(file_from, file_to)
return resp_result(data)
# 判断是否新提交图片，如果是则置任务状态为 NEW，否则为 PENDING
from ..engine_v1.task import TaskStat
if is_new_image:
task['task_stat'] = TaskStat.NEW
else:
#TODO: 在此通过后端逻辑判断是否为 VIP 功能，以及检查用户是否有 VIP 权限！（v2版本实现）
task['task_stat'] = TaskStat.PENDING
with current_app.app_context():
try:
#func_list = current_app.mm_client.shared_obj_func_lists('task_controller')
#current_app.logger.warn('task controller functions: {}'.format(func_list))
result = current_app.mm_client.shared_obj_func_call(
'task_controller', 'new_task_add', (task,))
current_app.logger.info('call new_task_add return: {}'.format(result))
if result['return'] == 'success' and result['data'] is not None:
data['task_id'] = result['data']['task_id']
except Exception as e:
current_app.logger.error(e)
else:
return resp_error("invalid method")
else:
return resp_error("invalid tk")
return resp_result(data)
def task_query():
"""查询任务状态
备注：为给引擎留出处理时间，可每秒钟重新查询，最多10次（可配置），由于不便在 flask 中使用阻塞式 sleep，需要在前端实现该逻辑。
"""
data = {}
result = None
task_id = request.args.get('task_id')
if task_id:
got_result_image = False
data['progress'] = 5
# 如果是请求 demo 图片处理结果的内定 task_id，则直接读取缓存并返回
if task_id == 'THE_CACHE_TASK_ID':
m_id = request.args.get('m_id')
img_hash = m_id.split('-')[-1]
if img_hash in current_app.config['IMG_HASH_APP_DEMO']: # 前端所用四个 demo 图片之一
current_app.logger.info('task_query() got request args: {}'.format(repr(request.args)))
func_group, func_class, func_name, func_param0 = __get_func_args_from_request_args(request)
cache_file = os.path.join(current_app.config['DEMO_CACHE_DIR'], img_hash, make_cache_file_name(func_class, func_group, func_name, func_param0))
current_app.logger.info('task_query() check cache file {}'.format(cache_file))
if os.path.isfile(cache_file):
result_image = read_local_image(cache_file)
if func_class == 'fc_fun':
# 演示图片，右下角加 logo 水印
result_image = image_add_watermark(result_image)
got_result_image = True
data['m_id'] = m_id
data['msg'] = json.dumps('got cache result.')
current_app.logger.info('task_query() loaded cache result for the demo image.')
# 如果没找到缓存，则向后端查询处理结果
if not got_result_image:
result = current_app.mm_client.shared_obj_func_call(
'task_controller', 'task_query', (task_id,))
current_app.logger.info('call task_query {} return: {}'
.format(task_id, result))
if result['return'] == 'success' and result['data'] is not None:
data['progress'] = result['data']['progress']
if data['progress'] == 100:
result_image = result['data']['image']
got_result_image = True
# 做超解析缩放为指定输出大小（已在引擎中处理，故此注释掉）
#out_img_w, out_img_h = current_app.config['OUTPUT_SIZE_FUN']
#with current_app.app_context():
# result_image = esrgan_adapter.process_image(result_image, current_app.esrgan, out_img_w, out_img_h)
if result['data']['func_class'] == 'fc_fun':
# 演示图片，右下角加 logo 水印
result_image = image_add_watermark(result_image)
data['m_id'] = result['data']['m_id']
data['msg'] = json.dumps(result['data']['msg'])
if got_result_image:
data['progress'] = 100
data['image'] = image_to_base64(result_image) # 转成 jpg base64 返回
return resp_result(data)