main.py 11 KB

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