Add sample files

.gitignore

@@ -1,4 +1,5 @@
 temps/*
 !temps/.gitkeep
 
+sample_files/images/*
 env

app.py

@@ -14,55 +14,88 @@ from PIL import Image
 
 import shlex
 import subprocess
+from file_picker import st_file_selector
 
 import os
 
+DEBUG = False
+
 def main():
     st.markdown("###")    
     uploaded_file = st.file_uploader('Upload a picture', type=['mp4', 'jpg', 'jpeg', 'png'], accept_multiple_files=False)
+    st.markdown("### or")
+    selected_file = st_file_selector(st, path='.\sample_files', key = 'selected', label = 'Choose a sample image/video')
+
     if uploaded_file: 
         random_id = uuid.uuid1()
+        base_folder = "temps"
         filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
         file_type = uploaded_file.type.split("/")[0]
+        filepath = f"{base_folder}/{filename}"
 
         if uploaded_file.type == 'video/mp4':
             with open(f"temps/{filename}", mode='wb') as f:
                 f.write(uploaded_file.read())
             st.video(uploaded_file)
-            pass
         else:
             img = Image.open(uploaded_file).convert('RGB')
             ext = uploaded_file.type.split("/")[-1]
             with open(f"temps/{filename}", mode='wb') as f:
                 f.write(uploaded_file.getbuffer())  
             st.image(img)
+    elif selected_file:
+        base_folder = "sample_files"
+        file_type = selected_file.split(".")[-1]
+        filename = selected_file.split("/")[-1]
+        filepath = f"{base_folder}/{selected_file}"
+
+        if file_type == 'mp4':
+            video_file = open(filepath, 'rb')
+            video_bytes = video_file.read()
+            st.video(video_bytes)
+        else:
+            image_file = open(filepath, 'rb')
+            image_bytes = image_file.read()
+            st.image(image_bytes)
+    else:
+        return
+
 
-        
-
-        with st.spinner(f'Processing {file_type}...'):
-            subprocess.run(shlex.split(f"python extract_video.py --device cuda --max_frames 50 --bs 2 --frame_interval 5 --confidence_threshold 0.997 --data_path temps/{filename}"))
-        st.text(f'1. Processing {file_type} ✅')
-        with st.spinner(f'Analyzing {file_type}...'):
-            pred = subprocess.run(shlex.split(f"python inference.py --weight weights/model_params_ffpp_c23.pickle --device cpu --image_folder temps/images/{filename}"), capture_output=True)
-        st.text(f'2. Analyzing {file_type} ✅')
-        
-        print(pred)
-        try:
-            fake_probability = float(pred.stdout.decode('utf-8').split('Mean prediction: ')[-1])
-            if fake_probability > 0.6:
-                st.error('  FAKE! ', icon="🚨")
-            else:
-                st.success("  REAL FOOTAGE! ", icon="✅")
-            st.text("fake probability {:.2f}".format(fake_probability))
-
-            os.remove(f"temps/{filename}")
-            folder_name = ".".join(filename.split(".")[:-1])
-            shutil.rmtree(f"temps/images/{folder_name}")
-        except:
-            st.text(pred.stdout.decode('utf-8'))
+    
+
+    with st.spinner(f'Processing {file_type}...'):
+        processing_stdout = subprocess.run(shlex.split(f"""python extract_video.py --device cpu --max_frames 50 --bs 2 --frame_interval 60 --confidence_threshold 0.997 --data_path "{filepath}" """), capture_output=True)
+    st.text(f'1. Processing {file_type} ✅')
+    with st.spinner(f'Analyzing {file_type}...'):
+        analyze_stdout = subprocess.run(shlex.split(f"""python inference.py --weight weights/model_params_ffpp_c23.pickle --device cpu --image_folder "{base_folder}/images/{filename}" """), capture_output=True)
+    st.text(f'2. Analyzing {file_type} ✅')
+
+    if len(os.listdir("{}/images/{}".format(base_folder, filename.split(".")[0]))) < 1:
+        st.text("No faces could be detected! 🚨")
+        return
+    
+    try:
+        fake_probability = float(analyze_stdout.stdout.decode('utf-8').split('Mean prediction: ')[-1])
+        if fake_probability > 0.6:
+            st.error('  FAKE! ', icon="🚨")
+        else:
+            st.success("  REAL FOOTAGE! ", icon="✅")
+        st.text("fake probability {:.2f}".format(fake_probability))
+
+        # os.remove(f"{base_folder}/{filename}")
+        folder_name = ".".join(filename.split(".")[:-1])
+        shutil.rmtree(f"{base_folder}/images/{folder_name}")
+    except Exception as e:
+        if DEBUG:
+            st.text(processing_stdout.stdout.decode('utf-8'))
+            st.text(analyze_stdout.stdout.decode('utf-8'))
 
             st.text("")
-            st.text(pred)
+            st.text(processing_stdout)
+            st.text(analyze_stdout)
+            st.write(e)
+        else:
+            st.text("Encountered a problem while analyzing video/image 🚨")
 
 
 def setup():

file_picker.py

@@ -0,0 +1,45 @@
+"""FilePicker for streamlit. 
+Still doesn't seem to be a good solution for a way to select files to process from the server Streamlit is running on.
+Here's a pretty functional solution. 
+Usage:
+```
+import streamlit as st
+from filepicker import st_file_selector
+tif_file = st_file_selector(st, key = 'tif', label = 'Choose tif file')
+```
+"""
+
+import os
+import streamlit as st
+
+def update_dir(key):
+    choice = st.session_state[key]
+    if os.path.isdir(os.path.join(st.session_state[key+'curr_dir'], choice)):
+        st.session_state[key+'curr_dir'] = os.path.normpath(os.path.join(st.session_state[key+'curr_dir'], choice))
+        files = sorted(os.listdir(st.session_state[key+'curr_dir']))
+        files.remove("images")
+        st.session_state[key+'files'] = files
+
+def st_file_selector(st_placeholder, path='.', label='Select a file/folder', key = 'selected'):
+    if key+'curr_dir' not in st.session_state:
+        base_path = '.' if path is None or path == '' else path
+        base_path = base_path if os.path.isdir(base_path) else os.path.dirname(base_path)
+        base_path = '.' if base_path is None or base_path is '' else base_path
+
+        files = sorted(os.listdir(base_path))
+        files.insert(0, 'Choose a file...')
+        files.remove("images")
+        st.session_state[key+'files'] = files
+        st.session_state[key+'curr_dir'] = base_path
+    else:
+        base_path = st.session_state[key+'curr_dir']
+
+    selected_file = st_placeholder.selectbox(label=label, 
+                                        options=st.session_state[key+'files'], 
+                                        key=key, 
+                                        on_change = lambda: update_dir(key))
+    
+    if selected_file == "Choose a file...":
+        return None
+
+    return selected_file

sample_files/jokowi deepfake - cupid.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83f057b648c76a1d1a92224767222cf9a74f8d8849ea4979a0201d0d8e14d906
+size 2163429

sample_files/messi_deepfake.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88ac25406e754c467b50686961fae2c7975eb816dfe8dc2974304ac9d57510ac
+size 2520688

sample_files/obama_real.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5fb30fd607ed6395c56644eec2dbe71e32354c69e68c52d667e25645b648050
+size 4856656