repair.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "from shutil import copyfile\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import models\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.transforms import InterpolationMode\n",
    "from utils.loader import environments, WeatherTransform\n",
    "\n",
    "from utils.common import setup_seed \n",
    "from tqdm import tqdm\n",
    "\n",
    "class DomainClassifier(nn.Module):\n",
    "    def __init__(self, domin_num=10) -> None:\n",
    "        super(DomainClassifier, self).__init__()\n",
    "        \n",
    "        # adatped from modern backbone, change last fc to adopt domain\n",
    "        self.net = models.resnet18(pretrained=True)\n",
    "        ftr_num = self.net.fc.in_features\n",
    "        self.net.fc = nn.Linear(ftr_num, domin_num)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class FrequencyFilter(nn.Module):\n",
    "    def __init__(self, h=256, w=256, channel=3) -> None:\n",
    "        super(FrequencyFilter, self).__init__()\n",
    "        self.filter = nn.Parameter(torch.ones(channel, h, w))\n",
    "\n",
    "    def forward(self, img):\n",
    "        fft_img = torch.fft.fftn(img, dim=(-2,-1))\n",
    "        amplitude = torch.abs(fft_img)\n",
    "        phase = torch.angle(fft_img)\n",
    "        amplitude_ = self.filter * amplitude \n",
    "        ret = torch.real(torch.fft.ifftn(torch.polar(amplitude_, phase), dim=(-2,-1)))\n",
    "        return ret\n",
    "    \n",
    "class MyDataset(datasets.ImageFolder):\n",
    "    def __init__(self, root, transform = None, target_transform = None, style='normal', h=256, w=256):\n",
    "        super().__init__(root, transform = transform, target_transform = target_transform)\n",
    "        self.h = h \n",
    "        self.w = w \n",
    "        self.weather = style\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img, target = super().__getitem__(index)\n",
    "        t_normal = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        t_aug = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=self.weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        t = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=self.weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        img_normal = t_normal(img)\n",
    "        img_aug = t_aug(img)\n",
    "        return img_aug, img_normal \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    "    \n",
    "def train(style='fog', batchsize=8, lr=0.001, num_epoch=60):\n",
    "    learnable_filter = FrequencyFilter()\n",
    "    data_set = MyDataset(os.path.join(os.getcwd(), 'University-Release', 'train', 'drone'), style=style)\n",
    "    data_loader = torch.utils.data.DataLoader(data_set, batch_size=batchsize, shuffle=True, num_workers=0, pin_memory=False)\n",
    "    optimizer = optim.Adam(learnable_filter.parameters(), lr=lr)\n",
    "    criterion = nn.MSELoss()\n",
    "\n",
    "    learnable_filter.cuda()\n",
    "\n",
    "    for epoch in range(1, num_epoch+1):\n",
    "        epoch_loss = 0\n",
    "        for img_aug, img in tqdm(data_loader):\n",
    "            img_aug, img = img_aug.cuda(), img.cuda() \n",
    "            outputs = learnable_filter(img_aug)\n",
    "            loss = criterion(outputs, img)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        epoch_loss /= len(data_loader)\n",
    "        print(f'epoch:{epoch} Loss:{epoch_loss:.4f}')\n",
    "\n",
    "        save_filename = '{}_{:03d}.pth'.format(style, epoch)\n",
    "        save_path = os.path.join(os.getcwd(), 'model', 'FFM', save_filename)\n",
    "        torch.save(learnable_filter.cpu().state_dict(), save_path)\n",
    "        learnable_filter.cuda()  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dark\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [07:55<00:00,  9.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [08:22<00:00,  9.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "rain\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [1:27:20<00:00,  1.11s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [1:26:03<00:00,  1.09s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "snow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [45:50<00:00,  1.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [45:50<00:00,  1.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "fog_rain\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [45:07<00:00,  1.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [44:56<00:00,  1.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "fog_snow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [29:08<00:00,  2.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [29:07<00:00,  2.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "rain_snow\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [1:10:44<00:00,  1.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [1:12:34<00:00,  1.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "light\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [11:37<00:00,  6.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [07:51<00:00, 10.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n",
      "wind\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [11:33<00:00,  6.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [11:35<00:00,  6.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for style in environments:\n",
    "    if style=='fog' or style=='normal':\n",
    "        continue\n",
    "    print(style)\n",
    "    train(style=style, num_epoch=2)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "from shutil import copyfile\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import models\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.transforms import InterpolationMode\n",
    "from utils.loader import environments, WeatherTransform\n",
    "\n",
    "from utils.common import setup_seed \n",
    "from tqdm import tqdm\n",
    "\n",
    "class DomainClassifier(nn.Module):\n",
    "    def __init__(self, domin_num=10) -> None:\n",
    "        super(DomainClassifier, self).__init__()\n",
    "        \n",
    "        # adatped from modern backbone, change last fc to adopt domain\n",
    "        self.net = models.resnet18(pretrained=True)\n",
    "        ftr_num = self.net.fc.in_features\n",
    "        self.net.fc = nn.Linear(ftr_num, domin_num)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class FrequencyFilter(nn.Module):\n",
    "    def __init__(self, h=256, w=256, channel=3) -> None:\n",
    "        super(FrequencyFilter, self).__init__()\n",
    "        self.filter = nn.Parameter(torch.ones(channel, h, w))\n",
    "\n",
    "    def forward(self, img):\n",
    "        fft_img = torch.fft.fftn(img, dim=(-2,-1))\n",
    "        amplitude = torch.abs(fft_img)\n",
    "        phase = torch.angle(fft_img)\n",
    "        amplitude_ = self.filter * amplitude \n",
    "        ret = torch.real(torch.fft.ifftn(torch.polar(amplitude_, phase), dim=(-2,-1)))\n",
    "        return ret\n",
    "    \n",
    "class MyDataset_mix(datasets.ImageFolder):\n",
    "    def __init__(self, root, transform = None, target_transform = None, h=256, w=256):\n",
    "        super().__init__(root, transform = transform, target_transform = target_transform)\n",
    "        self.h = h \n",
    "        self.w = w \n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img, target = super().__getitem__(index)\n",
    "        xx = [i for i in environments]\n",
    "        weather = np.random.choice(xx)\n",
    "        t_normal = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        t_aug = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        t = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        img_normal = t_normal(img)\n",
    "        img_aug = t_aug(img)\n",
    "        return img_aug, img_normal \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    "    \n",
    "def train(batchsize=8, lr=0.001, num_epoch=60):\n",
    "    learnable_filter = FrequencyFilter()\n",
    "    data_set = MyDataset_mix(os.path.join(os.getcwd(), 'University-Release', 'train', 'drone'))\n",
    "    data_loader = torch.utils.data.DataLoader(data_set, batch_size=batchsize, shuffle=True, num_workers=0, pin_memory=False)\n",
    "    optimizer = optim.Adam(learnable_filter.parameters(), lr=lr)\n",
    "    criterion = nn.MSELoss()\n",
    "\n",
    "    learnable_filter.cuda()\n",
    "\n",
    "    for epoch in range(1, num_epoch+1):\n",
    "        epoch_loss = 0\n",
    "        for img_aug, img in tqdm(data_loader):\n",
    "            img_aug, img = img_aug.cuda(), img.cuda() \n",
    "            outputs = learnable_filter(img_aug)\n",
    "            loss = criterion(outputs, img)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        epoch_loss /= len(data_loader)\n",
    "        print(f'epoch:{epoch} Loss:{epoch_loss:.4f}')\n",
    "\n",
    "        save_filename = '{}_{:03d}.pth'.format('mix', epoch)\n",
    "        save_path = os.path.join(os.getcwd(), 'model', 'FFM', save_filename)\n",
    "        torch.save(learnable_filter.cpu().state_dict(), save_path)\n",
    "        learnable_filter.cuda()  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [41:49<00:00,  1.89it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:1 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4732/4732 [37:13<00:00,  2.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:2 Loss:0.0000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "train(num_epoch=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use GAN to repair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\lyq\\.conda\\envs\\pytorch\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import math\n",
    "from shutil import copyfile\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import models\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.transforms import InterpolationMode\n",
    "from utils.loader import environments, WeatherTransform\n",
    "\n",
    "from utils.common import setup_seed \n",
    "from tqdm import tqdm\n",
    "\n",
    "# 生成器\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Generator, self).__init__()\n",
    "        # 定义U-Net或其他适合图像修复的网络结构\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        x = self.decoder(x)\n",
    "        return x\n",
    "\n",
    "# 判别器\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=1)\n",
    "            #nn.Sigmoid()\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(15*15,1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "    \n",
    "class MyDataset(datasets.ImageFolder):\n",
    "    def __init__(self, root, transform = None, target_transform = None, style='normal', h=256, w=256):\n",
    "        super().__init__(root, transform = transform, target_transform = target_transform)\n",
    "        self.h = h \n",
    "        self.w = w \n",
    "        self.weather = style\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img, target = super().__getitem__(index)\n",
    "        t_normal = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        t_aug = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=self.weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        t = transforms.Compose(\n",
    "                [\n",
    "                    transforms.Resize((self.h, self.w), interpolation=InterpolationMode.BICUBIC),\n",
    "                    WeatherTransform(aug=self.weather),\n",
    "                    transforms.ToTensor(),\n",
    "                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "                ]\n",
    "            )\n",
    "        \n",
    "        img_normal = t_normal(img)\n",
    "        img_aug = t_aug(img)\n",
    "        return img_aug, img_normal \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n",
    " \n",
    "\n",
    "    \n",
    "def train(model_dir = None, style='fog', batchsize=8, lr_g=0.0002,lr_d=0.0002, num_epoch=60):\n",
    "    if model_dir == None:\n",
    "        model_dir = os.path.join(os.getcwd(), 'model', 'GAN')\n",
    "    if not os.path.exists(model_dir):\n",
    "        os.makedirs(model_dir)\n",
    "\n",
    "    device = torch.device('cuda')\n",
    "    generator = Generator().to(device)\n",
    "    discriminator = Discriminator().to(device)\n",
    "\n",
    "    data_set = MyDataset(os.path.join(os.getcwd(), 'University-Release', 'train', 'drone'), style=style)\n",
    "    data_loader = torch.utils.data.DataLoader(data_set, batch_size=batchsize, shuffle=True, num_workers=0, pin_memory=False)\n",
    "\n",
    "    optimizer_g = optim.Adam(generator.parameters(), lr=lr_g, betas=(0.5, 0.999))\n",
    "    optimizer_d = optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.5, 0.999))\n",
    "    \n",
    "    criterion = nn.BCELoss()\n",
    "    l1_loss = nn.L1Loss()\n",
    "    \n",
    "    for epoch in range(1, num_epoch+1):\n",
    "        epoch_loss = 0\n",
    "        for i, (img_aug, img) in enumerate(data_loader):\n",
    "            img_aug, img_normal = img_aug.cuda(), img.cuda() \n",
    "            \n",
    "            # 训练判别器\n",
    "            optimizer_d.zero_grad()\n",
    "            \n",
    "            # 使用真实图像训练判别器\n",
    "            real_labels = torch.ones(img_normal.size(0), 1)\n",
    "            real_labels = real_labels.cuda()\n",
    "            outputs = discriminator(img_normal)\n",
    "            #print(outputs.shape, real_labels.shape)\n",
    "            d_loss_real = criterion(outputs, real_labels)\n",
    "            \n",
    "            # 使用生成图像训练判别器\n",
    "            fake_images = generator(img_aug)\n",
    "            fake_labels = torch.zeros(img_aug.size(0), 1)\n",
    "            fake_labels = fake_labels.cuda()\n",
    "            outputs = discriminator(fake_images.detach())\n",
    "            d_loss_fake = criterion(outputs, fake_labels)\n",
    "            \n",
    "            # 总的判别器损失\n",
    "            d_loss = d_loss_real + d_loss_fake\n",
    "            d_loss.backward()\n",
    "            optimizer_d.step()\n",
    "            \n",
    "            # 训练生成器\n",
    "            optimizer_g.zero_grad()\n",
    "            \n",
    "            # 生成器损失\n",
    "            outputs = discriminator(fake_images)\n",
    "            g_loss = criterion(outputs, real_labels) + l1_loss(fake_images, img_normal)\n",
    "            g_loss.backward()\n",
    "            optimizer_g.step()\n",
    "        \n",
    "            if (i+1) % 100 == 0:\n",
    "                print(f'Epoch [{epoch}/{num_epoch}], Step [{i+1}/{len(data_loader)}], '\n",
    "                    f'D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}')\n",
    "\n",
    "        save_filename = 'generator_{}_{:03d}.pth'.format(style, epoch)\n",
    "        save_path = os.path.join(model_dir, save_filename)\n",
    "        torch.save(generator.cpu().state_dict(), save_path)\n",
    "        generator.cuda()  \n",
    "\n",
    "        save_filename = 'discriminator_{}_{:03d}.pth'.format(style, epoch)\n",
    "        save_path = os.path.join(model_dir, save_filename)\n",
    "        torch.save(discriminator.cpu().state_dict(), save_path)\n",
    "        discriminator.cuda()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(style='fog_snow')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}