Spaces:
Sleeping
Sleeping
File size: 5,925 Bytes
f58f98d | 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 | import os
import gc
import sys
import torch
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score, precision_score, recall_score, confusion_matrix
from tqdm import tqdm
import matplotlib.pyplot as plt
import seaborn as sns
# Add your model directory to sys.path if needed
sys.path.append('dfdc_model')
from dfdc_model.kernel_utils import VideoReader, FaceExtractor
from dfdc_model.training.zoo.classifiers import DeepFakeClassifier
from enhanced_prediction import enhanced_predict_on_video
real_dir = r"C:\Users\babel\Downloads\archive\FaceForensics++_C23\original"
fake_dir = r"C:\Users\babel\Downloads\archive\FaceForensics++_C23\FaceSwap"
def cleanup_memory():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
def load_models():
models = []
weights_dir = "dfdc_model/weights"
model_files = [
"final_888_DeepFakeClassifier_tf_efficientnet_b7_ns_0_40",
]
device = "cuda" if torch.cuda.is_available() else "cpu"
for model_file in model_files:
model_path = os.path.join(weights_dir, model_file)
if os.path.exists(model_path):
model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns").to(device)
checkpoint = torch.load(model_path, map_location="cpu")
state_dict = checkpoint.get("state_dict", checkpoint)
model.load_state_dict({k.replace("module.", ""): v for k, v in state_dict.items()}, strict=True)
model.eval()
del checkpoint
if device == "cuda":
models.append(model.half())
else:
models.append(model)
return models
def initialize_preprocessing():
frames_per_video = 10 # Process 5 frames per video
video_reader = VideoReader()
video_read_fn = lambda x: video_reader.read_frames(x, num_frames=frames_per_video)
face_extractor = FaceExtractor(video_read_fn)
input_size = 380
return face_extractor, input_size, frames_per_video
def smart_prediction_strategy(predictions):
predictions = np.array(predictions)
if len(predictions) == 0:
return 0.5
fake_frames_50 = np.sum(predictions > 0.5)
fake_frames_60 = np.sum(predictions > 0.6)
fake_frames_70 = np.sum(predictions > 0.7)
total_frames = len(predictions)
fake_ratio = fake_frames_50 / total_frames
if fake_ratio > 0.5:
if fake_frames_70 > total_frames * 0.3:
return min(0.95, np.mean(predictions[predictions > 0.7]))
elif fake_frames_60 > total_frames * 0.4:
return min(0.9, np.mean(predictions[predictions > 0.6]))
else:
suspicious_mean = np.mean(predictions[predictions > 0.5])
overall_mean = np.mean(predictions)
return (suspicious_mean * 0.7 + overall_mean * 0.3)
elif fake_ratio > 0.3:
return min(0.7, np.mean(predictions))
else:
return np.mean(predictions)
def evaluate(dataset_csv, real_dir, fake_dir):
df = pd.read_csv(dataset_csv)
device = "cuda" if torch.cuda.is_available() else "cpu"
models = load_models()
face_extractor, input_size, frames_per_video = initialize_preprocessing()
y_trues = []
y_preds = []
loop = tqdm(df.iterrows(), total=len(df), desc="Evaluating")
for idx, row in loop:
# Determine source dir
filename = row["filename"]
if row["label"] == 0:
file_path = os.path.join(real_dir, filename)
else:
file_path = os.path.join(fake_dir, filename)
gt_label = row["label"]
try:
result = enhanced_predict_on_video(
face_extractor=face_extractor,
video_path=file_path,
batch_size=1, # Keep batch size at 1 for memory efficiency
input_size=input_size,
models=models,
strategy=smart_prediction_strategy,
apply_compression=False,
return_frame_predictions=False
)
pred_score = result["prediction"]
pred_label = 1 if pred_score > 0.5 else 0
y_trues.append(gt_label)
y_preds.append(pred_label)
except Exception as e:
print(f"Error processing {file_path}: {e}")
y_trues.append(gt_label)
y_preds.append(0.5) # neutral
cleanup_memory()
# Metrics
y_trues_np = np.array(y_trues)
y_preds_np = np.array(y_preds)
correct = y_preds_np != 0.5 # exclude neutral (errors)
acc = accuracy_score(y_trues_np[correct], y_preds_np[correct])
prec = precision_score(y_trues_np[correct], y_preds_np[correct])
rec = recall_score(y_trues_np[correct], y_preds_np[correct])
cm = confusion_matrix(y_trues_np[correct], y_preds_np[correct])
print(f"Accuracy: {acc:.4f}")
print(f"Precision: {prec:.4f}")
print(f"Recall: {rec:.4f}")
print("Confusion Matrix:")
print(cm)
# Plot confusion matrix and metrics, save as PNG
plt.figure(figsize=(7, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
plt.title('Confusion Matrix')
plt.xlabel('Predicted Label')
plt.ylabel('True Label')
# Add metrics as text box
metrics_text = f"Accuracy: {acc:.2f}\nPrecision: {prec:.2f}\nRecall: {rec:.2f}"
plt.gcf().text(0.7, 0.2, metrics_text, fontsize=12, bbox=dict(facecolor='white', alpha=0.7))
plt.tight_layout()
plt.savefig('confusion_matrix.png')
plt.close()
return {
"accuracy": acc,
"precision": prec,
"recall": rec,
"confusion_matrix": cm,
"total": len(df),
"evaluated": sum(correct)
}
if __name__ == "__main__":
# Usage
# Make sure 'evaluation_list.csv' is created as above
dataset_csv = "evaluation_list.csv"
stats = evaluate(dataset_csv, real_dir, fake_dir)
print(stats)
|