{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## References and Ackowledgements:\n- Architecture: CycleGAN\n- for this notebook, I have used the official pytorch implementation of CycleGAN and Pix2Pix papers, and modified it to fit the needs of this competition.\n- git repo link of official (CycleGAN + Pix2Pix) pytorch repo: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/\n- CycleGAN paper link: https://arxiv.org/abs/1703.10593\n- Also referred to this notebook to understand how to submit the result images for this competition: https://www.kaggle.com/alllancheung/mlproject1/notebook.\n- Remember to keep internet on while submitting as it needs to install dominate module.","metadata":{"id":"uXpepTYHsOwm"}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport glob\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\n\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        if filename.endswith(\".jpg\"):\n            print(os.path.join(dirname, filename))\n            break\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\nprint(torch.__version__)","metadata":{"execution":{"iopub.status.busy":"2022-02-04T13:57:27.484414Z","iopub.execute_input":"2022-02-04T13:57:27.484741Z","iopub.status.idle":"2022-02-04T13:58:14.694469Z","shell.execute_reply.started":"2022-02-04T13:57:27.484711Z","shell.execute_reply":"2022-02-04T13:58:14.693572Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\n'''Paths for colab notebook'''\n\n'''Paths for kaggle notebook'''\ninput_path = '/kaggle/input'\nkaggle_input_path = '/kaggle/input/gan-getting-started'\nkaggle_output_path = '/kaggle/working'\ncheckpoint_path = './checkpoints'\nresult_img_path = '/kaggle/working/results'\nimages_submit_path = '/kaggle/working/images'\n\n'''Kaggle dataset names'''\nkaggle_split_dataset_name = 'painter-colab-dataset-splits'\nkaggle_cycleGAN_checkpoint_dataset_name= 'painter-colab-cyclegan-chpt-e200-trainimg1K'\nkaggle_cycleGAN_checkpoint_dataset_path = '../input/painter-colab-cyclegan-chpt-e200-trainimg1k'","metadata":{"id":"yEdowb8xtwnV","execution":{"iopub.status.busy":"2022-02-04T13:56:48.720502Z","iopub.execute_input":"2022-02-04T13:56:48.720781Z","iopub.status.idle":"2022-02-04T13:56:48.727284Z","shell.execute_reply.started":"2022-02-04T13:56:48.72075Z","shell.execute_reply":"2022-02-04T13:56:48.726078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport glob\n\nimport glob\nimport random\nimport shutil\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\n\nfor dirname, _, filenames in os.walk(kaggle_input_path):\n    for filename in filenames:\n        if filename.endswith(\".jpg\"):\n            print(os.path.join(dirname, filename))\n            break\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\nprint(torch.__version__)","metadata":{"id":"cXInf1fPsOwr","outputId":"32746612-de28-4a17-897c-415521a8a9b7","execution":{"iopub.status.busy":"2022-02-04T13:56:53.367275Z","iopub.execute_input":"2022-02-04T13:56:53.367759Z","iopub.status.idle":"2022-02-04T13:56:54.531604Z","shell.execute_reply.started":"2022-02-04T13:56:53.36771Z","shell.execute_reply":"2022-02-04T13:56:54.530852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating trainA, trainB, valA, valB, testA, testB folders under the path /kaggle/output/working/","metadata":{"id":"UNCeD4L3sOws"}},{"cell_type":"code","source":"## Code to created input folders and fill them with images from /kaggle/input/gan_getting_started/ folder.\n\n\n# random.seed(1)\n\n# # train_ratio, val_ratio, test_ratio = 0.9, 0.05, 0.05\n# trainA_ratio, valA_ratio, testA_ratio = 0.14208582, 0.05, 0.05 # 0.9, 0.05, 0.05\n# trainB_ratio, valB_ratio, testB_ratio = 0.9, 0.05, 0.05\n# train_count, val_count, test_count = 0, 0, 0\n# trainA_count, valA_count, testA_count = 0, 0, 0\n# trainB_count, valB_count, testB_count = 0, 0, 0\n\n\n# src_A_folder = 'photo_jpg'\n# src_B_folder = 'monet_jpg'\n# A_photos=glob.glob(os.path.join(kaggle_input_path, src_A_folder, '*.jpg'))\n# B_photos=glob.glob(os.path.join(kaggle_input_path, src_B_folder, '*.jpg'))\n# A_photos_len = len(A_photos)\n# B_photos_len = len(B_photos)\n# print(A_photos_len)\n# print(B_photos_len)\n\n# trainA_count, valA_count = int(A_photos_len*trainA_ratio), int(A_photos_len*valA_ratio)\n# testA_count = A_photos_len - (trainA_count + valA_count)\n\n# trainB_count, valB_count = int(B_photos_len*trainB_ratio), int(B_photos_len*valB_ratio)\n# testB_count = B_photos_len - (trainB_count + valB_count)\n# print(\"trainA_count={}, valA_count={}, testA_count={}\".format(trainA_count, valA_count, testA_count))\n# print(\"trainB_count={}, valB_count={}, testB_count={}\".format(trainB_count, valB_count, testB_count))\n\n\n# trainA_folder = 'trainA'\n# trainB_folder = 'trainB'\n# trainA_path = os.path.join(kaggle_output_path, trainA_folder)\n# trainB_path = os.path.join(kaggle_output_path, trainB_folder)\n    \n# valA_folder = 'valA'\n# valB_folder = 'valB'\n# valA_path = os.path.join(kaggle_output_path, valA_folder)\n# valB_path = os.path.join(kaggle_output_path, valB_folder)\n\n# testA_folder = 'testA'\n# testB_folder = 'testB'\n# testA_path = os.path.join(kaggle_output_path, testA_folder)\n# testB_path = os.path.join(kaggle_output_path, testB_folder)\n\n# print(os.path.isdir(trainA_path))\n# if os.path.isdir(trainA_path):\n#     shutil.rmtree(trainA_path)\n# if not os.path.isdir(trainA_path):\n#     os.mkdir(trainA_path)\n    \n# if os.path.isdir(trainB_path):\n#     shutil.rmtree(trainB_path)\n# if not os.path.isdir(trainB_path):\n#     os.mkdir(trainB_path)\n    \n# if os.path.isdir(valA_path):\n#     shutil.rmtree(valA_path)\n# if not os.path.isdir(valA_path):\n#     os.mkdir(valA_path)\n    \n# if os.path.isdir(valB_path):\n#     shutil.rmtree(valB_path)\n# if not os.path.isdir(valB_path):\n#     os.mkdir(valB_path)\n    \n# if os.path.isdir(testA_path):\n#     shutil.rmtree(testA_path)\n# if not os.path.isdir(testA_path):\n#     os.mkdir(testA_path)\n\n# if os.path.isdir(testB_path):\n#     shutil.rmtree(testB_path)\n# if not os.path.isdir(testB_path):\n#     os.mkdir(testB_path)\n\n# print(len(glob.glob(os.path.join(kaggle_output_path, \"*\"))), end=\"\\n\\n\")\n# print(glob.glob(os.path.join(kaggle_output_path, \"*\")), end=\"\\n\\n\")\n\n# print(A_photos[:10], end=\"\\n\\n\")\n# print(B_photos[:10], end=\"\\n\\n\")","metadata":{"id":"MfQfDoJZsOwt","execution":{"iopub.status.busy":"2022-01-02T14:35:01.90047Z","iopub.execute_input":"2022-01-02T14:35:01.900742Z","iopub.status.idle":"2022-01-02T14:35:01.906272Z","shell.execute_reply.started":"2022-01-02T14:35:01.900713Z","shell.execute_reply":"2022-01-02T14:35:01.905495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# random.seed(1)\n\n# A_idx_arr = [i for i in range(0, A_photos_len)]\n# B_idx_arr = [i for i in range(0, B_photos_len)]\n\n# # print(A_idx_arr[:10], end=\"\\n\\n\")\n\n# def random_split(idx_arr, train_count, val_count):\n#     train_idx_arr = random.sample(idx_arr, k=train_count)\n#     val_idx_arr = random.sample(list(set.difference(set(idx_arr), set(train_idx_arr))), k=val_count)\n#     test_idx_arr = list(set.difference(set(idx_arr), set.union(set(train_idx_arr), set(val_idx_arr))))\n#     return train_idx_arr, val_idx_arr, test_idx_arr\n\n# def get_arr_by_idx(arr, idx_arr, result_arr=[]):\n#     result_arr=[]\n#     print(\"====================================================\")\n#     # print(arr[:5], end=\"\\n\\n\")\n#     for i in idx_arr:\n#         if i == 0:\n#             print(arr[i])\n#         result_arr.append(arr[i])\n#     # print(result_arr[:5], end=\"\\n\\n\")\n#     return result_arr\n    \n# def copy_files_arr_to_split_folder(filepath_arr, split_folder_path):\n#     for filepath in filepath_arr:\n#         filename = os.path.basename(filepath)\n#         target_path = os.path.join(split_folder_path, filename)\n#         shutil.copy(filepath, target_path)","metadata":{"id":"efaQAzQJsOwu","execution":{"iopub.status.busy":"2022-01-02T14:35:02.900074Z","iopub.execute_input":"2022-01-02T14:35:02.900735Z","iopub.status.idle":"2022-01-02T14:35:02.905467Z","shell.execute_reply.started":"2022-01-02T14:35:02.900699Z","shell.execute_reply":"2022-01-02T14:35:02.904715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## Copying images into split folders\n\n\n# import shutil\n\n# random.seed(1)\n\n# # trainA_idx_arr = random.sample(A_idx_arr, k=trainA_count)\n# # valA_idx_arr = random.sample(list(set.difference(set(A_idx_arr), set(trainA_idx_arr))), k=valA_count)\n# # testA_idx_arr = list(set.difference(set(A_idx_arr), set.union(set(trainA_idx_arr), set(valA_idx_arr))))\n\n# # trainB_idx_arr = random.sample(B_idx_arr, k=trainB_count)\n# # valB_idx_arr = random.sample(list(set.difference(set(B_idx_arr), set(trainB_idx_arr))), k=valB_count)\n# # testB_idx_arr = list(set.difference(set(B_idx_arr), set.union(set(trainB_idx_arr), set(valB_idx_arr))))\n\n# trainA_idx_arr, valA_idx_arr, testA_idx_arr = random_split(A_idx_arr, trainA_count, valA_count)\n# trainB_idx_arr, valB_idx_arr, testB_idx_arr = random_split(B_idx_arr, trainB_count, valB_count)\n                     \n# print(trainA_idx_arr[:5], end=\"\\n\\n\")\n# print(valA_idx_arr[:5], end=\"\\n\\n\")\n# # print(testA_idx_arr[:5], end=\"\\n\\n\")\n\n# print(trainB_idx_arr[:5], end=\"\\n\\n\")\n# print(valB_idx_arr[:5], end=\"\\n\\n\")\n# # print(testB_idx_arr[:5], end=\"\\n\\n\")\n\n# trainA_photos, valA_photos, testA_photos = [], [], []\n# trainB_monets, valB_monets, testB_monets = [], [], []\n\n# # for i in trainA_idx_arr:\n# #     train_A_photos.append(A_photos[i])\n# # for i in valA_idx_arr:\n# #     val_A_photos.append(A_photos[i])\n# # for i in testA_idx_arr:\n# #     test_A_photos.append(A_photos[i])\n\n# # print(A_photos[:5], end=\"\\n\\n\")\n# # print(B_photos[:5], end=\"\\n\\n\")\n    \n# trainA_photos = get_arr_by_idx(A_photos, trainA_idx_arr)\n# valA_photos = get_arr_by_idx(A_photos, valA_idx_arr)\n# # testA_photos = get_arr_by_idx(A_photos, testA_idx_arr)\n# testA_photos = A_photos\n\n# trainB_monets = get_arr_by_idx(B_photos, trainB_idx_arr) \n# valB_monets = get_arr_by_idx(B_photos, valB_idx_arr)\n# # testB_monets = get_arr_by_idx(B_photos, testB_idx_arr)\n# testB_monets = B_photos\n    \n# print(\"len(trainA_photos)={}\".format(len(trainA_photos)))\n# print(\"len(valA_photos)={}\".format(len(valA_photos)))\n# print(\"len(testA_photos)={}\".format(len(testA_photos)))\n# print(\"len(trainB_monets)={}\".format(len(trainB_monets)))\n# print(\"len(valB_monets)={}\".format(len(valB_monets)))\n# print(\"len(testB_monets)={}\".format(len(testB_monets)))\n\n# print(trainA_photos[:5], end=\"\\n\\n\")\n\n# print(valA_photos[:5], end=\"\\n\\n\")\n\n# print(testA_photos[:5], end=\"\\n\\n\")\n\n# print(trainB_monets[:5], end=\"\\n\\n\")\n\n# print(valB_monets[:5], end=\"\\n\\n\")\n\n# print(testB_monets[:5], end=\"\\n\\n\")\n\n# copy_files_arr_to_split_folder(trainA_photos, trainA_path)\n# copy_files_arr_to_split_folder(valA_photos, valA_path)\n# copy_files_arr_to_split_folder(testA_photos, testA_path)\n\n# copy_files_arr_to_split_folder(trainB_monets, trainB_path)\n# copy_files_arr_to_split_folder(valB_monets, valB_path)\n# copy_files_arr_to_split_folder(testB_monets, testB_path)","metadata":{"id":"rOr-gw49sOwv","execution":{"iopub.status.busy":"2022-01-02T14:35:03.705455Z","iopub.execute_input":"2022-01-02T14:35:03.705885Z","iopub.status.idle":"2022-01-02T14:35:03.71342Z","shell.execute_reply.started":"2022-01-02T14:35:03.705851Z","shell.execute_reply":"2022-01-02T14:35:03.712703Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Copy the data split used in google colab code into the /kaggle/working folder","metadata":{}},{"cell_type":"code","source":"def copy_split_folders(input_folder, output_folder, split_arr=['trainA', 'trainB', 'valA', 'valB', 'testA', 'testB']):\n    for split in split_arr:\n        dataset_src_path = os.path.join(input_path, kaggle_split_dataset_name, split, \"*\")\n        inner_split_folder_arr = glob.glob(dataset_src_path)\n        for inner_split_folder in inner_split_folder_arr:\n            shutil.copytree(inner_split_folder, output_folder+os.sep+split)","metadata":{"execution":{"iopub.status.busy":"2022-01-02T14:35:07.115559Z","iopub.execute_input":"2022-01-02T14:35:07.11613Z","iopub.status.idle":"2022-01-02T14:35:07.122236Z","shell.execute_reply.started":"2022-01-02T14:35:07.116091Z","shell.execute_reply":"2022-01-02T14:35:07.121343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# dataset_src_path = os.path.join(input_path, kaggle_split_dataset_name, \"trainA\", \"*\") # , 'trainA.zip'\n# path = glob.glob(dataset_src_path)[0]\n# print(path)\n# print(os.path.isdir(path))\n# print(kaggle_split_dataset_name)\n# print(kaggle_output_path)\n\ncopy_split_folders(input_path, kaggle_output_path)","metadata":{"execution":{"iopub.status.busy":"2022-01-02T14:35:08.519564Z","iopub.execute_input":"2022-01-02T14:35:08.519828Z","iopub.status.idle":"2022-01-02T14:35:28.921729Z","shell.execute_reply.started":"2022-01-02T14:35:08.519799Z","shell.execute_reply":"2022-01-02T14:35:28.92095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Print dataset split sizes","metadata":{}},{"cell_type":"code","source":"path=os.path.join(kaggle_output_path, \"trainA\", \"*.jpg\")\nprint(\"len(trainA)={}\".format(len(glob.glob(path))))\n# path=\"/kaggle/working/valA/*.jpg\"\npath=os.path.join(kaggle_output_path, \"valA\", \"*.jpg\")\nprint(\"len(valA)={}\".format(len(glob.glob(path))))\n# path=\"/kaggle/working/testA/*.jpg\"\npath=os.path.join(kaggle_output_path, \"testA\", \"*.jpg\")\nnum_test_sample=len(glob.glob(path))\nprint(\"len(testA)={}\".format(num_test_sample))\n\n# path=\"/kaggle/working/trainB/*.jpg\"\npath=os.path.join(kaggle_output_path, \"trainB\", \"*.jpg\")\nprint(\"len(trainB)={}\".format(len(glob.glob(path))))\n# path=\"/kaggle/working/valB/*.jpg\"\npath=os.path.join(kaggle_output_path, \"valB\", \"*.jpg\")\nprint(\"len(valB)={}\".format(len(glob.glob(path))))\n# path=\"/kaggle/working/testB/*.jpg\"\npath=os.path.join(kaggle_output_path, \"testB\", \"*.jpg\")\nprint(\"len(testB)={}\".format(len(glob.glob(path))))\n","metadata":{"id":"y4W5AO1EsOwv","outputId":"2a68907f-e445-4ea1-d873-310709560e80","execution":{"iopub.status.busy":"2022-01-02T14:35:43.128795Z","iopub.execute_input":"2022-01-02T14:35:43.129075Z","iopub.status.idle":"2022-01-02T14:35:43.17179Z","shell.execute_reply.started":"2022-01-02T14:35:43.129024Z","shell.execute_reply":"2022-01-02T14:35:43.17111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Copying latest checkpoint files from Datasource path into path: /kaggle/working/checkpoints/\\<experiment_name\\>","metadata":{"id":"qVBWZlIhsOww"}},{"cell_type":"code","source":"## test\n\ndata_path = os.path.join(input_path, kaggle_split_dataset_name, 'trainA', 'trainA')\nchkpt_path = os.path.join(input_path, kaggle_cycleGAN_checkpoint_dataset_name)\n\nprint(data_path)\nprint(chkpt_path)\nprint(len(glob.glob(os.path.join(data_path, '*'))))\nprint(len(glob.glob(os.path.join(kaggle_cycleGAN_checkpoint_dataset_path, '*'))))","metadata":{"execution":{"iopub.status.busy":"2022-01-02T14:35:53.656761Z","iopub.execute_input":"2022-01-02T14:35:53.657016Z","iopub.status.idle":"2022-01-02T14:35:53.670188Z","shell.execute_reply.started":"2022-01-02T14:35:53.656986Z","shell.execute_reply":"2022-01-02T14:35:53.669271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# input_ckpt_folder = kaggle_cycleGAN_checkpoint_dataset_name\n# src_ckpt_folderpath = os.path.join(input_path, input_ckpt_folder)\n\nsrc_ckpt_folderpath = kaggle_cycleGAN_checkpoint_dataset_path\ndest_ckpt_folderpath = os.path.join(kaggle_output_path, \"checkpoints\", \"photo_to_monet\")\n\nif not os.path.isdir(os.path.join(kaggle_output_path, \"checkpoints\")):\n    os.mkdir(os.path.join(kaggle_output_path, \"checkpoints\"))\n    if not os.path.isdir(os.path.join(kaggle_output_path, \"checkpoints\", \"photo_to_monet\")):\n        os.mkdir(os.path.join(kaggle_output_path, \"checkpoints\", \"photo_to_monet\"))\n\nprint(glob.glob(input_path+os.sep+'*'))\nsrc_ckpt_filepath_arr = glob.glob(os.path.join(src_ckpt_folderpath, \"*\"))\n\nprint(src_ckpt_folderpath)\nprint(len(src_ckpt_filepath_arr))\n\nfor src_ckpt_file in src_ckpt_filepath_arr:\n    filename = os.path.basename(src_ckpt_file)\n    print(filename)\n    dest_ckpt_file = os.path.join(dest_ckpt_folderpath, filename)\n    shutil.copy(src_ckpt_file, dest_ckpt_file)","metadata":{"id":"HgADua6bsOww","execution":{"iopub.status.busy":"2022-01-02T14:36:03.303941Z","iopub.execute_input":"2022-01-02T14:36:03.304482Z","iopub.status.idle":"2022-01-02T14:36:05.810457Z","shell.execute_reply.started":"2022-01-02T14:36:03.304444Z","shell.execute_reply":"2022-01-02T14:36:05.809723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils (path: /util/util.py, /util/image_pool.py)","metadata":{"id":"gUS7mUWosOwx"}},{"cell_type":"code","source":"##=================================================\n## /util/util.py\n##=================================================\n\n\n\"\"\"This module contains simple helper functions \"\"\"\nfrom __future__ import print_function\nimport torch\nimport numpy as np\nfrom PIL import Image\nimport os\n\n\ndef tensor2im(input_image, imtype=np.uint8):\n    \"\"\"\"Converts a Tensor array into a numpy image array.\n\n    Parameters:\n        input_image (tensor) --  the input image tensor array\n        imtype (type)        --  the desired type of the converted numpy array\n    \"\"\"\n    if not isinstance(input_image, np.ndarray):\n        if isinstance(input_image, torch.Tensor):  # get the data from a variable\n            image_tensor = input_image.data\n        else:\n            return input_image\n        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array\n        if image_numpy.shape[0] == 1:  # grayscale to RGB\n            image_numpy = np.tile(image_numpy, (3, 1, 1))\n        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling\n    else:  # if it is a numpy array, do nothing\n        image_numpy = input_image\n    return image_numpy.astype(imtype)\n\n\ndef diagnose_network(net, name='network'):\n    \"\"\"Calculate and print the mean of average absolute(gradients)\n\n    Parameters:\n        net (torch network) -- Torch network\n        name (str) -- the name of the network\n    \"\"\"\n    mean = 0.0\n    count = 0\n    for param in net.parameters():\n        if param.grad is not None:\n            mean += torch.mean(torch.abs(param.grad.data))\n            count += 1\n    if count > 0:\n        mean = mean / count\n    print(name)\n    print(mean)\n\n\ndef save_image(image_numpy, image_path, aspect_ratio=1.0):\n    \"\"\"Save a numpy image to the disk\n\n    Parameters:\n        image_numpy (numpy array) -- input numpy array\n        image_path (str)          -- the path of the image\n    \"\"\"\n\n    image_pil = Image.fromarray(image_numpy)\n    h, w, _ = image_numpy.shape\n\n    if aspect_ratio > 1.0:\n        image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)\n    if aspect_ratio < 1.0:\n        image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)\n    image_pil.save(image_path)\n\n\ndef print_numpy(x, val=True, shp=False):\n    \"\"\"Print the mean, min, max, median, std, and size of a numpy array\n\n    Parameters:\n        val (bool) -- if print the values of the numpy array\n        shp (bool) -- if print the shape of the numpy array\n    \"\"\"\n    x = x.astype(np.float64)\n    if shp:\n        print('shape,', x.shape)\n    if val:\n        x = x.flatten()\n        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (\n            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))\n\n\ndef mkdirs(paths):\n    \"\"\"create empty directories if they don't exist\n\n    Parameters:\n        paths (str list) -- a list of directory paths\n    \"\"\"\n    if isinstance(paths, list) and not isinstance(paths, str):\n        for path in paths:\n            mkdir(path)\n    else:\n        mkdir(paths)\n\n\ndef mkdir(path):\n    \"\"\"create a single empty directory if it didn't exist\n\n    Parameters:\n        path (str) -- a single directory path\n    \"\"\"\n    if not os.path.exists(path):\n        os.makedirs(path)\n","metadata":{"id":"3a718jWRsOwx","execution":{"iopub.status.busy":"2022-01-02T14:36:12.619656Z","iopub.execute_input":"2022-01-02T14:36:12.619919Z","iopub.status.idle":"2022-01-02T14:36:12.637208Z","shell.execute_reply.started":"2022-01-02T14:36:12.61989Z","shell.execute_reply":"2022-01-02T14:36:12.636371Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /util/image_pool.py\n##=================================================\n\n\nimport random\nimport torch\n\n\nclass ImagePool():\n    \"\"\"This class implements an image buffer that stores previously generated images.\n\n    This buffer enables us to update discriminators using a history of generated images\n    rather than the ones produced by the latest generators.\n    \"\"\"\n\n    def __init__(self, pool_size):\n        \"\"\"Initialize the ImagePool class\n\n        Parameters:\n            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created\n        \"\"\"\n        self.pool_size = pool_size\n        if self.pool_size > 0:  # create an empty pool\n            self.num_imgs = 0\n            self.images = []\n\n    def query(self, images):\n        \"\"\"Return an image from the pool.\n\n        Parameters:\n            images: the latest generated images from the generator\n\n        Returns images from the buffer.\n\n        By 50/100, the buffer will return input images.\n        By 50/100, the buffer will return images previously stored in the buffer,\n        and insert the current images to the buffer.\n        \"\"\"\n        if self.pool_size == 0:  # if the buffer size is 0, do nothing\n            return images\n        return_images = []\n        for image in images:\n            image = torch.unsqueeze(image.data, 0)\n            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer\n                self.num_imgs = self.num_imgs + 1\n                self.images.append(image)\n                return_images.append(image)\n            else:\n                p = random.uniform(0, 1)\n                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer\n                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive\n                    tmp = self.images[random_id].clone()\n                    self.images[random_id] = image\n                    return_images.append(tmp)\n                else:       # by another 50% chance, the buffer will return the current image\n                    return_images.append(image)\n        return_images = torch.cat(return_images, 0)   # collect all the images and return\n        return return_images\n","metadata":{"id":"4FSR0V1LsOwy","execution":{"iopub.status.busy":"2022-01-02T14:36:14.023462Z","iopub.execute_input":"2022-01-02T14:36:14.023851Z","iopub.status.idle":"2022-01-02T14:36:14.0334Z","shell.execute_reply.started":"2022-01-02T14:36:14.023815Z","shell.execute_reply":"2022-01-02T14:36:14.032675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config (path: /configs/base_config.py, /config/train_config.py)","metadata":{"id":"AyrMjbuVsOwz"}},{"cell_type":"code","source":"##=================================================\n## /configs/base_config.py\n##=================================================\n\nimport argparse\nimport os\n# from util import util\nimport torch\n# import models\n# import data\n\n\nclass BaseConfigs():\n    \"\"\"This class defines configs used during both training and test time.\n\n    It also implements several helper functions such as parsing, printing, and saving the configs.\n    It also gathers additional configs defined in <modify_commandline_configs> functions in both dataset class and model class.\n    \"\"\"\n\n    def __init__(self):\n        \"\"\"Reset the class; indicates the class hasn't been initailized\"\"\"\n        self.initialized = False\n\n    def initialize(self, parser):\n        \"\"\"Define the common configs that are used in both training and test.\"\"\"\n        # basic parameters\n        parser.add_argument('--dataroot', type=str, default=kaggle_output_path, help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')\n        parser.add_argument('--name', type=str, default='photo_to_monet', help='name of the experiment. It decides where to store samples and models')\n        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')\n        parser.add_argument('--checkpoints_dir', type=str, default=checkpoint_path, help='models are saved here')\n        # model parameters\n        parser.add_argument('--model', type=str, default='cycle_gan', help='chooses which model to use. [cycle_gan | pix2pix | test | colorization]')\n        parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')\n        parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels: 3 for RGB and 1 for grayscale')\n        parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')\n        parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')\n        parser.add_argument('--netD', type=str, default='basic', help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')\n        parser.add_argument('--netG', type=str, default='resnet_9blocks', help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')\n        parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')\n        parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')\n        parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')\n        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')\n        parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')\n        # dataset parameters\n        parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]')\n        parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA')\n        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')\n        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')\n        parser.add_argument('--batch_size', type=int, default=1, help='input batch size')\n        parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')\n        parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')\n        parser.add_argument('--max_dataset_size', type=int, default=float(\"inf\"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')\n        parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')\n        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')\n        parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML')\n        # additional parameters\n        parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')\n        parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')\n        parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')\n        parser.add_argument('--suffix', default='', type=str, help='customized suffix: config.name = config.name + suffix: e.g., {model}_{netG}_size{load_size}')\n        self.initialized = True\n        return parser\n\n    def gather_configs(self):\n        \"\"\"Initialize our parser with basic configs(only once).\n        Add additional model-specific and dataset-specific configs.\n        These configs are defined in the <modify_commandline_configs> function\n        in model and dataset classes.\n        \"\"\"\n        if not self.initialized:  # check if it has been initialized\n            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n            parser = self.initialize(parser)\n\n        # get the basic configs\n        config, _ = parser.parse_known_args()\n\n        # modify model-related parser configs\n        model_name = config.model\n        # model_config_setter = models.get_config_setter(model_name)\n        # model_config_setter = get_model_config_setter(model_name)\n        # parser = model_config_setter(parser, self.isTrain)\n        parser = CycleGANModel.modify_commandline_configs(parser, self.isTrain)\n        config, _ = parser.parse_known_args()  # parse again with new defaults\n\n        # modify dataset-related parser configs\n        dataset_name = config.dataset_mode\n        # dataset_config_setter = data.get_config_setter(dataset_name)\n        # dataset_config_setter = get_data_config_setter(dataset_name)\n        # parser = dataset_config_setter(parser, self.isTrain)\n        parser = UnalignedDataset.modify_commandline_configs(parser, self.isTrain)\n\n        # save and return the parser\n        self.parser = parser\n        print(parser)\n        return parser.parse_args([])\n\n    def print_configs(self, config):\n        \"\"\"Print and save configs\n\n        It will print both current configs and default values(if different).\n        It will save configs into a text file / [checkpoints_dir] / config.txt\n        \"\"\"\n        message = ''\n        message += '----------------- Configs ---------------\\n'\n        for k, v in sorted(vars(config).items()):\n            comment = ''\n            default = self.parser.get_default(k)\n            if v != default:\n                comment = '\\t[default: %s]' % str(default)\n            message += '{:>25}: {:<30}{}\\n'.format(str(k), str(v), comment)\n        message += '----------------- End -------------------'\n        print(message)\n\n        # save to the disk\n        expr_dir = os.path.join(config.checkpoints_dir, config.name)\n        mkdirs(expr_dir)\n        file_name = os.path.join(expr_dir, '{}_config.txt'.format(config.phase))\n        with open(file_name, 'wt') as config_file:\n            config_file.write(message)\n            config_file.write('\\n')\n\n    def parse(self):\n        \"\"\"Parse our configs, create checkpoints directory suffix, and set up gpu device.\"\"\"\n        config = self.gather_configs()\n        config.isTrain = self.isTrain   # train or test\n\n        # process config.suffix\n        if config.suffix:\n            suffix = ('_' + config.suffix.format(**vars(config))) if config.suffix != '' else ''\n            config.name = config.name + suffix\n\n        self.print_configs(config)\n\n        # set gpu ids\n        # if torch.cuda.is_available() == True:\n        str_ids = config.gpu_ids.split(',')\n        config.gpu_ids = []\n        for str_id in str_ids:\n            id = int(str_id)\n            if id >= 0:\n                config.gpu_ids.append(id)\n        if len(config.gpu_ids) > 0:\n            torch.cuda.set_device(config.gpu_ids[0])\n        # else:\n        #     pass\n\n        self.config = config\n        return self.config\n","metadata":{"id":"0HE2NVufsOwz","execution":{"iopub.status.busy":"2022-01-02T14:36:15.438889Z","iopub.execute_input":"2022-01-02T14:36:15.439427Z","iopub.status.idle":"2022-01-02T14:36:15.46779Z","shell.execute_reply.started":"2022-01-02T14:36:15.439385Z","shell.execute_reply":"2022-01-02T14:36:15.46708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /config/train_config.py\n##=================================================\n\n# from .base_config import BaseConfigs\n\n\nclass TrainConfigs(BaseConfigs):\n    \"\"\"This class includes training configs.\n\n    It also includes shared configs defined in BaseConfigs.\n    \"\"\"\n\n    def initialize(self, parser):\n        parser = BaseConfigs.initialize(self, parser)\n        # visdom and HTML visualization parameters\n        parser.add_argument('--display_freq', type=int, default=400, help='frequency of showing training results on screen')\n        parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.')\n        parser.add_argument('--display_id', type=int, default=0, help='window id of the web display')\n        parser.add_argument('--display_server', type=str, default=\"http://localhost\", help='visdom server of the web display')\n        parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is \"main\")')\n        parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')\n        parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')\n        parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')\n        parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [config.checkpoints_dir]/[config.name]/web/')\n        # network saving and loading parameters\n        parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results') # 5000\n        parser.add_argument('--save_epoch_freq', type=int, default=1, help='frequency of saving checkpoints at the end of epochs') # 100\n        parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration')\n        parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')\n        parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...') # 1, LAST_EPOCH+1\n        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')\n        # training parameters\n        parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs with the initial learning rate') # 100\n        parser.add_argument('--n_epochs_decay', type=int, default=100, help='number of epochs to linearly decay learning rate to zero') # 100\n        parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')\n        parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')\n        parser.add_argument('--gan_mode', type=str, default='lsgan', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')\n        parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')\n        parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]')\n        parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')\n\n        self.isTrain = True\n        return parser\n","metadata":{"id":"DU3h5pcRsOw0","execution":{"iopub.status.busy":"2022-01-02T14:36:16.430554Z","iopub.execute_input":"2022-01-02T14:36:16.431144Z","iopub.status.idle":"2022-01-02T14:36:16.445632Z","shell.execute_reply.started":"2022-01-02T14:36:16.43111Z","shell.execute_reply":"2022-01-02T14:36:16.444853Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CycleGAN Model (path: /models/\\_\\_init\\_\\_.py, /models/networks.py, /models/base_model.py, /models/cycle_gan_model.py)","metadata":{"id":"7TLynN0FsOw0"}},{"cell_type":"code","source":"##=================================================\n## /models/__init__.py\n##=================================================\n\n\"\"\"This package contains modules related to objective functions, optimizations, and network architectures.\n\nTo add a custom model class called 'dummy', you need to add a file called 'dummy_model.py' and define a subclass DummyModel inherited from BaseModel.\nYou need to implement the following five functions:\n    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, config).\n    -- <set_input>:                     unpack data from dataset and apply preprocessing.\n    -- <forward>:                       produce intermediate results.\n    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.\n    -- <modify_commandline_configs>:    (optionally) add model-specific configs and set default configs.\n\nIn the function <__init__>, you need to define four lists:\n    -- self.loss_names (str list):          specify the training losses that you want to plot and save.\n    -- self.model_names (str list):         define networks used in our training.\n    -- self.visual_names (str list):        specify the images that you want to display and save.\n    -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an usage.\n\nNow you can use the model class by specifying flag '--model dummy'.\nSee our template model class 'template_model.py' for more details.\n\"\"\"\n\nimport importlib\n# from models.base_model import BaseModel\n\n\ndef find_model_using_name(model_name):\n    \"\"\"Import the module \"models/[model_name]_model.py\".\n\n    In the file, the class called DatasetNameModel() will\n    be instantiated. It has to be a subclass of BaseModel,\n    and it is case-insensitive.\n    \"\"\"\n    # model_filename = \"models.\" + model_name + \"_model\"\n    model_filename =  model_name + \"_model\"\n    modellib = importlib.import_module(model_filename)\n    model = None\n    target_model_name = model_name.replace('_', '') + 'model'\n    for name, cls in modellib.__dict__.items():\n        if name.lower() == target_model_name.lower() \\\n           and issubclass(cls, BaseModel):\n            model = cls\n\n    if model is None:\n        print(\"In %s.py, there should be a subclass of BaseModel with class name that matches %s in lowercase.\" % (model_filename, target_model_name))\n        exit(0)\n\n    return model\n\n\ndef get_model_config_setter(model_name):\n    \"\"\"Return the static method <modify_commandline_configs> of the model class.\"\"\"\n    # model_class = find_model_using_name(model_name)\n    model_class = CycleGANModel\n    return model_class.modify_commandline_configs\n\n\ndef create_model(config):\n    \"\"\"Create a model given the config.\n\n    This function warps the class CustomDatasetDataLoader.\n    This is the main interface between this package and 'train.py'/'test.py'\n    \"\"\"\n    model = find_model_using_name(config.model)\n    instance = model(config)\n    print(\"model [%s] was created\" % type(instance).__name__)\n    return instance\n","metadata":{"id":"yQnEHgdOsOw1","execution":{"iopub.status.busy":"2022-01-02T14:36:18.145822Z","iopub.execute_input":"2022-01-02T14:36:18.146145Z","iopub.status.idle":"2022-01-02T14:36:18.167976Z","shell.execute_reply.started":"2022-01-02T14:36:18.146109Z","shell.execute_reply":"2022-01-02T14:36:18.167301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /models/networks.py\n##=================================================\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport functools\nfrom torch.optim import lr_scheduler\n\n\n###############################################################################\n# Helper Functions\n###############################################################################\n\n\nclass Identity(nn.Module):\n    def forward(self, x):\n        return x\n\n\ndef get_norm_layer(norm_type='instance'):\n    \"\"\"Return a normalization layer\n\n    Parameters:\n        norm_type (str) -- the name of the normalization layer: batch | instance | none\n\n    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).\n    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.\n    \"\"\"\n    if norm_type == 'batch':\n        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)\n    elif norm_type == 'instance':\n        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)\n    elif norm_type == 'none':\n        def norm_layer(x): return Identity()\n    else:\n        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)\n    return norm_layer\n\n\ndef get_scheduler(optimizer, config):\n    \"\"\"Return a learning rate scheduler\n\n    Parameters:\n        optimizer          -- the optimizer of the network\n        config (config class) -- stores all the experiment flags; needs to be a subclass of BaseConfigs．　\n                              config.lr_policy is the name of learning rate policy: linear | step | plateau | cosine\n\n    For 'linear', we keep the same learning rate for the first <config.n_epochs> epochs\n    and linearly decay the rate to zero over the next <config.n_epochs_decay> epochs.\n    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.\n    See https://pytorch.org/docs/stable/optim.html for more details.\n    \"\"\"\n    if config.lr_policy == 'linear':\n        def lambda_rule(epoch):\n            lr_l = 1.0 - max(0, epoch + config.epoch_count - config.n_epochs) / float(config.n_epochs_decay + 1)\n            return lr_l\n        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)\n    elif config.lr_policy == 'step':\n        scheduler = lr_scheduler.StepLR(optimizer, step_size=config.lr_decay_iters, gamma=0.1)\n    elif config.lr_policy == 'plateau':\n        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)\n    elif config.lr_policy == 'cosine':\n        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.n_epochs, eta_min=0)\n    else:\n        return NotImplementedError('learning rate policy [%s] is not implemented', config.lr_policy)\n    return scheduler\n\n\ndef init_weights(net, init_type='normal', init_gain=0.02):\n    \"\"\"Initialize network weights.\n\n    Parameters:\n        net (network)   -- network to be initialized\n        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal\n        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.\n\n    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might\n    work better for some applications. Feel free to try yourself.\n    \"\"\"\n    def init_func(m):  # define the initialization function\n        classname = m.__class__.__name__\n        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n            if init_type == 'normal':\n                init.normal_(m.weight.data, 0.0, init_gain)\n            elif init_type == 'xavier':\n                init.xavier_normal_(m.weight.data, gain=init_gain)\n            elif init_type == 'kaiming':\n                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n            elif init_type == 'orthogonal':\n                init.orthogonal_(m.weight.data, gain=init_gain)\n            else:\n                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)\n            if hasattr(m, 'bias') and m.bias is not None:\n                init.constant_(m.bias.data, 0.0)\n        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.\n            init.normal_(m.weight.data, 1.0, init_gain)\n            init.constant_(m.bias.data, 0.0)\n\n    print('initialize network with %s' % init_type)\n    net.apply(init_func)  # apply the initialization function <init_func>\n\n\ndef init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):\n    \"\"\"Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights\n    Parameters:\n        net (network)      -- the network to be initialized\n        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal\n        gain (float)       -- scaling factor for normal, xavier and orthogonal.\n        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2\n\n    Return an initialized network.\n    \"\"\"\n    if len(gpu_ids) > 0:\n        assert(torch.cuda.is_available())\n        net.to(gpu_ids[0])\n        net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs\n    init_weights(net, init_type, init_gain=init_gain)\n    return net\n\n\ndef define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):\n    \"\"\"Create a generator\n\n    Parameters:\n        input_nc (int) -- the number of channels in input images\n        output_nc (int) -- the number of channels in output images\n        ngf (int) -- the number of filters in the last conv layer\n        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128\n        norm (str) -- the name of normalization layers used in the network: batch | instance | none\n        use_dropout (bool) -- if use dropout layers.\n        init_type (str)    -- the name of our initialization method.\n        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.\n        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2\n\n    Returns a generator\n\n    Our current implementation provides two types of generators:\n        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)\n        The original U-Net paper: https://arxiv.org/abs/1505.04597\n\n        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)\n        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.\n        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).\n\n\n    The generator has been initialized by <init_net>. It uses RELU for non-linearity.\n    \"\"\"\n    net = None\n    norm_layer = get_norm_layer(norm_type=norm)\n\n    if netG == 'resnet_9blocks':\n        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)\n    elif netG == 'resnet_6blocks':\n        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)\n    elif netG == 'unet_128':\n        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)\n    elif netG == 'unet_256':\n        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)\n    else:\n        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)\n    return init_net(net, init_type, init_gain, gpu_ids)\n\n\ndef define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):\n    \"\"\"Create a discriminator\n\n    Parameters:\n        input_nc (int)     -- the number of channels in input images\n        ndf (int)          -- the number of filters in the first conv layer\n        netD (str)         -- the architecture's name: basic | n_layers | pixel\n        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'\n        norm (str)         -- the type of normalization layers used in the network.\n        init_type (str)    -- the name of the initialization method.\n        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.\n        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2\n\n    Returns a discriminator\n\n    Our current implementation provides three types of discriminators:\n        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.\n        It can classify whether 70×70 overlapping patches are real or fake.\n        Such a patch-level discriminator architecture has fewer parameters\n        than a full-image discriminator and can work on arbitrarily-sized images\n        in a fully convolutional fashion.\n\n        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator\n        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)\n\n        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.\n        It encourages greater color diversity but has no effect on spatial statistics.\n\n    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.\n    \"\"\"\n    net = None\n    norm_layer = get_norm_layer(norm_type=norm)\n\n    if netD == 'basic':  # default PatchGAN classifier\n        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)\n    elif netD == 'n_layers':  # more configs\n        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)\n    elif netD == 'pixel':     # classify if each pixel is real or fake\n        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)\n    else:\n        raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)\n    return init_net(net, init_type, init_gain, gpu_ids)\n\n\n##############################################################################\n# Classes\n##############################################################################\nclass GANLoss(nn.Module):\n    \"\"\"Define different GAN objectives.\n\n    The GANLoss class abstracts away the need to create the target label tensor\n    that has the same size as the input.\n    \"\"\"\n\n    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):\n        \"\"\" Initialize the GANLoss class.\n\n        Parameters:\n            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.\n            target_real_label (bool) - - label for a real image\n            target_fake_label (bool) - - label of a fake image\n\n        Note: Do not use sigmoid as the last layer of Discriminator.\n        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.\n        \"\"\"\n        super(GANLoss, self).__init__()\n        self.register_buffer('real_label', torch.tensor(target_real_label))\n        self.register_buffer('fake_label', torch.tensor(target_fake_label))\n        self.gan_mode = gan_mode\n        if gan_mode == 'lsgan':\n            self.loss = nn.MSELoss()\n        elif gan_mode == 'vanilla':\n            self.loss = nn.BCEWithLogitsLoss()\n        elif gan_mode in ['wgangp']:\n            self.loss = None\n        else:\n            raise NotImplementedError('gan mode %s not implemented' % gan_mode)\n\n    def get_target_tensor(self, prediction, target_is_real):\n        \"\"\"Create label tensors with the same size as the input.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            A label tensor filled with ground truth label, and with the size of the input\n        \"\"\"\n\n        if target_is_real:\n            target_tensor = self.real_label\n        else:\n            target_tensor = self.fake_label\n        return target_tensor.expand_as(prediction)\n\n    def __call__(self, prediction, target_is_real):\n        \"\"\"Calculate loss given Discriminator's output and grount truth labels.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction output from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            the calculated loss.\n        \"\"\"\n        if self.gan_mode in ['lsgan', 'vanilla']:\n            target_tensor = self.get_target_tensor(prediction, target_is_real)\n            loss = self.loss(prediction, target_tensor)\n        elif self.gan_mode == 'wgangp':\n            if target_is_real:\n                loss = -prediction.mean()\n            else:\n                loss = prediction.mean()\n        return loss\n\n\ndef cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):\n    \"\"\"Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028\n\n    Arguments:\n        netD (network)              -- discriminator network\n        real_data (tensor array)    -- real images\n        fake_data (tensor array)    -- generated images from the generator\n        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')\n        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].\n        constant (float)            -- the constant used in formula ( ||gradient||_2 - constant)^2\n        lambda_gp (float)           -- weight for this loss\n\n    Returns the gradient penalty loss\n    \"\"\"\n    if lambda_gp > 0.0:\n        if type == 'real':   # either use real images, fake images, or a linear interpolation of two.\n            interpolatesv = real_data\n        elif type == 'fake':\n            interpolatesv = fake_data\n        elif type == 'mixed':\n            alpha = torch.rand(real_data.shape[0], 1, device=device)\n            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)\n            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)\n        else:\n            raise NotImplementedError('{} not implemented'.format(type))\n        interpolatesv.requires_grad_(True)\n        disc_interpolates = netD(interpolatesv)\n        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,\n                                        grad_outputs=torch.ones(disc_interpolates.size()).to(device),\n                                        create_graph=True, retain_graph=True, only_inputs=True)\n        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data\n        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp        # added eps\n        return gradient_penalty, gradients\n    else:\n        return 0.0, None\n\n\nclass ResnetGenerator(nn.Module):\n    \"\"\"Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.\n\n    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)\n    \"\"\"\n\n    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):\n        \"\"\"Construct a Resnet-based generator\n\n        Parameters:\n            input_nc (int)      -- the number of channels in input images\n            output_nc (int)     -- the number of channels in output images\n            ngf (int)           -- the number of filters in the last conv layer\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers\n            n_blocks (int)      -- the number of ResNet blocks\n            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero\n        \"\"\"\n        assert(n_blocks >= 0)\n        super(ResnetGenerator, self).__init__()\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        model = [nn.ReflectionPad2d(3),\n                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),\n                 norm_layer(ngf),\n                 nn.ReLU(True)]\n\n        n_downsampling = 2\n        for i in range(n_downsampling):  # add downsampling layers\n            mult = 2 ** i\n            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),\n                      norm_layer(ngf * mult * 2),\n                      nn.ReLU(True)]\n\n        mult = 2 ** n_downsampling\n        for i in range(n_blocks):       # add ResNet blocks\n\n            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]\n\n        for i in range(n_downsampling):  # add upsampling layers\n            mult = 2 ** (n_downsampling - i)\n            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),\n                                         kernel_size=3, stride=2,\n                                         padding=1, output_padding=1,\n                                         bias=use_bias),\n                      norm_layer(int(ngf * mult / 2)),\n                      nn.ReLU(True)]\n        model += [nn.ReflectionPad2d(3)]\n        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]\n        model += [nn.Tanh()]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input):\n        \"\"\"Standard forward\"\"\"\n        return self.model(input)\n\n\nclass ResnetBlock(nn.Module):\n    \"\"\"Define a Resnet block\"\"\"\n\n    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n        \"\"\"Initialize the Resnet block\n\n        A resnet block is a conv block with skip connections\n        We construct a conv block with build_conv_block function,\n        and implement skip connections in <forward> function.\n        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf\n        \"\"\"\n        super(ResnetBlock, self).__init__()\n        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)\n\n    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n        \"\"\"Construct a convolutional block.\n\n        Parameters:\n            dim (int)           -- the number of channels in the conv layer.\n            padding_type (str)  -- the name of padding layer: reflect | replicate | zero\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers.\n            use_bias (bool)     -- if the conv layer uses bias or not\n\n        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))\n        \"\"\"\n        conv_block = []\n        p = 0\n        if padding_type == 'reflect':\n            conv_block += [nn.ReflectionPad2d(1)]\n        elif padding_type == 'replicate':\n            conv_block += [nn.ReplicationPad2d(1)]\n        elif padding_type == 'zero':\n            p = 1\n        else:\n            raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n\n        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]\n        if use_dropout:\n            conv_block += [nn.Dropout(0.5)]\n\n        p = 0\n        if padding_type == 'reflect':\n            conv_block += [nn.ReflectionPad2d(1)]\n        elif padding_type == 'replicate':\n            conv_block += [nn.ReplicationPad2d(1)]\n        elif padding_type == 'zero':\n            p = 1\n        else:\n            raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]\n\n        return nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        \"\"\"Forward function (with skip connections)\"\"\"\n        out = x + self.conv_block(x)  # add skip connections\n        return out\n\n\nclass UnetGenerator(nn.Module):\n    \"\"\"Create a Unet-based generator\"\"\"\n\n    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):\n        \"\"\"Construct a Unet generator\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            output_nc (int) -- the number of channels in output images\n            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,\n                                image of size 128x128 will become of size 1x1 # at the bottleneck\n            ngf (int)       -- the number of filters in the last conv layer\n            norm_layer      -- normalization layer\n\n        We construct the U-Net from the innermost layer to the outermost layer.\n        It is a recursive process.\n        \"\"\"\n        super(UnetGenerator, self).__init__()\n        # construct unet structure\n        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer\n        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters\n            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)\n        # gradually reduce the number of filters from ngf * 8 to ngf\n        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer\n\n    def forward(self, input):\n        \"\"\"Standard forward\"\"\"\n        return self.model(input)\n\n\nclass UnetSkipConnectionBlock(nn.Module):\n    \"\"\"Defines the Unet submodule with skip connection.\n        X -------------------identity----------------------\n        |-- downsampling -- |submodule| -- upsampling --|\n    \"\"\"\n\n    def __init__(self, outer_nc, inner_nc, input_nc=None,\n                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):\n        \"\"\"Construct a Unet submodule with skip connections.\n\n        Parameters:\n            outer_nc (int) -- the number of filters in the outer conv layer\n            inner_nc (int) -- the number of filters in the inner conv layer\n            input_nc (int) -- the number of channels in input images/features\n            submodule (UnetSkipConnectionBlock) -- previously defined submodules\n            outermost (bool)    -- if this module is the outermost module\n            innermost (bool)    -- if this module is the innermost module\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers.\n        \"\"\"\n        super(UnetSkipConnectionBlock, self).__init__()\n        self.outermost = outermost\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n        if input_nc is None:\n            input_nc = outer_nc\n        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,\n                             stride=2, padding=1, bias=use_bias)\n        downrelu = nn.LeakyReLU(0.2, True)\n        downnorm = norm_layer(inner_nc)\n        uprelu = nn.ReLU(True)\n        upnorm = norm_layer(outer_nc)\n\n        if outermost:\n            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1)\n            down = [downconv]\n            up = [uprelu, upconv, nn.Tanh()]\n            model = down + [submodule] + up\n        elif innermost:\n            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1, bias=use_bias)\n            down = [downrelu, downconv]\n            up = [uprelu, upconv, upnorm]\n            model = down + up\n        else:\n            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1, bias=use_bias)\n            down = [downrelu, downconv, downnorm]\n            up = [uprelu, upconv, upnorm]\n\n            if use_dropout:\n                model = down + [submodule] + up + [nn.Dropout(0.5)]\n            else:\n                model = down + [submodule] + up\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        if self.outermost:\n            return self.model(x)\n        else:   # add skip connections\n            return torch.cat([x, self.model(x)], 1)\n\n\nclass NLayerDiscriminator(nn.Module):\n    \"\"\"Defines a PatchGAN discriminator\"\"\"\n\n    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):\n        \"\"\"Construct a PatchGAN discriminator\n\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            ndf (int)       -- the number of filters in the last conv layer\n            n_layers (int)  -- the number of conv layers in the discriminator\n            norm_layer      -- normalization layer\n        \"\"\"\n        super(NLayerDiscriminator, self).__init__()\n        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        kw = 4\n        padw = 1\n        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]\n        nf_mult = 1\n        nf_mult_prev = 1\n        for n in range(1, n_layers):  # gradually increase the number of filters\n            nf_mult_prev = nf_mult\n            nf_mult = min(2 ** n, 8)\n            sequence += [\n                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),\n                norm_layer(ndf * nf_mult),\n                nn.LeakyReLU(0.2, True)\n            ]\n\n        nf_mult_prev = nf_mult\n        nf_mult = min(2 ** n_layers, 8)\n        sequence += [\n            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n            norm_layer(ndf * nf_mult),\n            nn.LeakyReLU(0.2, True)\n        ]\n\n        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map\n        self.model = nn.Sequential(*sequence)\n\n    def forward(self, input):\n        \"\"\"Standard forward.\"\"\"\n        return self.model(input)\n\n\nclass PixelDiscriminator(nn.Module):\n    \"\"\"Defines a 1x1 PatchGAN discriminator (pixelGAN)\"\"\"\n\n    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):\n        \"\"\"Construct a 1x1 PatchGAN discriminator\n\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            ndf (int)       -- the number of filters in the last conv layer\n            norm_layer      -- normalization layer\n        \"\"\"\n        super(PixelDiscriminator, self).__init__()\n        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        self.net = [\n            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),\n            nn.LeakyReLU(0.2, True),\n            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),\n            norm_layer(ndf * 2),\n            nn.LeakyReLU(0.2, True),\n            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]\n\n        self.net = nn.Sequential(*self.net)\n\n    def forward(self, input):\n        \"\"\"Standard forward.\"\"\"\n        return self.net(input)\n","metadata":{"id":"FsNjQ8gUsOw1","execution":{"iopub.status.busy":"2022-01-02T14:36:19.290748Z","iopub.execute_input":"2022-01-02T14:36:19.291004Z","iopub.status.idle":"2022-01-02T14:36:19.380346Z","shell.execute_reply.started":"2022-01-02T14:36:19.290974Z","shell.execute_reply":"2022-01-02T14:36:19.379626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /models/base_model.py\n##=================================================\n\nimport os\nimport torch\nfrom collections import OrderedDict\nfrom abc import ABC, abstractmethod\n# from . import networks\nfrom time import time\n\n\nclass BaseModel(ABC):\n    \"\"\"This class is an abstract base class (ABC) for models.\n    To create a subclass, you need to implement the following five functions:\n        -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, config).\n        -- <set_input>:                     unpack data from dataset and apply preprocessing.\n        -- <forward>:                       produce intermediate results.\n        -- <optimize_parameters>:           calculate losses, gradients, and update network weights.\n        -- <modify_commandline_configs>:    (optionally) add model-specific configs and set default configs.\n    \"\"\"\n\n    def __init__(self, config):\n        \"\"\"Initialize the BaseModel class.\n\n        Parameters:\n            config (Config class)-- stores all the experiment flags; needs to be a subclass of BaseConfigs\n\n        When creating your custom class, you need to implement your own initialization.\n        In this function, you should first call <BaseModel.__init__(self, config)>\n        Then, you need to define four lists:\n            -- self.loss_names (str list):          specify the training losses that you want to plot and save.\n            -- self.model_names (str list):         define networks used in our training.\n            -- self.visual_names (str list):        specify the images that you want to display and save.\n            -- self.optimizers (optimizer list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.\n            -- self.log_file\n        \"\"\"\n        self.config = config\n        self.gpu_ids = config.gpu_ids\n        self.isTrain = config.isTrain\n        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU\n        self.save_dir = os.path.join(config.checkpoints_dir, config.name)  # save all the checkpoints to save_dir\n        if config.preprocess != 'scale_width':  # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.\n            torch.backends.cudnn.benchmark = True\n        self.loss_names = []\n        self.model_names = []\n        self.visual_names = []\n        self.optimizers = []\n        self.image_paths = []\n        self.metric = 0  # used for learning rate policy 'plateau'\n        self.log_name = os.path.join(config.checkpoints_dir, config.name, 'loss_log.txt')\n        with open(self.log_name, \"a\") as log_file:\n            now = time.strftime(\"%c\")\n            log_file.write('================ Training Loss (%s) ================\\n' % now)\n\n    @staticmethod\n    def modify_commandline_configs(parser, is_train):\n        \"\"\"Add new model-specific configs, and rewrite default values for existing configs.\n\n        Parameters:\n            parser          -- original config parser\n            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific configs.\n\n        Returns:\n            the modified parser.\n        \"\"\"\n        return parser\n\n    @abstractmethod\n    def set_input(self, input):\n        \"\"\"Unpack input data from the dataloader and perform necessary pre-processing steps.\n\n        Parameters:\n            input (dict): includes the data itself and its metadata information.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def forward(self):\n        \"\"\"Run forward pass; called by both functions <optimize_parameters> and <test>.\"\"\"\n        pass\n\n    @abstractmethod\n    def forward_test_kaggle(self):\n        \"\"\"Run forward pass to run only netG_A; called by function <test>.\"\"\"\n        pass\n\n    @abstractmethod\n    def optimize_parameters(self):\n        \"\"\"Calculate losses, gradients, and update network weights; called in every training iteration\"\"\"\n        pass\n\n    def setup(self, config):\n        \"\"\"Load and print networks; create schedulers\n\n        Parameters:\n            config (Config class) -- stores all the experiment flags; needs to be a subclass of BaseConfigs\n        \"\"\"\n        if self.isTrain:\n            self.schedulers = [get_scheduler(optimizer, config) for optimizer in self.optimizers]\n        if not self.isTrain or config.continue_train:\n            load_suffix = 'iter_%d' % config.load_iter if config.load_iter > 0 else config.epoch\n            self.load_networks(load_suffix)\n        self.print_networks(config.verbose)\n\n    def eval(self):\n        \"\"\"Make models eval mode during test time\"\"\"\n        for name in self.model_names:\n            if isinstance(name, str):\n                net = getattr(self, 'net' + name)\n                net.eval()\n\n    def test(self):\n        \"\"\"Forward function used in test time.\n\n        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop\n        It also calls <compute_visuals> to produce additional visualization results\n        \"\"\"\n        with torch.no_grad():\n            # self.forward()\n            self.forward_test_kaggle()\n            self.compute_visuals()\n\n    def compute_visuals(self):\n        \"\"\"Calculate additional output images for visdom and HTML visualization\"\"\"\n        pass\n\n    def get_image_paths(self):\n        \"\"\" Return image paths that are used to load current data\"\"\"\n        return self.image_paths\n\n    def update_learning_rate(self):\n        \"\"\"Update learning rates for all the networks; called at the end of every epoch\"\"\"\n        old_lr = self.optimizers[0].param_groups[0]['lr']\n        for scheduler in self.schedulers:\n            if self.config.lr_policy == 'plateau':\n                scheduler.step(self.metric)\n            else:\n                scheduler.step()\n\n        lr = self.optimizers[0].param_groups[0]['lr']\n        print('learning rate %.7f -> %.7f' % (old_lr, lr))\n\n    def get_current_visuals(self):\n        \"\"\"Return visualization images. train.py will display these images with visdom, and save the images to a HTML\"\"\"\n        visual_ret = OrderedDict()\n        for name in self.visual_names:\n            if isinstance(name, str):\n                visual_ret[name] = getattr(self, name)\n        return visual_ret\n\n    def get_current_losses(self):\n        \"\"\"Return traning losses / errors. train.py will print out these errors on console, and save them to a file\"\"\"\n        errors_ret = OrderedDict()\n        for name in self.loss_names:\n            if isinstance(name, str):\n                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number\n        return errors_ret\n\n    def save_networks(self, epoch):\n        \"\"\"Save all the networks to the disk.\n\n        Parameters:\n            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)\n        \"\"\"\n        for name in self.model_names:\n            if isinstance(name, str):\n                save_filename = '%s_net_%s.pth' % (epoch, name)\n                save_path = os.path.join(self.save_dir, save_filename)\n                net = getattr(self, 'net' + name)\n\n                if len(self.gpu_ids) > 0 and torch.cuda.is_available():\n                    torch.save(net.module.cpu().state_dict(), save_path)\n                    net.cuda(self.gpu_ids[0])\n                else:\n                    torch.save(net.cpu().state_dict(), save_path)\n\n    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):\n        \"\"\"Fix InstanceNorm checkpoints incompatibility (prior to 0.4)\"\"\"\n        key = keys[i]\n        if i + 1 == len(keys):  # at the end, pointing to a parameter/buffer\n            if module.__class__.__name__.startswith('InstanceNorm') and \\\n                    (key == 'running_mean' or key == 'running_var'):\n                if getattr(module, key) is None:\n                    state_dict.pop('.'.join(keys))\n            if module.__class__.__name__.startswith('InstanceNorm') and \\\n               (key == 'num_batches_tracked'):\n                state_dict.pop('.'.join(keys))\n        else:\n            self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)\n\n    def load_networks(self, epoch):\n        \"\"\"Load all the networks from the disk.\n\n        Parameters:\n            epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)\n        \"\"\"\n        for name in self.model_names:\n            if isinstance(name, str):\n                load_filename = '%s_net_%s.pth' % (epoch, name)\n                load_path = os.path.join(self.save_dir, load_filename)\n                net = getattr(self, 'net' + name)\n                if isinstance(net, torch.nn.DataParallel):\n                    net = net.module\n                print('loading the model from %s' % load_path)\n                # if you are using PyTorch newer than 0.4 (e.g., built from\n                # GitHub source), you can remove str() on self.device\n                state_dict = torch.load(load_path, map_location=str(self.device))\n                if hasattr(state_dict, '_metadata'):\n                    del state_dict._metadata\n\n                # patch InstanceNorm checkpoints prior to 0.4\n                for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop\n                    self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))\n                net.load_state_dict(state_dict)\n\n    def print_networks(self, verbose):\n        \"\"\"Print the total number of parameters in the network and (if verbose) network architecture\n\n        Parameters:\n            verbose (bool) -- if verbose: print the network architecture\n        \"\"\"\n        print('---------- Networks initialized -------------')\n        for name in self.model_names:\n            if isinstance(name, str):\n                net = getattr(self, 'net' + name)\n                num_params = 0\n                for param in net.parameters():\n                    num_params += param.numel()\n                if verbose:\n                    print(net)\n                print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))\n        print('-----------------------------------------------')\n\n    def set_requires_grad(self, nets, requires_grad=False):\n        \"\"\"Set requies_grad=Fasle for all the networks to avoid unnecessary computations\n        Parameters:\n            nets (network list)   -- a list of networks\n            requires_grad (bool)  -- whether the networks require gradients or not\n        \"\"\"\n        if not isinstance(nets, list):\n            nets = [nets]\n        for net in nets:\n            if net is not None:\n                for param in net.parameters():\n                    param.requires_grad = requires_grad\n\n        # losses: same format as |losses| of plot_current_losses\n    def print_current_losses(self, epoch, iters, losses, t_comp, t_data):\n        \"\"\"print current losses on console; also save the losses to the disk\n\n        Parameters:\n            epoch (int) -- current epoch\n            iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)\n            losses (OrderedDict) -- training losses stored in the format of (name, float) pairs\n            t_comp (float) -- computational time per data point (normalized by batch_size)\n            t_data (float) -- data loading time per data point (normalized by batch_size)\n        \"\"\"\n        message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, iters, t_comp, t_data)\n        for k, v in losses.items():\n            message += '%s: %.3f ' % (k, v)\n\n        print(message)  # print the message\n        with open(self.log_name, \"a\") as log_file:\n            log_file.write('%s\\n' % message)  # save the message","metadata":{"id":"pQegwMHIsOw4","execution":{"iopub.status.busy":"2022-01-02T14:36:19.977459Z","iopub.execute_input":"2022-01-02T14:36:19.977799Z","iopub.status.idle":"2022-01-02T14:36:20.01671Z","shell.execute_reply.started":"2022-01-02T14:36:19.977761Z","shell.execute_reply":"2022-01-02T14:36:20.015968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /models/cycle_gan_model.py\n##=================================================\n\n\nimport torch\nimport itertools\n# from util.image_pool import ImagePool\n# from .base_model import BaseModel\n# from . import networks\n\n\nclass CycleGANModel(BaseModel):\n    \"\"\"\n    This class implements the CycleGAN model, for learning image-to-image translation without paired data.\n\n    The model training requires '--dataset_mode unaligned' dataset.\n    By default, it uses a '--netG resnet_9blocks' ResNet generator,\n    a '--netD basic' discriminator (PatchGAN introduced by pix2pix),\n    and a least-square GANs objective ('--gan_mode lsgan').\n\n    CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf\n    \"\"\"\n    @staticmethod\n    def modify_commandline_configs(parser, is_train=True):\n        \"\"\"Add new dataset-specific configs, and rewrite default values for existing configs.\n\n        Parameters:\n            parser          -- original config parser\n            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific configs.\n\n        Returns:\n            the modified parser.\n\n        For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.\n        A (source domain), B (target domain).\n        Generators: G_A: A -> B; G_B: B -> A.\n        Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.\n        Forward cycle loss:  lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)\n        Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)\n        Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 \"Photo generation from paintings\" in the paper)\n        Dropout is not used in the original CycleGAN paper.\n        \"\"\"\n        parser.set_defaults(no_dropout=True)  # default CycleGAN did not use dropout\n        if is_train:\n            parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')\n            parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')\n            parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')\n            print(parser)\n        return parser\n\n    def __init__(self, config):\n        \"\"\"Initialize the CycleGAN class.\n\n        Parameters:\n            config (Config class)-- stores all the experiment flags; needs to be a subclass of BaseConfigs\n        \"\"\"\n        BaseModel.__init__(self, config)\n        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>\n        self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B']\n        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>\n        visual_names_A = ['real_A', 'fake_B', 'rec_A']\n        visual_names_B = ['real_B', 'fake_A', 'rec_B']\n        if self.isTrain and self.config.lambda_identity > 0.0:  # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)\n            visual_names_A.append('idt_B')\n            visual_names_B.append('idt_A')\n\n        self.visual_names = visual_names_A + visual_names_B  # combine visualizations for A and B\n        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.\n        if self.isTrain:\n            self.model_names = ['G_A', 'G_B', 'D_A', 'D_B']\n        else:  # during test time, only load Gs\n            self.model_names = ['G_A', 'G_B']\n\n        # define networks (both Generators and discriminators)\n        # The naming is different from those used in the paper.\n        # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)\n        self.netG_A = define_G(config.input_nc, config.output_nc, config.ngf, config.netG, config.norm,\n                                        not config.no_dropout, config.init_type, config.init_gain, self.gpu_ids)\n        self.netG_B = define_G(config.output_nc, config.input_nc, config.ngf, config.netG, config.norm,\n                                        not config.no_dropout, config.init_type, config.init_gain, self.gpu_ids)\n\n        if self.isTrain:  # define discriminators\n            self.netD_A = define_D(config.output_nc, config.ndf, config.netD,\n                                            config.n_layers_D, config.norm, config.init_type, config.init_gain, self.gpu_ids)\n            self.netD_B = define_D(config.input_nc, config.ndf, config.netD,\n                                            config.n_layers_D, config.norm, config.init_type, config.init_gain, self.gpu_ids)\n\n        if self.isTrain:\n            if config.lambda_identity > 0.0:  # only works when input and output images have the same number of channels\n                assert(config.input_nc == config.output_nc)\n            self.fake_A_pool = ImagePool(config.pool_size)  # create image buffer to store previously generated images\n            self.fake_B_pool = ImagePool(config.pool_size)  # create image buffer to store previously generated images\n            # define loss functions\n            self.criterionGAN = GANLoss(config.gan_mode).to(self.device)  # define GAN loss.\n            self.criterionCycle = torch.nn.L1Loss()\n            self.criterionIdt = torch.nn.L1Loss()\n            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.\n            self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=config.lr, betas=(config.beta1, 0.999))\n            self.optimizer_D = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=config.lr, betas=(config.beta1, 0.999))\n            self.optimizers.append(self.optimizer_G)\n            self.optimizers.append(self.optimizer_D)\n\n    def set_input(self, input):\n        \"\"\"Unpack input data from the dataloader and perform necessary pre-processing steps.\n\n        Parameters:\n            input (dict): include the data itself and its metadata information.\n\n        The config 'direction' can be used to swap domain A and domain B.\n        \"\"\"\n        AtoB = self.config.direction == 'AtoB'\n        self.real_A = input['A' if AtoB else 'B'].to(self.device)\n        self.real_B = input['B' if AtoB else 'A'].to(self.device)\n        self.image_paths = input['A_paths' if AtoB else 'B_paths']\n\n    def forward(self):\n        \"\"\"Run forward pass; called by both functions <optimize_parameters> and <test>.\"\"\"\n        self.fake_B = self.netG_A(self.real_A)  # G_A(A)\n        self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))\n        self.fake_A = self.netG_B(self.real_B)  # G_B(B)\n        self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))\n\n    def forward_test_kaggle(self):\n        \"\"\"Run forward pass; called by both function <test>.\"\"\"\n        self.visual_names = ['fake_B']\n        self.fake_B = self.netG_A(self.real_A)  # G_A(A)\n\n    def backward_D_basic(self, netD, real, fake):\n        \"\"\"Calculate GAN loss for the discriminator\n\n        Parameters:\n            netD (network)      -- the discriminator D\n            real (tensor array) -- real images\n            fake (tensor array) -- images generated by a generator\n\n        Return the discriminator loss.\n        We also call loss_D.backward() to calculate the gradients.\n        \"\"\"\n        # Real\n        pred_real = netD(real)\n        loss_D_real = self.criterionGAN(pred_real, True)\n        # Fake\n        pred_fake = netD(fake.detach())\n        loss_D_fake = self.criterionGAN(pred_fake, False)\n        # Combined loss and calculate gradients\n        loss_D = (loss_D_real + loss_D_fake) * 0.5\n        loss_D.backward()\n        return loss_D\n\n    def backward_D_A(self):\n        \"\"\"Calculate GAN loss for discriminator D_A\"\"\"\n        fake_B = self.fake_B_pool.query(self.fake_B)\n        self.loss_D_A = self.backward_D_basic(self.netD_A, self.real_B, fake_B)\n\n    def backward_D_B(self):\n        \"\"\"Calculate GAN loss for discriminator D_B\"\"\"\n        fake_A = self.fake_A_pool.query(self.fake_A)\n        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)\n\n    def backward_G(self):\n        \"\"\"Calculate the loss for generators G_A and G_B\"\"\"\n        lambda_idt = self.config.lambda_identity\n        lambda_A = self.config.lambda_A\n        lambda_B = self.config.lambda_B\n        # Identity loss\n        if lambda_idt > 0:\n            # G_A should be identity if real_B is fed: ||G_A(B) - B||\n            self.idt_A = self.netG_A(self.real_B)\n            self.loss_idt_A = self.criterionIdt(self.idt_A, self.real_B) * lambda_B * lambda_idt\n            # G_B should be identity if real_A is fed: ||G_B(A) - A||\n            self.idt_B = self.netG_B(self.real_A)\n            self.loss_idt_B = self.criterionIdt(self.idt_B, self.real_A) * lambda_A * lambda_idt\n        else:\n            self.loss_idt_A = 0\n            self.loss_idt_B = 0\n\n        # GAN loss D_A(G_A(A))\n        self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)\n        # GAN loss D_B(G_B(B))\n        self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)\n        # Forward cycle loss || G_B(G_A(A)) - A||\n        self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A\n        # Backward cycle loss || G_A(G_B(B)) - B||\n        self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B\n        # combined loss and calculate gradients\n        self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_cycle_A + self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B\n        self.loss_G.backward()\n\n    def optimize_parameters(self):\n        \"\"\"Calculate losses, gradients, and update network weights; called in every training iteration\"\"\"\n        # forward\n        self.forward()      # compute fake images and reconstruction images.\n        # G_A and G_B\n        self.set_requires_grad([self.netD_A, self.netD_B], False)  # Ds require no gradients when optimizing Gs\n        self.optimizer_G.zero_grad()  # set G_A and G_B's gradients to zero\n        self.backward_G()             # calculate gradients for G_A and G_B\n        self.optimizer_G.step()       # update G_A and G_B's weights\n        # D_A and D_B\n        self.set_requires_grad([self.netD_A, self.netD_B], True)\n        self.optimizer_D.zero_grad()   # set D_A and D_B's gradients to zero\n        self.backward_D_A()      # calculate gradients for D_A\n        self.backward_D_B()      # calculate graidents for D_B\n        self.optimizer_D.step()  # update D_A and D_B's weights","metadata":{"id":"unwr9wUKsOw5","execution":{"iopub.status.busy":"2022-01-02T14:36:20.759247Z","iopub.execute_input":"2022-01-02T14:36:20.759657Z","iopub.status.idle":"2022-01-02T14:36:20.794796Z","shell.execute_reply.started":"2022-01-02T14:36:20.759621Z","shell.execute_reply":"2022-01-02T14:36:20.793946Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset and Dataloader classes (path: /data/image_folder.py, /data/base_dataset.py, /data/\\_\\_init\\_\\_.py, /data/unaligned_dataset.py)","metadata":{"id":"x5lsJ4sisOw6"}},{"cell_type":"code","source":"##=================================================\n## /data/image_folder.py\n##=================================================\n\n\"\"\"A modified image folder class\n\nWe modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)\nso that this class can load images from both current directory and its subdirectories.\n\"\"\"\n\nimport torch.utils.data as data\n\nfrom PIL import Image\nimport os\n\nIMG_EXTENSIONS = [\n    '.jpg', '.JPG', '.jpeg', '.JPEG',\n    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n    '.tif', '.TIF', '.tiff', '.TIFF',\n]\n\n\ndef is_image_file(filename):\n    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)\n\n\ndef make_dataset(dir, max_dataset_size=float(\"inf\")):\n    images = []\n    assert os.path.isdir(dir), '%s is not a valid directory' % dir\n\n    for root, _, fnames in sorted(os.walk(dir)):\n        for fname in fnames:\n            if is_image_file(fname):\n                path = os.path.join(root, fname)\n                images.append(path)\n    return images[:min(max_dataset_size, len(images))]\n\n\ndef default_loader(path):\n    return Image.open(path).convert('RGB')\n\n\nclass ImageFolder(data.Dataset):\n\n    def __init__(self, root, transform=None, return_paths=False,\n                 loader=default_loader):\n        imgs = make_dataset(root)\n        if len(imgs) == 0:\n            raise(RuntimeError(\"Found 0 images in: \" + root + \"\\n\"\n                               \"Supported image extensions are: \" + \",\".join(IMG_EXTENSIONS)))\n\n        self.root = root\n        self.imgs = imgs\n        self.transform = transform\n        self.return_paths = return_paths\n        self.loader = loader\n\n    def __getitem__(self, index):\n        path = self.imgs[index]\n        img = self.loader(path)\n        if self.transform is not None:\n            img = self.transform(img)\n        if self.return_paths:\n            return img, path\n        else:\n            return img\n\n    def __len__(self):\n        return len(self.imgs)\n\n","metadata":{"id":"xp3lYXBIsOw6","execution":{"iopub.status.busy":"2022-01-02T14:36:22.953759Z","iopub.execute_input":"2022-01-02T14:36:22.954007Z","iopub.status.idle":"2022-01-02T14:36:22.966179Z","shell.execute_reply.started":"2022-01-02T14:36:22.953978Z","shell.execute_reply":"2022-01-02T14:36:22.965298Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /data/base_dataset.py\n##=================================================\n\n\n\"\"\"This module implements an abstract base class (ABC) 'BaseDataset' for datasets.\n\nIt also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.\n\"\"\"\nimport random\nimport numpy as np\nimport torch.utils.data as data\nfrom PIL import Image\nimport torchvision.transforms as transforms\nfrom abc import ABC, abstractmethod\n\n\nclass BaseDataset(data.Dataset, ABC):\n    \"\"\"This class is an abstract base class (ABC) for datasets.\n\n    To create a subclass, you need to implement the following four functions:\n    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, config).\n    -- <__len__>:                       return the size of dataset.\n    -- <__getitem__>:                   get a data point.\n    -- <modify_commandline_configs>:    (optionally) add dataset-specific configs and set default configs.\n    \"\"\"\n\n    def __init__(self, config):\n        \"\"\"Initialize the class; save the configs in the class\n\n        Parameters:\n            config (Config class)-- stores all the experiment flags; needs to be a subclass of BaseConfigs\n        \"\"\"\n        self.config = config\n        self.root = config.dataroot\n\n    @staticmethod\n    def modify_commandline_configs(parser, is_train):\n        \"\"\"Add new dataset-specific configs, and rewrite default values for existing configs.\n\n        Parameters:\n            parser          -- original config parser\n            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific configs.\n\n        Returns:\n            the modified parser.\n        \"\"\"\n        return parser\n\n    @abstractmethod\n    def __len__(self):\n        \"\"\"Return the total number of images in the dataset.\"\"\"\n        return 0\n\n    @abstractmethod\n    def __getitem__(self, index):\n        \"\"\"Return a data point and its metadata information.\n\n        Parameters:\n            index - - a random integer for data indexing\n\n        Returns:\n            a dictionary of data with their names. It ususally contains the data itself and its metadata information.\n        \"\"\"\n        pass\n\n\ndef get_params(config, size):\n    w, h = size\n    new_h = h\n    new_w = w\n    if config.preprocess == 'resize_and_crop':\n        new_h = new_w = config.load_size\n    elif config.preprocess == 'scale_width_and_crop':\n        new_w = config.load_size\n        new_h = config.load_size * h // w\n\n    x = random.randint(0, np.maximum(0, new_w - config.crop_size))\n    y = random.randint(0, np.maximum(0, new_h - config.crop_size))\n\n    flip = random.random() > 0.5\n\n    return {'crop_pos': (x, y), 'flip': flip}\n\n\ndef get_transform(config, params=None, grayscale=False, method=Image.BICUBIC, convert=True):\n    transform_list = []\n    if grayscale:\n        transform_list.append(transforms.Grayscale(1))\n    if 'resize' in config.preprocess:\n        osize = [config.load_size, config.load_size]\n        transform_list.append(transforms.Resize(osize, method))\n    elif 'scale_width' in config.preprocess:\n        transform_list.append(transforms.Lambda(lambda img: __scale_width(img, config.load_size, config.crop_size, method)))\n\n    if 'crop' in config.preprocess:\n        if params is None:\n            transform_list.append(transforms.RandomCrop(config.crop_size))\n        else:\n            transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], config.crop_size)))\n\n    if config.preprocess == 'none':\n        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))\n\n    if not config.no_flip:\n        if params is None:\n            transform_list.append(transforms.RandomHorizontalFlip())\n        elif params['flip']:\n            transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))\n\n    if convert:\n        transform_list += [transforms.ToTensor()]\n        if grayscale:\n            transform_list += [transforms.Normalize((0.5,), (0.5,))]\n        else:\n            transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n    return transforms.Compose(transform_list)\n\n\ndef __make_power_2(img, base, method=Image.BICUBIC):\n    ow, oh = img.size\n    h = int(round(oh / base) * base)\n    w = int(round(ow / base) * base)\n    if h == oh and w == ow:\n        return img\n\n    __print_size_warning(ow, oh, w, h)\n    return img.resize((w, h), method)\n\n\ndef __scale_width(img, target_size, crop_size, method=Image.BICUBIC):\n    ow, oh = img.size\n    if ow == target_size and oh >= crop_size:\n        return img\n    w = target_size\n    h = int(max(target_size * oh / ow, crop_size))\n    return img.resize((w, h), method)\n\n\ndef __crop(img, pos, size):\n    ow, oh = img.size\n    x1, y1 = pos\n    tw = th = size\n    if (ow > tw or oh > th):\n        return img.crop((x1, y1, x1 + tw, y1 + th))\n    return img\n\n\ndef __flip(img, flip):\n    if flip:\n        return img.transpose(Image.FLIP_LEFT_RIGHT)\n    return img\n\n\ndef __print_size_warning(ow, oh, w, h):\n    \"\"\"Print warning information about image size(only print once)\"\"\"\n    if not hasattr(__print_size_warning, 'has_printed'):\n        print(\"The image size needs to be a multiple of 4. \"\n              \"The loaded image size was (%d, %d), so it was adjusted to \"\n              \"(%d, %d). This adjustment will be done to all images \"\n              \"whose sizes are not multiples of 4\" % (ow, oh, w, h))\n        __print_size_warning.has_printed = True\n","metadata":{"id":"GD4gbqxNsOw6","execution":{"iopub.status.busy":"2022-01-02T14:36:23.585603Z","iopub.execute_input":"2022-01-02T14:36:23.585847Z","iopub.status.idle":"2022-01-02T14:36:23.881189Z","shell.execute_reply.started":"2022-01-02T14:36:23.585818Z","shell.execute_reply":"2022-01-02T14:36:23.880481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /data/__init__.py\n##=================================================\n\n\n\"\"\"This package includes all the modules related to data loading and preprocessing\n\n To add a custom dataset class called 'dummy', you need to add a file called 'dummy_dataset.py' and define a subclass 'DummyDataset' inherited from BaseDataset.\n You need to implement four functions:\n    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, config).\n    -- <__len__>:                       return the size of dataset.\n    -- <__getitem__>:                   get a data point from data loader.\n    -- <modify_commandline_configs>:    (optionally) add dataset-specific configs and set default configs.\n\nNow you can use the dataset class by specifying flag '--dataset_mode dummy'.\nSee our template dataset class 'template_dataset.py' for more details.\n\"\"\"\nimport importlib\nimport torch.utils.data\n# from data.base_dataset import BaseDataset\n\n\ndef find_dataset_using_name(dataset_name):\n    \"\"\"Import the module \"data/[dataset_name]_dataset.py\".\n\n    In the file, the class called DatasetNameDataset() will\n    be instantiated. It has to be a subclass of BaseDataset,\n    and it is case-insensitive.\n    \"\"\"\n    dataset_filename = \"data.\" + dataset_name + \"_dataset\"\n    datasetlib = importlib.import_module(dataset_filename)\n\n    dataset = None\n    target_dataset_name = dataset_name.replace('_', '') + 'dataset'\n    for name, cls in datasetlib.__dict__.items():\n        if name.lower() == target_dataset_name.lower() \\\n           and issubclass(cls, BaseDataset):\n            dataset = cls\n\n    if dataset is None:\n        raise NotImplementedError(\"In %s.py, there should be a subclass of BaseDataset with class name that matches %s in lowercase.\" % (dataset_filename, target_dataset_name))\n\n    return dataset\n\n\ndef get_data_config_setter(dataset_name):\n    \"\"\"Return the static method <modify_commandline_configs> of the dataset class.\"\"\"\n    # dataset_class = find_dataset_using_name(dataset_name)\n    dataset_class = UnalignedDataset\n    return dataset_class.modify_commandline_configs\n\n\ndef create_dataset(config):\n    \"\"\"Create a dataset given the config.\n\n    This function wraps the class CustomDatasetDataLoader.\n        This is the main interface between this package and 'train.py'/'test.py'\n    \"\"\"\n    data_loader = CustomDatasetDataLoader(config)\n    dataset = data_loader.load_data()\n    return dataset\n\n\nclass CustomDatasetDataLoader():\n    \"\"\"Wrapper class of Dataset class that performs multi-threaded data loading\"\"\"\n\n    def __init__(self, config):\n        \"\"\"Initialize this class\n\n        Step 1: create a dataset instance given the name [dataset_mode]\n        Step 2: create a multi-threaded data loader.\n        \"\"\"\n        self.config = config\n        # dataset_class = find_dataset_using_name(config.dataset_mode)\n        dataset_class = UnalignedDataset\n        self.dataset = dataset_class(config)\n        print(\"dataset [%s] was created\" % type(self.dataset).__name__)\n        self.dataloader = torch.utils.data.DataLoader(\n            self.dataset,\n            batch_size=config.batch_size,\n            shuffle=not config.serial_batches,\n            num_workers=int(config.num_threads))\n\n    def load_data(self):\n        return self\n\n    def __len__(self):\n        \"\"\"Return the number of data in the dataset\"\"\"\n        return min(len(self.dataset), self.config.max_dataset_size)\n\n    def __iter__(self):\n        \"\"\"Return a batch of data\"\"\"\n        for i, data in enumerate(self.dataloader):\n            if i * self.config.batch_size >= self.config.max_dataset_size:\n                break\n            yield data","metadata":{"id":"jkyYFH6ssOw7","execution":{"iopub.status.busy":"2022-01-02T14:36:24.258405Z","iopub.execute_input":"2022-01-02T14:36:24.259131Z","iopub.status.idle":"2022-01-02T14:36:24.272288Z","shell.execute_reply.started":"2022-01-02T14:36:24.259084Z","shell.execute_reply":"2022-01-02T14:36:24.271579Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /data/unaligned_dataset.py\n##=================================================\n\nimport os\n# from data.base_dataset import BaseDataset, get_transform\n# from data.image_folder import make_dataset\nfrom PIL import Image\nimport random\n\n\nclass UnalignedDataset(BaseDataset):\n    \"\"\"\n    This dataset class can load unaligned/unpaired datasets.\n\n    It requires two directories to host training images from domain A '/path/to/data/trainA'\n    and from domain B '/path/to/data/trainB' respectively.\n    You can train the model with the dataset flag '--dataroot /path/to/data'.\n    Similarly, you need to prepare two directories:\n    '/path/to/data/testA' and '/path/to/data/testB' during test time.\n    \"\"\"\n\n    def __init__(self, config):\n        \"\"\"Initialize this dataset class.\n\n        Parameters:\n            config (Config class) -- stores all the experiment flags; needs to be a subclass of BaseConfigs\n        \"\"\"\n        BaseDataset.__init__(self, config)\n        self.dir_A = os.path.join(config.dataroot, config.phase + 'A')  # create a path '/path/to/data/trainA'\n        self.dir_B = os.path.join(config.dataroot, config.phase + 'B')  # create a path '/path/to/data/trainB'\n\n        self.A_paths = sorted(make_dataset(self.dir_A, config.max_dataset_size))   # load images from '/path/to/data/trainA'\n        self.B_paths = sorted(make_dataset(self.dir_B, config.max_dataset_size))    # load images from '/path/to/data/trainB'\n        self.A_size = len(self.A_paths)  # get the size of dataset A\n        self.B_size = len(self.B_paths)  # get the size of dataset B\n        btoA = self.config.direction == 'BtoA'\n        input_nc = self.config.output_nc if btoA else self.config.input_nc       # get the number of channels of input image\n        output_nc = self.config.input_nc if btoA else self.config.output_nc      # get the number of channels of output image\n        self.transform_A = get_transform(self.config, grayscale=(input_nc == 1))\n        self.transform_B = get_transform(self.config, grayscale=(output_nc == 1))\n\n    def __getitem__(self, index):\n        \"\"\"Return a data point and its metadata information.\n\n        Parameters:\n            index (int)      -- a random integer for data indexing\n\n        Returns a dictionary that contains A, B, A_paths and B_paths\n            A (tensor)       -- an image in the input domain\n            B (tensor)       -- its corresponding image in the target domain\n            A_paths (str)    -- image paths\n            B_paths (str)    -- image paths\n        \"\"\"\n        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range\n        if self.config.serial_batches:   # make sure index is within then range\n            index_B = index % self.B_size\n        else:   # randomize the index for domain B to avoid fixed pairs.\n            index_B = random.randint(0, self.B_size - 1)\n        B_path = self.B_paths[index_B]\n        A_img = Image.open(A_path).convert('RGB')\n        B_img = Image.open(B_path).convert('RGB')\n        # apply image transformation\n        A = self.transform_A(A_img)\n        B = self.transform_B(B_img)\n\n        return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}\n\n    def __len__(self):\n        \"\"\"Return the total number of images in the dataset.\n\n        As we have two datasets with potentially different number of images,\n        we take a maximum of\n        \"\"\"\n        return max(self.A_size, self.B_size)\n\n","metadata":{"id":"nJeLbNXqsOw7","execution":{"iopub.status.busy":"2022-01-02T14:36:25.963136Z","iopub.execute_input":"2022-01-02T14:36:25.963568Z","iopub.status.idle":"2022-01-02T14:36:25.976932Z","shell.execute_reply.started":"2022-01-02T14:36:25.963533Z","shell.execute_reply":"2022-01-02T14:36:25.976269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Main Train-loop (path: train.py)","metadata":{"id":"V8-vIm9-sOw7"}},{"cell_type":"code","source":"##=================================================\n## train.py\n##=================================================\n\n\nimport time\n# from configs.train_config import TrainConfigs\n# from data import create_dataset\n# from models import create_model\n# from models.cycle_gan_model import CycleGANModel\n# from util.visualizer import Visualizer\n\n#TODO: replace visualizer by creating new print_losses() function\n\n# if __name__ == '__main__':\ndef train():\n    config_obj = TrainConfigs().parse()   # get training configs\n    dataset = create_dataset(config_obj)  # create a dataset given config.dataset_mode and other configs\n    \n    dataset_size = len(dataset)    # get the number of images in the dataset.\n    print('The number of training images = %d' % dataset_size)\n\n    # model = create_model(config)      # create a model given config.model and other configs\n    model = CycleGANModel(config_obj)\n    model.setup(config_obj)               # regular setup: load and print networks; create schedulers\n    # visualizer = Visualizer(config_obj)   # create a visualizer that display/save images and plots\n    total_iters = 0                # the total number of training iterations\n\n    for epoch in range(config_obj.epoch_count, config_obj.n_epochs + config_obj.n_epochs_decay + 1):    # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>\n        epoch_start_time = time.time()  # timer for entire epoch\n        iter_data_time = time.time()    # timer for data loading per iteration\n        epoch_iter = 0                  # the number of training iterations in current epoch, reset to 0 every epoch\n        # visualizer.reset()              # reset the visualizer: make sure it saves the results to HTML at least once every epoch\n        # model.update_learning_rate()    # update learning rates in the beginning of every epoch.\n        for i, data in enumerate(dataset):  # inner loop within one epoch\n            t_data = 0\n            iter_start_time = time.time()  # timer for computation per iteration\n            if total_iters % config_obj.print_freq == 0:\n                t_data = iter_start_time - iter_data_time\n\n            total_iters += config_obj.batch_size\n            epoch_iter += config_obj.batch_size\n            model.set_input(data)         # unpack data from dataset and apply preprocessing\n            model.optimize_parameters()   # calculate loss functions, get gradients, update network weights\n\n            # if total_iters % config_obj.display_freq == 0:   # display images on visdom and save images to a HTML file\n            #     save_result = total_iters % config_obj.update_html_freq == 0\n            #     model.compute_visuals()\n            #     visualizer.display_current_results(model.get_current_visuals(), epoch, save_result)\n\n            if total_iters % config_obj.print_freq == 0:    # print training losses and save logging information to the disk\n                losses = model.get_current_losses()\n                t_comp = (time.time() - iter_start_time) / config_obj.batch_size\n                model.print_current_losses(epoch, epoch_iter, losses, t_comp, t_data)\n                # if config_obj.display_id > 0:\n                #     visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, losses)\n\n            if total_iters % config_obj.save_latest_freq == 0:   # cache our latest model every <save_latest_freq> iterations\n                print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters))\n                save_suffix = 'iter_%d' % total_iters if config_obj.save_by_iter else 'latest'\n                model.save_networks(save_suffix)\n\n            iter_data_time = time.time()\n            \n        model.update_learning_rate()  ## Added by Nandita, updating LR at the end of every epoch instead of beginning.\n        if epoch % config_obj.save_epoch_freq == 0:              # cache our model every <save_epoch_freq> epochs\n            print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters))\n            model.save_networks('latest')\n            model.save_networks(epoch)\n\n        print('End of epoch %d / %d \\t Time Taken: %d sec' % (epoch, config_obj.n_epochs + config_obj.n_epochs_decay, time.time() - epoch_start_time))","metadata":{"id":"-VrMAVuhsOw8","execution":{"iopub.status.busy":"2022-01-02T14:36:27.489857Z","iopub.execute_input":"2022-01-02T14:36:27.490269Z","iopub.status.idle":"2022-01-02T14:36:27.502826Z","shell.execute_reply.started":"2022-01-02T14:36:27.490231Z","shell.execute_reply":"2022-01-02T14:36:27.501864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # !python train.py\n\n# import sys\n\n# if __name__ == '__main__':\n#   if '-f' in sys.argv:\n#     sys.argv.remove('-f')\n#   train()","metadata":{"id":"4llxIqxKsOw8","outputId":"2513dbaa-821e-4960-80a9-ab1599ce0916","execution":{"iopub.status.busy":"2022-01-02T14:36:28.591071Z","iopub.execute_input":"2022-01-02T14:36:28.591538Z","iopub.status.idle":"2022-01-02T14:36:28.595457Z","shell.execute_reply.started":"2022-01-02T14:36:28.591503Z","shell.execute_reply":"2022-01-02T14:36:28.594733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%tb","metadata":{"id":"J0I3gUT4sOw8","execution":{"iopub.status.busy":"2022-01-02T14:36:29.370388Z","iopub.execute_input":"2022-01-02T14:36:29.370858Z","iopub.status.idle":"2022-01-02T14:36:29.375679Z","shell.execute_reply.started":"2022-01-02T14:36:29.37082Z","shell.execute_reply":"2022-01-02T14:36:29.374966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Some utils for Inference (/util/html.py, /util/visualizer.py)","metadata":{"id":"t1RaU2Q-sOw8"}},{"cell_type":"code","source":"!pip install dominate","metadata":{"id":"xh4LxCZnsOw8","outputId":"e7a2564c-d4e4-4666-d847-c6ac745c82ea","execution":{"iopub.status.busy":"2022-01-02T14:36:33.382935Z","iopub.execute_input":"2022-01-02T14:36:33.383476Z","iopub.status.idle":"2022-01-02T14:36:43.590345Z","shell.execute_reply.started":"2022-01-02T14:36:33.383439Z","shell.execute_reply":"2022-01-02T14:36:43.589526Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /util/html.py\n##=================================================\n\nimport dominate\nfrom dominate.tags import meta, h3, table, tr, td, p, a, img, br\nimport os\n\n\nclass HTML:\n    \"\"\"This HTML class allows us to save images and write texts into a single HTML file.\n\n     It consists of functions such as <add_header> (add a text header to the HTML file),\n     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).\n     It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API.\n    \"\"\"\n\n    def __init__(self, web_dir, title, refresh=0):\n        \"\"\"Initialize the HTML classes\n\n        Parameters:\n            web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/\n            title (str)   -- the webpage name\n            refresh (int) -- how often the website refresh itself; if 0; no refreshing\n        \"\"\"\n        self.title = title\n        self.web_dir = web_dir\n        self.img_dir = os.path.join(self.web_dir, 'images')\n        if not os.path.exists(self.web_dir):\n            os.makedirs(self.web_dir)\n        if not os.path.exists(self.img_dir):\n            os.makedirs(self.img_dir)\n\n        self.doc = dominate.document(title=title)\n        if refresh > 0:\n            with self.doc.head:\n                meta(http_equiv=\"refresh\", content=str(refresh))\n\n    def get_image_dir(self):\n        \"\"\"Return the directory that stores images\"\"\"\n        return self.img_dir\n\n    def add_header(self, text):\n        \"\"\"Insert a header to the HTML file\n\n        Parameters:\n            text (str) -- the header text\n        \"\"\"\n        with self.doc:\n            h3(text)\n\n    def add_images(self, ims, txts, links, width=400):\n        \"\"\"add images to the HTML file\n\n        Parameters:\n            ims (str list)   -- a list of image paths\n            txts (str list)  -- a list of image names shown on the website\n            links (str list) --  a list of hyperref links; when you click an image, it will redirect you to a new page\n        \"\"\"\n        self.t = table(border=1, style=\"table-layout: fixed;\")  # Insert a table\n        self.doc.add(self.t)\n        with self.t:\n            with tr():\n                for im, txt, link in zip(ims, txts, links):\n                    with td(style=\"word-wrap: break-word;\", halign=\"center\", valign=\"top\"):\n                        with p():\n                            with a(href=os.path.join('images', link)):\n                                img(style=\"width:%dpx\" % width, src=os.path.join('images', im))\n                            br()\n                            p(txt)\n\n    def save(self):\n        \"\"\"save the current content to the HMTL file\"\"\"\n        html_file = '%s/index.html' % self.web_dir\n        f = open(html_file, 'wt')\n        f.write(self.doc.render())\n        f.close()\n\n\nif __name__ == '__main__':  # we show an example usage here.\n    html = HTML('web/', 'test_html')\n    html.add_header('hello world')\n\n    ims, txts, links = [], [], []\n    for n in range(4):\n        ims.append('image_%d.png' % n)\n        txts.append('text_%d' % n)\n        links.append('image_%d.png' % n)\n    html.add_images(ims, txts, links)\n    html.save()","metadata":{"id":"LHE6waFwsOw9","execution":{"iopub.status.busy":"2022-01-02T14:37:49.624706Z","iopub.execute_input":"2022-01-02T14:37:49.624982Z","iopub.status.idle":"2022-01-02T14:37:49.6523Z","shell.execute_reply.started":"2022-01-02T14:37:49.624949Z","shell.execute_reply":"2022-01-02T14:37:49.651616Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /util/visualizer.py\n##=================================================\n\nimport numpy as np\nimport os\nimport sys\nimport ntpath\nimport time\n# from . import util, html\nfrom subprocess import Popen, PIPE\n\n\nif sys.version_info[0] == 2:\n    VisdomExceptionBase = Exception\nelse:\n    VisdomExceptionBase = ConnectionError\n\n\ndef save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):\n    \"\"\"Save images to the disk.\n\n    Parameters:\n        webpage (the HTML class) -- the HTML webpage class that stores these imaegs (see html.py for more details)\n        visuals (OrderedDict)    -- an ordered dictionary that stores (name, images (either tensor or numpy) ) pairs\n        image_path (str)         -- the string is used to create image paths\n        aspect_ratio (float)     -- the aspect ratio of saved images\n        width (int)              -- the images will be resized to width x width\n\n    This function will save images stored in 'visuals' to the HTML file specified by 'webpage'.\n    \"\"\"\n    image_dir = webpage.get_image_dir()\n    short_path = ntpath.basename(image_path[0])\n    name = os.path.splitext(short_path)[0]\n\n    webpage.add_header(name)\n    ims, txts, links = [], [], []\n\n    for label, im_data in visuals.items():\n        im = tensor2im(im_data)\n        image_name = '%s_%s.png' % (name, label)\n        save_path = os.path.join(image_dir, image_name)\n        save_image(im, save_path, aspect_ratio=aspect_ratio)\n        ims.append(image_name)\n        txts.append(label)\n        links.append(image_name)\n    webpage.add_images(ims, txts, links, width=width)\n","metadata":{"id":"pOncjZxxsOw9","execution":{"iopub.status.busy":"2022-01-02T14:37:50.320834Z","iopub.execute_input":"2022-01-02T14:37:50.321228Z","iopub.status.idle":"2022-01-02T14:37:50.332189Z","shell.execute_reply.started":"2022-01-02T14:37:50.321193Z","shell.execute_reply":"2022-01-02T14:37:50.331475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## TestConfigs","metadata":{"id":"_3RJiPTFsOw9"}},{"cell_type":"code","source":"import os\nimport shutil\n\nif os.path.isdir(result_img_path):\n    shutil.rmtree(result_img_path)\nif not os.path.isdir(result_img_path):\n    os.mkdir(result_img_path)","metadata":{"id":"2qi2e8KHsOw9","execution":{"iopub.status.busy":"2022-01-02T14:37:52.026626Z","iopub.execute_input":"2022-01-02T14:37:52.026881Z","iopub.status.idle":"2022-01-02T14:37:52.031729Z","shell.execute_reply.started":"2022-01-02T14:37:52.026851Z","shell.execute_reply":"2022-01-02T14:37:52.030818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"##=================================================\n## /configs/test_config.py\n##=================================================\n\n\nclass TestConfigs(BaseConfigs):\n    \"\"\"This class includes test configs.\n\n    It also includes shared configs defined in BaseConfigs.\n    \"\"\"\n\n    def initialize(self, parser):\n        parser = BaseConfigs.initialize(self, parser)  # define shared configs\n        parser.add_argument('--results_dir', type=str, default=result_img_path, help='saves results here.')\n        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')\n        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')\n        # Dropout and Batchnorm has different behavioir during training and test.\n        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')\n        parser.add_argument('--num_test', type=int, default=num_test_sample, help='how many test images to run') # 50\n        # rewrite devalue values\n        parser.set_defaults(model='test')\n        # To avoid cropping, the load_size should be the same as crop_size\n        parser.set_defaults(load_size=parser.get_default('crop_size'))\n        self.isTrain = False\n        return parser\n","metadata":{"id":"svn4m4WtsOw9","execution":{"iopub.status.busy":"2022-01-02T14:37:53.029754Z","iopub.execute_input":"2022-01-02T14:37:53.030297Z","iopub.status.idle":"2022-01-02T14:37:53.038222Z","shell.execute_reply.started":"2022-01-02T14:37:53.030258Z","shell.execute_reply":"2022-01-02T14:37:53.037468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference code (path: infer.py)","metadata":{"id":"kGGXZp4esOw-"}},{"cell_type":"code","source":"##=================================================\n## test.py\n##=================================================\n\n\n\"\"\"General-purpose test script for image-to-image translation.\n\nOnce you have trained your model with train.py, you can use this script to test the model.\nIt will load a saved model from '--checkpoints_dir' and save the results to '--results_dir'.\n\nIt first creates model and dataset given the config. It will hard-code some parameters.\nIt then runs inference for '--num_test' images and save results to an HTML file.\n\nExample (You need to train models first or download pre-trained models from our website):\n    Test a CycleGAN model (both sides):\n        python test.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan\n\n    Test a CycleGAN model (one side only):\n        python test.py --dataroot datasets/horse2zebra/testA --name horse2zebra_pretrained --model test --no_dropout\n\n    The config '--model test' is used for generating CycleGAN results only for one side.\n    This config will automatically set '--dataset_mode single', which only loads the images from one set.\n    On the contrary, using '--model cycle_gan' requires loading and generating results in both directions,\n    which is sometimes unnecessary. The results will be saved at ./results/.\n    Use '--results_dir <directory_path_to_save_result>' to specify the results directory.\n\n    Test a pix2pix model:\n        python test.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA\n\nSee configs/base_configs.py and configs/test_configs.py for more test configs.\nSee training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md\nSee frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md\n\"\"\"\nimport os\n# from configs.test_config import TestConfigs\n# from data import create_dataset\n# from models import create_model\n# from util.visualizer import save_images\n# from util import html\n\n\ndef test():\n    config = TestConfigs().parse()  # get test configs\n    # hard-code some parameters for test\n    config.num_threads = 0   # test code only supports num_threads = 1\n    config.batch_size = 1    # test code only supports batch_size = 1\n    config.serial_batches = True  # disable data shuffling; comment this line if results on randomly chosen images are needed.\n    config.no_flip = True    # no flip; comment this line if results on flipped images are needed.\n    config.display_id = -1   # no visdom display; the test code saves the results to a HTML file.\n    dataset = create_dataset(config)  # create a dataset given config.dataset_mode and other configs\n    model = CycleGANModel(config)      # create a model given config.model and other configs\n    model.setup(config)               # regular setup: load and print networks; create schedulers\n    # create a website\n    web_dir = os.path.join(config.results_dir, config.name, '{}_{}'.format(config.phase, config.epoch))  # define the website directory\n    if config.load_iter > 0:  # load_iter is 0 by default\n        web_dir = '{:s}_iter{:d}'.format(web_dir, config.load_iter)\n    print('creating web directory', web_dir)\n    webpage = HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (config.name, config.phase, config.epoch))\n    # test with eval mode. This only affects layers like batchnorm and dropout.\n    # For [pix2pix]: we use batchnorm and dropout in the original pix2pix. You can experiment it with and without eval() mode.\n    # For [CycleGAN]: It should not affect CycleGAN as CycleGAN uses instancenorm without dropout.\n    if config.eval:\n        model.eval()\n    for i, data in enumerate(dataset):\n        if i >= config.num_test:  # only apply our model to config.num_test images.\n            break\n        model.set_input(data)  # unpack data from data loader\n        model.test()           # run inference\n        visuals = model.get_current_visuals()  # get image results\n        img_path = model.get_image_paths()     # get image paths\n        if i % 5 == 0:  # save images to an HTML file\n            print('processing (%04d)-th image... %s' % (i, img_path))\n        save_images(webpage, visuals, img_path, aspect_ratio=config.aspect_ratio, width=config.display_winsize)\n    webpage.save()  # save the HTML\n","metadata":{"id":"c9_xjcpasOw-","execution":{"iopub.status.busy":"2022-01-02T14:37:54.346817Z","iopub.execute_input":"2022-01-02T14:37:54.347452Z","iopub.status.idle":"2022-01-02T14:37:54.366913Z","shell.execute_reply.started":"2022-01-02T14:37:54.347412Z","shell.execute_reply":"2022-01-02T14:37:54.365878Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# trainA_path = os.path.join(kaggle_output_path, \"trainA\")\n# os.path.isdir(trainA_path)\n# glob.glob(os.path.join(kaggle_output_path, \"*\"))","metadata":{"id":"3FBZp2wAsOw-","execution":{"iopub.status.busy":"2022-01-02T14:37:55.337822Z","iopub.execute_input":"2022-01-02T14:37:55.3381Z","iopub.status.idle":"2022-01-02T14:37:55.341265Z","shell.execute_reply.started":"2022-01-02T14:37:55.338063Z","shell.execute_reply":"2022-01-02T14:37:55.340539Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !python test.py\n\nimport sys\n\nif __name__ == '__main__':\n    if '-f' in sys.argv:\n        sys.argv.remove('-f')\n    test()\n\ntest()","metadata":{"id":"a5yV2al4sOw-","outputId":"4b8c2b8f-a49e-4daa-dcac-2940f313c956","execution":{"iopub.status.busy":"2022-01-02T14:37:56.040605Z","iopub.execute_input":"2022-01-02T14:37:56.041221Z","iopub.status.idle":"2022-01-02T14:48:31.84438Z","shell.execute_reply.started":"2022-01-02T14:37:56.041181Z","shell.execute_reply":"2022-01-02T14:48:31.843625Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Copy results into /kaggle/working/images.zip","metadata":{"id":"EBrolmMTsOw_"}},{"cell_type":"code","source":"import shutil\nimport glob\nimport os.path\n\nsrc_path = os.path.join(result_img_path, 'photo_to_monet', 'test_latest', 'images', '*_fake_B.png')\n# src_path = os.path.join(\"/kaggle/working/results/photo_to_monet/test_latest/images\",\"*_fake_B.png\")\n\n# dest_folder = os.path.join(\"/kaggle/working/images\")\ndest_folder = os.path.join(images_submit_path)\n\nif not os.path.isdir(dest_folder):\n    os.mkdir(dest_folder)\n\nsrc_files = glob.glob(src_path)\nfor file in src_files:\n    filename = os.path.basename(file)\n    dest_filepath = os.path.join(dest_folder, filename)\n    shutil.copy(file, dest_filepath)\n\n# shutil.make_archive('/kaggle/working/images/', 'zip', '/kaggle/working/images')\n# print(f\"Generated samples: {len([name for name in os.listdir('/kaggle/working/images/') if os.path.isfile(os.path.join('/kaggle/working/images/', name))])}\")\n\nimages_submit_folder = images_submit_path + os.sep\nshutil.make_archive(images_submit_folder, 'zip', images_submit_path)\nprint(f\"Generated samples: {len([name for name in os.listdir(images_submit_folder) if os.path.isfile(os.path.join(images_submit_folder, name))])}\")","metadata":{"id":"MmtC8GlJsOw_","outputId":"fdd01ecf-6776-456a-b745-5ee91198d5a3","execution":{"iopub.status.busy":"2022-01-02T14:48:36.804499Z","iopub.execute_input":"2022-01-02T14:48:36.804748Z","iopub.status.idle":"2022-01-02T14:49:08.653012Z","shell.execute_reply.started":"2022-01-02T14:48:36.804721Z","shell.execute_reply":"2022-01-02T14:49:08.652279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## [OPTIONAL] Delete all other folders inside /kaggle/working/","metadata":{"id":"D-M680hHsOw_"}},{"cell_type":"code","source":"# optional","metadata":{"id":"kz1GNRZasOw_"},"execution_count":null,"outputs":[]}]}