# !/usr/bin/python3 # Built-in libraries import gc from pathlib import Path import json import os from aiohttp import web import aiohttp_cors import time import numpy as np from PIL import Image, ImageOps import requests import urllib3 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) from io import BytesIO from multiprocessing import Process, cpu_count from socket import SOL_SOCKET, SO_REUSEADDR, socket import cv2 from libs.app_tool import AppTool from InsPy import Remover class AiMattingV2: def __init__(self): self.app_tool = AppTool('ps_matting.log') self.BASE_DIR = Path(__file__).parent self.ZK_CONF = self.app_tool.zk() self.hostname = os.uname()[1] self.logger = self.app_tool.logger self.inspy_remover = None def __init_seg_net__(self): if isinstance(self.inspy_remover, Remover) is False: self.inspy_remover = Remover() #self.inspy_remover = Remover(fast=True, jit=True) def __load_image__(self, data, mode = None): """ Loads an image file for other processing :param data: Path to image file or PIL image :return: image tensor, original pil image """ if isinstance(data, str) or isinstance(data, Path): try: self.logger.info("StartFetch ImgUrl %s", data) is_url = self.app_tool.uri_validator(data) if is_url is True: headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36'} response = requests.get(data, verify=False, timeout=10, headers=headers) image = Image.open(BytesIO(response.content)) # Load image if there is a path else: image = Image.open(BytesIO(data)) # Load image if there is a path self.logger.info("EndFetch ImgUrl %s", data) except IOError: self.logger.error('Cannot retrieve image. Please check file: ' + str(data)) return False elif isinstance(data, bytes): image = Image.open(BytesIO(data)) else: image = data if mode is not None: #image = ImageOps.exif_transpose(image) mode = mode.upper() image = image.convert(mode) return image def __process_extend__(self, image, extend = None): try: #image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA) """ if isinstance(image, Image.Image) is False: return False """ # 将 PIL 图像转换为 OpenCV 格式 image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA) # 如果图像是RGB模式 image_x = image.shape[1] image_y = image.shape[0] #总面积 total_area = image_x * image_y # 提取透明通道 alpha_channel = image[:, :, 3] alpha_channel_cp = alpha_channel #非透明区域 total_alpha_channel_area = cv2.countNonZero(alpha_channel) #抠图面积小于10则任务抠图失败 alpha_precent = round((total_alpha_channel_area / total_area) * 100, 2) if alpha_precent <= 10: self.logger.info("process v3 fail alpha_precent is %s", alpha_precent) return 'bad' #半透明（0 ~ 50） 大于抠图面积的 30% drop alpha_mask = ((alpha_channel >= 1) & (alpha_channel <= 50)).astype(np.uint8) valid_alpha_count = np.count_nonzero(alpha_mask) valid_alpha_ratio = round((valid_alpha_count / total_alpha_channel_area) * 100, 2) if valid_alpha_ratio >= 10: self.logger.info("process v3 fail valid_alpha_ratio is %s", valid_alpha_ratio) return 'bad' # 连通组件分析 num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(alpha_channel) tl = num_labels - 1 #去噪点 面积数小于5 且 且抠图面积大于40% if num_labels <= 5 and alpha_precent >= 40: for label in range(1, num_labels): area = stats[label, cv2.CC_STAT_AREA] #区域面积占比 area_ratio = round((area / total_area) * 100, 2) #半透明区域面积 region_mask = (labels == label).astype(np.uint8) # 计算当前区域半透明像素相对于该区域的比例 total_region_pixels = np.count_nonzero(region_mask) valid_alpha_count = np.count_nonzero(alpha_mask * region_mask) alpha_ratio = round(valid_alpha_count / total_region_pixels, 2) * 100 #独立区域小于2% 为噪点 if area_ratio < 2: alpha_channel[labels == label] = 0 tl = tl - 1 #独立区域2-10%，半透明区域大于20% 为噪点 elif (area_ratio >= 2 and area_ratio <= 10) and alpha_ratio >= 20: alpha_channel[labels == label] = 0 tl = tl - 1 #面积数大于5 或 抠图面积小于40% elif num_labels > 5 or alpha_precent < 40: for label in range(1, num_labels): area = stats[label, cv2.CC_STAT_AREA] #区域面积占比 area_ratio = round((area / total_alpha_channel_area) * 100, 2) #半透明区域面积 region_mask = (labels == label).astype(np.uint8) # 计算当前区域半透明像素相对于该区域的比例 total_region_pixels = np.count_nonzero(region_mask) valid_alpha_count = np.count_nonzero(alpha_mask * region_mask) alpha_ratio = round(valid_alpha_count / total_region_pixels, 2) * 100 #独立区域小于10%抠图面积 为噪点 if area_ratio < 10: alpha_channel[labels == label] = 0 tl = tl - 1 #独立区域10-20%，半透明区域大于20% 为噪点 elif (area_ratio >= 10 and area_ratio <= 30) and alpha_ratio >= 35: alpha_channel[labels == label] = 0 tl = tl - 1 #经过噪点处理重新计算面积 if tl < (num_labels - 1): #非透明区域 total_alpha_channel_area = cv2.countNonZero(alpha_channel) #抠图面积小于10则任务抠图失败 alpha_precent = round((total_alpha_channel_area / total_area) * 100, 2) if alpha_precent <= 10: return 'bad' image[:, :, 3] = alpha_channel_cp if tl == 0 else alpha_channel if 'center' in extend: # 找到非透明区域的边界框 alpha_channel = image[:, :, 3] non_transparent_pixels = np.where(alpha_channel != 0) min_row, max_row = np.min(non_transparent_pixels[0]), np.max(non_transparent_pixels[0]) min_col, max_col = np.min(non_transparent_pixels[1]), np.max(non_transparent_pixels[1]) # 计算需要平移的像素数量 image_center_x = image.shape[1] // 2 image_center_y = image.shape[0] // 2 shift_x = image_center_x - (max_col + min_col) // 2 shift_y = image_center_y - (max_row + min_row) // 2 # 创建新的图像，大小与原图像相同 new_image = np.zeros_like(image) # 计算平移后的非透明区域位置 new_min_row = min_row + shift_y new_max_row = max_row + shift_y new_min_col = min_col + shift_x new_max_col = max_col + shift_x # 将原图像的非透明区域放置到新图像的相应位置 new_image[new_min_row:new_max_row, new_min_col:new_max_col] = image[min_row:max_row, min_col:max_col] image = new_image alpha_channel = image[:, :, 3] result = cv2.bitwise_and(image, image, mask=alpha_channel) return result except Exception as e: self.logger.info("porcess image error msg: %s", e) print(e) return False def matting_by_md(self, img_url, extend=None): # 构造 POST 请求参数 params = { 'imgUrls': img_url, 'extend': extend, 'sign': self.app_tool.build_post_sign(img_url) } try: # 发送 POST 请求，设置 timeout 参数为 30 秒 response = requests.post('http://ay-ai-ps2-intra.chengji-inc.com/matting', data=params, timeout=10) # 获取返回的 JSON 结果 json_result = response.json() json_result['seg'] = 'v3 - ' + json_result['seg'] return json_result except requests.exceptions.Timeout: self.logger.info("matting_by_md image %s timeout", img_url) except requests.exceptions.RequestException as e: self.logger.info("matting_by_md image %s error", img_url) return self.app_tool.fail_result("matting image [%s] fail, please try again", img_url) def do_matting_by_url(self, img_url, extend=None): if self.app_tool.uri_validator(img_url) is False: return self.app_tool.fail_result("file is not img url") self.__init_seg_net__() start = time.time() self.logger.info("start matting imgurl list and imgurl is %s", img_url) pil_img = self.__load_image__(img_url) if isinstance(pil_img, Image.Image) is False: return self.app_tool.fail_result("fetch image data error") image = self.inspy_remover.process(pil_img, 'rgba') process_image = self.__process_extend__(image, extend=extend) if process_image is None: return self.app_tool.fail_result("process image error") if isinstance(process_image, str) is True and process_image == 'bad': md_matting_ret = self.matting_by_md(img_url, extend=extend) is_fail_result = self.app_tool.is_fail_result(md_matting_ret) self.logger.info("process_extend retrun bad and imgurl is %s and is_fail_result is [%s] - md_matting_ret [%s]", img_url, is_fail_result, md_matting_ret['result']) if is_fail_result: return md_matting_ret else: image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA) # 如果图像是RGB模式 process_image = image img_data = cv2.imencode('.png', process_image)[1].tobytes() date_h = time.strftime("%Y-%m-%d/%H", time.localtime()) oss_base_path = 'matting/' + date_h + '/' ret = self.app_tool.save_image_to_oss_by_bytes(img_byte_str=img_data, img_url=img_url, oss_base_path=oss_base_path, img_format='png') cost = (time.time() - start) * 1000 ret['cost_time'] = cost ret['seg'] = 'v3' self.logger.info("end matting upoad oss end cost time is %s and ", cost) gc.collect() return ret def do_matting(self, img_urls, extend = None): if isinstance(img_urls, str): ret = self.do_matting_by_url(img_url=img_urls, extend=extend) return ret elif isinstance(img_urls , list): ret_list = [] for img_url in img_urls: ret = self.do_matting_by_url(img_url=img_url, extend=extend) ret_list.append(ret) return ret_list async def matting(self, request): load1, load5, load15 = os.getloadavg() max_lod = 4 * cpu_cnt if load1 > max_lod: error_log = "server is busy now hostname {}, load1 {}".format(self.hostname, load1) self.logger.error(error_log) resp = json.dumps(self.app_tool.fail_result(error_log)) return web.Response(text = resp, status=503) form = await request.post() resp = json.dumps(self.app_tool.fail_result("empty avgs")) if request.method == 'POST': img_urls_str = None if 'imgUrls' not in form else form['imgUrls'] sign = None if 'sign' not in form else form['sign'] extend = 'denoise' if 'extend' not in form else form['extend'] is_self_ip = self.app_tool.is_self_ip(request) try: self.logger.info("start process and extend is %s", extend) is_vaild = self.app_tool.verify_post_sign(sign, img_urls_str) self.logger.info("__verify_post_sign is %s", is_vaild) if is_vaild is False and is_self_ip is False: return web.Response(text = json.dumps(self.app_tool.fail_result("verify sign fail"))) data_json = self.app_tool.check_is_json(img_urls_str) if data_json is False: data_json = img_urls_str self.logger.info("process_matting start and data_json is %s", data_json) process_resp = self.do_matting(data_json, extend=extend) self.logger.info("process_matting end") return web.Response(text = json.dumps(process_resp)) except ValueError as e: return web.Response(text = resp) else: return web.Response(text = resp) async def del_object(self, request): form = await request.post() resp = json.dumps( { "result": "fail", "reason": "格式解析失败" } ) if request.method == 'POST': object_key = form['objectKey'] sign = form['sign'] try: is_vaild = self.app_tool.verify_post_sign(sign, object_key) self.logger.info("del objet obj key is %s and vaild is %s", object_key, is_vaild) if is_vaild is False: return web.Response(text = json.dumps( { "result": "fail", "reason": "非法的请求参数" } )) resp = self.app_tool.del_oss_object(object_key) resp = json.dumps(resp) return web.Response(text = resp) except ValueError as e: return web.Response(text = resp) else: return web.Response(text = resp) def serve_multiple(self, app, workers): sock = socket() sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1) sock.bind(('0.0.0.0', 8910)) sock.set_inheritable(True) processes = [] for i in range(workers): process = Process(target=web.run_app, name=f'worker-{i}', kwargs=dict(app=app, sock=sock)) process.daemon = True process.start() processes.append(process) try: for process in processes: process.join() except KeyboardInterrupt: pass finally: for process in processes: process.terminate() sock.close() def main_multi(self, process_cnt): app = web.Application(client_max_size=1024**2*4) app.router.add_get('/matting', self.matting, name="matting") app.router.add_post('/matting', self.matting, name="matting") app.router.add_get('/del_object', self.del_object, name="del_object") app.router.add_post('/del_object', self.del_object, name="del_object") cors = aiohttp_cors.setup(app, defaults={ "*": aiohttp_cors.ResourceOptions( allow_credentials=True, expose_headers="*", allow_headers="*", ) }) self.serve_multiple(app, process_cnt) if __name__ == '__main__': cpu_cnt = cpu_count() AiPs = AiMattingV2() AiPs.main_multi(cpu_cnt)