{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":4117,"databundleVersionId":46665}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport subprocess\nimport re\nimport time\nimport torch\nimport gc\nimport shutil\nfrom PIL import Image\nfrom collections import Counter\nfrom torchvision import transforms, models\nfrom torch.utils.data import DataLoader, Dataset\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.utils.class_weight import compute_sample_weight\nfrom sklearn.metrics import log_loss, classification_report, confusion_matrix\nimport seaborn as sns\nimport matplotlib.pyplot as plt\nimport xgboost as xgb\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\n# --- 1. CẤU HÌNH HỆ THỐNG ---\nTRAIN_7Z = '/kaggle/input/competitions/malware-classification/train.7z'\nTEST_7Z = '/kaggle/input/competitions/malware-classification/test.7z'\nTRAIN_LABELS = '/kaggle/input/competitions/malware-classification/trainLabels.csv'\nTEMP_DIR = '/kaggle/working/temp_data'\nIMG_DIR = '/kaggle/working/train_imgs_png'  # Thư mục lưu ảnh tạm trên ổ đĩa\nSUBMISSION_FILE = 'submission_final.csv'\n\nos.makedirs(TEMP_DIR, exist_ok=True)\nos.makedirs(IMG_DIR, exist_ok=True)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nopcodes = ['jmp', 'mov', 'retf', 'push', 'pop', 'add', 'sub', 'mul', 'div', 'and', 'or', 'xor', 'cmp', 'test', 'call']\nkeywords = ['db', 'dw', 'dd', 'offset', 'api', '.dll', 'std::']\nn_asm_features = len(opcodes) + len(keywords)\n\n# --- 2. HÀM TRÍCH XUẤT ĐẶC TRƯNG AN TOÀN ---\ndef process_and_save_features(fid):\n    b_path = os.path.join(TEMP_DIR, f\"{fid}.bytes\")\n    a_path = os.path.join(TEMP_DIR, f\"{fid}.asm\")\n    img_save_path = os.path.join(IMG_DIR, f\"{fid}.png\")\n    asm_v = [0 for _ in range(n_asm_features)]\n    \n    # 1. Trích xuất file bytes -> Lưu thành file ảnh PNG vật lý để tiết kiệm RAM\n    if os.path.exists(b_path):\n        try:\n            with open(b_path, 'r') as f:\n                content = f.read()\n            hex_values = [int(x, 16) for x in content.split() if len(x) == 2 and x != '??']\n            if len(hex_values) > 0:\n                d = np.array(hex_values, dtype=np.uint8)\n                width = 256\n                height = len(d) // width\n                if height > 0:\n                    img = d[:height * width].reshape((height, width))\n                    Image.fromarray(img).save(img_save_path)\n        except Exception as e:\n            print(f\"Error processing bytes {fid}: {e}\")\n\n    # Fallback tạo ảnh trống nếu file bytes lỗi để không làm sập DataLoader\n    if not os.path.exists(img_save_path):\n        Image.fromarray(np.zeros((224, 224), dtype=np.uint8)).save(img_save_path)\n\n    # 2. Trích xuất file ASM lấy đặc trưng cho XGBoost\n    if os.path.exists(a_path):\n        try:\n            if os.path.getsize(a_path) < 150 * 1024 * 1024:\n                opcode_counts = Counter()\n                keyword_counts = Counter()\n                with open(a_path, 'r', encoding='latin-1') as f:\n                    for line in f:\n                        parts = line.strip().split()\n                        if len(parts) > 1:\n                            token = parts[1].lower()\n                            if token in opcodes: opcode_counts[token] += 1\n                            if token in keywords: keyword_counts[token] += 1\n                total = sum(opcode_counts.values()) + 1\n                asm_v = ([opcode_counts.get(o, 0) / total for o in opcodes] +\n                         [keyword_counts.get(k, 0) / total for k in keywords])\n        except Exception as e:\n            print(f\"Error processing asm {fid}: {e}\")\n            \n    return asm_v\n\n# --- 3. BẮT ĐẦU TRÍCH XUẤT TẬP TRAIN ---\nprint(\">>> BẮT ĐẦU TRÍCH XUẤT DỮ LIỆU\")\nlabels_df = pd.read_csv(TRAIN_LABELS)\ntrain_ids = labels_df['Id'].values\nX_asm_all, Y_all, valid_train_ids = [], [], []\nT_BATCH = 100\nlabel_map = dict(zip(labels_df['Id'], labels_df['Class']))\n\nfor i in range(0, len(train_ids), T_BATCH):\n    batch = train_ids[i : i+T_BATCH]\n    with open('train_list.txt', 'w') as f:\n        for fid in batch:\n            f.write(f\"train/{fid}.bytes\\ntrain/{fid}.asm\\n\")\n            \n    subprocess.run(f\"7z e {TRAIN_7Z} -o{TEMP_DIR} -i@train_list.txt -y > /dev/null\", shell=True)\n\n    for fid in batch:\n        if fid not in label_map: continue\n        asm = process_and_save_features(fid)\n        \n        X_asm_all.append(asm)\n        Y_all.append(label_map[fid] - 1)\n        valid_train_ids.append(fid)\n    \n    shutil.rmtree(TEMP_DIR)\n    os.makedirs(TEMP_DIR, exist_ok=True)\n    \n    if (i+T_BATCH) % 1000 == 0:\n        print(f\"Tiến độ: {min(i+T_BATCH, len(train_ids))} file\")\n\nX_asm_all = np.array(X_asm_all, dtype=np.float32)\nY_all = np.array(Y_all).astype(np.int64)\nvalid_train_ids = np.array(valid_train_ids)\n\n# --- 4. CHIA DỮ LIỆU 80% TRAIN - 20% VALIDATION ---\nprint(\">>> ĐANG CHIA DỮ LIỆU 80/20\")\nidx_train, idx_val = train_test_split(\n    np.arange(len(Y_all)), test_size=0.2, random_state=42, stratify=Y_all\n)\n\nY_train, Y_val = Y_all[idx_train], Y_all[idx_val]\nX_asm_train, X_asm_val = X_asm_all[idx_train], X_asm_all[idx_val]\nids_train, ids_val = valid_train_ids[idx_train], valid_train_ids[idx_val]\n\n# Giải phóng bớt mảng gốc trên RAM\ndel X_asm_all, Y_all\ngc.collect()\n\n# --- 5. HUẤN LUYỆN MODEL XGBOOST ---\nprint(\">>> ĐANG LUYỆN XGBOOST CÓ CÂN BẰNG TRỌNG SỐ (PHƯƠNG ÁN 1)\")\nxgb_sample_weights = compute_sample_weight(class_weight='balanced', y=Y_train)\n\nxgb_model = xgb.XGBClassifier(\n    device='cuda' if torch.cuda.is_available() else 'cpu', \n    n_estimators=800, max_depth=6, learning_rate=0.05,\n    subsample=0.8, colsample_bytree=0.8, objective='multi:softprob',\n    num_class=9, tree_method='hist'\n)\n# Đảm bảo đầu vào là NumPy Array chuẩn trên CPU để XGBoost tự chuyển đổi thiết bị hợp lệ\nxgb_model.fit(X_asm_train, Y_train, sample_weight=xgb_sample_weights)\n\n# --- 6. HUẤN LUYỆN MODEL RESNET THÍCH ỨNG ---\nprint(\">>> ĐANG LUYỆN RESNET VỚI WEIGHTED CROSS ENTROPY LOSS (PHƯƠNG ÁN 2)\")\nclass LazyImageDataset(Dataset):\n    def __init__(self, fids, labels, img_dir=IMG_DIR):\n        self.fids = fids\n        self.labels = labels\n        self.img_dir = img_dir\n        self.transform = transforms.Compose([\n            transforms.Resize((224, 224)),\n            transforms.Grayscale(num_output_channels=3),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n        ])\n\n    def __len__(self):\n        return len(self.labels)\n\n    def __getitem__(self, idx):\n        img_path = os.path.join(self.img_dir, f\"{self.fids[idx]}.png\")\n        img = Image.open(img_path).convert('RGB')\n        return self.transform(img), int(self.labels[idx])\n\nmodel_resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)\nmodel_resnet.fc = nn.Linear(512, 9)\nmodel_resnet.to(device)\n\ntrain_loader = DataLoader(LazyImageDataset(ids_train, Y_train), batch_size=32, shuffle=True, num_workers=2, pin_memory=True)\nval_loader = DataLoader(LazyImageDataset(ids_val, Y_val), batch_size=32, shuffle=False, num_workers=2, pin_memory=True)\n\noptimizer = optim.Adam(model_resnet.parameters(), lr=1e-4)\nclass_counts = np.bincount(Y_train)\nclass_weights = 1.0 / np.sqrt(class_counts)\nclass_weights = class_weights / np.sum(class_weights) * 9\ncriterion = nn.CrossEntropyLoss(weight=torch.FloatTensor(class_weights).to(device))\n\nfor epoch in range(10):\n    model_resnet.train()\n    train_loss = 0\n    for imgs, labels in train_loader:\n        imgs, labels = imgs.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n        optimizer.zero_grad()\n        loss = criterion(model_resnet(imgs), labels)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n    print(f\"Epoch {epoch+1} | Train Loss: {train_loss/len(train_loader):.4f}\")\n\n# --- 7. ĐÁNH GIÁ & TÌM TRỌNG SỐ ENSEMBLE TỐI ƯU ---\nmodel_resnet.eval()\nres_val = []\nwith torch.no_grad():\n    for imgs, _ in val_loader:\n        res_val.append(F.softmax(model_resnet(imgs.to(device)), dim=1).cpu().numpy())\nres_val = np.vstack(res_val)\n\n# Dự đoán XGBoost sử dụng mảng NumPy trên CPU\nxgb_val = xgb_model.predict_proba(X_asm_val)\n\n# Chuẩn hóa lại xác suất\nres_val = res_val / res_val.sum(axis=1, keepdims=True)\nxgb_val = xgb_val / xgb_val.sum(axis=1, keepdims=True)\n\n# Tìm trọng số tối ưu\nbest_loss, best_w = float(\"inf\"), 0.5\nfor w in np.linspace(0, 1, 21):\n    combined = w * res_val + (1 - w) * xgb_val\n    loss = log_loss(Y_val, combined)\n    if loss < best_loss:\n        best_loss, best_w = loss, w\n\nprint(f\"Best weight (ResNet ratio): {best_w:.2f} | Best Log Loss: {best_loss:.4f}\")\n\n# In báo cáo kết quả đánh giá chi tiết\npred = np.argmax(best_w * res_val + (1 - best_w) * xgb_val, axis=1)\nprint(classification_report(Y_val, pred))\n\n# Dọn dẹp ổ đĩa tập train trước khi xử lý tập test\nshutil.rmtree(IMG_DIR)\nos.makedirs(IMG_DIR, exist_ok=True)\ngc.collect()\ntorch.cuda.empty_cache()\n\n# --- 8. DỰ ĐOÁN TẬP TEST THỰC TẾ ---\nprint(\">>> ĐANG DỰ ĐOÁN TẬP TEST\")\noutput = subprocess.check_output([\"7z\", \"l\", TEST_7Z]).decode()\nraw_test = [line.split()[-1] for line in output.splitlines() if \".bytes\" in line]\ntest_ids = [os.path.basename(f).replace('.bytes', '') for f in raw_test]\n\nif os.path.exists(SUBMISSION_FILE): os.remove(SUBMISSION_FILE)\n\ntransform_test = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.Grayscale(num_output_channels=3),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n])\n\nfor i in range(0, len(test_ids), 100):\n    batch = test_ids[i : i+100]\n    with open('test_batch.txt', 'w') as f:\n        for fid in batch: f.write(f\"test/{fid}.bytes\\ntest/{fid}.asm\\n\")\n    \n    subprocess.run(f\"7z e {TEST_7Z} -o{TEMP_DIR} -i@test_batch.txt -y > /dev/null\", shell=True)\n    \n    imgs_list = []\n    asm_features = []\n    valid_ids = []\n    \n    for fid in batch:\n        asm_f = process_and_save_features(fid)\n        img_path = os.path.join(IMG_DIR, f\"{fid}.png\")\n        \n        if os.path.exists(img_path):\n            img = Image.open(img_path).convert('RGB')\n            imgs_list.append(transform_test(img).unsqueeze(0))\n            asm_features.append(asm_f)\n            valid_ids.append(fid)\n            \n    if len(imgs_list) == 0: continue\n    \n    # Chạy mô hình ResNet\n    imgs_tensor = torch.cat(imgs_list).to(device)\n    with torch.no_grad():\n        res_batch = F.softmax(model_resnet(imgs_tensor), dim=1).cpu().numpy()\n    \n    # Chạy mô hình XGBoost (Đưa vào NumPy array float32 chuẩn)\n    xgb_batch = xgb_model.predict_proba(np.array(asm_features, dtype=np.float32))\n    \n    batch_probs = []\n    for idx, fid in enumerate(valid_ids):\n        p_resnet = res_batch[idx]\n        p_xgb = xgb_batch[idx]\n        \n        p_resnet /= (p_resnet.sum() + 1e-7)\n        p_xgb /= (p_xgb.sum() + 1e-7)\n        \n        combined = best_w * p_resnet + (1 - best_w) * p_xgb\n        batch_probs.append([fid] + combined.tolist())\n    \n    cols = ['Id'] + [f'Prediction{j}' for j in range(1, 10)]\n    pd.DataFrame(batch_probs, columns=cols).to_csv(SUBMISSION_FILE, mode='a', header=not os.path.exists(SUBMISSION_FILE), index=False)\n    \n    # Dọn dẹp bộ nhớ đệm sau mỗi Batch để tránh đầy ổ cứng Kaggle\n    shutil.rmtree(TEMP_DIR)\n    shutil.rmtree(IMG_DIR)\n    os.makedirs(TEMP_DIR, exist_ok=True)\n    os.makedirs(IMG_DIR, exist_ok=True)\n    \n    if (i+100) % 500 == 0: \n        print(f\"Tiến độ Test: {min(i+100, len(test_ids))}/{len(test_ids)}\")\n        gc.collect()\n        torch.cuda.empty_cache()\n\nprint(f\"File nộp bài hoàn chỉnh thành công: {SUBMISSION_FILE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-07-19T05:53:05.804626Z","iopub.execute_input":"2026-07-19T05:53:05.805044Z","iopub.status.idle":"2026-07-19T05:55:12.457373Z","shell.execute_reply.started":"2026-07-19T05:53:05.804997Z","shell.execute_reply":"2026-07-19T05:55:12.456152Z"}},"outputs":[],"execution_count":null}]}