2 Commits 13cb3e6e36 ... 0576009a6f

Autor SHA1 Mensaje Fecha
  xuzuoyun 0576009a6f 'a' hace 3 años
  xuzuoyun 5526db410d '1' hace 3 años
Se han modificado 2 ficheros con 35 adiciones y 15 borrados
  1. BIN
      __pycache__/main.cpython-36.pyc
  2. 35 15
      main.py

BIN
__pycache__/main.cpython-36.pyc


+ 35 - 15
main.py

@@ -27,6 +27,7 @@ from database.database import insert,update,query,Task
 from datetime import datetime
 from fastapi.middleware.cors import CORSMiddleware
 import hashlib
+import traceback
 
 threadPool = ThreadPoolExecutor(max_workers=4)
 def lcm(a, b): return abs(a * b) / fractions.gcd(a, b) if a and b else 0
@@ -173,6 +174,14 @@ 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/retry')
+def retry(item:Item):
+    #插入数据库
+    task = query(item.taskId,None,None,None)[0]
+    item.inputVideoUrl = task.input_video_url
+    item.inputImageUrl = task.input_image_url
+    threadPool.submit(videoSwap,item).add_done_callback(swapFinish)
+    return {'code':0,'data':{'taskId': task.id}}
 
 @app.post('/jeecg-boot/task/single')
 def single(item:Item):
@@ -194,10 +203,11 @@ def single(item:Item):
 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)
+    if(task.status == 'processing'):
+        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)
@@ -229,6 +239,7 @@ def videoSwap(opt):
         open(image_path,'wb').write(image_file.content)
         task = query(opt.taskId,None,None,None)[0]
         task.status = 'downloaded'
+        print('downloaded')
         update(task)
         #md5判断视频和图片是否生成过
         opt.pic_a_path = image_path
@@ -240,11 +251,13 @@ def videoSwap(opt):
     except:
         task = query(opt.taskId,None,None,None)[0]
         task.status = 'download_error'
+        print('download_error')
         update(task)
         del_file(opt.base_path)
 
     task = query(opt.taskId,None,None,None)[0]
     task.status = 'processing'
+    print('processing')
     task.start_time =  datetime.now()
     update(task)
     try:
@@ -256,12 +269,14 @@ def videoSwap(opt):
         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))
+        app.prepare(ctx_id= 0, det_thresh=0.6, det_size=(320,320))
         
         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)
+            print(app.get(img_a_whole,crop_size))
             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)
@@ -285,19 +300,24 @@ def videoSwap(opt):
 
             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:
+            
+        output_video_md5 = getmd5(opt.output_path)
         task = query(opt.taskId,None,None,None)[0]
-        task.status = 'process_error'
+        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)
-    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}
+        return {'taskId':opt.taskId,'outputUrl':url}
+    except Exception as ex:
+        traceback.print_exc()
+        task = query(opt.taskId,None,None,None)[0]
+        task.status = 'process_error'
+        print('process_error')
+        update(task)
+        #del_file(opt.base_path)
+        return {'taskId':opt.taskId} 
+    
 
 
 if __name__ == '__main__':