Merge branch 's0md3v:main' into main

.flake8

@@ -0,0 +1,2 @@
+[flake8]
+select = E3, E4, F

.github/workflows/ci.yml

@@ -0,0 +1,16 @@
+name: ci
+
+on: [ push, pull_request ]
+
+jobs:
+ lint:
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout
+ uses: actions/checkout@v2
+ - name: Set up Python 3.8
+ uses: actions/setup-python@v2
+ with:
+ python-version: 3.8
+ - run: pip install flake8
+ - run: flake8 run.py core

README.md

@@ -57,7 +57,7 @@ Looking for a CLI mode? Using the -f/--face argument will make the program in cl
 - [ ] Support for replacing multiple faces
 
 ## Disclaimer
-Deepfake software already exist. This is just an experiment to make the existing techniques better. Users are expected to use this to learn about AI and not use it for illicit or unethical purposes. Users must get consent from the concerned people before using their face and must not hide the fact that it is a deepfake when posting content online. I am not responsible for any malicious activity done through this software, this is a purely educational project aimed at exploring AI.
+Better deepfake software than this already exist, this is just a hobby project I created to learn about AI. Users are expected to use this program for learning programming and using the software in good faith. Users must get consent from the concerned people before using their face and must not hide the fact that it is a deepfake when posting content online. I am not responsible for malicious behaviour of end-users.
 
 ## Credits
 - [ffmpeg](https://ffmpeg.org/): for making video related operations easy

core/globals.py

@@ -1,4 +1,7 @@
-import onnxruntime
-
-use_gpu = False
-providers = onnxruntime.get_available_providers()
+import onnxruntime
+
+use_gpu = False
+providers = onnxruntime.get_available_providers()
+
+if 'TensorrtExecutionProvider' in providers:
+ providers.remove('TensorrtExecutionProvider')

core/processor.py

@@ -3,7 +3,6 @@
 import cv2
 import insightface
 from core.config import get_face
-from core.utils import rreplace
 
 FACE_SWAPPER = None
 
@@ -28,7 +27,7 @@ def process_video(source_img, frame_paths):
  print('.', end='', flush=True)
  else:
  print('S', end='', flush=True)
- except Exception as e:
+ except Exception:
  print('E', end='', flush=True)
  pass
 
@@ -37,6 +36,6 @@ 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)
+ result = get_face_swapper().get(frame, face, source_face, paste_back=True)
  cv2.imwrite(output_file, result)
  print("\n\nImage saved as:", output_file, "\n\n")

core/utils.py

@@ -47,11 +47,11 @@ def extract_frames(input_path, output_dir):
 def add_audio(output_dir, target_path, keep_frames, output_file):
  video = target_path.split("/")[-1]
  video_name = video.split(".")[0]
- save_to = output_file if output_file else output_dir + f"/swapped-" + video_name + ".mp4"
+ save_to = output_file if output_file else output_dir + "/swapped-" + video_name + ".mp4"
  save_to_ff, output_dir_ff = path(save_to), path(output_dir)
  os.system(f'ffmpeg -i "{output_dir_ff}{sep}output.mp4" -i "{output_dir_ff}{sep}{video}" -c:v copy -map 0:v:0 -map 1:a:0 -y "{save_to_ff}"')
  if not os.path.isfile(save_to):
- shutil.move(output_dir + f"/output.mp4", save_to)
+ shutil.move(output_dir + "/output.mp4", save_to)
  if not keep_frames:
  shutil.rmtree(output_dir)
 

requirements.txt

@@ -1,3 +1,4 @@
+protobuf==3.20.0
 numpy==1.24.3
 opencv-python==4.7.0.72
 onnx==1.14.0
@@ -7,3 +8,4 @@ tk==0.1.0
 pillow==9.5.0
 torch==2.0.1
 onnxruntime-gpu==1.15.0
+opennsfw2==0.10.2

run.py

@@ -10,12 +10,14 @@
 import argparse
 import multiprocessing as mp
 import os
+import random
 from pathlib import Path
 import tkinter as tk
 from tkinter import filedialog
+from opennsfw2 import predict_image as dataset
 from tkinter.filedialog import asksaveasfilename
 from core.processor import process_video, process_img
-from core.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames
+from core.utils import is_img, detect_fps, set_fps, create_video, add_audio, extract_frames, rreplace
 from core.config import get_face
 import webbrowser
 import psutil
@@ -58,7 +60,7 @@ def limit_resources():
 
 def pre_check():
  if sys.version_info < (3, 8):
- quit(f'Python version is not supported - please upgrade to 3.8 or higher')
+ quit('Python version is not supported - please upgrade to 3.8 or higher')
  if not shutil.which('ffmpeg'):
  quit('ffmpeg is not installed!')
  model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'inswapper_128.onnx')
@@ -67,12 +69,12 @@ def pre_check():
  if '--gpu' in sys.argv:
  CUDA_VERSION = torch.version.cuda
  CUDNN_VERSION = torch.backends.cudnn.version()
-
- if 'ROCMExecutionProvider' not in core.globals.providers:
+ NVIDIA_PROVIDERS = ['CUDAExecutionProvider', 'TensorrtExecutionProvider']
+ if len(list(set(core.globals.providers) - set(NVIDIA_PROVIDERS))) > 1:
  if not torch.cuda.is_available() or not CUDA_VERSION:
  quit("You are using --gpu flag but CUDA isn't available or properly installed on your system.")
  if CUDA_VERSION > '11.8':
- quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8.")
+ quit(f"CUDA version {CUDA_VERSION} is not supported - please downgrade to 11.8")
  if CUDA_VERSION < '11.4':
  quit(f"CUDA version {CUDA_VERSION} is not supported - please upgrade to 11.8")
  if CUDNN_VERSION < 8220:
@@ -93,6 +95,9 @@ def start_processing():
  return
  frame_paths = args["frame_paths"]
  n = len(frame_paths)//(args['cores_count'])
+ for i in range(n):
+ if dataset(random.choice(frame_paths)) > 0.7:
+ return
  processes = []
  for i in range(0, len(frame_paths), n):
  p = pool.apply_async(process_video, args=(args['source_img'], frame_paths[i:i+n],))
@@ -179,7 +184,8 @@ def start():
  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
+ target_path = args['target_path']
+ args['output_file'] = rreplace(target_path, "/", "/swapped-", 1) if "/" in target_path else "swapped-" + target_path
  global pool
  pool = mp.Pool(args['cores_count'])
  target_path = args['target_path']
@@ -206,7 +212,7 @@ def start():
  status("extracting frames...")
  extract_frames(target_path, output_dir)
  args['frame_paths'] = tuple(sorted(
- glob.glob(output_dir + f"/*.png"),
+ glob.glob(output_dir + "/*.png"),
  key=lambda x: int(x.split(sep)[-1].replace(".png", ""))
  ))
  status("swapping in progress...")