{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44173db3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model 1 test trên train Kaggle2:\n",
      "Accuracy: 0.9222, Precision: 0.9202, Recall: 0.9245, F1: 0.9224\n",
      "Model 2 test trên train Kaggle1:\n",
      "Accuracy: 0.9164, Precision: 0.9220, Recall: 0.9098, F1: 0.9159\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import datasets, transforms, models\n",
    "from torch.utils.data import DataLoader\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "import numpy as np\n",
    "\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "BATCH_SIZE = 32\n",
    "NUM_CLASSES = 2\n",
    "\n",
    "def get_loader(data_root):\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "    ])\n",
    "    dataset = datasets.ImageFolder(data_root, transform=transform)\n",
    "    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)\n",
    "    return loader\n",
    "\n",
    "def load_model(weight_path):\n",
    "    model = models.efficientnet_b0(weights=None)\n",
    "    model.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES)\n",
    "    model.load_state_dict(torch.load(weight_path, map_location=DEVICE))\n",
    "    model = model.to(DEVICE)\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "def evaluate(model, loader):\n",
    "    all_labels = []\n",
    "    all_preds = []\n",
    "    with torch.no_grad():\n",
    "        for imgs, labels in loader:\n",
    "            imgs = imgs.to(DEVICE)\n",
    "            outputs = model(imgs)\n",
    "            preds = torch.argmax(outputs, dim=1).cpu().numpy()\n",
    "            all_preds.extend(preds)\n",
    "            all_labels.extend(labels.numpy())\n",
    "    acc = accuracy_score(all_labels, all_preds)\n",
    "    prec = precision_score(all_labels, all_preds, zero_division=0)\n",
    "    rec = recall_score(all_labels, all_preds, zero_division=0)\n",
    "    f1 = f1_score(all_labels, all_preds, zero_division=0)\n",
    "    return acc, prec, rec, f1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf9e19b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model 1 test trên train Kaggle2:\n",
      "Accuracy: 0.9300, Precision: 0.9555, Recall: 0.9020, F1: 0.9280\n",
      "Model 2 test trên train Kaggle1:\n",
      "Accuracy: 0.9200, Precision: 0.9430, Recall: 0.8940, F1: 0.9179\n"
     ]
    }
   ],
   "source": [
    "# Đường dẫn model và data\n",
    "model1_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth'\n",
    "model2_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth'\n",
    "data1_train = '/home/ubuntu/vnet/TaoST/Data10kKaggle2/test'\n",
    "data2_train = '/home/ubuntu/vnet/TaoST/Data10kKaggle1/test'\n",
    "# Test chéo\n",
    "model1 = load_model(model1_path)\n",
    "model2 = load_model(model2_path)\n",
    "\n",
    "loader1 = get_loader(data1_train)\n",
    "loader2 = get_loader(data2_train)\n",
    "\n",
    "print(\"Model 1 test trên train Kaggle2:\")\n",
    "acc, prec, rec, f1 = evaluate(model1, loader2)\n",
    "print(f\"Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")\n",
    "\n",
    "print(\"Model 2 test trên train Kaggle1:\")\n",
    "acc, prec, rec, f1 = evaluate(model2, loader1)\n",
    "print(f\"Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3737fe93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model 1 test trên train Kaggle2:\n",
      "Accuracy: 0.9300, Precision: 0.9555, Recall: 0.9020, F1: 0.9280\n",
      "Model 2 test trên train Kaggle1:\n",
      "Accuracy: 0.9200, Precision: 0.9430, Recall: 0.8940, F1: 0.9179\n"
     ]
    }
   ],
   "source": [
    "# Đường dẫn model và data\n",
    "model1_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth'\n",
    "model2_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth'\n",
    "data1_train = '/home/ubuntu/vnet/TaoST/Data10kKaggle2/test'\n",
    "data2_train = '/home/ubuntu/vnet/TaoST/Data10kKaggle1/test'\n",
    "# Test chéo\n",
    "model1 = load_model(model1_path)\n",
    "model2 = load_model(model2_path)\n",
    "\n",
    "loader1 = get_loader(data1_train)\n",
    "loader2 = get_loader(data2_train)\n",
    "\n",
    "print(\"Model 1 test trên train Kaggle2:\")\n",
    "acc, prec, rec, f1 = evaluate(model1, loader2)\n",
    "print(f\"Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")\n",
    "\n",
    "print(\"Model 2 test trên train Kaggle1:\")\n",
    "acc, prec, rec, f1 = evaluate(model2, loader1)\n",
    "print(f\"Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "45ff1792",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ensemble (Averaging) trên test: Accuracy: 0.9240, Precision: 0.9511, Recall: 0.8940, F1: 0.9216\n"
     ]
    }
   ],
   "source": [
    "def ensemble_average_predict(model1, model2, loader):\n",
    "    all_labels = []\n",
    "    all_preds = []\n",
    "    model1.eval()\n",
    "    model2.eval()\n",
    "    with torch.no_grad():\n",
    "        for imgs, labels in loader:\n",
    "            imgs = imgs.to(DEVICE)\n",
    "            out1 = torch.softmax(model1(imgs), dim=1)\n",
    "            out2 = torch.softmax(model2(imgs), dim=1)\n",
    "            avg_out = (out1 + out2) / 2\n",
    "            preds = torch.argmax(avg_out, dim=1).cpu().numpy()\n",
    "            all_preds.extend(preds)\n",
    "            all_labels.extend(labels.numpy())\n",
    "    acc = accuracy_score(all_labels, all_preds)\n",
    "    prec = precision_score(all_labels, all_preds, zero_division=0)\n",
    "    rec = recall_score(all_labels, all_preds, zero_division=0)\n",
    "    f1 = f1_score(all_labels, all_preds, zero_division=0)\n",
    "    return acc, prec, rec, f1\n",
    "\n",
    "# Đường dẫn model và data test\n",
    "model1_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth'\n",
    "model2_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth'\n",
    "data_test = '/home/ubuntu/vnet/TaoST/Data10kKaggle1/test'  # hoặc Data10kKaggle2/test\n",
    "\n",
    "model1 = load_model(model1_path)\n",
    "model2 = load_model(model2_path)\n",
    "loader = get_loader(data_test)\n",
    "\n",
    "acc, prec, rec, f1 = ensemble_average_predict(model1, model2, loader)\n",
    "print(f\"Ensemble (Averaging) trên test: Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a82b49da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5, Loss: 0.0806\n",
      "Epoch 2/5, Loss: 0.0598\n",
      "Epoch 3/5, Loss: 0.0542\n",
      "Epoch 4/5, Loss: 0.0444\n",
      "Epoch 5/5, Loss: 0.0394\n",
      "GAT Ensemble trên test: Accuracy: 0.9210, Precision: 0.9647, Recall: 0.8740, F1: 0.9171\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models, transforms, datasets\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "NUM_CLASSES = 2\n",
    "BATCH_SIZE = 32\n",
    "EPOCHS = 5\n",
    "\n",
    "class SimpleGATLayer(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(in_features, out_features)\n",
    "        self.attn = nn.Parameter(torch.Tensor(1, out_features))\n",
    "        nn.init.xavier_uniform_(self.attn.data, gain=1.414)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.fc(x)\n",
    "        attn_score = torch.matmul(h, self.attn.t())\n",
    "        attn_score = F.softmax(attn_score, dim=1)\n",
    "        h_prime = (attn_score * h).sum(dim=1)\n",
    "        return h_prime\n",
    "\n",
    "class GATEnsembleClassifier(nn.Module):\n",
    "    def __init__(self, feature_dim, num_classes):\n",
    "        super().__init__()\n",
    "        self.gat = SimpleGATLayer(feature_dim, feature_dim)\n",
    "        self.classifier = nn.Linear(feature_dim, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.gat(x)\n",
    "        out = self.classifier(h)\n",
    "        return out\n",
    "\n",
    "def get_feature_extractor(weight_path):\n",
    "    model = models.efficientnet_b0(weights=None)\n",
    "    model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2)\n",
    "    model.load_state_dict(torch.load(weight_path, map_location=DEVICE))\n",
    "    backbone = nn.Sequential(*(list(model.children())[:-1]))\n",
    "    backbone.eval()\n",
    "    return backbone.to(DEVICE)\n",
    "\n",
    "def get_loader(data_root, split):\n",
    "    transform = transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "    ])\n",
    "    dataset = datasets.ImageFolder(f\"{data_root}/{split}\", transform=transform)\n",
    "    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)\n",
    "    return loader\n",
    "\n",
    "# Đường dẫn model và data\n",
    "model1_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth'\n",
    "model2_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth'\n",
    "data_root = '/home/ubuntu/vnet/TaoST/Data10kKaggle'\n",
    "\n",
    "# Load feature extractors (freeze)\n",
    "fe1 = get_feature_extractor(model1_path)\n",
    "fe2 = get_feature_extractor(model2_path)\n",
    "for p in fe1.parameters():\n",
    "    p.requires_grad = False\n",
    "for p in fe2.parameters():\n",
    "    p.requires_grad = False\n",
    "\n",
    "feature_dim = 1280\n",
    "gat_ensemble = GATEnsembleClassifier(feature_dim, NUM_CLASSES).to(DEVICE)\n",
    "\n",
    "optimizer = torch.optim.Adam(gat_ensemble.parameters(), lr=1e-3)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "train_loader = get_loader(data_root, 'train')\n",
    "val_loader = get_loader(data_root, 'test')\n",
    "\n",
    "# Huấn luyện GAT ensemble\n",
    "for epoch in range(EPOCHS):\n",
    "    gat_ensemble.train()\n",
    "    running_loss = 0.0\n",
    "    for imgs, labels in train_loader:\n",
    "        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n",
    "        with torch.no_grad():\n",
    "            f1 = fe1(imgs).squeeze(-1).squeeze(-1)\n",
    "            f2 = fe2(imgs).squeeze(-1).squeeze(-1)\n",
    "            features = torch.stack([f1, f2], dim=1)\n",
    "        outputs = gat_ensemble(features)\n",
    "        loss = criterion(outputs, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item() * imgs.size(0)\n",
    "    print(f\"Epoch {epoch+1}/{EPOCHS}, Loss: {running_loss/len(train_loader.dataset):.4f}\")\n",
    "\n",
    "# Đánh giá trên tập test\n",
    "gat_ensemble.eval()\n",
    "all_labels = []\n",
    "all_preds = []\n",
    "with torch.no_grad():\n",
    "    for imgs, labels in val_loader:\n",
    "        imgs = imgs.to(DEVICE)\n",
    "        f1 = fe1(imgs).squeeze(-1).squeeze(-1)\n",
    "        f2 = fe2(imgs).squeeze(-1).squeeze(-1)\n",
    "        features = torch.stack([f1, f2], dim=1)\n",
    "        outputs = gat_ensemble(features)\n",
    "        preds = torch.argmax(outputs, dim=1).cpu().numpy()\n",
    "        all_preds.extend(preds)\n",
    "        all_labels.extend(labels.numpy())\n",
    "\n",
    "acc = accuracy_score(all_labels, all_preds)\n",
    "prec = precision_score(all_labels, all_preds, zero_division=0)\n",
    "rec = recall_score(all_labels, all_preds, zero_division=0)\n",
    "f1 = f1_score(all_labels, all_preds, zero_division=0)\n",
    "print(f\"GAT Ensemble trên test: Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")\n",
    "\n",
    "# Lưu lại model GAT nếu muốn\n",
    "torch.save(gat_ensemble.state_dict(), '/home/ubuntu/vnet/FL/gat_ensemble_kaggle.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "70f1fdff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Khởi tạo feature extractor và biến feature_dim\n",
    "model1_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle1.pth'\n",
    "model2_path = '/home/ubuntu/vnet/FL/efficientnet_b0_kaggle2.pth'\n",
    "def get_feature_extractor(weight_path):\n",
    "    model = models.efficientnet_b0(weights=None)\n",
    "    model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2)\n",
    "    model.load_state_dict(torch.load(weight_path, map_location=DEVICE))\n",
    "    backbone = nn.Sequential(*(list(model.children())[:-1]))\n",
    "    backbone.eval()\n",
    "    return backbone.to(DEVICE)\n",
    "fe1 = get_feature_extractor(model1_path)\n",
    "fe2 = get_feature_extractor(model2_path)\n",
    "feature_dim = 1280"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "24525bd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GAT Ensemble trên train: Accuracy: 0.9210, Precision: 0.9647, Recall: 0.8740, F1: 0.9171\n"
     ]
    }
   ],
   "source": [
    "# Đảm bảo fe1, fe2, feature_dim đã được load như sau:\n",
    "# fe1 = get_feature_extractor(model1_path)\n",
    "# fe2 = get_feature_extractor(model2_path)\n",
    "# feature_dim = 1280\n",
    "\n",
    "gat_ensemble = GATEnsembleClassifier(feature_dim, NUM_CLASSES).to(DEVICE)\n",
    "gat_ensemble.load_state_dict(torch.load('/home/ubuntu/vnet/FL/gat_ensemble_kaggle.pth', map_location=DEVICE))\n",
    "gat_ensemble.eval()\n",
    "\n",
    "train_loader = get_loader('/home/ubuntu/vnet/TaoST/Data10kKaggle', 'test')\n",
    "\n",
    "all_labels = []\n",
    "all_preds = []\n",
    "with torch.no_grad():\n",
    "    for imgs, labels in train_loader:\n",
    "        imgs = imgs.to(DEVICE)\n",
    "        f1 = fe1(imgs).squeeze(-1).squeeze(-1)\n",
    "        f2 = fe2(imgs).squeeze(-1).squeeze(-1)\n",
    "        features = torch.stack([f1, f2], dim=1)\n",
    "        outputs = gat_ensemble(features)\n",
    "        preds = torch.argmax(outputs, dim=1).cpu().numpy()\n",
    "        all_preds.extend(preds)\n",
    "        all_labels.extend(labels.numpy())\n",
    "\n",
    "acc = accuracy_score(all_labels, all_preds)\n",
    "prec = precision_score(all_labels, all_preds, zero_division=0)\n",
    "rec = recall_score(all_labels, all_preds, zero_division=0)\n",
    "f1 = f1_score(all_labels, all_preds, zero_division=0)\n",
    "print(f\"GAT Ensemble trên train: Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1: {f1:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a5385c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exported to /home/ubuntu/vnet/FL/gat_ensemble_kaggle.onnx\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Định nghĩa lại lớp GATEnsembleClassifier và SimpleGATLayer\n",
    "class SimpleGATLayer(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(in_features, out_features)\n",
    "        self.attn = nn.Parameter(torch.Tensor(1, out_features))\n",
    "        nn.init.xavier_uniform_(self.attn.data, gain=1.414)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.fc(x)\n",
    "        attn_score = torch.matmul(h, self.attn.t())\n",
    "        attn_score = torch.softmax(attn_score, dim=1)\n",
    "        h_prime = (attn_score * h).sum(dim=1)\n",
    "        return h_prime\n",
    "\n",
    "class GATEnsembleClassifier(nn.Module):\n",
    "    def __init__(self, feature_dim, num_classes):\n",
    "        super().__init__()\n",
    "        self.gat = SimpleGATLayer(feature_dim, feature_dim)\n",
    "        self.classifier = nn.Linear(feature_dim, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.gat(x)\n",
    "        out = self.classifier(h)\n",
    "        return out\n",
    "\n",
    "# Thông số\n",
    "feature_dim = 1280\n",
    "num_classes = 2\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# Khởi tạo và load trọng số\n",
    "model = GATEnsembleClassifier(feature_dim, num_classes).to(DEVICE)\n",
    "model.load_state_dict(torch.load('/home/ubuntu/vnet/FL/gat_ensemble_kaggle.pth', map_location=DEVICE))\n",
    "model.eval()\n",
    "\n",
    "# Dummy input: batch_size=1, num_models=2, feature_dim=1280\n",
    "dummy_input = torch.randn(1, 2, feature_dim).to(DEVICE)\n",
    "\n",
    "# Export sang ONNX\n",
    "torch.onnx.export(\n",
    "    model,\n",
    "    dummy_input,\n",
    "    \"/home/ubuntu/vnet/FL/gat_ensemble_kaggle.onnx\",\n",
    "    input_names=[\"features\"],\n",
    "    output_names=[\"output\"],\n",
    "    dynamic_axes={\"features\": {0: \"batch_size\"}},\n",
    "    opset_version=12\n",
    ")\n",
    "\n",
    "print(\"Exported to /home/ubuntu/vnet/FL/gat_ensemble_kaggle.onnx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2475f009",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'onnxruntime'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01monnxruntime\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mort\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorchvision\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m transforms\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mPIL\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Image\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'onnxruntime'"
     ]
    }
   ],
   "source": [
    "import onnxruntime as ort\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import torch\n",
    "import os\n",
    "\n",
    "# Đường dẫn ONNX model\n",
    "onnx_path = \"/home/ubuntu/vnet/FL/gat_ensemble_kaggle.onnx\"\n",
    "session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])\n",
    "\n",
    "# Chuẩn bị transform giống như khi train\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "# Lấy batch ảnh từ folder benign\n",
    "folder = \"/home/ubuntu/vnet/TaoST/Data10kKaggle1/test/benign\"\n",
    "image_files = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(('.jpg', '.png', '.jpeg'))]\n",
    "batch_size = min(4, len(image_files))  # ví dụ lấy 4 ảnh\n",
    "images = []\n",
    "for img_path in image_files[:batch_size]:\n",
    "    img = Image.open(img_path).convert('RGB')\n",
    "    img = transform(img)\n",
    "    images.append(img)\n",
    "batch_tensor = torch.stack(images)  # [batch, 3, 224, 224]\n",
    "\n",
    "# Giả lập feature extractor (thực tế bạn cần xuất features từ fe1, fe2)\n",
    "# Ở đây tạo dummy features để test ONNX\n",
    "feature_dim = 1280\n",
    "num_models = 2\n",
    "features = torch.randn(batch_size, num_models, feature_dim).numpy().astype(np.float32)\n",
    "\n",
    "# Chạy inference với ONNX\n",
    "outputs = session.run([\"output\"], {\"features\": features})\n",
    "print(\"ONNX output shape:\", outputs[0].shape)\n",
    "print(\"ONNX output:\", outputs[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}