Compare commits
No commits in common. "master" and "master" have entirely different histories.
|
|
@ -1,14 +0,0 @@
|
|||
## :rocket: 背景
|
||||
|
||||
近年来，开源软件供应链遭受持续的软件投毒和恶意代码攻击，造成了无法估计的损失。例如，Apache Log4j2远程代码执行漏洞被认为是近10年最严重的漏洞之一，攻击者可以在目标服务器上执行任意代码和嗅探系统信息。网络安全专家认为Log4j 中的远程代码执行漏洞可能需要数月甚至数年时间才能得到妥善解决。受Log4J漏洞影响组件包括Apache的Struts2、Solr、Druid、Flink等，Github上60,644个开源项目发布321,094软件存在风险。因此，当前急需智能化技术辅助降低漏洞风险，提高漏洞工程能力，减少漏洞损失。
|
||||
|
||||
|
||||
## :checkered_flag: 比赛要求
|
||||
本项赛事共设计4个赛题，参赛团队选择其中一个完成即可，最终评奖将结合作品完成质量、创新性、实用性等多个维度进行综合评选。
|
||||
|
||||
参赛作品要求在官方竞赛平台“GitLink（确实开源）”提交，包括算法代码、README文件、技术报告以及可以展示算法性能的Docker镜像。具体要求参见赛事网站的“参赛指南”。
|
||||
|
||||
注：推荐使用开源大模型。
|
||||
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,102 @@
|
|||
# 基于大语言模型的软件漏洞补丁分类研究
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
├── Dockerfile
|
||||
├── README.md
|
||||
├── config.yaml
|
||||
├── requirements.txt
|
||||
├── scripts
|
||||
│ ├── evaluate_model.py
|
||||
│ ├── generate_predictions.py
|
||||
│ ├── run_all.sh
|
||||
│ └── train_model.py
|
||||
└── vulnpatch
|
||||
├── __init__.py
|
||||
├── modeling
|
||||
│ ├── __init__.py
|
||||
│ ├── patch_classifier.py
|
||||
│ └── qwen_adapter.py
|
||||
├── preprocessing
|
||||
│ ├── __init__.py
|
||||
│ └── data_cleaner.py
|
||||
├── utils
|
||||
│ ├── __init__.py
|
||||
│ └── metrics.py
|
||||
└── visualization
|
||||
├── __init__.py
|
||||
└── plot_generator.py
|
||||
```
|
||||
|
||||
## 环境配置
|
||||
|
||||
本项目使用 Docker 来确保环境的一致性和可复现性。请按照以下步骤配置环境：
|
||||
|
||||
1. **确保您的系统已安装 Docker 和 NVIDIA Container Toolkit**（用于 GPU 支持）。
|
||||
|
||||
2. **克隆项目目录**：
|
||||
|
||||
```
|
||||
cd vulnerability-patch-classification
|
||||
```
|
||||
|
||||
3. **构建 Docker 镜像**：
|
||||
|
||||
```
|
||||
docker build -t vulnpatch_classifier .
|
||||
```
|
||||
|
||||
这个过程可能需要一些时间，因为它会安装所有必要的依赖项。
|
||||
|
||||
## 运行代码
|
||||
|
||||
为了运行代码并复现报告中的结果，请按照以下步骤操作：
|
||||
|
||||
1. **准备数据**：
|
||||
|
||||
将训练数据、验证数据和测试数据分别放在项目根目录下的 `data/train.csv`、`data/valid.csv` 和 `data/test.csv` 中。
|
||||
|
||||
2. **修改配置**：
|
||||
|
||||
编辑 `config.yaml` 文件，确保所有路径都正确设置，特别注意以下配置项：
|
||||
|
||||
```
|
||||
yaml复制代码train_data_path: '/app/data/train.csv'
|
||||
val_data_path: '/app/data/valid.csv'
|
||||
test_data_path: '/app/data/test.csv'
|
||||
model_output_path: '/app/outputs/model_config.yaml'
|
||||
train_predictions_path: '/app/outputs/train_predictions.csv'
|
||||
val_predictions_path: '/app/outputs/val_predictions.csv'
|
||||
test_predictions_path: '/app/outputs/test_predictions.csv'
|
||||
final_predictions_path: '/app/outputs/final_predictions.csv'
|
||||
prob_dist_plot_path: '/app/outputs/probability_distribution.png'
|
||||
roc_curve_plot_path: '/app/outputs/roc_curve.png'
|
||||
pr_curve_plot_path: '/app/outputs/pr_curve.png'
|
||||
```
|
||||
|
||||
3. **运行 Docker 容器**：
|
||||
|
||||
```
|
||||
docker run --gpus all -v $(pwd):/app vulnpatch_classifier
|
||||
```
|
||||
|
||||
这个命令会自动执行 `scripts/run_all.sh` 脚本，该脚本依次运行训练、评估和预测过程。
|
||||
|
||||
4. **查看结果**：
|
||||
|
||||
运行完成后，您可以在 `outputs` 目录下找到所有的输出文件，包括：
|
||||
|
||||
- 模型配置文件（`model_config.yaml`）
|
||||
- 训练、验证和测试集的预测结果（CSV 文件）
|
||||
- 概率分布图、ROC 曲线和 PR 曲线（PNG 文件）
|
||||
|
||||
5. **分析结果**：
|
||||
|
||||
使用 `scripts/evaluate_model.py` 脚本来分析测试集的性能指标。您可以通过以下命令查看详细的评估结果：
|
||||
|
||||
```
|
||||
docker run --gpus all -v $(pwd):/app vulnpatch_classifier python scripts/evaluate_model.py
|
||||
```
|
||||
|
||||
这将输出测试集上的召回率、F1 分数、NDCG@1 和 NDCG@5 等指标。
|
||||
Binary file not shown.
|
|
@ -0,0 +1,26 @@
|
|||
train_data_path: 'datasets/raw/new_train.csv'
|
||||
valid_data_path: 'datasets/raw/new_valid.csv'
|
||||
test_data_path: 'datasets/raw/new_test.csv'
|
||||
predict_data_path: 'datasets/raw/new_test.csv'
|
||||
|
||||
train_predictions_path: 'outputs/predictions/train_predictions.csv'
|
||||
val_predictions_path: 'outputs/predictions/val_predictions.csv'
|
||||
test_predictions_path: 'outputs/predictions/test_predictions.csv'
|
||||
final_predictions_path: 'outputs/predictions/final_predictions.csv'
|
||||
|
||||
model_output_path: 'outputs/models/model_config.yaml'
|
||||
|
||||
prob_dist_plot_path: 'outputs/figures/probability_distribution.png'
|
||||
roc_curve_plot_path: 'outputs/figures/roc_curve.png'
|
||||
pr_curve_plot_path: 'outputs/figures/precision_recall_curve.png'
|
||||
|
||||
model_name: 'Qwen/Qwen2.5-14B-Instruct'
|
||||
max_new_tokens: 256
|
||||
|
||||
random_seed: 42
|
||||
test_size: 0.2
|
||||
|
||||
use_cuda: true
|
||||
half_precision: true
|
||||
|
||||
log_level: 'INFO'
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-devel
|
||||
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
cmake \
|
||||
git \
|
||||
curl \
|
||||
vim \
|
||||
wget \
|
||||
ca-certificates \
|
||||
libjpeg-dev \
|
||||
libpng-dev \
|
||||
librdmacm1 \
|
||||
libibverbs1 \
|
||||
ibverbs-providers \
|
||||
tzdata \
|
||||
libgl1-mesa-glx \
|
||||
libglib2.0-0 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN curl -sS https://dl.yarnpkg.com/debian/pubkey.gpg | apt-key add -
|
||||
RUN echo "deb https://dl.yarnpkg.com/debian/ stable main" | tee /etc/apt/sources.list.d/yarn.list
|
||||
RUN apt-get update && apt-get install -y yarn
|
||||
|
||||
RUN curl -fsSL https://deb.nodesource.com/setup_16.x | bash -
|
||||
RUN apt-get install -y nodejs
|
||||
|
||||
RUN wget https://github.com/mikefarah/yq/releases/download/v4.16.2/yq_linux_amd64 -O /usr/bin/yq && chmod +x /usr/bin/yq
|
||||
|
||||
RUN git clone https://github.com/facebook/zstd.git \
|
||||
&& cd zstd \
|
||||
&& make -j$(nproc) \
|
||||
&& make install \
|
||||
&& cd .. \
|
||||
&& rm -rf zstd
|
||||
|
||||
RUN pip install --no-cache-dir --upgrade pip
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY requirements.txt .
|
||||
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
RUN git clone https://github.com/NVIDIA/apex \
|
||||
&& cd apex \
|
||||
&& pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ \
|
||||
&& cd .. \
|
||||
&& rm -rf apex
|
||||
|
||||
RUN conda install -c conda-forge -y gperftools
|
||||
|
||||
ENV LD_PRELOAD=/opt/conda/lib/libtcmalloc.so
|
||||
|
||||
RUN mkdir -p /root/.cache/huggingface
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN chmod +x scripts/*.sh
|
||||
|
||||
ENV PYTHONPATH=/app
|
||||
|
||||
CMD ["bash", "scripts/run_all.sh"]
|
||||
|
|
@ -0,0 +1,91 @@
|
|||
torch==2.1.2
|
||||
transformers==4.36.2
|
||||
datasets==2.16.1
|
||||
accelerate==0.25.0
|
||||
peft==0.7.1
|
||||
bitsandbytes==0.41.3
|
||||
scipy==1.11.4
|
||||
scikit-learn==1.3.2
|
||||
pandas==2.1.4
|
||||
numpy==1.26.2
|
||||
matplotlib==3.8.2
|
||||
seaborn==0.13.1
|
||||
plotly==5.18.0
|
||||
dash==2.14.2
|
||||
streamlit==1.29.0
|
||||
fastapi==0.105.0
|
||||
uvicorn==0.25.0
|
||||
gunicorn==21.2.0
|
||||
pyyaml==6.0.1
|
||||
tqdm==4.66.1
|
||||
joblib==1.3.2
|
||||
pytest==7.4.3
|
||||
black==23.12.1
|
||||
isort==5.13.2
|
||||
mypy==1.7.1
|
||||
flake8==6.1.0
|
||||
pre-commit==3.6.0
|
||||
wandb==0.16.1
|
||||
mlflow==2.9.2
|
||||
optuna==3.4.0
|
||||
ray[tune]==2.8.1
|
||||
hydra-core==1.3.2
|
||||
dvc==3.33.3
|
||||
great-expectations==0.18.7
|
||||
evidently==0.4.12
|
||||
shap==0.44.0
|
||||
lime==0.2.0.1
|
||||
alibi==0.9.4
|
||||
fairlearn==0.9.1
|
||||
imbalanced-learn==0.11.0
|
||||
catboost==1.2.2
|
||||
lightgbm==4.1.0
|
||||
xgboost==2.0.2
|
||||
opencv-python==4.8.1.78
|
||||
pillow==10.1.0
|
||||
albumentations==1.3.1
|
||||
torchvision==0.16.2
|
||||
torchtext==0.16.2
|
||||
torchaudio==2.1.2
|
||||
transformers[sentencepiece]==4.36.2
|
||||
tokenizers==0.15.0
|
||||
sacremoses==0.1.1
|
||||
nltk==3.8.1
|
||||
spacy==3.7.2
|
||||
gensim==4.3.2
|
||||
networkx==3.2.1
|
||||
neo4j==5.14.1
|
||||
psycopg2-binary==2.9.9
|
||||
sqlalchemy==2.0.23
|
||||
alembic==1.13.1
|
||||
fastapi-users[sqlalchemy]==12.1.2
|
||||
pydantic==2.5.2
|
||||
pydantic-settings==2.1.0
|
||||
python-jose==3.3.0
|
||||
passlib==1.7.4
|
||||
bcrypt==4.1.2
|
||||
httpx==0.26.0
|
||||
aiohttp==3.9.1
|
||||
requests==2.31.0
|
||||
beautifulsoup4==4.12.2
|
||||
selenium==4.16.0
|
||||
playwright==1.40.0
|
||||
scrapy==2.11.0
|
||||
celery==5.3.6
|
||||
redis==5.0.1
|
||||
pymongo==4.6.1
|
||||
motor==3.3.2
|
||||
elasticsearch==8.11.1
|
||||
kafka-python==2.0.2
|
||||
confluent-kafka==2.3.0
|
||||
apache-airflow==2.7.3
|
||||
dagster==1.5.13
|
||||
prefect==2.14.7
|
||||
dask[complete]==2023.12.1
|
||||
vaex==4.17.0
|
||||
polars==0.19.19
|
||||
pyarrow==14.0.2
|
||||
fastparquet==2023.10.1
|
||||
numba==0.58.1
|
||||
cupy-cuda12x==12.3.0
|
||||
jax[cuda12_pip]==0.4.20
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
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()
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
import yaml
|
||||
from vulnpatch.preprocessing.data_cleaner import DataCleaner
|
||||
from vulnpatch.modeling.patch_classifier import PatchClassifier
|
||||
|
||||
def load_config():
|
||||
with open('config.yaml', 'r') as file:
|
||||
return yaml.safe_load(file)
|
||||
|
||||
def main():
|
||||
config = load_config()
|
||||
|
||||
data_cleaner = DataCleaner(config['predict_data_path'])
|
||||
predict_data = data_cleaner.prepare_for_model()
|
||||
|
||||
classifier = PatchClassifier()
|
||||
|
||||
predictions = classifier.classify_patches(predict_data)
|
||||
|
||||
output_predictions = predictions[['commit_id', 'cve_id', 'probability']]
|
||||
classifier.save_predictions(output_predictions, config['final_predictions_path'])
|
||||
|
||||
print(f"预测结果已保存到 {config['final_predictions_path']}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
#!/bin/bash
|
||||
|
||||
set -e
|
||||
|
||||
export CUDA_VISIBLE_DEVICES=0,1
|
||||
export OMP_NUM_THREADS=8
|
||||
export MKL_NUM_THREADS=8
|
||||
|
||||
mkdir -p logs outputs/models outputs/predictions outputs/figures
|
||||
|
||||
python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; AutoTokenizer.from_pretrained('Qwen/Qwen2.5-14B-Instruct', trust_remote_code=True); AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-14B-Instruct', trust_remote_code=True)"
|
||||
|
||||
python scripts/preprocess_data.py
|
||||
|
||||
python scripts/train_model.py
|
||||
|
||||
python scripts/evaluate_model.py
|
||||
|
||||
python scripts/generate_predictions.py
|
||||
|
||||
python scripts/generate_report.py
|
||||
|
||||
find . -name "*.pyc" -exec rm -f {} \;
|
||||
find . -name "__pycache__" -exec rm -rf {} \;
|
||||
|
||||
tar -czvf outputs.tar.gz outputs
|
||||
|
||||
sleep 300
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
import yaml
|
||||
from vulnpatch.preprocessing.data_cleaner import DataCleaner
|
||||
from vulnpatch.modeling.patch_classifier import PatchClassifier
|
||||
from vulnpatch.utils.metrics import calculate_metrics, optimal_threshold
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
from sklearn.model_selection import train_test_split
|
||||
|
||||
def load_config():
|
||||
with open('config.yaml', 'r') as file:
|
||||
return yaml.safe_load(file)
|
||||
|
||||
def main():
|
||||
config = load_config()
|
||||
|
||||
data_cleaner = DataCleaner(config['train_data_path'])
|
||||
prepared_data = data_cleaner.prepare_for_model()
|
||||
|
||||
train_data, val_data = train_test_split(prepared_data, test_size=0.2, random_state=42)
|
||||
|
||||
classifier = PatchClassifier()
|
||||
|
||||
train_results = classifier.classify_patches(train_data)
|
||||
classifier.save_predictions(train_results, config['train_predictions_path'])
|
||||
|
||||
val_results = classifier.classify_patches(val_data)
|
||||
classifier.save_predictions(val_results, config['val_predictions_path'])
|
||||
|
||||
threshold = optimal_threshold(val_results['label'], val_results['probability'])
|
||||
|
||||
val_metrics = calculate_metrics(
|
||||
val_results['label'],
|
||||
(val_results['probability'] >= threshold).astype(int),
|
||||
val_results['probability']
|
||||
)
|
||||
|
||||
print("验证集指标:")
|
||||
for metric, value in val_metrics.items():
|
||||
print(f"{metric}: {value:.4f}")
|
||||
|
||||
print(f"最佳阈值: {threshold:.4f}")
|
||||
|
||||
with open(config['model_output_path'], 'w') as f:
|
||||
yaml.dump({'threshold': threshold}, f)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
from vulnpatch.modeling.qwen_adapter import QwenAdapter
|
||||
from typing import List, Dict
|
||||
import pandas as pd
|
||||
|
||||
class PatchClassifier:
|
||||
def __init__(self):
|
||||
self.qwen_adapter = QwenAdapter()
|
||||
|
||||
def classify_patches(self, data: List[Dict]) -> pd.DataFrame:
|
||||
results = []
|
||||
for item in data:
|
||||
probability = self.qwen_adapter.classify_patch(item['desc_cve'], item['diff_code'])
|
||||
results.append({
|
||||
'commit_id': item['commit_id'],
|
||||
'cve_id': item['cve_id'],
|
||||
'probability': probability,
|
||||
'label': item['label']
|
||||
})
|
||||
return pd.DataFrame(results)
|
||||
|
||||
def save_predictions(self, predictions: pd.DataFrame, output_path: str):
|
||||
predictions.to_csv(output_path, index=False)
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
import torch
|
||||
|
||||
class QwenAdapter:
|
||||
def __init__(self, model_name: str = "Qwen/Qwen2.5-14B-Instruct"):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).half().cuda()
|
||||
self.model.eval()
|
||||
|
||||
def generate_response(self, prompt: str) -> str:
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda")
|
||||
with torch.no_grad():
|
||||
generated = self.model.generate(**inputs, max_new_tokens=100)
|
||||
response = self.tokenizer.decode(generated[0], skip_special_tokens=True)
|
||||
return response.split("Human:")[0].strip()
|
||||
|
||||
def classify_patch(self, cve_desc: str, patch_diff: str) -> float:
|
||||
prompt = f"Given the following CVE description and patch diff, determine the probability (0-1) that this patch is related to the described vulnerability:\n\nCVE Description: {cve_desc}\n\nPatch Diff: {patch_diff}\n\nProbability:"
|
||||
response = self.generate_response(prompt)
|
||||
try:
|
||||
probability = float(response)
|
||||
return max(0, min(1, probability))
|
||||
except ValueError:
|
||||
return 0.5
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
import pandas as pd
|
||||
import json
|
||||
from typing import Dict, List
|
||||
|
||||
class DataCleaner:
|
||||
def __init__(self, file_path: str):
|
||||
self.df = pd.read_csv(file_path)
|
||||
|
||||
def clean_data(self) -> pd.DataFrame:
|
||||
self.df['diff_code'] = self.df['diff_code'].apply(json.loads)
|
||||
self.df['cwe'] = self.df['cwe'].apply(eval)
|
||||
self.df['commit_mess'] = self.df['commit_mess'].fillna('')
|
||||
self.df['desc_cve'] = self.df['desc_cve'].fillna('')
|
||||
return self.df
|
||||
|
||||
def prepare_for_model(self) -> List[Dict]:
|
||||
cleaned_df = self.clean_data()
|
||||
prepared_data = []
|
||||
for _, row in cleaned_df.iterrows():
|
||||
prepared_data.append({
|
||||
'commit_id': row['commit_id'],
|
||||
'cve_id': row['cve_id'],
|
||||
'diff_code': json.dumps(row['diff_code']),
|
||||
'commit_mess': row['commit_mess'],
|
||||
'cwe': ','.join(map(str, row['cwe'])),
|
||||
'desc_cve': row['desc_cve'],
|
||||
'label': row['label']
|
||||
})
|
||||
return prepared_data
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
import numpy as np
|
||||
from sklearn.metrics import recall_score, f1_score, ndcg_score
|
||||
|
||||
def calculate_metrics(y_true, y_pred, y_score):
|
||||
recall = recall_score(y_true, y_pred)
|
||||
f1 = f1_score(y_true, y_pred)
|
||||
ndcg_1 = ndcg_score(y_true.reshape(1, -1), y_score.reshape(1, -1), k=1)
|
||||
ndcg_5 = ndcg_score(y_true.reshape(1, -1), y_score.reshape(1, -1), k=5)
|
||||
|
||||
return {
|
||||
'Recall': recall,
|
||||
'F1-Score': f1,
|
||||
'NDCG@1': ndcg_1,
|
||||
'NDCG@5': ndcg_5
|
||||
}
|
||||
|
||||
def optimal_threshold(y_true, y_score):
|
||||
thresholds = np.linspace(0, 1, 100)
|
||||
f1_scores = [f1_score(y_true, y_score >= threshold) for threshold in thresholds]
|
||||
return thresholds[np.argmax(f1_scores)]
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
import matplotlib.pyplot as plt
|
||||
import seaborn as sns
|
||||
import pandas as pd
|
||||
|
||||
class PlotGenerator:
|
||||
def __init__(self, results_df: pd.DataFrame):
|
||||
self.results_df = results_df
|
||||
plt.style.use('seaborn')
|
||||
|
||||
def plot_probability_distribution(self, output_path: str):
|
||||
plt.figure(figsize=(10, 6))
|
||||
sns.histplot(data=self.results_df, x='probability', hue='label', kde=True, element='step')
|
||||
plt.title('Distribution of Patch Relevance Probabilities')
|
||||
plt.xlabel('Probability')
|
||||
plt.ylabel('Count')
|
||||
plt.legend(['Non-relevant', 'Relevant'])
|
||||
plt.savefig(output_path)
|
||||
plt.close()
|
||||
|
||||
def plot_roc_curve(self, output_path: str):
|
||||
from sklearn.metrics import roc_curve, auc
|
||||
fpr, tpr, _ = roc_curve(self.results_df['label'], self.results_df['probability'])
|
||||
roc_auc = auc(fpr, tpr)
|
||||
|
||||
plt.figure(figsize=(8, 6))
|
||||
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
|
||||
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
|
||||
plt.xlim([0.0, 1.0])
|
||||
plt.ylim([0.0, 1.05])
|
||||
plt.xlabel('False Positive Rate')
|
||||
plt.ylabel('True Positive Rate')
|
||||
plt.title('Receiver Operating Characteristic (ROC) Curve')
|
||||
plt.legend(loc="lower right")
|
||||
plt.savefig(output_path)
|
||||
plt.close()
|
||||
|
||||
def plot_precision_recall_curve(self, output_path: str):
|
||||
from sklearn.metrics import precision_recall_curve, average_precision_score
|
||||
precision, recall, _ = precision_recall_curve(self.results_df['label'], self.results_df['probability'])
|
||||
avg_precision = average_precision_score(self.results_df['label'], self.results_df['probability'])
|
||||
|
||||
plt.figure(figsize=(8, 6))
|
||||
plt.step(recall, precision, color='b', alpha=0.2, where='post')
|
||||
plt.fill_between(recall, precision, step='post', alpha=0.2, color='b')
|
||||
plt.xlabel('Recall')
|
||||
plt.ylabel('Precision')
|
||||
plt.ylim([0.0, 1.05])
|
||||
plt.xlim([0.0, 1.0])
|
||||
plt.title(f'Precision-Recall curve: AP={avg_precision:.2f}')
|
||||
plt.savefig(output_path)
|
||||
plt.close()
|
||||
Loading…
Reference in New Issue