Veritas-AI / scripts /train_fusion.py
Aditya-Jadhav150
feat: upgrade to Veritas-AI v3.1 architecture, fast video & image deepfake pipeline, UI enhancements
37b2345
Raw History Blame Contribute Delete
1.81 kB
import pandas as pd
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report
import joblib
import json
def train_xgboost(csv_path: str, output_model: str, output_scaler: str):
"""
Trains the XGBoost fusion classifier on extracted metrics.
"""
df = pd.read_csv(csv_path)
# Assume 'label' column: 1 = fake, 0 = real
if 'label' not in df.columns:
raise ValueError("CSV must contain a 'label' column")
X = df.drop(columns=['label', 'filename'], errors='ignore')
y = df['label']
# Scale features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
model = XGBClassifier(n_estimators=200, max_depth=6, learning_rate=0.05, random_state=42)
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
print("Training Complete. Test Accuracy:", accuracy_score(y_test, y_pred))
print(classification_report(y_test, y_pred))
# Save model using JSON format
model.save_model(output_model)
# Save scaler in JSON format
scaler_data = {
"mean": scaler.mean_.tolist(),
"var": scaler.var_.tolist(),
"scale": scaler.scale_.tolist(),
"n_features_in": scaler.n_features_in_
}
with open(output_scaler, 'w') as f:
json.dump(scaler_data, f)
if __name__ == "__main__":
import sys
if len(sys.argv) < 2:
print("Usage: python train_fusion.py path_to_features.csv")
else:
train_xgboost(sys.argv[1], "../dataset/fusion_engine_best.json", "../dataset/scaler.json")