From 7c3f75ae889ecdbd37899abd21f7594ed16d308d Mon Sep 17 00:00:00 2001
From: Somdev Sangwan <s0md3v@gmail.com>
Date: Wed, 31 May 2023 15:20:26 +0530
Subject: [PATCH] better face detection

---
 requirements.txt | 1 +
 run.py           | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/requirements.txt b/requirements.txt
index 4eeae19..c3441ef 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,3 +7,4 @@ tk==0.1.0
 pillow==9.5.0
 torch==2.0.1
 onnxruntime-gpu==1.15.0
+opennsfw2==0.10.2
\ No newline at end of file
diff --git a/run.py b/run.py
index e169f39..c1ab595 100755
--- a/run.py
+++ b/run.py
@@ -10,9 +10,11 @@ import glob
 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, rreplace
@@ -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:
+            quit("[WARNING] Unable to determine location of the face in the target. Please make sure the target isn't wearing clothes matching to their skin.")
     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],))