main.py 9.1 KB

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