# !/usr/bin/python3 # Built-in libraries from pathlib import Path import json import os from aiohttp import web import aiohttp_cors import time import numpy as np from cairosvg import svg2png import base64 from multiprocessing import Process, cpu_count from socket import SOL_SOCKET, SO_REUSEADDR, socket import cv2 from libs.model_manager import ModelManager from libs.schema import Config import torch try: torch._C._jit_override_can_fuse_on_cpu(False) torch._C._jit_override_can_fuse_on_gpu(False) torch._C._jit_set_texpr_fuser_enabled(False) torch._C._jit_set_nvfuser_enabled(False) except: pass from libs.helper import ( load_img, numpy_to_bytes, resize_max_size, ) from libs.app_tool import AppTool device = 'cuda' if torch.cuda.is_available() else 'cpu' class InpaintServer: def __init__(self): self.app_tool = AppTool('ai_inpaint.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.matting_modle = None self._init_modle() def _init_modle(self): self.inpaint_model = ModelManager('lama', device=torch.device(device)) async def del_object(self, request): form = await request.post() resp = self.app_tool.fail_result("格式解析失败") 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(self.app_tool.fail_result("非法的请求参数"))) resp = self.app_tool.del_oss_object(object_key) 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)) def do_inpaint_V2(self, img_bytes, mask_bytes): # RGB origin_image_bytes = img_bytes image, alpha_channel = load_img(origin_image_bytes) original_shape = image.shape interpolation = cv2.INTER_CUBIC size_limit = max(image.shape) config = Config( ldm_steps=50, hd_strategy='Resize', hd_strategy_crop_margin=128, hd_strategy_crop_trigger_size=2048, hd_strategy_resize_limit=size_limit, ) #self.logger.info(f"Origin image shape: {original_shape}") self.logger.info(f"Origin image shape: {original_shape}") image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation) self.logger.info(f"Resized image shape: {image.shape}") mask, _ = load_img(mask_bytes, gray=True) mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation) start = time.time() self.logger.info(f"process start time: {start * 1000}ms") res_np_img = self.inpaint_model(image, mask, config) self.logger.info(f"process inpaint time: {(time.time() - start) * 1000}ms") torch.cuda.empty_cache() if alpha_channel is not None: if alpha_channel.shape[:2] != res_np_img.shape[:2]: alpha_channel = cv2.resize( alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0]) ) res_np_img = np.concatenate( (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1 ) ext = self.app_tool.get_image_ext(origin_image_bytes) fin_img_bytes_str = numpy_to_bytes(res_np_img, ext) date_h = time.strftime("%Y-%m-%d/%H", time.localtime()) oss_base_path = 'inpaint/' + date_h + '/' img_path = self.app_tool.uuid() resp = self.app_tool.save_image_to_oss_by_bytes(fin_img_bytes_str, img_path, oss_base_path, ext) self.logger.info(f"process enbd time: {(time.time() - start) * 1000}ms") return resp """ fin_img_bytes = io.BytesIO(fin_img_bytes_str) image = Image.open(fin_img_bytes) image.show('test.png) """ 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_V2(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)) def serve_multiple(self, app, workers): sock = socket() sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1) sock.bind(('0.0.0.0', 8902)) 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) cors = aiohttp_cors.setup(app, defaults={ "*": aiohttp_cors.ResourceOptions( allow_credentials=True, expose_headers="*", allow_headers="*", ) }) resource = cors.add(app.router.add_resource("/inpaint")) cors.add(resource.add_route("GET", self.inpaint)) cors.add(resource.add_route("POST", self.inpaint)) self.serve_multiple(app, process_cnt) def main(self): app = web.Application() cors = aiohttp_cors.setup(app, defaults={ "*": aiohttp_cors.ResourceOptions( allow_credentials=True, expose_headers="*", allow_headers="*", ) }) resource = cors.add(app.router.add_resource("/inpaint")) cors.add(resource.add_route("GET", self.inpaint)) cors.add(resource.add_route("POST", self.inpaint)) web.run_app(app, port=8902) if __name__ == '__main__': IS = InpaintServer() if device == 'cuda': IS.main() else: process_cnt = cpu_count() IS.main_multi(process_cnt)