[update] - thêm các tệp LocFileFcg.py và thongkeNodeAndEdge.py, cập nhật hàm data_vizualization trong output.py để hỗ trợ vẽ PCA cho dữ liệu benign và malware

Fix by: pikconlonton

Author: SuZeAI

core/addition/LocFileFcg.py

@@ -0,0 +1,27 @@
+import os
+
+# Định nghĩa đường dẫn của các thư mục và tệp
+folder_b = '../../data_storage/processed/fcg_prune'  # Thư mục chứa các file .fcg
+file_a = '../..data_storage/processed/constant/hash_benign2.txt'  # Tệp chứa danh sách các tên file không đuôi
+output_file = '../..data_storage/metric_pruning/fcg_prun_benign.txt'  # Tệp lưu đường dẫn của các file tìm thấy
+
+# Đọc các tên file từ A.txt vào một danh sách
+with open(file_a, 'r') as f:
+    file_names = [line.strip() for line in f.readlines()]
+
+# Lấy tất cả các file trong folder B có đuôi .fcg
+fcg_files = [f for f in os.listdir(folder_b) if f.endswith('.fcg')]
+
+# Lọc ra những file trong folder B mà tên không có đuôi trùng với những tên trong A.txt
+matched_files = []
+for fcg in fcg_files:
+    file_name_without_extension = os.path.splitext(fcg)[0]  # Bỏ đuôi .fcg
+    if file_name_without_extension in file_names:
+        matched_files.append(os.path.join(folder_b, fcg))  # Lưu đường dẫn đầy đủ
+
+# Lưu đường dẫn của các file tìm thấy vào C.txt
+with open(output_file, 'w') as f:
+    for file_path in matched_files:
+        f.write(file_path + '\n')
+
+print(f"Đã lưu các đường dẫn file vào tệp: {output_file}")

core/addition/thongkeNodeAndEdge.py

@@ -0,0 +1,82 @@
+import dgl
+import networkx as nx
+import numpy as np
+
+malware_fcg = []
+malware_fcg_prune = []
+benign_fcg = []
+benign_fcg_prune = []
+
+# Đọc danh sách các file từ các tệp
+with open('/workspace/Graph_reduce_Contrastive_Learning_ADM-master/core/metric_pruning/fcg_malware.txt', 'r') as f:
+    malware_fcg = [line.strip() for line in f.readlines()]
+
+with open('/workspace/Graph_reduce_Contrastive_Learning_ADM-master/core/metric_pruning/fcg_benign.txt', 'r') as f:
+    benign_fcg = [line.strip() for line in f.readlines()]
+
+with open('/workspace/Graph_reduce_Contrastive_Learning_ADM-master/core/metric_pruning/fcg_prun_malware.txt', 'r') as f:
+    malware_fcg_prune = [line.strip() for line in f.readlines()]
+
+with open('/workspace/Graph_reduce_Contrastive_Learning_ADM-master/core/metric_pruning/fcg_prun_benign.txt', 'r') as f:
+    benign_fcg_prune = [line.strip() for line in f.readlines()]
+
+# Hàm đọc file .fcg
+def loadGraph(path):
+    g = dgl.load_graphs(path)[0][0]
+    return g
+
+def edge_node_count(graph):
+
+    n_nodes = []
+    n_edges = []
+    n_nodes.append(len(graph.nodes()))
+    tmp = np.array(graph.edges()).tolist()
+    n_edges.append(len(tmp[0]))
+    return n_nodes, n_edges
+    
+# Khởi tạo các biến đếm
+sum_node_malware = 0
+sum_edge_malware = 0
+sum_node_benign = 0
+sum_edge_benign = 0
+sum_node_malware_prun = 0
+sum_edge_malware_prun = 0
+sum_node_benign_prun = 0
+sum_edge_benign_prun = 0
+# Tính tổng số node và cạnh cho malware
+for fcg in malware_fcg:
+    g = loadGraph(fcg)
+    nodes, edges = edge_node_count(g)
+    sum_node_malware += nodes[0]
+    sum_edge_malware += edges[0]
+
+for fcg in benign_fcg:
+    g = loadGraph(fcg)
+    nodes, edges = edge_node_count(g)
+    sum_node_benign += nodes[0]
+    sum_edge_benign += edges[0]
+
+for fcg in malware_fcg_prune:
+    g = loadGraph(fcg)
+    nodes, edges = edge_node_count(g)
+    sum_node_malware_prun += nodes[0]
+    sum_edge_malware_prun += edges[0]
+
+for fcg in benign_fcg_prune:
+    g = loadGraph(fcg)
+    nodes, edges = edge_node_count(g)
+    sum_node_benign_prun += nodes[0]
+    sum_edge_benign_prun += edges[0]
+
+# Ghi kết quả vào file A.txt
+with open('/workspace/Graph_reduce_Contrastive_Learning_ADM-master/core/metric_pruning/LD.txt', 'w') as f:
+    f.write(f"Avg node malware: {sum_node_malware / len(malware_fcg)}\n")
+    f.write(f"Avg edge malware: {sum_edge_malware / len(malware_fcg)}\n")
+    f.write(f"Avg node benign: {sum_node_benign / len(benign_fcg)}\n")
+    f.write(f"Avg edge benign: {sum_edge_benign / len(benign_fcg)}\n")
+    f.write(f"Avg node malware_prune: {sum_node_malware_prun / len(malware_fcg_prune)}\n")
+    f.write(f"Avg edge malware_prune: {sum_edge_malware_prun / len(malware_fcg_prune)}\n")
+    f.write(f"Avg node benign_prune: {sum_node_benign_prun / len(benign_fcg_prune)}\n")
+    f.write(f"Avg edge benign_prune: {sum_edge_benign_prun / len(benign_fcg_prune)}\n")
+
+

