from typing import Optional from fastapi import FastAPI from concurrent.futures import ThreadPoolExecutor from numpy.core.records import array from pydantic import BaseModel import cv2 from sqlalchemy.sql.elements import Null from sqlalchemy.sql.expression import null import torch import fractions import numpy as np from PIL import Image from datetime import date from torch._C import Node import torch.nn.functional as F from torchvision import transforms from models.models import create_model from options.test_options import TestOptions from insightface_func.face_detect_crop_single import Face_detect_crop from util.videoswap import video_swap import requests import uuid import os import uvicorn from util.cos_util import cos_upload from database.database import insert,update,query,Task from datetime import datetime from fastapi.middleware.cors import CORSMiddleware import hashlib threadPool = ThreadPoolExecutor(max_workers=4) def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0 transformer = transforms.Compose([ transforms.ToTensor(), #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) transformer_Arcface = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) def get_args_from_json(json_file_path, args_dict): import json summary_filename = json_file_path with open(summary_filename) as f: summary_dict = json.load(fp=f) for key in summary_dict.keys(): args_dict[key] = summary_dict[key] return args_dict def getmd5(file): m = hashlib.md5() with open(file,'rb') as f: for line in f: m.update(line) md5code = m.hexdigest() print(md5code) return md5code class Item(BaseModel): isTrain:bool = False use_mask:bool = False name:str = 'people' Arc_path:str = './arcface_model/arcface_checkpoint.tar' pic_a_path:str = '' pic_b_path:str = '' video_path:str = '' pic_specific_path:str = './crop_224/zrf.jpg' multisepcific_dir:str = './demo_file/multispecific' output_path:str = './output' temp_path:str = './temp_results' id_thres:float = 0.03 no_simswaplogo:bool = True model:str = 'pix2pixHD' gpu_ids:str = '0' checkpoints_dir:str = './checkpoints' norm:str = 'batch' use_dropout:bool = False data_type:int = 32 verbose:bool = False fp16:bool = False local_rank:int = 0 batchSize:int = 8 loadSize:int = 1024 fineSize:int = 512 label_nc:int = 0 input_nc:int = 3 output_nc:int = 3 dataroot:str = './datasets/cityscapes/' resize_or_crop:str = 'scale_width' serial_batches:bool = False no_flip:bool = False nThreads:int = 2 max_dataset_size:int = float("inf") display_winsize:int = 512 tf_log:bool = False netG:str = 'global' latent_size:int = 512 ngf:int = 64 n_downsample_global:int = 3 n_blocks_global:int = 6 n_block_local:int = 3 n_local_enhancers:int = 1 niter_fix_global:int = 0 no_instance:bool = False instance_feat:bool = False label_feat:bool = False feat_num:int = 3 load_features:bool = False n_downsample_E:int = 4 net:int = 16 n_clusters:int = 10 image_size:int = 224 norg_G:str = 'spectralspadesyncbatch3x3' semantic_nc:int = 3 ntest:int = float("inf") results_dir:str = './results/' aspect_ratio:float = 1.0 phase:str = 'test' which_epoch:str = 'latest' how_many:int = 50 cluster_path:str = 'features_clusered_010.npy' use_encoded_image:bool = False export_onnx:str = '' engine:str = '' onnx:str = '' inputVideoUrl:str = '' inputImageUrl:str = '' videoMd5:str = '' imageMd5:str = '' output_file_name:str = '' taskId:int = 0 base_path:str = '' createBy:str = '' class QueryItem(BaseModel): id:Optional[int]=None videoMd5:Optional[str]=None imageMd5:Optional[str]=None createBy:Optional[str]=None app = FastAPI() origins = [ "http://192.168.1.34", "http://192.168.1.34:8000", "http://192.168.1.105", "http://192.168.1.105:3000", "http://111.206.86.186", "http://111.206.86.186:3000", "http://adsp.tjyourong.com.cn", "http://adsp.tjyourong.com.cn:3000" ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.get('/') def index(): return {'message': '你已经正确创建 FastApi 服务!'} @app.post('/jeecg-boot/task/query') def task_query(queryItem:QueryItem): task = query(queryItem.id,queryItem.videoMd5,queryItem.imageMd5,queryItem.createBy) return {'code':0,'data':task} @app.post('/jeecg-boot/task/single') def single(item:Item): #插入数据库 old_task = query(None,item.videoMd5,item.imageMd5,None) if len(old_task) > 0: return {'code':-1,'data':old_task} uid = str(uuid.uuid4()) suid = ''.join(uid.split('-')) video_input = item.inputVideoUrl image_input = item.inputImageUrl task = Task(input_video_url = video_input,input_image_url=image_input,status='waiting',input_video_md5=item.videoMd5,input_image_md5=item.imageMd5,create_by=item.createBy) task = insert(task) item.taskId = task.id threadPool.submit(videoSwap,item).add_done_callback(swapFinish) return {'code':0,'data':{'taskId': task.id}} def swapFinish(res): print('solute',res.result()) task = query(res.result()['taskId'],None,None,None)[0] task.status = 'finished' task.finish_time = datetime.now() task.output_video_url = res.result()['outputUrl'] update(task) def del_file(path): ls = os.listdir(path) for i in ls: c_path = os.path.join(path, i) if os.path.isdir(c_path): del_file(c_path) else: os.remove(c_path) def videoSwap(opt): task = query(opt.taskId,None,None,None)[0] task.status = 'downloading' update(task) try: uid = str(uuid.uuid4()) suid = ''.join(uid.split('-')) video_input = opt.inputVideoUrl image_input = opt.inputImageUrl base_path = '/data/swap_file_temp/'+suid+'/' # base_path = './'+suid+'/' video_path = base_path + suid + os.path.splitext(video_input)[-1] image_path = base_path + suid + os.path.splitext(image_input)[-1] out_path = base_path + suid + 'out' + os.path.splitext(video_input)[-1] temp_path = base_path + 'temp/' os.makedirs(temp_path) video_file = requests.get(video_input) image_file = requests.get(image_input) open(video_path,'wb').write(video_file.content) open(image_path,'wb').write(image_file.content) task = query(opt.taskId,None,None,None)[0] task.status = 'downloaded' update(task) #md5判断视频和图片是否生成过 opt.pic_a_path = image_path opt.video_path = video_path opt.output_path = out_path opt.temp_path = temp_path opt.base_path = base_path opt.output_file_name = suid + os.path.splitext(video_input)[-1] except: task = query(opt.taskId,None,None,None)[0] task.status = 'download_error' update(task) del_file(opt.base_path) task = query(opt.taskId,None,None,None)[0] task.status = 'processing' task.start_time = datetime.now() update(task) try: start_epoch, epoch_iter = 1, 0 crop_size = 224 # opt_dict = vars(opt) # args = get_args_from_json(item,opt_dict) torch.nn.Module.dump_patches = True model = create_model(opt) model.eval() app = Face_detect_crop(name='antelope', root='./insightface_func/models') app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640)) with torch.no_grad(): pic_a = opt.pic_a_path # img_a = Image.open(pic_a).convert('RGB') img_a_whole = cv2.imread(pic_a) img_a_align_crop, _ = app.get(img_a_whole,crop_size) img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB)) img_a = transformer_Arcface(img_a_align_crop_pil) img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2]) # pic_b = opt.pic_b_path # img_b_whole = cv2.imread(pic_b) # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size) # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB)) # img_b = transformer(img_b_align_crop_pil) # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2]) # convert numpy to tensor img_id = img_id.cuda() # img_att = img_att.cuda() #create latent id img_id_downsample = F.interpolate(img_id, scale_factor=0.5) latend_id = model.netArc(img_id_downsample) latend_id = F.normalize(latend_id, p=2, dim=1) video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\ no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask) except: task = query(opt.taskId,None,None,None)[0] task.status = 'process_error' update(task) del_file(opt.base_path) output_video_md5 = getmd5(opt.output_path) task = query(opt.taskId,None,None,None)[0] task.status = 'uploading' task.output_video_md5 = output_video_md5 update(task) url = cos_upload(opt.output_path,'faceswap/'+datetime.now().strftime('%Y%m%d')+'/'+output_video_md5+'.mp4') del_file(opt.base_path) return {'taskId':opt.taskId,'outputUrl':url} if __name__ == '__main__': uvicorn.run(app='main:app', host="0.0.0.0", port=8000, reload=True, debug=True) #gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker #线上启动命令