{"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":[{"sourceId":21755,"databundleVersionId":1475600,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":1913658,"sourceType":"datasetVersion","datasetId":1036526},{"sourceId":3589214,"sourceType":"datasetVersion","datasetId":2143015}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import cv2\nimport os\n\ninput_folder = \"/kaggle/input/landscape-image-colorization/landscape Images/color\"\noutput_sketch_folder = \"sketch_image\"\noutput_color_folder = \"color_image\"  # Thư mục chứa ảnh gốc\n\n# Tạo thư mục đầu ra nếu chưa có\nos.makedirs(output_sketch_folder, exist_ok=True)\nos.makedirs(output_color_folder, exist_ok=True)\n\n# Sắp xếp danh sách file theo thứ tự tên\nfile_list = sorted(os.listdir(input_folder))  \n\nfor filename in file_list:\n    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  \n        img_path = os.path.join(input_folder, filename)\n        img = cv2.imread(img_path)\n\n        if img is None:\n            print(f\"Lỗi khi đọc ảnh: {filename}\")\n            continue\n\n        # Chuyển ảnh sang grayscale và tạo sketch\n        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n        sketch = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, \n                                       cv2.THRESH_BINARY, 9, 10)      \n        \n        # Lưu ảnh sketch\n        sketch_output_path = os.path.join(output_sketch_folder, filename)\n        cv2.imwrite(sketch_output_path, sketch)\n\n        # Lưu ảnh màu gốc\n        color_output_path = os.path.join(output_color_folder, filename)\n        cv2.imwrite(color_output_path, img)\n\nprint(\"Done.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:10:30.919653Z","iopub.execute_input":"2025-03-10T15:10:30.919863Z","iopub.status.idle":"2025-03-10T15:11:53.222357Z","shell.execute_reply.started":"2025-03-10T15:10:30.919844Z","shell.execute_reply":"2025-03-10T15:11:53.221473Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(len(os.listdir('/kaggle/working/sketch_image')))\nprint(len(os.listdir('/kaggle/working/color_image')))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:11:53.225340Z","iopub.execute_input":"2025-03-10T15:11:53.225533Z","iopub.status.idle":"2025-03-10T15:11:53.238299Z","shell.execute_reply.started":"2025-03-10T15:11:53.225516Z","shell.execute_reply":"2025-03-10T15:11:53.237568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport random\nfrom PIL import Image\nfrom torch.utils.data import Dataset,DataLoader,Subset\nfrom torchvision.transforms import ToTensor,Compose,Resize,Normalize\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nclass MyDataset(Dataset):\n    def __init__(self, gray_root, color_root, gray_transform=None, color_transform=None):\n        self.gray_images_path = [os.path.join(gray_root, i) for i in os.listdir(gray_root)]\n        self.color_images_path = [os.path.join(color_root, i) for i in os.listdir(color_root)]\n        self.gray_transform = gray_transform\n        self.color_transform = color_transform\n        print(len(self.gray_images_path))\n        print(len(self.color_images_path))\n\n    def __len__(self):\n        return len(self.color_images_path)\n\n    def __getitem__(self, item):\n        gray = cv2.imread(self.gray_images_path[item])\n        color = cv2.imread(self.color_images_path[item])\n        color = cv2.cvtColor(color, cv2.COLOR_BGR2RGB)\n        gray = Image.fromarray(gray)\n        color = Image.fromarray(color)\n        if self.color_transform:\n            color = self.color_transform(color)\n        if self.gray_transform:\n            gray = self.gray_transform(gray)\n        return gray, color\n\nif __name__ == '__main__':\n    data = MyDataset(gray_root='/kaggle/working/sketch_image',\n                     color_root='/kaggle/working/color_image',\n                     color_transform=Compose([\n                         Resize((128, 128)),\n                         ToTensor(),\n                         Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n                     ]),\n                     gray_transform=Compose([\n                         Resize((128, 128)),\n                         ToTensor(),\n                         Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n                     ]))\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\n        plt.show()\n        break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:16:04.427048Z","iopub.execute_input":"2025-03-10T15:16:04.427426Z","iopub.status.idle":"2025-03-10T15:16:05.451311Z","shell.execute_reply.started":"2025-03-10T15:16:04.427398Z","shell.execute_reply":"2025-03-10T15:16:05.450487Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\n\nclass Gen(nn.Module):\n    def __init__(self,in_image_channel,out_image_channel,feature_g=64):\n        super().__init__()\n        self.down1 = nn.Sequential(\n            nn.Conv2d(in_image_channel,feature_g,4,2,1),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n        self.down2 = self.down_block_conv(feature_g,feature_g*2,4,2,1)\n        self.down3 = self.down_block_conv(feature_g*2, feature_g*4, 4, 2, 1)\n        self.down4 = self.down_block_conv(feature_g*4, feature_g*8, 4, 2, 1)\n        self.down5 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n        self.down6 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n        self.down7 = self.down_block_conv(feature_g*8, feature_g*8, 4, 2, 1)\n\n        self.bottom = nn.Sequential(\n            nn.Conv2d(feature_g*8, feature_g*8, 4, 2, 1),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n\n        self.up1 = self.up_block_conv(feature_g*8,feature_g*8,4,  2,  1,drop_out=True)\n        self.up2 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1,drop_out=True)\n        self.up3 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1,drop_out=True)\n        self.up4 = self.up_block_conv(feature_g*8*2,feature_g*8,4, 2, 1)\n        self.up5 = self.up_block_conv(feature_g*8*2,feature_g*4,4, 2, 1)\n        self.up6 = self.up_block_conv(feature_g*4*2,feature_g*2,4, 2, 1)\n        self.up7 = self.up_block_conv(feature_g*2*2,feature_g,4, 2, 1)\n\n        self.final = nn.Sequential(\n            nn.ConvTranspose2d(feature_g,out_image_channel,4,2,1),\n            nn.Tanh()\n        )\n\n    def up_block_conv(self,in_channel,out_channel,kernel,stride,padding,drop_out=False):\n        return nn.Sequential(\n            nn.ConvTranspose2d(in_channel,out_channel,kernel,stride,padding,bias=False),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.Dropout(0.5) if drop_out else nn.Identity(),\n            nn.ReLU(inplace=True)\n        )\n\n    def down_block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel, out_channel, kernel, stride, padding,padding_mode='reflect',bias=False),\n            nn.InstanceNorm2d(out_channel, affine=True),\n            nn.LeakyReLU(0.2,inplace=True)\n        )\n\n    def forward(self, x):\n        d1 = self.down1(x)\n        d2 = self.down2(d1)\n        d3 = self.down3(d2)\n        d4 = self.down4(d3)\n        d5 = self.down5(d4)\n        d6 = self.down6(d5)\n        d7 = self.down7(d6)\n\n        bottom = self.bottom(d7)\n\n        u1 = self.up1(bottom)\n        u2 = self.up2(torch.cat([u1, d7], 1))\n        u3 = self.up3(torch.cat([u2, d6], 1))\n        u4 = self.up4(torch.cat([u3, d5], 1))\n        u5 = self.up5(torch.cat([u4, d4], 1))\n        u6 = self.up6(torch.cat([u5, d3], 1))\n        u7 = self.up7(torch.cat([u6, d2], 1))\n        final = self.final(u7)\n\n        return final\n\nif __name__ == '__main__':\n    model = Gen(1,3,64)\n    print(summary(model,input_size=(1,1,256,256)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:16:12.260324Z","iopub.execute_input":"2025-03-10T15:16:12.260658Z","iopub.status.idle":"2025-03-10T15:16:13.676346Z","shell.execute_reply.started":"2025-03-10T15:16:12.260633Z","shell.execute_reply":"2025-03-10T15:16:13.675455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchinfo import summary\n\nclass Dis(nn.Module):\n    def __init__(self,image_channel=3,feature_d=64):\n        super().__init__()\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(image_channel,feature_d,4,2,1),\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n        self.conv2 = self.block_conv(feature_d,feature_d*2,4,2,1)\n        self.conv3 = self.block_conv(feature_d*2, feature_d*4, 4, 2, 1)\n        self.conv4 = self.block_conv(feature_d*4, feature_d*8, 4, 1, 1)\n        self.conv5 = nn.Sequential(\n            nn.Conv2d(feature_d*8,1,4,1,1),\n            #nn.Sigmoid()\n        )\n\n    def block_conv(self,in_channel,out_channel,kernel,stride,padding):\n        return nn.Sequential(\n            nn.Conv2d(in_channel,out_channel,kernel_size=kernel,stride=stride,padding=padding,bias=False),\n            nn.InstanceNorm2d(out_channel,affine=True),\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n\n    def forward(self,x,return_feature=False):\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.conv3(x)\n        f_x = self.conv4(x)\n        x = self.conv5(f_x)\n        if return_feature:\n            return x,f_x\n        return x\n\nif __name__ == '__main__':\n    model = Dis(image_channel=3,feature_d=64)\n    print(summary(model,input_size=(32,3,256,256)))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:16:17.554161Z","iopub.execute_input":"2025-03-10T15:16:17.554474Z","iopub.status.idle":"2025-03-10T15:16:17.717606Z","shell.execute_reply.started":"2025-03-10T15:16:17.554452Z","shell.execute_reply":"2025-03-10T15:16:17.716830Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nimport torch\nimport torchvision.utils\n# from Dis import Dis\n# from Gen import Gen\nfrom torch.utils.data import DataLoader, Dataset, Subset\n# from dataset import MyDataset\nfrom torchvision.transforms import ToTensor, Resize, Normalize, RandomCrop\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=10)\n    parser.add_argument('--batch_size', '-b', type=int, default=8)\n    parser.add_argument('--learning_rate', '-lr', type=float, default=0.0002)\n    parser.add_argument('--ratio_train_size', '-r', type=float, default=0.2)\n    parser.add_argument('--val_size', '-v', type=int, default=16)\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\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\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    writer = SummaryWriter(arg.log)\n\n    data = MyDataset(gray_root='/kaggle/working/sketch_image',\n                     color_root=\"/kaggle/working/color_image\",\n                     color_transform=Compose([\n                         Resize((286, 286)),\n                         RandomCrop(256),\n                         ToTensor(),\n                         Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n                     ]),\n                     gray_transform=Compose([\n                         Resize((286, 286)),\n                         RandomCrop(256),\n                         ToTensor(),\n                         Normalize(mean=0.5, std=0.5)\n                     ]))\n\n    train,val = split_data(data,arg.ratio_train_size,arg.val_size)\n    print(len(train))\n    print(len(val))\n    train_loader = DataLoader(train, batch_size=arg.batch_size, shuffle=True)\n    val_loader = DataLoader(val, batch_size=arg.val_size, shuffle=False)\n\n    gen_a = Gen(3,3)\n    gen_b = Gen(3,3)\n    dis_a = Dis(3)\n    dis_b = Dis(3)\n\n    gen_a.to(device)\n    gen_b.to(device)\n    dis_a.to(device)\n    dis_b.to(device)\n\n    def init_weights(m):\n        if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n            #nn.init.kaiming_normal_(m.weight)\n            nn.init.normal_(m.weight.data, 0.0, 0.02)\n\n    gen_a.apply(init_weights)\n    gen_b.apply(init_weights)\n    dis_a.apply(init_weights)\n    dis_b.apply(init_weights)\n\n    op_g_a = optim.Adam(gen_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_g_b = optim.Adam(gen_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_d_a = optim.Adam(dis_a.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n    op_d_b = optim.Adam(dis_b.parameters(), lr=arg.learning_rate, betas=(0.5, 0.999))\n\n    bce = nn.BCEWithLogitsLoss()\n    #mse = nn.MSELoss()\n    l1 = nn.L1Loss()\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        dis_a.load_state_dict(checkpoint['dis_a'])\n        dis_b.load_state_dict(checkpoint['dis_b'])\n        gen_a.load_state_dict(checkpoint['gen_a'])\n        gen_b.load_state_dict(checkpoint['gen_b'])\n        op_g_a.load_state_dict(checkpoint['op_g_a'])\n        op_g_b.load_state_dict(checkpoint['op_g_b'])\n        op_d_a.load_state_dict(checkpoint['op_d_a'])\n        op_d_b.load_state_dict(checkpoint['op_d_b'])\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.val_size\n\n    for epoch in range(start_epoch, arg.epochs):\n        dis_a.train()\n        dis_b.train()\n        gen_a.train()\n        gen_b.train()\n\n        mean_loss_da = 0.0\n        mean_loss_db = 0.0\n        mean_loss_ga = 0.0\n        mean_loss_gb = 0.0\n\n        for raw_image, true_image in tqdm(train_loader):\n            raw_image = raw_image.to(device)\n            true_image = true_image.to(device)\n            out_g1 = gen_a(raw_image)\n            out_g2 = gen_b(true_image)\n\n            out_d1a = dis_a(true_image)\n            loss_d1a = bce(out_d1a,torch.ones_like(out_d1a)*0.9)\n            out_d2a = dis_a(out_g1.detach())\n            loss_d2a = bce(out_d2a, torch.zeros_like(out_d2a)+0.1)\n            loss_da = (loss_d1a + loss_d2a)/2\n            mean_loss_da += loss_da\n\n            op_d_a.zero_grad()\n            loss_da.backward()\n            op_d_a.step()\n\n            out_d1b = dis_b(raw_image)\n            loss_d1b = bce(out_d1b, torch.ones_like(out_d1b))\n            out_d2b = dis_b(out_g2.detach())\n            loss_d2b = bce(out_d2b, torch.zeros_like(out_d2b))\n            loss_db = (loss_d1b + loss_d2b)/2\n            mean_loss_db += loss_db\n\n            op_d_b.zero_grad()\n            loss_db.backward()\n            op_d_b.step()\n\n            out_da = dis_a(out_g1)\n            loss_ga = bce(out_da, torch.ones_like(out_da))\n\n            cycle_a = gen_a(out_g2.detach())\n            iden_a = gen_a(true_image)\n            loss_iden_a = l1(iden_a,true_image)\n            loss_cycle_a = l1(cycle_a,true_image)\n            loss_ga_final = loss_ga + 5 * loss_cycle_a + 10 * loss_iden_a\n            mean_loss_ga += loss_ga_final\n\n            op_g_a.zero_grad()\n            loss_ga_final.backward()\n            op_g_a.step()\n\n            out_db = dis_b(out_g2)\n            loss_gb = bce(out_db, torch.ones_like(out_db))\n\n            cycle_b = gen_b(out_g1.detach())\n            iden_b = gen_b(raw_image)\n            loss_cycle_b = l1(cycle_b,raw_image)\n            loss_iden_b = l1(iden_b,raw_image)\n            loss_gb_final = loss_gb + 5 * loss_cycle_b + 10 * loss_iden_b\n            mean_loss_gb += loss_gb_final\n\n            op_g_b.zero_grad()\n            loss_gb_final.backward()\n            op_g_b.step()\n\n        print('Epoch {}/{} -Da: Loss: {} | -Ga: Loss: {} | -Db: Loss: {} | -Gb: Loss: {}\\n'.format(\n                                                                epoch + 1,arg.epochs,\n                                                                mean_loss_da / train_iters,\n                                                                mean_loss_ga / train_iters,\n                                                                mean_loss_db / train_iters,\n                                                                mean_loss_gb / train_iters))\n\n        gen_a.eval()\n        gen_b.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_ga = gen_a(raw_image)\n                out_gb = gen_b(out_ga)\n                true_image = torchvision.utils.make_grid(true_image[:arg.val_size], normalize=True)\n                gen_image = torchvision.utils.make_grid(out_ga[:arg.val_size], normalize=True)\n                raw_image = torchvision.utils.make_grid(raw_image[:arg.val_size], normalize=True)\n                cycle_image = torchvision.utils.make_grid(out_gb[:arg.val_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_image('cycle image', cycle_image, epoch)\n                writer.add_scalar('loss da', mean_loss_da / train_iters, epoch)\n                writer.add_scalar('loss ga', mean_loss_ga / train_iters, epoch)\n                writer.add_scalar('loss db', mean_loss_db / train_iters, epoch)\n                writer.add_scalar('loss gb', mean_loss_gb / train_iters, epoch)\n\n        checkpoint = {\n            'epoch': epoch + 1,\n            'gen_a': gen_a.state_dict(),\n            'dis_a': dis_a.state_dict(),\n            'gen_b': gen_b.state_dict(),\n            'dis_b': dis_b.state_dict(),\n            'op_g_a': op_g_a.state_dict(),\n            'op_d_a': op_d_a.state_dict(),\n            'op_g_b': op_g_b.state_dict(),\n            'op_d_b': op_d_b.state_dict(),\n            'best_loss': best_loss\n        }\n        torch.save(checkpoint, '{}/last.pt'.format(arg.trained))\n        if ((mean_loss_ga / train_iters) + (mean_loss_gb / train_iters))/2 < best_loss:\n            best_loss = ((mean_loss_ga / train_iters) + (mean_loss_gb / train_iters))/2\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-10T15:19:47.409957Z","iopub.execute_input":"2025-03-10T15:19:47.410351Z","iopub.status.idle":"2025-03-10T15:39:27.444627Z","shell.execute_reply.started":"2025-03-10T15:19:47.410325Z","shell.execute_reply":"2025-03-10T15:39:27.443638Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport torch\nfrom argparse import ArgumentParser\n#from Gen import Gen\nimport matplotlib.pyplot as plt\nfrom torchvision.transforms import Compose,Resize,Normalize,ToTensor\nfrom torch.utils.data import DataLoader\n\ndef get_arg():\n    parser = ArgumentParser(description='Test Cycle GAN')\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    \nif __name__ == '__main__':\n    arg = get_arg()\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    gen_a = Gen(3,3)\n    gen_b = Gen(3,3)\n\n    checkpoint = torch.load('/kaggle/working/trained/last.pt',\n                            map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()),\n                            weights_only=True)\n    gen_a.load_state_dict(checkpoint['gen_a'])\n    gen_b.load_state_dict(checkpoint['gen_b'])\n\n    mean = [0.5, 0.5, 0.5]\n    std = [0.5, 0.5, 0.5]\n\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    #data = MyDataset(root=day2night_root, transform=transform)\n    data = MyDataset(gray_root='/kaggle/working/sketch_image',\n                     color_root=\"/kaggle/working/color_image\",\n                     color_transform=Compose([\n                         Resize((256,256)),\n                         ToTensor(),\n                         Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n                     ]),\n                     gray_transform=Compose([\n                         Resize((256,256)),\n                         ToTensor(),\n                         Normalize(mean=0.5, std=0.5)\n                     ]))\n\n    loader = DataLoader(dataset=data,\n                             batch_size=16,\n                             shuffle=True,\n                             num_workers=4)\n\n    show_iter = 0\n    for in_image, out_image in loader:\n        show_iter += 1\n        fix, ax = plt.subplots(2, 4, figsize=(15, 6))\n        ax = ax.flatten()\n\n        with torch.no_grad():\n            out_a = gen_a(in_image).detach().cpu()\n            #out_b = gen_b(out_a).detach().cpu()\n        \n        ori_image = (in_image.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        gen_image_a = (out_a.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        #gen_image_b = (out_b.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        #out_image = (out_image.numpy().transpose(0, 2, 3, 1) + 1) * 127.5\n        print(ori_image.shape,gen_image_a.shape)\n\n        for value in range(0, len(ax), 2):\n            ax[value].imshow(ori_image[value].astype(np.uint8))\n            ax[value + 1].imshow(gen_image_a[value].astype(np.uint8))\n            #ax[value + 2].imshow([value].astype(np.uint8))\n\n            ax[value].set_title('ori image')\n            ax[value + 1].set_title('paint image')\n            #ax[value + 2].set_title('cycle image')\n            ax[value].axis('off')\n            ax[value + 1].axis('off')\n            #ax[value + 2].axis('off')\n\n        plt.show()\n        if show_iter == 3:\n            break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-10T15:47:55.112443Z","iopub.execute_input":"2025-03-10T15:47:55.112830Z","iopub.status.idle":"2025-03-10T15:48:04.654320Z","shell.execute_reply.started":"2025-03-10T15:47:55.112803Z","shell.execute_reply":"2025-03-10T15:48:04.653384Z"}},"outputs":[],"execution_count":null}]}