123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304 |
- 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://api.tjyourong.com.cn",
- "http://api.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 #线上启动命令
|