{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":21755,"databundleVersionId":1475600},{"sourceType":"datasetVersion","sourceId":2684635,"datasetId":1634071,"databundleVersionId":2729088},{"sourceType":"datasetVersion","sourceId":6124479,"datasetId":3450357,"databundleVersionId":6203197}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport random\nimport cv2\nfrom PIL import Image\nfrom torch.utils.data import Dataset,DataLoader\nfrom torchvision.transforms import ToTensor,Compose,Normalize,Resize\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nclass MyDataset(Dataset):\n    def __init__(self, raw_root,true_root, raw_transform=None, style_transform=None):\n        self.raw_images_path = [os.path.join(raw_root, i) for i in os.listdir(raw_root)]\n        self.true_images_path = [os.path.join(true_root,i) for i in os.listdir(true_root)]\n        self.raw_transform = raw_transform\n        self.style_transform = style_transform\n\n    def __len__(self):\n        return len(self.raw_images_path)\n\n    def __getitem__(self, item):\n        raw_image = cv2.imread(self.raw_images_path[item])\n        raw_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)\n        raw_image = Image.fromarray(raw_image)\n\n        true_image = cv2.imread(self.true_images_path[item % len(self.true_images_path)])\n        true_image = cv2.cvtColor(true_image, cv2.COLOR_BGR2RGB)\n        true_image = Image.fromarray(true_image)\n        if (self.raw_transform):\n            raw_image = self.raw_transform(raw_image)\n        if (self.style_transform):\n            true_image = self.style_transform(true_image)\n\n        return raw_image, true_image\n\nif __name__ == '__main__':\n    data = MyDataset(raw_root='/kaggle/input/gan-getting-started/photo_jpg',\n                     true_root =\"/kaggle/input/neural-style-transfer/Style Images\",\n                     raw_transform=Compose([\n                         ToTensor(),\n                        Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])]),\n                     style_transform=Compose([\n                         Resize((256,256)),\n                         ToTensor(),\n                        Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])\n                     ]))\n    data_loader = DataLoader(dataset=data,\n                             batch_size=32,\n                             shuffle=True,\n                             num_workers=4)\n\n    for i, j in data_loader:\n        fix, ax = plt.subplots(2, 8, figsize=(15, 6))\n        ax = ax.flatten()\n        i = (i.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        j = (j.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        print(i.shape, j.shape)\n\n        for value in range(0, len(ax), 2):\n            ax[value].imshow(i[value].astype(np.uint8))\n            ax[value + 1].imshow(j[value].astype(np.uint8))\n            ax[value].axis('off')\n            ax[value + 1].axis('off')\n            ax[value].set_title('ori image')\n            ax[value+1].set_title('paint style')\n\n        plt.show()\n        break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-04T02:43:12.777157Z","iopub.execute_input":"2025-03-04T02:43:12.777471Z","iopub.status.idle":"2025-03-04T02:43:25.262585Z","shell.execute_reply.started":"2025-03-04T02:43:12.777438Z","shell.execute_reply":"2025-03-04T02:43:25.261603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass Net(torch.nn.Module):\n    def __init__(self):\n        super(Net, self).__init__()\n        self.model = nn.Sequential(\n            ConvBlock(3, 32, kernel_size=9, stride=1),\n            ConvBlock(32, 64, kernel_size=3, stride=2),\n            ConvBlock(64, 128, kernel_size=3, stride=2),\n            ResidualBlock(128),\n            ResidualBlock(128),\n            ResidualBlock(128),\n            ResidualBlock(128),\n            ResidualBlock(128),\n            ConvBlock(128, 64, kernel_size=3, upsample=True),\n            ConvBlock(64, 32, kernel_size=3, upsample=True),\n            ConvBlock(32, 3, kernel_size=9, stride=1, normalize=False, relu=False),\n        )\n\n    def forward(self, x):\n        return self.model(x)\n\n\n\"\"\" Components of Transformer Net \"\"\"\nclass ResidualBlock(torch.nn.Module):\n    def __init__(self, channels):\n        super(ResidualBlock, self).__init__()\n        self.block = nn.Sequential(\n            ConvBlock(channels, channels, kernel_size=3, stride=1, normalize=True, relu=True),\n            ConvBlock(channels, channels, kernel_size=3, stride=1, normalize=True, relu=False),\n        )\n\n    def forward(self, x):\n        return self.block(x) + x\n\n\nclass ConvBlock(torch.nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, stride=1, upsample=False, normalize=True, relu=True):\n        super(ConvBlock, self).__init__()\n        self.upsample = upsample\n        self.block = nn.Sequential(\n            nn.ReflectionPad2d(kernel_size // 2), nn.Conv2d(in_channels, out_channels, kernel_size, stride)\n        )\n        self.norm = nn.InstanceNorm2d(out_channels, affine=True) if normalize else None\n        self.relu = relu\n\n    def forward(self, x):\n        if self.upsample:\n            x = F.interpolate(x, scale_factor=2)\n        x = self.block(x)\n        if self.norm is not None:\n            x = self.norm(x)\n        if self.relu:\n            x = F.relu(x)\n        return x\n\n# Kiểm tra mô hình\nif __name__ == '__main__':\n    model = Net()\n    sample_input = torch.randn(1, 3, 256, 256)  \n    output = model(sample_input)\n    print(output.shape) ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-04T02:43:25.263714Z","iopub.execute_input":"2025-03-04T02:43:25.264290Z","iopub.status.idle":"2025-03-04T02:43:25.720595Z","shell.execute_reply.started":"2025-03-04T02:43:25.264250Z","shell.execute_reply":"2025-03-04T02:43:25.719590Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchvision.models import vgg16\n\ndef gram_matrix(tensor):\n    B, C, H, W = tensor.shape\n    tensor = tensor.view(B, C, H * W)\n    return torch.bmm(tensor, tensor.transpose(1, 2)) / (C * H * W)\n\nclass PerceptualLoss(nn.Module):\n    def __init__(self, device, content_weight=1, style_weight=1000):\n        super(PerceptualLoss, self).__init__()\n        vgg = vgg16(pretrained=True).features.to(device).eval()\n\n        self.layers = {\n            '3': 'relu1_2',\n            '8': 'relu2_2',\n            '15': 'relu3_3',\n            '22': 'relu4_3',\n        }\n        self.feature_extractor = vgg\n        for param in self.feature_extractor.parameters():\n            param.requires_grad = False\n\n        self.loss_fn = nn.MSELoss()\n        self.content_weight = content_weight\n        self.style_weight = style_weight\n\n        self.mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)\n        self.std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)\n\n    def forward(self, pred, target, style_target=None):\n        pred_features = {}\n        target_features = {}\n        style_features = {}\n\n        x_pred = (pred + 1) / 2 \n        x_pred = (x_pred - self.mean) / self.std\n        x_target = (target + 1) / 2 \n        x_target = (x_target - self.mean) / self.std\n        if style_target is not None:\n            x_style = (style_target + 1) / 2\n            x_style = (x_style - self.mean) / self.std\n\n        for i, layer in enumerate(self.feature_extractor.children()):\n            x_pred = layer(x_pred)\n            x_target = layer(x_target)\n            if style_target is not None:\n                x_style = layer(x_style)\n        \n            if str(i) in self.layers:  # So sánh với key dạng string\n                pred_features[self.layers[str(i)]] = x_pred\n                target_features[self.layers[str(i)]] = x_target\n                if style_target is not None:\n                    style_features[self.layers[str(i)]] = x_style\n\n        content_loss = self.loss_fn(pred_features['relu3_3'], target_features['relu3_3'])\n\n        style_loss = 0.0\n        if style_target is not None:\n            for key in style_features.keys():\n                if key in pred_features and key in style_features:\n                    style_loss += self.loss_fn(gram_matrix(pred_features[key]), gram_matrix(style_features[key]))\n\n        return self.content_weight * content_loss + self.style_weight * style_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-04T02:43:25.721625Z","iopub.execute_input":"2025-03-04T02:43:25.721985Z","iopub.status.idle":"2025-03-04T02:43:25.735029Z","shell.execute_reply.started":"2025-03-04T02:43:25.721957Z","shell.execute_reply":"2025-03-04T02:43:25.734067Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nimport tensorflow as tf\nos.environ[\"TF_ENABLE_ONEDNN_OPTS\"] = \"0\"\nimport torch\nimport torchvision.utils\nfrom torch.utils.data import DataLoader, Dataset, Subset\nfrom torchvision.transforms import ToTensor, Resize, Normalize, RandomCrop, RandomHorizontalFlip, ColorJitter, GaussianBlur\nfrom torchvision.transforms.v2 import Compose\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.tensorboard import SummaryWriter\nfrom argparse import ArgumentParser\nimport random\n\n\ndef get_arg():\n    parser = ArgumentParser(description='Train')\n    parser.add_argument('--epochs', '-e', type=int, default=50)\n    parser.add_argument('--batch_size', '-b', type=int, default=1)\n    parser.add_argument('--learning_rate', '-lr', type=float, default=0.001)\n    parser.add_argument('--ratio_train_size', '-rt', type=float, default=0.1)\n    parser.add_argument('--val_size', '-r', type=int, default=16)\n    parser.add_argument('--content_lambda', '-cl', type=float, default=1)\n    parser.add_argument('--style_lambda', '-sl', type=float, default=1000)\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--trained', '-t', type=str, default='trained')\n    parser.add_argument('--log', '-l', type=str, default='tensorboard')\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef split_data(data, ratio_train_size, val_size, seed=42):\n    random.seed(seed)\n    index_data = list(range(len(data)))\n    random.shuffle(index_data)\n    train_size = int(ratio_train_size * len(index_data))\n    train_index = index_data[:train_size]\n    val_index = index_data[train_size:train_size + val_size]\n\n    train = Subset(data, train_index)\n    val = Subset(data, val_index)\n    print(f'Training set size: {len(train)}')\n    print(f'Validation set size: {len(val)}')\n    return train, val\n\ndef train(arg):\n    arg = get_arg()\n    if not os.path.isdir(arg.trained):\n        os.mkdir(arg.trained)\n    if os.path.isdir(arg.log):\n        shutil.rmtree(arg.log)\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    writer = SummaryWriter(arg.log)\n\n    raw_transform = Compose([\n        Resize(286),\n        RandomCrop(256),\n        ToTensor(),\n        Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n    ])\n\n    style_transform = Compose([\n        Resize(286),\n        RandomCrop(256),\n        ToTensor(),\n        Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n    ])\n\n    data = MyDataset(raw_root='/kaggle/input/gan-getting-started/photo_jpg',\n                     true_root =\"/kaggle/input/neural-style-transfer/Style Images\",\n                     raw_transform=raw_transform,\n                     style_transform=style_transform)\n\n    train,val = split_data(data,arg.ratio_train_size,arg.val_size)\n    train_loader = DataLoader(train, batch_size=arg.batch_size, shuffle=True,num_workers=0,drop_last=True)\n    val_loader = DataLoader(val, batch_size=arg.val_size, shuffle=False,num_workers=0)\n\n    if arg.checkpoint:\n        checkpoint = torch.load(arg.checkpoint,map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),weights_only=True)\n        start_epoch = checkpoint['epoch']\n        model.load_state_dict(checkpoint['model'])\n        op.load_state_dict(checkpoint['op'])\n        best_loss = checkpoint['best_loss']\n    else:\n        start_epoch = 0\n        best_loss = float('inf')\n\n    train_iters = len(train) // arg.batch_size\n    val_iters = len(val) // arg.batch_size\n\n    def init_weights(m,mode='kaiming'):\n        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n            if mode=='standard':\n                nn.init.normal_(m.weight.data, 0.0, 0.02)\n            else:\n                nn.init.kaiming_normal_(m.weight)\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    \n    model = Net().to(device)\n    model.apply(init_weights)\n    op = torch.optim.Adam(model.parameters(), lr=arg.learning_rate)\n    \n    _ , style_image = train.__getitem__(2)\n    plt.imshow(np.clip((np.transpose(style_image, (1, 2, 0)) + 1) * 127.5, 0, 255).to(torch.uint8))\n    plt.axis(\"off\")\n    plt.show()\n    \n    style_image = style_image.repeat(arg.batch_size, 1, 1, 1).to(device)\n    perceptual_loss = PerceptualLoss(device,1,1000000)\n\n    for epoch in range(start_epoch, arg.epochs):\n        model.train()\n        mean_loss = 0.0\n\n        for i, (raw_image, _) in enumerate(tqdm(train_loader)):\n            raw_image = raw_image.to(device)\n            gen_image = model(raw_image)\n            loss = perceptual_loss(gen_image, raw_image, style_image)\n            mean_loss += loss.item() \n            \n            op.zero_grad()\n            loss.backward()\n            op.step()\n\n        print('Epoch {}/{} - loss: {:.4f}'.format(\n            epoch + 1, arg.epochs,\n            mean_loss / len(train_loader)\n        ))\n\n        model.eval()\n        for raw_image, true_image in val_loader:\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            with torch.no_grad():\n                out_g = model(raw_image)\n                true_image = torchvision.utils.make_grid(true_image[:arg.batch_size], normalize=True)\n                gen_image = torchvision.utils.make_grid(out_g[:arg.batch_size], normalize=True)\n                raw_image = torchvision.utils.make_grid(raw_image[:arg.batch_size], normalize=True)\n                writer.add_image('true image', true_image, epoch)\n                writer.add_image('gen image', gen_image, epoch)\n                writer.add_image('raw image', raw_image, epoch)\n                writer.add_scalar('loss', mean_loss / train_iters, epoch)\n\n        checkpoint = {\n            'epoch': epoch + 1,\n            'model': model.state_dict(),\n            'op': op.state_dict(),\n            'best_loss': best_loss\n        }\n        torch.save(checkpoint, '{}/last.pt'.format(arg.trained))\n        if mean_loss/train_iters < best_loss:\n            best_loss = mean_loss/train_iters\n            torch.save(checkpoint, '{}/best.pt'.format(arg.trained))\n\nif __name__ == '__main__':\n    arg = get_arg()\n    train(arg)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-04T02:43:25.737235Z","iopub.execute_input":"2025-03-04T02:43:25.737862Z","iopub.status.idle":"2025-03-04T03:24:18.222859Z","shell.execute_reply.started":"2025-03-04T02:43:25.737822Z","shell.execute_reply":"2025-03-04T03:24:18.221658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport torch\nfrom argparse import ArgumentParser\n#from gan.pic2pic.G import Gen\nimport matplotlib.pyplot as plt\nfrom torchvision.transforms import Compose, Normalize, ToTensor\nfrom torch.utils.data import DataLoader\n#from torchmetrics.image.fid import FrechetInceptionDistance\nfrom tqdm import tqdm\n\n\ndef get_arg():\n    parser = ArgumentParser(description='Test')\n    parser.add_argument('--checkpoint', '-c', type=str, default=None)\n    parser.add_argument('--image', '-i', type=str, default=None)\n    parser.add_argument('--threshold', '-t', type=float, default=0.0)\n    args, unknown = parser.parse_known_args()\n    return args\n\ndef upscale(img, factor=10):\n    h, w = img.shape[:2]\n    return cv2.resize(img, (w * factor, h * factor), interpolation=cv2.INTER_NEAREST)\n\nif __name__ == '__main__':\n    arg = get_arg()\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    #device = 'cpu'\n    model = Net().to(device)\n\n    checkpoint = torch.load('trained/last.pt',\n                            map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),weights_only=True)\n    model.load_state_dict(checkpoint['model'])\n\n    mean = [0.5, 0.5, 0.5]\n    std = [0.5, 0.5, 0.5]\n    std = [s if s > 0 else 1e-6 for s in std]\n\n    transform = Compose([\n        ToTensor(),\n        Normalize(mean, std)\n    ])\n\n    data = MyDataset(\n        raw_root=\"/kaggle/input/gan-getting-started/photo_jpg\",\n        true_root=\"/kaggle/input/gan-getting-started/monet_jpg\",\n        raw_transform=transform,\n        style_transform=transform\n    )\n    loader = DataLoader(dataset=data, batch_size=16, shuffle=True)\n    print(len(data))\n\n    model.eval()\n\n    # fid = FrechetInceptionDistance(feature=2048).to(device)\n    # with torch.no_grad():\n    #     for in_image, out_image in tqdm(loader):\n    #         gen_image = gen_a(in_image.to(device))\n\n    #         gen_image = ((gen_image + 1) * 127.5).to(torch.uint8)\n    #         out_image = ((out_image + 1) * 127.5).to(torch.uint8)\n    #         if gen_image.shape[-1] == 3:\n    #             gen_image = gen_image.permute(0, 3, 1, 2)\n    #             out_image = out_image.permute(0, 3, 1, 2)\n\n    #         fid.update(gen_image.to(device), real=False)\n    #         fid.update(out_image.to(device), real=True)\n\n    # fid_score = fid.compute().item()\n    # print(f\"FID score: {fid_score:.2f}\")\n    # exit()\n\n    show_iter = 0\n    for in_image, out_image in loader:\n        show_iter += 1\n        fig, ax = plt.subplots(2, 4, figsize=(10,5))\n        ax = ax.flatten()\n\n        with torch.no_grad():\n            out = model(in_image.to(device)).detach().cpu()\n\n        ori_image = (in_image.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        paint_image = (out_image.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        gen_image = (out.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n\n        # Phóng to ảnh\n        ori_image = np.array([upscale(img) for img in ori_image])\n        paint_image = np.array([upscale(img) for img in paint_image])\n        gen_image = np.array([upscale(img) for img in gen_image])\n\n        for value in range(0, len(ax), 2):\n            if value < 16:\n                ax[value].imshow(ori_image[value].astype(np.uint8), interpolation='nearest')\n                ax[value + 1].imshow(gen_image[value].astype(np.uint8), interpolation='nearest')\n\n                ax[value].set_title('in')\n                ax[value + 1].set_title('gen')\n\n                ax[value].axis('off')\n                ax[value + 1].axis('off')\n\n        plt.show()\n        if show_iter == 3:\n            break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-04T03:27:02.886228Z","iopub.execute_input":"2025-03-04T03:27:02.886609Z","iopub.status.idle":"2025-03-04T03:27:20.621177Z","shell.execute_reply.started":"2025-03-04T03:27:02.886576Z","shell.execute_reply":"2025-03-04T03:27:20.619970Z"}},"outputs":[],"execution_count":null}]}