File size: 8,251 Bytes
2fbad05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 |
import streamlit as st
from PIL import Image
from models import mesonet, mesoinception, fal_detector, local_detector
from utils.visualization import display_results
from utils.preprocessing import preprocess_image, preprocess_video
from utils.preprocessing import preprocess_image, generate_local_image
import numpy as np
import cv2
import tempfile
import os
# Initialize session state for tabs and uploaded files
if "active_tab" not in st.session_state:
st.session_state["active_tab"] = "Face Photoshop Detection"
if "uploaded_file" not in st.session_state:
st.session_state["uploaded_file"] = None
# Load models
models = {
"MesoNet": mesonet.load_mesonet("models/weights/Meso4_DF.h5"),
"MesoInception": mesoinception.load_mesonetInception("models/weights/MesoInception_DF.h5"),
"Photoshop FALdetector Global": fal_detector.load_fal_detector("models/weights/global.pth"),
"Photoshop FALdetector Local": local_detector.load_local_detector("models/weights/local.pth", gpu_id=-1),
}
st.title("DeepSAIF")
# Create tabs for different functionalities
tab1, tab2, tab3 = st.tabs(["Face Photoshop Detection", "DeepFake Detection for Images", "DeepFake Detection for Videos"])
# Tab 1: Photoshop Detection
with tab1:
if st.session_state["active_tab"] != "Face Photoshop Detection":
st.session_state["uploaded_file"] = None
st.session_state["active_tab"] = "Face Photoshop Detection"
st.header("Face Photoshop Detection")
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="photoshop")
if uploaded_file:
st.session_state["uploaded_file"] = uploaded_file
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
local_image = generate_local_image(image)
# Run inference on all models
results = {}
for model_name, model in models.items():
if model_name == "Photoshop FALdetector Global":
results[model_name] = fal_detector.predict_fal_detector(model, image)
elif model_name == "Photoshop FALdetector Local":
heatmap_path, prediction = local_detector.predict_and_generate_heatmap(model, image)
if heatmap_path:
# Display the heatmap using Streamlit
st.image(heatmap_path, caption=f"Heatmap for {model_name}", use_container_width=True)
# Delete the temporary heatmap file after display
os.remove(heatmap_path)
os.remove('cropped_input.jpg')
os.remove('warped.jpg')
else:
st.error(f"Failed to generate heatmap for {model_name}")
results[model_name] = prediction
# elif model_name == "Global Classifier":
# results[model_name] = global_classifier.classify_fake(model, image)
# Display results
display_results(results)
# Tab 2: DeepFake Detection for Images
with tab2:
if st.session_state["active_tab"] != "DeepFake Detection for Images":
st.session_state["uploaded_file"] = None
st.session_state["active_tab"] = "DeepFake Detection for Images"
st.header("DeepFake Detection for Images")
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="deepfake_image")
if uploaded_file:
st.session_state["uploaded_file"] = uploaded_file
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
local_image = generate_local_image(image)
# Preprocess the image
# preprocessed_image = preprocess_image(uploaded_file)
# Run inference on all models
results = {}
for model_name, model in models.items():
if model_name == "MesoNet":
results[model_name] = mesonet.predict_mesonet(model, image)
elif model_name == "MesoInception":
results[model_name] = mesoinception.predict_mesonetInception(model, image)
# Display results
display_results(results)
def confident_strategy(pred, t=0.8):
"""
Implements the confident averaging strategy for predictions.
Args:
pred (list[float]): List of predictions for each frame.
t (float): Threshold for high-confidence fake detection.
Returns:
float: Final confidence score for the video.
"""
if len(pred) == 0:
return np.nan
pred = np.array(pred)
sz = len(pred)
fakes = np.count_nonzero(pred > t)
if fakes > sz // 2.5 and fakes > 11:
return np.mean(pred[pred > t])
elif np.count_nonzero(pred < 0.2) > 0.9 * sz:
return np.mean(pred[pred < 0.2])
else:
return np.mean(pred)
# Tab 3: DeepFake Detection for Videos
with tab3:
if st.session_state["active_tab"] != "DeepFake Detection for Videos":
st.session_state["uploaded_file"] = None
st.session_state["active_tab"] = "DeepFake Detection for Videos"
st.header("DeepFake Detection for Videos")
uploaded_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"], key="deepfake_video")
if uploaded_file:
st.session_state["uploaded_file"] = uploaded_file
with st.spinner("Processing video..."):
# Save uploaded file to a temporary location
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
temp_file.write(uploaded_file.getbuffer())
video_path = temp_file.name
try:
# Test video accessibility
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
st.error("Failed to open video file.")
else:
st.success("Video file opened successfully!")
# Extract frames from video
frames = preprocess_video(video_path, frame_count=32)
if len(frames) == 0:
st.error("Failed to extract frames from the video.")
else:
# st.success(f"Extracted {len(frames)} frames.")
# for frame in frames[:5]: # Display first 5 frames
# st.image(frame, caption="Extracted Frame")
# Dictionary to store model predictions
model_results = {
"MesoNet": [],
"Photoshop FALdetector Global": []
}
# Iterate over frames and make predictions for each model
for frame in frames:
preprocessed_frame = preprocess_image(frame) # Preprocess frame
local_image = generate_local_image(preprocessed_frame)
# Predictions for MesoNet and Photoshop FALdetector Global
model_results["MesoNet"].append(
mesonet.predict_mesonet(models["MesoNet"], preprocessed_frame)
)
model_results["Photoshop FALdetector Global"].append(
fal_detector.predict_fal_detector(models["Photoshop FALdetector Global"], local_image)
)
# Apply the confident averaging strategy for each model
final_results = {}
for model_name, predictions in model_results.items():
final_results[model_name] = confident_strategy(predictions)
# Display results
st.write("### Video Analysis Results")
display_results(final_results)
# Optionally show detailed frame predictions per model
if st.checkbox("Show Detailed Frame Predictions"):
for model_name, predictions in model_results.items():
st.write(f"### Predictions for {model_name}")
st.bar_chart(predictions)
finally:
# Clean up temporary file
os.remove(video_path)
|