# !/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 cv2 import numpy as np from multiprocessing import Process, cpu_count from socket import SOL_SOCKET, SO_REUSEADDR, socket from cairosvg import svg2png import base64 from PIL import Image, ImageOps import requests import urllib3 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) from io import BytesIO from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from modelscope.outputs import OutputKeys from libs.app_tool import AppTool cpu_cnt = cpu_count() class MsPs: def __init__(self): self.app_tool = AppTool('ms_ps.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.universal_matting = None self.image_inpainting = None def _init_pipeline(self, action): if action == 'matting' and self.universal_matting is None: self.universal_matting = pipeline(Tasks.universal_matting, model='damo/cv_unet_universal-matting') if action == 'inpaint' and self.image_inpainting is None: self.image_inpainting = pipeline(Tasks.image_inpainting, model='damo/cv_fft_inpainting_lama') 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(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: # 将 PIL 图像转换为 OpenCV 格式 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 image #半透明（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 image # 连通组件分析 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 image 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) return False 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_pipeline(action='matting') start = time.time() self.logger.info("start matting imgurl list and imgurl is %s", img_url) pil_img = self.__load_image__(img_url, mode='rgb') if isinstance(pil_img, Image.Image) is False: return self.app_tool.fail_result("image load fail") result = self.universal_matting(pil_img) image = result[OutputKeys.OUTPUT_IMG] image = self.__process_extend__(image, extend=extend) if image is None: return self.app_tool.fail_result("process error") img_data = cv2.imencode('.png', 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'] = 'damo' 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 def do_inpaint_by_betys(self, image_data, mask_data): self._init_pipeline(action='inpaint') start = time.time() pil_img = self.__load_image__(image_data, 'rgb') pil_mask = self.__load_image__(mask_data, mode='L') if isinstance(pil_img, Image.Image) is False or isinstance(pil_mask, Image.Image) is False: return self.app_tool.fail_result("data error") result = self.image_inpainting({ 'img':pil_img, 'mask':pil_mask, }) image_ndarray = result[OutputKeys.OUTPUT_IMG] img_format = self.app_tool.get_image_ext(image_data) img_data = cv2.imencode('.' + img_format, image_ndarray)[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='bytest', oss_base_path=oss_base_path, img_format=img_format) cost = (time.time() - start) * 1000 ret['cost_time'] = cost ret['seg'] = 'damo' self.logger.info("end matting upoad oss end cost time is %s and ", cost) gc.collect() return ret 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 inpaint(self, request): form = await request.post() resp = self.app_tool.fail_result("无效的参数!!") if request.method == 'POST': try: img_bytes = None if 'imgBytes' not in form else form['imgBytes'] mask_data = None if 'maskSvg' not in form else form['maskSvg'] sign = None if 'sign' not in form else form['sign'] is_vaild = self.app_tool.verify_post_sign(sign, img_bytes) is_self_ip = self.app_tool.is_self_ip(request) self.logger.info("__verify_post_sign is %s and sign is %s" % (is_vaild, sign)) if is_vaild is False and is_self_ip is False: return web.Response(text = json.dumps(self.app_tool.fail_result("sign检验失败"))) image_data = base64.b64decode(img_bytes) mask_data = svg2png(bytestring=mask_data) resp = self.do_inpaint_by_betys(image_data, mask_data) return web.Response(text=json.dumps(resp)) except ValueError as e: return web.Response(text=json.dumps(resp)) else: return web.Response(text=json.dumps(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('/inpaint', self.inpaint, name="inpaint") app.router.add_post('/inpaint', self.inpaint, name="inpaint") 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__': #process_cnt = 1 if cpu_cnt == 1 else int(cpu_cnt / 2) MsPs = MsPs() MsPs.main_multi(cpu_cnt)