core/output/output.py

@@ -87,25 +87,52 @@ def list_score(types: List[str], y_preds, y_trues, path_save_metrics):
         write_txt(paths["metrics"], metrics)
 
     @staticmethod
-    def data_vizualization(preds, y_trues, path_save_metrics):
-        
+    def data_vizualization(preds, y_trues, path_save_metrics, name_img_visual):
         os.makedirs(path_save_metrics, exist_ok=True)
+
         pca = PCA(n_components=2)
         vectors_2d = pca.fit_transform(preds)
+
+        y_trues = np.array(y_trues)
+
         plt.figure(figsize=(8, 6))
-        
-        plt.scatter(vectors_2d[:, 0], vectors_2d[:, 1], c='blue', edgecolors='black')
-        for i, (x, y) in enumerate(vectors_2d):
-            plt.text(x, y, f"{y_trues[i]}", fontsize=8, color='red', ha='right')
 
+        # Benign (label 0)
+        idx_0 = y_trues == 0
+        plt.scatter(
+            vectors_2d[idx_0, 0],
+            vectors_2d[idx_0, 1],
+            c='steelblue',
+            label='Benign (0)',
+            edgecolors='black',
+            alpha=0.7
+        )
+
+        # Malware (label 1)
+        idx_1 = y_trues == 1
+        plt.scatter(
+            vectors_2d[idx_1, 0],
+            vectors_2d[idx_1, 1],
+            c='darkorange',
+            label='Malware (1)',
+            edgecolors='black',
+            alpha=0.7
+        )
+
+        # Trục & tiêu đề
         plt.axhline(0, color='gray', linewidth=0.5)
         plt.axvline(0, color='gray', linewidth=0.5)
         plt.xlabel("PCA Component 1")
         plt.ylabel("PCA Component 2")
-        plt.title("2D Visualization of High-Dimensional Vectors")
+        plt.title("PCA Visualization of Benign and Malware Samples")
         plt.grid(True)
-        path_save_image = f"{path_save_metrics}/pca_cl.png"
+
+        # Chú thích ở góc trên bên phải
+        plt.legend(loc='upper right', title='Class Labels')
+
+        path_save_image = f"{path_save_metrics}/{name_img_visual}.png"
         plt.savefig(path_save_image, dpi=300, bbox_inches='tight')
+        plt.close()
 
 # Mã thử nghiệm
 # pred = [1, 1, 0]

core/trainer/trainer.py

@@ -2,6 +2,7 @@
 import os
 import pathlib
 import time
+import numpy as np
 import torch
 import logging
 from omegaconf import OmegaConf
@@ -130,6 +131,7 @@ def test_engine(config, test_dataloader):
         progress_loader = tqdm(test_dataloader, file=open(os.devnull, "w"))
         preds = []
         trusts = []
+        h_list = []
 
         # time start
         start_time = time.time()
@@ -146,7 +148,8 @@ def test_engine(config, test_dataloader):
                 ans_batch = ans.argmax(dim=1).long()
                 preds.extend(ans_batch.cpu().tolist())
                 trusts.extend(labels.cpu().tolist())
-                
+                h_list.extend(h.cpu().numpy())  # lưu embedding
+        embeddings = np.array(h_list, dtype=np.float32)
         end_time = time.time()
         elapsed_time = end_time - start_time
 
@@ -163,6 +166,8 @@ def test_engine(config, test_dataloader):
                                y_preds=preds,
                                y_trues=trusts,
                                path_save_metrics=config.training.path_save_metrics)
+        logger.info(f"Plot points of data")
+        ScoreOutput.data_vizualization(embeddings, trusts, config.training.path_save_metrics, "pca_benign_malware")
         logger.info(f"Completed Testing model with output:\n\
                         \tPath to model: {config.training.path_save_metrics}")