main.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304
  1. from typing import Optional
  2. from fastapi import FastAPI
  3. from concurrent.futures import ThreadPoolExecutor
  4. from numpy.core.records import array
  5. from pydantic import BaseModel
  6. import cv2
  7. from sqlalchemy.sql.elements import Null
  8. from sqlalchemy.sql.expression import null
  9. import torch
  10. import fractions
  11. import numpy as np
  12. from PIL import Image
  13. from datetime import date
  14. from torch._C import Node
  15. import torch.nn.functional as F
  16. from torchvision import transforms
  17. from models.models import create_model
  18. from options.test_options import TestOptions
  19. from insightface_func.face_detect_crop_single import Face_detect_crop
  20. from util.videoswap import video_swap
  21. import requests
  22. import uuid
  23. import os
  24. import uvicorn
  25. from util.cos_util import cos_upload
  26. from database.database import insert,update,query,Task
  27. from datetime import datetime
  28. from fastapi.middleware.cors import CORSMiddleware
  29. import hashlib
  30. threadPool = ThreadPoolExecutor(max_workers=4)
  31. def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
  32. transformer = transforms.Compose([
  33. transforms.ToTensor(),
  34. #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  35. ])
  36. transformer_Arcface = transforms.Compose([
  37. transforms.ToTensor(),
  38. transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  39. ])
  40. def get_args_from_json(json_file_path, args_dict):
  41. import json
  42. summary_filename = json_file_path
  43. with open(summary_filename) as f:
  44. summary_dict = json.load(fp=f)
  45. for key in summary_dict.keys():
  46. args_dict[key] = summary_dict[key]
  47. return args_dict
  48. def getmd5(file):
  49. m = hashlib.md5()
  50. with open(file,'rb') as f:
  51. for line in f:
  52. m.update(line)
  53. md5code = m.hexdigest()
  54. print(md5code)
  55. return md5code
  56. class Item(BaseModel):
  57. isTrain:bool = False
  58. use_mask:bool = False
  59. name:str = 'people'
  60. Arc_path:str = './arcface_model/arcface_checkpoint.tar'
  61. pic_a_path:str = ''
  62. pic_b_path:str = ''
  63. video_path:str = ''
  64. pic_specific_path:str = './crop_224/zrf.jpg'
  65. multisepcific_dir:str = './demo_file/multispecific'
  66. output_path:str = './output'
  67. temp_path:str = './temp_results'
  68. id_thres:float = 0.03
  69. no_simswaplogo:bool = True
  70. model:str = 'pix2pixHD'
  71. gpu_ids:str = '0'
  72. checkpoints_dir:str = './checkpoints'
  73. norm:str = 'batch'
  74. use_dropout:bool = False
  75. data_type:int = 32
  76. verbose:bool = False
  77. fp16:bool = False
  78. local_rank:int = 0
  79. batchSize:int = 8
  80. loadSize:int = 1024
  81. fineSize:int = 512
  82. label_nc:int = 0
  83. input_nc:int = 3
  84. output_nc:int = 3
  85. dataroot:str = './datasets/cityscapes/'
  86. resize_or_crop:str = 'scale_width'
  87. serial_batches:bool = False
  88. no_flip:bool = False
  89. nThreads:int = 2
  90. max_dataset_size:int = float("inf")
  91. display_winsize:int = 512
  92. tf_log:bool = False
  93. netG:str = 'global'
  94. latent_size:int = 512
  95. ngf:int = 64
  96. n_downsample_global:int = 3
  97. n_blocks_global:int = 6
  98. n_block_local:int = 3
  99. n_local_enhancers:int = 1
  100. niter_fix_global:int = 0
  101. no_instance:bool = False
  102. instance_feat:bool = False
  103. label_feat:bool = False
  104. feat_num:int = 3
  105. load_features:bool = False
  106. n_downsample_E:int = 4
  107. net:int = 16
  108. n_clusters:int = 10
  109. image_size:int = 224
  110. norg_G:str = 'spectralspadesyncbatch3x3'
  111. semantic_nc:int = 3
  112. ntest:int = float("inf")
  113. results_dir:str = './results/'
  114. aspect_ratio:float = 1.0
  115. phase:str = 'test'
  116. which_epoch:str = 'latest'
  117. how_many:int = 50
  118. cluster_path:str = 'features_clusered_010.npy'
  119. use_encoded_image:bool = False
  120. export_onnx:str = ''
  121. engine:str = ''
  122. onnx:str = ''
  123. inputVideoUrl:str = ''
  124. inputImageUrl:str = ''
  125. videoMd5:str = ''
  126. imageMd5:str = ''
  127. output_file_name:str = ''
  128. taskId:int = 0
  129. base_path:str = ''
  130. createBy:str = ''
  131. class QueryItem(BaseModel):
  132. id:Optional[int]=None
  133. videoMd5:Optional[str]=None
  134. imageMd5:Optional[str]=None
  135. createBy:Optional[str]=None
  136. app = FastAPI()
  137. origins = [
  138. "http://192.168.1.34",
  139. "http://192.168.1.34:8000",
  140. "http://192.168.1.105",
  141. "http://192.168.1.105:3000",
  142. "http://111.206.86.186",
  143. "http://111.206.86.186:3000",
  144. "http://adsp.tjyourong.com.cn",
  145. "http://adsp.tjyourong.com.cn:3000"
  146. ]
  147. app.add_middleware(
  148. CORSMiddleware,
  149. allow_origins=origins,
  150. allow_credentials=True,
  151. allow_methods=["*"],
  152. allow_headers=["*"],
  153. )
  154. @app.get('/')
  155. def index():
  156. return {'message': '你已经正确创建 FastApi 服务!'}
  157. @app.post('/jeecg-boot/task/query')
  158. def task_query(queryItem:QueryItem):
  159. task = query(queryItem.id,queryItem.videoMd5,queryItem.imageMd5,queryItem.createBy)
  160. return {'code':0,'data':task}
  161. @app.post('/jeecg-boot/task/single')
  162. def single(item:Item):
  163. #插入数据库
  164. old_task = query(None,item.videoMd5,item.imageMd5,None)
  165. if len(old_task) > 0:
  166. return {'code':-1,'data':old_task}
  167. uid = str(uuid.uuid4())
  168. suid = ''.join(uid.split('-'))
  169. video_input = item.inputVideoUrl
  170. image_input = item.inputImageUrl
  171. 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)
  172. task = insert(task)
  173. item.taskId = task.id
  174. threadPool.submit(videoSwap,item).add_done_callback(swapFinish)
  175. return {'code':0,'data':{'taskId': task.id}}
  176. def swapFinish(res):
  177. print('solute',res.result())
  178. task = query(res.result()['taskId'],None,None,None)[0]
  179. task.status = 'finished'
  180. task.finish_time = datetime.now()
  181. task.output_video_url = res.result()['outputUrl']
  182. update(task)
  183. def del_file(path):
  184. ls = os.listdir(path)
  185. for i in ls:
  186. c_path = os.path.join(path, i)
  187. if os.path.isdir(c_path):
  188. del_file(c_path)
  189. else:
  190. os.remove(c_path)
  191. def videoSwap(opt):
  192. task = query(opt.taskId,None,None,None)[0]
  193. task.status = 'downloading'
  194. update(task)
  195. try:
  196. uid = str(uuid.uuid4())
  197. suid = ''.join(uid.split('-'))
  198. video_input = opt.inputVideoUrl
  199. image_input = opt.inputImageUrl
  200. base_path = '/data/swap_file_temp/'+suid+'/'
  201. # base_path = './'+suid+'/'
  202. video_path = base_path + suid + os.path.splitext(video_input)[-1]
  203. image_path = base_path + suid + os.path.splitext(image_input)[-1]
  204. out_path = base_path + suid + 'out' + os.path.splitext(video_input)[-1]
  205. temp_path = base_path + 'temp/'
  206. os.makedirs(temp_path)
  207. video_file = requests.get(video_input)
  208. image_file = requests.get(image_input)
  209. open(video_path,'wb').write(video_file.content)
  210. open(image_path,'wb').write(image_file.content)
  211. task = query(opt.taskId,None,None,None)[0]
  212. task.status = 'downloaded'
  213. update(task)
  214. #md5判断视频和图片是否生成过
  215. opt.pic_a_path = image_path
  216. opt.video_path = video_path
  217. opt.output_path = out_path
  218. opt.temp_path = temp_path
  219. opt.base_path = base_path
  220. opt.output_file_name = suid + os.path.splitext(video_input)[-1]
  221. except:
  222. task = query(opt.taskId,None,None,None)[0]
  223. task.status = 'download_error'
  224. update(task)
  225. del_file(opt.base_path)
  226. task = query(opt.taskId,None,None,None)[0]
  227. task.status = 'processing'
  228. task.start_time = datetime.now()
  229. update(task)
  230. try:
  231. start_epoch, epoch_iter = 1, 0
  232. crop_size = 224
  233. # opt_dict = vars(opt)
  234. # args = get_args_from_json(item,opt_dict)
  235. torch.nn.Module.dump_patches = True
  236. model = create_model(opt)
  237. model.eval()
  238. app = Face_detect_crop(name='antelope', root='./insightface_func/models')
  239. app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(640,640))
  240. with torch.no_grad():
  241. pic_a = opt.pic_a_path
  242. # img_a = Image.open(pic_a).convert('RGB')
  243. img_a_whole = cv2.imread(pic_a)
  244. img_a_align_crop, _ = app.get(img_a_whole,crop_size)
  245. img_a_align_crop_pil = Image.fromarray(cv2.cvtColor(img_a_align_crop[0],cv2.COLOR_BGR2RGB))
  246. img_a = transformer_Arcface(img_a_align_crop_pil)
  247. img_id = img_a.view(-1, img_a.shape[0], img_a.shape[1], img_a.shape[2])
  248. # pic_b = opt.pic_b_path
  249. # img_b_whole = cv2.imread(pic_b)
  250. # img_b_align_crop, b_mat = app.get(img_b_whole,crop_size)
  251. # img_b_align_crop_pil = Image.fromarray(cv2.cvtColor(img_b_align_crop,cv2.COLOR_BGR2RGB))
  252. # img_b = transformer(img_b_align_crop_pil)
  253. # img_att = img_b.view(-1, img_b.shape[0], img_b.shape[1], img_b.shape[2])
  254. # convert numpy to tensor
  255. img_id = img_id.cuda()
  256. # img_att = img_att.cuda()
  257. #create latent id
  258. img_id_downsample = F.interpolate(img_id, scale_factor=0.5)
  259. latend_id = model.netArc(img_id_downsample)
  260. latend_id = F.normalize(latend_id, p=2, dim=1)
  261. video_swap(opt.video_path, latend_id, model, app, opt.output_path,temp_results_dir=opt.temp_path,\
  262. no_simswaplogo=opt.no_simswaplogo,use_mask=opt.use_mask)
  263. except:
  264. task = query(opt.taskId,None,None,None)[0]
  265. task.status = 'process_error'
  266. update(task)
  267. del_file(opt.base_path)
  268. output_video_md5 = getmd5(opt.output_path)
  269. task = query(opt.taskId,None,None,None)[0]
  270. task.status = 'uploading'
  271. task.output_video_md5 = output_video_md5
  272. update(task)
  273. url = cos_upload(opt.output_path,'faceswap/'+datetime.now().strftime('%Y%m%d')+'/'+output_video_md5+'.mp4')
  274. del_file(opt.base_path)
  275. return {'taskId':opt.taskId,'outputUrl':url}
  276. if __name__ == '__main__':
  277. uvicorn.run(app='main:app', host="0.0.0.0", port=8000, reload=True, debug=True)
  278. #gunicorn main:app -w 4 -k uvicorn.workers.UvicornWorker #线上启动命令