admin/LC_NeoRefacer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replaced ProgressBar with the standard tqdm library for enhanced user experience #3

Browse Source
This commit is contained in:
Xavi Vinaixa
2023-06-05 09:04:02 +02:00
parent f0a73b4914
commit 50ab0e1fbe
3 changed files with 7 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
refacer.py
View File

@@ -9,7 +9,7 @@ from arcface_onnx import ArcFaceONNX
import os.path as osp import os.path as osp
import os import os
from pathlib import Path from pathlib import Path
import progressbar as pb from tqdm import tqdm
import ffmpeg import ffmpeg
import random import random
@@ -31,13 +31,6 @@ class Refacer:
self.rec_app.prepare(0) self.rec_app.prepare(0)
self.face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True, providers=['CoreMLExecutionProvider','CUDAExecutionProvider']) self.face_swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True, providers=['CoreMLExecutionProvider','CUDAExecutionProvider'])
self.pb_widgets =[
pb.Percentage(),
pb.Bar(),
' ', pb.SimpleProgress(),
' ', pb.ETA(),
' ', pb.FileTransferSpeed(unit='it')
]
def __prepare_faces(self, faces): def __prepare_faces(self, faces):
replacements=[] replacements=[]
@@ -59,15 +52,12 @@ class Refacer:
out.run() out.run()
return new_path return new_path
def reface(self, video_path, faces, gr_progress=None): def reface(self, video_path, faces):
output_video_path = os.path.join('out',Path(video_path).name) output_video_path = os.path.join('out',Path(video_path).name)
replacement_faces=self.__prepare_faces(faces) replacement_faces=self.__prepare_faces(faces)
cap = cv2.VideoCapture(video_path) cap = cv2.VideoCapture(video_path)
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
bar = pb.ProgressBar(maxval=total,widgets=self.pb_widgets)
bar.start()
fps = cap.get(cv2.CAP_PROP_FPS) fps = cap.get(cv2.CAP_PROP_FPS)
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -76,16 +66,13 @@ class Refacer:
fourcc = cv2.VideoWriter_fourcc(*'mp4v') fourcc = cv2.VideoWriter_fourcc(*'mp4v')
output = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) output = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
pbar = tqdm(total=total_frames)
while cap.isOpened(): while cap.isOpened():
flag, frame = cap.read() flag, frame = cap.read()
if flag and len(frame)>0: if flag and len(frame)>0:
pos_frame = cap.get(cv2.CAP_PROP_POS_FRAMES) pos_frame = cap.get(cv2.CAP_PROP_POS_FRAMES)
bar.update(pos_frame) pbar.update(pos_frame)
if gr_progress is not None:
gr_progress((pos_frame,total),desc=f"Processing frame {pos_frame} of {total}")
faces = self.face_app.get(frame) faces = self.face_app.get(frame)
res = frame.copy() res = frame.copy()
@@ -100,7 +87,6 @@ class Refacer:
else: else:
break break
bar.finish()
cap.release() cap.release()
output.release() output.release()

2
requirements-GPU.txt
View File

@@ -6,5 +6,5 @@ onnx==1.14.0
onnxruntime_gpu==1.15.0 onnxruntime_gpu==1.15.0
opencv_python==4.7.0.72 opencv_python==4.7.0.72
opencv_python_headless==4.7.0.72 opencv_python_headless==4.7.0.72
progressbar33==2.4
scikit-image==0.20.0 scikit-image==0.20.0
tqdm

2
requirements.txt
View File

@@ -6,5 +6,5 @@ onnx==1.14.0
onnxruntime==1.15.0 onnxruntime==1.15.0
opencv_python==4.7.0.72 opencv_python==4.7.0.72
opencv_python_headless==4.7.0.72 opencv_python_headless==4.7.0.72
progressbar33==2.4
scikit-image==0.20.0 scikit-image==0.20.0
tqdm