forked from Eshe/competition-vd
43 lines
1.4 KiB
Python
Executable File
43 lines
1.4 KiB
Python
Executable File
import yaml
|
|
from vulnpatch.preprocessing.data_cleaner import DataCleaner
|
|
from vulnpatch.modeling.patch_classifier import PatchClassifier
|
|
from vulnpatch.utils.metrics import calculate_metrics
|
|
from vulnpatch.visualization.plot_generator import PlotGenerator
|
|
|
|
def load_config():
|
|
with open('config.yaml', 'r') as file:
|
|
return yaml.safe_load(file)
|
|
|
|
def main():
|
|
config = load_config()
|
|
|
|
with open(config['model_output_path'], 'r') as f:
|
|
model_config = yaml.safe_load(f)
|
|
threshold = model_config['threshold']
|
|
|
|
data_cleaner = DataCleaner(config['test_data_path'])
|
|
test_data = data_cleaner.prepare_for_model()
|
|
|
|
classifier = PatchClassifier()
|
|
|
|
test_results = classifier.classify_patches(test_data)
|
|
classifier.save_predictions(test_results, config['test_predictions_path'])
|
|
|
|
test_metrics = calculate_metrics(
|
|
test_results['label'],
|
|
(test_results['probability'] >= threshold).astype(int),
|
|
test_results['probability']
|
|
)
|
|
|
|
print("测试集指标:")
|
|
for metric, value in test_metrics.items():
|
|
print(f"{metric}: {value:.4f}")
|
|
|
|
plot_generator = PlotGenerator(test_results)
|
|
plot_generator.plot_probability_distribution(config['prob_dist_plot_path'])
|
|
plot_generator.plot_roc_curve(config['roc_curve_plot_path'])
|
|
plot_generator.plot_precision_recall_curve(config['pr_curve_plot_path'])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |