{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "63726b83-b39e-426c-8227-de05e6e2093a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11.0\n",
      "0.12.0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader,Dataset\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.utils\n",
    "import numpy as np\n",
    "import random\n",
    "from PIL import Image\n",
    "import os\n",
    "import PIL.ImageOps \n",
    "print(torch.__version__)  #1.1.0\n",
    "print(torchvision.__version__)  #0.3.0\n",
    "\n",
    "\n",
    "#定义一些超参\n",
    "train_batch_size = 20        #训练时batch_size\n",
    "train_number_epochs = 50     #训练的epoch\n",
    "\n",
    "def imshow(img,text=None,should_save=False): \n",
    "    #展示一幅tensor图像，输入是(C,H,W)\n",
    "    npimg = img.numpy() #将tensor转为ndarray\n",
    "#     print(npimg.shape)\n",
    "    plt.axis(\"off\")\n",
    "    if text:\n",
    "        plt.text(75, 8, text, style='italic',fontweight='bold',\n",
    "            bbox={'facecolor':'white', 'alpha':0.8, 'pad':10})\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0))) #转换为(H,W,C)\n",
    "    plt.show()    \n",
    "\n",
    "def show_plot(iteration,loss):\n",
    "    #绘制损失变化图\n",
    "    plt.plot(iteration,loss)\n",
    "    plt.title(\"Loss Function Line\")\n",
    "    plt.xlabel('iteration')\n",
    "    plt.ylabel('loss')\n",
    "    plt.show()\n",
    "def file_name(file_dir):\n",
    "    for root,dirs,files in os.walk(file_dir):\n",
    "        return files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b7aa0620-35e5-460b-955e-9c9fe216343b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "122 108\n"
     ]
    }
   ],
   "source": [
    "#重写Dataset类，__getitem__(self,index)每次返回(img1, img2, 01标签)\n",
    "class SiameseNetworkDataset(Dataset):\n",
    "    \n",
    "    def __init__(self,imageFolderDataset1,imageFolderDataset2,transform=None,should_invert=True):\n",
    "        self.imageFolderDataset1 = imageFolderDataset1\n",
    "        self.imageFolderDataset2 = imageFolderDataset2\n",
    "        self.transform = transform\n",
    "        self.should_invert = should_invert\n",
    "        \n",
    "    def __getitem__(self,index):\n",
    "        img0_tuple = random.choice(self.imageFolderDataset1.imgs) #10个套装中任选一\n",
    "        should_get_same_class = random.randint(0,1) #保证同类样本约占一半\n",
    "        if should_get_same_class:\n",
    "            while True:\n",
    "                #直到找到同一类别\n",
    "                img1_tuple = random.choice(self.imageFolderDataset2.imgs)\n",
    "                if img0_tuple[1]==img1_tuple[1]:\n",
    "                    break\n",
    "\n",
    "        else:\n",
    "            while True:\n",
    "                #直到找到非同一类别\n",
    "                img1_tuple = random.choice(self.imageFolderDataset2.imgs)\n",
    "                if img0_tuple[1] != img1_tuple[1]:\n",
    "                    break\n",
    "\n",
    "\n",
    "        img0 = Image.open(img0_tuple[0])\n",
    "        img1 = Image.open(img1_tuple[0])\n",
    "        img0 = img0.convert(\"RGB\")\n",
    "        img1 = img1.convert(\"RGB\")\n",
    "        \n",
    "        if self.should_invert:\n",
    "            img0 = PIL.ImageOps.invert(img0)\n",
    "            img1 = PIL.ImageOps.invert(img1)\n",
    "\n",
    "        if self.transform is not None:\n",
    "            img0 = self.transform(img0)\n",
    "            img1 = self.transform(img1)\n",
    "        \n",
    "        return img0, img1, torch.from_numpy(np.array([int(img1_tuple[1]!=img0_tuple[1])],dtype=np.float32))\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.imageFolderDataset1.imgs)+len(self.imageFolderDataset1.imgs)\n",
    "    \n",
    "    \n",
    "    \n",
    "#定义文件dataset\n",
    "training_dir1 = \"./data2/train/1/\"  #上衣训练集地址\n",
    "folder_dataset1 = torchvision.datasets.ImageFolder(root=training_dir1)\n",
    "\n",
    "training_dir2 = \"./data2/train/2/\"  #上衣训练集地址\n",
    "folder_dataset2 = torchvision.datasets.ImageFolder(root=training_dir2)\n",
    "\n",
    "#定义图像dataset\n",
    "transform = transforms.Compose([transforms.Resize((100,100)),\n",
    "                                transforms.ToTensor()])\n",
    "siamese_dataset = SiameseNetworkDataset(imageFolderDataset1=folder_dataset1,\n",
    "                                        imageFolderDataset2=folder_dataset2,\n",
    "                                        transform=transform,\n",
    "                                        should_invert=False)\n",
    "\n",
    "#定义图像dataloader\n",
    "train_dataloader = DataLoader(siamese_dataset,\n",
    "                            shuffle=True,\n",
    "                            batch_size=train_batch_size)\n",
    "print(len(folder_dataset1.imgs),len(folder_dataset2.imgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8efd097d-a66b-4e43-9f8e-07a043c8a9a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 206, 818)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 50932 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1.]\n",
      " [1.]\n",
      " [1.]\n",
      " [0.]\n",
      " [0.]\n",
      " [1.]\n",
      " [0.]\n",
      " [1.]]\n"
     ]
    }
   ],
   "source": [
    "vis_dataloader = DataLoader(siamese_dataset,\n",
    "                        shuffle=True,\n",
    "                        batch_size=8)\n",
    "example_batch = next(iter(vis_dataloader)) #生成一批图像\n",
    "#其中example_batch[0] 维度为torch.Size([8, 1, 100, 100])\n",
    "concatenated = torch.cat((example_batch[0],example_batch[1]),0) \n",
    "imshow(torchvision.utils.make_grid(concatenated, nrow=8))\n",
    "print(example_batch[2].numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "af8c0e61-7c8b-4299-a43c-789c8092bcdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#搭建模型\n",
    "class SiameseNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.cnn1 = nn.Sequential(\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(3, 4, kernel_size=3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(4),\n",
    "            \n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(4, 8, kernel_size=3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(8),\n",
    "\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(8, 8, kernel_size=3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.BatchNorm2d(8),\n",
    "        )\n",
    "\n",
    "        self.fc1 = nn.Sequential(\n",
    "            nn.Linear(8*100*100, 500),\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.Linear(500, 500),\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.Linear(500, 16))\n",
    "\n",
    "    def forward_once(self, x):\n",
    "        output = self.cnn1(x)\n",
    "        output = output.view(output.size()[0], -1)\n",
    "        output = self.fc1(output)\n",
    "        return output\n",
    "\n",
    "    def forward(self, input1, input2):\n",
    "        output1 = self.forward_once(input1)\n",
    "        output2 = self.forward_once(input2)\n",
    "        return output1, output2\n",
    "    \n",
    "    \n",
    "#自定义ContrastiveLoss\n",
    "class ContrastiveLoss(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Contrastive loss function.\n",
    "    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, margin=2.0):\n",
    "        super(ContrastiveLoss, self).__init__()\n",
    "        self.margin = margin\n",
    "\n",
    "    def forward(self, output1, output2, label):\n",
    "        euclidean_distance = F.pairwise_distance(output1, output2, keepdim = True)\n",
    "        loss_contrastive = torch.mean((1-label) * torch.pow(euclidean_distance, 2) +\n",
    "                                      (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))\n",
    "\n",
    "        return loss_contrastive\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "1e552d57-5702-4b01-9132-68e529e92e6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19452 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#net = SiameseNetwork().cuda() #定义模型且移至GPU\n",
    "net = SiameseNetwork()\n",
    "criterion = ContrastiveLoss() #定义损失函数\n",
    "optimizer = optim.Adam(net.parameters(), lr = 0.0005) #定义优化器\n",
    "\n",
    "counter = []\n",
    "loss_history = [] \n",
    "iteration_number = 0\n",
    "\n",
    "\n",
    "#开始训练\n",
    "for epoch in range(0, train_number_epochs):\n",
    "    for i, data in enumerate(train_dataloader, 0):\n",
    "        img0, img1 , label = data\n",
    "        #img0维度为torch.Size([32, 3, 100, 100])，32是batch，label为torch.Size([32, 3])\n",
    "        #img0, img1 , label = img0.cuda(), img1.cuda(), label.cuda() #数据移至GPU\n",
    "        optimizer.zero_grad()\n",
    "        output1,output2 = net(img0, img1)\n",
    "        loss_contrastive = criterion(output1, output2, label)\n",
    "        loss_contrastive.backward()\n",
    "        optimizer.step()\n",
    "        if i % 10 == 0 :\n",
    "            iteration_number +=10\n",
    "            counter.append(iteration_number)\n",
    "            loss_history.append(loss_contrastive.item())\n",
    "    \n",
    "show_plot(counter, loss_history)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "2998d236-f018-4449-909f-96d058094d99",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55496 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 67804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 61740 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 58684 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 51528 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60980 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64312 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60048 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 54464 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "#定义测试的dataset和dataloader\n",
    "\n",
    "#定义文件dataset\n",
    "testing_dir1 = \"./data2/test/1/\"  #上衣测试集地址\n",
    "folder_dataset_test1 = torchvision.datasets.ImageFolder(root=testing_dir1)\n",
    "\n",
    "testing_dir2 = \"./data2/test/2/\"  #上衣测试集地址\n",
    "folder_dataset_test2 = torchvision.datasets.ImageFolder(root=testing_dir2)\n",
    "\n",
    "#定义图像dataset\n",
    "transform_test = transforms.Compose([transforms.Resize((100,100)), \n",
    "                                     transforms.ToTensor()])\n",
    "siamese_dataset_test = SiameseNetworkDataset(imageFolderDataset1=folder_dataset_test1,\n",
    "                                             imageFolderDataset2=folder_dataset_test2,\n",
    "                                        transform=transform_test,\n",
    "                                        should_invert=False)\n",
    "\n",
    "#定义图像dataloader\n",
    "test_dataloader = DataLoader(siamese_dataset_test,\n",
    "                            shuffle=True,\n",
    "                            batch_size=1)\n",
    "\n",
    "\n",
    "#生成对比图像\n",
    "dataiter = iter(test_dataloader)\n",
    "x0,_,_ = next(dataiter)\n",
    "\n",
    "for i in range(10):\n",
    "    _,x1,label2 = next(dataiter)\n",
    "    concatenated = torch.cat((x0,x1),0)\n",
    "#     output1,output2 = net(x0.cuda(),x1.cuda())\n",
    "    output1,output2 = net(x0,x1)\n",
    "    euclidean_distance = F.pairwise_distance(output1, output2)\n",
    "    imshow(torchvision.utils.make_grid(concatenated),'Dissimilarity: {:.2f}'.format(euclidean_distance.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4031cb0a-6fec-4050-8a68-7479bd8e1a74",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18a1a725-bff2-460a-a141-f40612ade747",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "886ba55b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
