properly name output file when its an image

This commit is contained in:
Somdev Sangwan 2023-05-30 19:07:22 +05:30
parent 03ac74edac
commit 90e2ba0fc7
2 changed files with 10 additions and 6 deletions

View File

@ -28,11 +28,10 @@ def process_video(source_img, frame_paths):
pass
def process_img(source_img, target_path):
def process_img(source_img, target_path, output_file):
frame = cv2.imread(target_path)
face = get_face(frame)
source_face = get_face(cv2.imread(source_img))
result = face_swapper.get(frame, face, source_face, paste_back=True)
target_path = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-"+target_path
print(target_path)
cv2.imwrite(target_path, result)
cv2.imwrite(output_file, result)
print("\n\nImage saved as:", output_file, "\n\n")

9
run.py
View File

@ -123,7 +123,10 @@ def toggle_keep_frames():
def save_file():
args['output_file'] = asksaveasfilename(initialfile='output.mp4', defaultextension=".mp4", filetypes=[("All Files","*.*"),("Videos","*.mp4")])
filename, ext = 'output.mp4', '.mp4'
if is_img(args['target_path']):
filename, ext = 'output.png', '.png'
args['output_file'] = asksaveasfilename(initialfile=filename, defaultextension=ext, filetypes=[("All Files","*.*"),("Videos","*.mp4")])
def status(string):
@ -142,6 +145,8 @@ def start():
elif not args['target_path'] or not os.path.isfile(args['target_path']):
print("\n[WARNING] Please select a video/image to swap face in.")
return
if not args['output_file']:
args['output_file'] = rreplace(args['target_path'], "/", "/swapped-", 1) if "/" in target_path else "swapped-"+target_path
global pool
pool = mp.Pool(psutil.cpu_count()-1)
target_path = args['target_path']
@ -150,7 +155,7 @@ def start():
print("\n[WARNING] No face detected in source image. Please try with another one.\n")
return
if is_img(target_path):
process_img(args['source_img'], target_path)
process_img(args['source_img'], target_path, args['output_file'])
status("swap successful!")
return
video_name = target_path.split("/")[-1].split(".")[0]