{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":71750,"databundleVersionId":7863244,"sourceType":"competition"},{"sourceId":9263711,"sourceType":"datasetVersion","datasetId":5605497}],"dockerImageVersionId":30763,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Решение ноутбука основано на репозитории  https://github.com/xinntao/Real-ESRGAN, в котором приведены инструменты для работы с ESRGAN. Поэтому часть функций запускается, как скрипт из отдельного файла, но все такие запускаемые файлы формируются в данном ноутбуке для наглядности решения.","metadata":{}},{"cell_type":"code","source":"%%capture\n!pip install basicsr\n!pip install torch","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:26:41.748415Z","iopub.execute_input":"2024-08-28T06:26:41.749340Z","iopub.status.idle":"2024-08-28T06:27:14.767046Z","shell.execute_reply.started":"2024-08-28T06:26:41.749295Z","shell.execute_reply":"2024-08-28T06:27:14.765837Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def change_code_in_env(path: str, old: str, new: str) -> None:\n    \"\"\"\n    Функция для адаптации basicsr для работы с kaggle в части исключения ошибок импорта\n    пакетов.\n    \"\"\"\n    with open (path, 'r') as f:\n        old_data = f.read()\n    new_data = old_data.replace(old, new)\n    with open (path, 'w') as f:\n        f.write(new_data)\n        \nchange_code_in_env(\n    '/opt/conda/lib/python3.10/site-packages/basicsr/data/degradations.py',\n    'from torchvision.transforms.functional_tensor import rgb_to_grayscale',\n    'from torchvision.transforms.functional import rgb_to_grayscale'\n    )","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:27:14.768997Z","iopub.execute_input":"2024-08-28T06:27:14.769332Z","iopub.status.idle":"2024-08-28T06:27:14.776177Z","shell.execute_reply.started":"2024-08-28T06:27:14.769295Z","shell.execute_reply":"2024-08-28T06:27:14.775310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import argparse\nimport glob\nimport os\nimport zipfile\n\nfrom IPython.display import FileLink","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:27:14.777337Z","iopub.execute_input":"2024-08-28T06:27:14.777645Z","iopub.status.idle":"2024-08-28T06:27:14.786775Z","shell.execute_reply.started":"2024-08-28T06:27:14.777604Z","shell.execute_reply":"2024-08-28T06:27:14.785895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def zip_dir(directory = os.curdir, file_name = 'directory.zip'):\n    \"\"\"\n    Функция для выгрузки папок с результатами работы на kaggle.\n    \n    \"\"\"\n    os.chdir(directory)\n    zip_ref = zipfile.ZipFile(file_name, mode='w')\n    for folder, _, files in os.walk(directory):\n        for file in files:\n            if file_name in file:\n                pass\n            else:\n                zip_ref.write(os.path.join(folder, file))\n    return FileLink(file_name)","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:27:14.788517Z","iopub.execute_input":"2024-08-28T06:27:14.788840Z","iopub.status.idle":"2024-08-28T06:27:14.798241Z","shell.execute_reply.started":"2024-08-28T06:27:14.788807Z","shell.execute_reply":"2024-08-28T06:27:14.797104Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\n\n# загрузка предобученных моделей\n!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /kaggle/working/\n!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P /kaggle/working/","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:15.011655Z","iopub.execute_input":"2024-08-27T20:05:15.011946Z","iopub.status.idle":"2024-08-27T20:05:18.504953Z","shell.execute_reply.started":"2024-08-27T20:05:15.011916Z","shell.execute_reply":"2024-08-27T20:05:18.503621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# объявление констант\nINPUT_FOLDERS = ['/kaggle/input/super-resolution-in-video-games/train/hr', '/kaggle/input/super-resolution-in-video-games/train/lr']\nROOT_FOLDERS = [None, None]\nMETA_INFO_PATH = '/kaggle/working/meta_info_DIV2K_sub_pair.txt'","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:18.506616Z","iopub.execute_input":"2024-08-27T20:05:18.507366Z","iopub.status.idle":"2024-08-27T20:05:18.512558Z","shell.execute_reply.started":"2024-08-27T20:05:18.507318Z","shell.execute_reply":"2024-08-27T20:05:18.511603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_meta_info(input_folders, root_folders, meta_info_path):\n    \"\"\"\n    Функция, реализующая создание .txt файла с сопоставлением lr и hr имен снимков.\n    \"\"\"\n    os.makedirs(os.path.dirname(meta_info_path), exist_ok=True)\n    for i in range(2):\n        if input_folders[i].endswith('/'):\n            input_folders[i] = input_folders[i][:-1]\n        if root_folders[i] is None:\n            root_folders[i] = os.path.dirname(input_folders[i])\n\n    with open(meta_info_path, 'w') as txt_file:\n        # Scan images\n        img_paths_gt = sorted(glob.glob(os.path.join(input_folders[0], '*')))\n        img_paths_lq = sorted(glob.glob(os.path.join(input_folders[1], '*')))\n\n        assert len(img_paths_gt) == len(img_paths_lq), ('GT folder and LQ folder should have the same length, but got '\n                                                        f'{len(img_paths_gt)} and {len(img_paths_lq)}.')\n\n        for img_path_gt, img_path_lq in zip(img_paths_gt, img_paths_lq):\n            img_name_gt = os.path.relpath(img_path_gt, root_folders[0])\n            img_name_lq = os.path.relpath(img_path_lq, root_folders[1])\n            print(f'{img_name_gt}, {img_name_lq}')\n            txt_file.write(f'{img_name_gt}, {img_name_lq}\\n')","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:18.513707Z","iopub.execute_input":"2024-08-27T20:05:18.513999Z","iopub.status.idle":"2024-08-27T20:05:18.524779Z","shell.execute_reply.started":"2024-08-27T20:05:18.513964Z","shell.execute_reply":"2024-08-27T20:05:18.523968Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%capture\ncreate_meta_info(INPUT_FOLDERS, ROOT_FOLDERS, META_INFO_PATH)","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:18.525975Z","iopub.execute_input":"2024-08-27T20:05:18.526299Z","iopub.status.idle":"2024-08-27T20:05:21.209646Z","shell.execute_reply.started":"2024-08-27T20:05:18.526268Z","shell.execute_reply":"2024-08-27T20:05:21.208838Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Создадим файл конфигураций для работы с моделью. В котором содержатся все необходимые настройки для обучения модели.","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/clustering_details.yaml\n\n# general settings\nname: finetune_RealESRGANx4plus_400k_pairdata\nmodel_type: RealESRGANModel\nscale: 4\nnum_gpu: auto\nmanual_seed: 0\n\n# USM the ground-truth\nl1_gt_usm: True\npercep_gt_usm: True\ngan_gt_usm: False\n\nhigh_order_degradation: False # do not use the high-order degradation generation process\n\n# dataset and data loader settings\ndatasets:\n  train:\n    name: DIV2K\n    type: RealESRGANPairedDataset\n    dataroot_gt: /kaggle/input/super-resolution-in-video-games/train\n    dataroot_lq: /kaggle/input/super-resolution-in-video-games/train\n    meta_info: /kaggle/working/meta_info_DIV2K_sub_pair.txt\n    io_backend:\n      type: disk\n\n    gt_size: 256\n    use_hflip: True\n    use_rot: False\n\n    # data loader\n    use_shuffle: true\n    num_worker_per_gpu: 5\n    batch_size_per_gpu: 12\n    dataset_enlarge_ratio: 1\n    prefetch_mode: ~\n\n# network structures\nnetwork_g:\n  type: RRDBNet\n  num_in_ch: 3\n  num_out_ch: 3\n  num_feat: 64\n  num_block: 23\n  num_grow_ch: 32\n\nnetwork_d:\n  type: UNetDiscriminatorSN\n  num_in_ch: 3\n  num_feat: 64\n  skip_connection: True\n\n# path\npath:\n  # use the pre-trained Real-ESRNet model\n  pretrain_network_g: /kaggle/working/RealESRGAN_x4plus.pth\n  param_key_g: params_ema\n  strict_load_g: true\n  pretrain_network_d: /kaggle/working/RealESRGAN_x4plus_netD.pth\n  param_key_d: params\n  strict_load_d: true\n  resume_state: ~\n\n# training settings\ntrain:\n  ema_decay: 0.999\n  optim_g:\n    type: Adam\n    lr: !!float 1e-4\n    weight_decay: 0\n    betas: [0.9, 0.99]\n  optim_d:\n    type: Adam\n    lr: !!float 1e-4\n    weight_decay: 0\n    betas: [0.9, 0.99]\n\n  scheduler:\n    type: MultiStepLR\n    milestones: [400000]\n    gamma: 0.5\n\n  total_iter: 50000\n  warmup_iter: -1  # no warm up\n\n  # losses\n  pixel_opt:\n    type: L1Loss\n    loss_weight: 1.0\n    reduction: mean\n  # perceptual loss (content and style losses)\n  perceptual_opt:\n    type: PerceptualLoss\n    layer_weights:\n      # before relu\n      'conv1_2': 0.1\n      'conv2_2': 0.1\n      'conv3_4': 1\n      'conv4_4': 1\n      'conv5_4': 1\n    vgg_type: vgg19\n    use_input_norm: true\n    perceptual_weight: !!float 1.0\n    style_weight: 0\n    range_norm: false\n    criterion: l1\n  # gan loss\n  gan_opt:\n    type: GANLoss\n    gan_type: vanilla\n    real_label_val: 1.0\n    fake_label_val: 0.0\n    loss_weight: !!float 1e-1\n\n  net_d_iters: 1\n  net_d_init_iters: 0\n\n# logging settings\nlogger:\n  print_freq: 100\n  save_checkpoint_freq: !!float 5e3\n  use_tb_logger: true\n  wandb:\n    project: ~\n    resume_id: ~\n\n# dist training settings\ndist_params:\n  backend: nccl\n  port: 29500","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:21.210993Z","iopub.execute_input":"2024-08-27T20:05:21.211611Z","iopub.status.idle":"2024-08-27T20:05:21.219496Z","shell.execute_reply.started":"2024-08-27T20:05:21.211568Z","shell.execute_reply":"2024-08-27T20:05:21.218555Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Реализуем создание скрипта для обучения модели и запустим его.","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/train.py\n\nimport os.path as osp\nfrom basicsr.train import train_pipeline\n\nif __name__ == '__main__':\n    root_path = osp.abspath(osp.join('', osp.pardir, osp.pardir, 'kaggle/working'))\n    train_pipeline(root_path)","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:21.223299Z","iopub.execute_input":"2024-08-27T20:05:21.223660Z","iopub.status.idle":"2024-08-27T20:05:21.233546Z","shell.execute_reply.started":"2024-08-27T20:05:21.223628Z","shell.execute_reply":"2024-08-27T20:05:21.232725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!python /kaggle/working/train.py \\\n-opt /kaggle/working/clustering_details.yaml \\\n--auto_resume","metadata":{"execution":{"iopub.status.busy":"2024-08-27T20:05:21.234469Z","iopub.execute_input":"2024-08-27T20:05:21.234781Z","iopub.status.idle":"2024-08-28T04:23:34.059787Z","shell.execute_reply.started":"2024-08-27T20:05:21.234750Z","shell.execute_reply":"2024-08-28T04:23:34.057781Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Обучение было прервано искусственно, принимая во внимание длительность его реализации. В связи с этим и возникла ошибка `KeyboardInterrupt`. Далее реализован скрипт для инференса модели.","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/infer.py\n\nimport argparse\nimport cv2\nimport glob\nimport os\nfrom basicsr.archs.rrdbnet_arch import RRDBNet\nfrom basicsr.utils.download_util import load_file_from_url\n\n# from realesrgan import RealESRGANer\n# from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n\nimport cv2\nimport math\nimport numpy as np\nimport os\nimport queue\nimport threading\nimport torch\nfrom basicsr.utils.download_util import load_file_from_url\nfrom torch.nn import functional as F\n\nROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\n\n\nclass RealESRGANer():\n    \"\"\"A helper class for upsampling images with RealESRGAN.\n\n    Args:\n        scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.\n        model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).\n        model (nn.Module): The defined network. Default: None.\n        tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop\n            input images into tiles, and then process each of them. Finally, they will be merged into one image.\n            0 denotes for do not use tile. Default: 0.\n        tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.\n        pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.\n        half (float): Whether to use half precision during inference. Default: False.\n    \"\"\"\n\n    def __init__(self,\n                 scale,\n                 model_path,\n                 dni_weight=None,\n                 model=None,\n                 tile=0,\n                 tile_pad=10,\n                 pre_pad=10,\n                 half=False,\n                 device=None,\n                 gpu_id=None):\n        self.scale = scale\n        self.tile_size = tile\n        self.tile_pad = tile_pad\n        self.pre_pad = pre_pad\n        self.mod_scale = None\n        self.half = half\n\n        # initialize model\n        if gpu_id:\n            self.device = torch.device(\n                f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device\n        else:\n            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device\n\n        if isinstance(model_path, list):\n            # dni\n            assert len(model_path) == len(dni_weight), 'model_path and dni_weight should have the save length.'\n            loadnet = self.dni(model_path[0], model_path[1], dni_weight)\n        else:\n            # if the model_path starts with https, it will first download models to the folder: weights\n            if model_path.startswith('https://'):\n                model_path = load_file_from_url(\n                    url=model_path, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)\n            loadnet = torch.load(model_path, map_location=torch.device('cpu'))\n\n        # prefer to use params_ema\n        if 'params_ema' in loadnet:\n            keyname = 'params_ema'\n        else:\n            keyname = 'params'\n        model.load_state_dict(loadnet[keyname], strict=True)\n\n        model.eval()\n        self.model = model.to(self.device)\n        if self.half:\n            self.model = self.model.half()\n\n    def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'):\n        \"\"\"Deep network interpolation.\n\n        ``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition``\n        \"\"\"\n        net_a = torch.load(net_a, map_location=torch.device(loc))\n        net_b = torch.load(net_b, map_location=torch.device(loc))\n        for k, v_a in net_a[key].items():\n            net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k]\n        return net_a\n\n    def pre_process(self, img):\n        \"\"\"Pre-process, such as pre-pad and mod pad, so that the images can be divisible\n        \"\"\"\n        img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()\n        self.img = img.unsqueeze(0).to(self.device)\n        if self.half:\n            self.img = self.img.half()\n\n        # pre_pad\n        if self.pre_pad != 0:\n            self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')\n        # mod pad for divisible borders\n        if self.scale == 2:\n            self.mod_scale = 2\n        elif self.scale == 1:\n            self.mod_scale = 4\n        if self.mod_scale is not None:\n            self.mod_pad_h, self.mod_pad_w = 0, 0\n            _, _, h, w = self.img.size()\n            if (h % self.mod_scale != 0):\n                self.mod_pad_h = (self.mod_scale - h % self.mod_scale)\n            if (w % self.mod_scale != 0):\n                self.mod_pad_w = (self.mod_scale - w % self.mod_scale)\n            self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')\n\n    def process(self):\n        # model inference\n        self.output = self.model(self.img)\n\n    def tile_process(self):\n        \"\"\"It will first crop input images to tiles, and then process each tile.\n        Finally, all the processed tiles are merged into one images.\n\n        Modified from: https://github.com/ata4/esrgan-launcher\n        \"\"\"\n        batch, channel, height, width = self.img.shape\n        output_height = height * self.scale\n        output_width = width * self.scale\n        output_shape = (batch, channel, output_height, output_width)\n\n        # start with black image\n        self.output = self.img.new_zeros(output_shape)\n        tiles_x = math.ceil(width / self.tile_size)\n        tiles_y = math.ceil(height / self.tile_size)\n\n        # loop over all tiles\n        for y in range(tiles_y):\n            for x in range(tiles_x):\n                # extract tile from input image\n                ofs_x = x * self.tile_size\n                ofs_y = y * self.tile_size\n                # input tile area on total image\n                input_start_x = ofs_x\n                input_end_x = min(ofs_x + self.tile_size, width)\n                input_start_y = ofs_y\n                input_end_y = min(ofs_y + self.tile_size, height)\n\n                # input tile area on total image with padding\n                input_start_x_pad = max(input_start_x - self.tile_pad, 0)\n                input_end_x_pad = min(input_end_x + self.tile_pad, width)\n                input_start_y_pad = max(input_start_y - self.tile_pad, 0)\n                input_end_y_pad = min(input_end_y + self.tile_pad, height)\n\n                # input tile dimensions\n                input_tile_width = input_end_x - input_start_x\n                input_tile_height = input_end_y - input_start_y\n                tile_idx = y * tiles_x + x + 1\n                input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]\n\n                # upscale tile\n                try:\n                    with torch.no_grad():\n                        output_tile = self.model(input_tile)\n                except RuntimeError as error:\n                    print('Error', error)\n                print(f'\\tTile {tile_idx}/{tiles_x * tiles_y}')\n\n                # output tile area on total image\n                output_start_x = input_start_x * self.scale\n                output_end_x = input_end_x * self.scale\n                output_start_y = input_start_y * self.scale\n                output_end_y = input_end_y * self.scale\n\n                # output tile area without padding\n                output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale\n                output_end_x_tile = output_start_x_tile + input_tile_width * self.scale\n                output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale\n                output_end_y_tile = output_start_y_tile + input_tile_height * self.scale\n\n                # put tile into output image\n                self.output[:, :, output_start_y:output_end_y,\n                            output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,\n                                                                       output_start_x_tile:output_end_x_tile]\n\n    def post_process(self):\n        # remove extra pad\n        if self.mod_scale is not None:\n            _, _, h, w = self.output.size()\n            self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]\n        # remove prepad\n        if self.pre_pad != 0:\n            _, _, h, w = self.output.size()\n            self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]\n        return self.output\n\n    @torch.no_grad()\n    def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):\n        h_input, w_input = img.shape[0:2]\n        # img: numpy\n        img = img.astype(np.float32)\n        if np.max(img) > 256:  # 16-bit image\n            max_range = 65535\n            print('\\tInput is a 16-bit image')\n        else:\n            max_range = 255\n        img = img / max_range\n        if len(img.shape) == 2:  # gray image\n            img_mode = 'L'\n            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)\n        elif img.shape[2] == 4:  # RGBA image with alpha channel\n            img_mode = 'RGBA'\n            alpha = img[:, :, 3]\n            img = img[:, :, 0:3]\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            if alpha_upsampler == 'realesrgan':\n                alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)\n        else:\n            img_mode = 'RGB'\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        # ------------------- process image (without the alpha channel) ------------------- #\n        self.pre_process(img)\n        if self.tile_size > 0:\n            self.tile_process()\n        else:\n            self.process()\n        output_img = self.post_process()\n        output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n        output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))\n        if img_mode == 'L':\n            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)\n\n        # ------------------- process the alpha channel if necessary ------------------- #\n        if img_mode == 'RGBA':\n            if alpha_upsampler == 'realesrgan':\n                self.pre_process(alpha)\n                if self.tile_size > 0:\n                    self.tile_process()\n                else:\n                    self.process()\n                output_alpha = self.post_process()\n                output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()\n                output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))\n                output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)\n            else:  # use the cv2 resize for alpha channel\n                h, w = alpha.shape[0:2]\n                output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)\n\n            # merge the alpha channel\n            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)\n            output_img[:, :, 3] = output_alpha\n\n        # ------------------------------ return ------------------------------ #\n        if max_range == 65535:  # 16-bit image\n            output = (output_img * 65535.0).round().astype(np.uint16)\n        else:\n            output = (output_img * 255.0).round().astype(np.uint8)\n\n        if outscale is not None and outscale != float(self.scale):\n            output = cv2.resize(\n                output, (\n                    int(w_input * outscale),\n                    int(h_input * outscale),\n                ), interpolation=cv2.INTER_LANCZOS4)\n\n        return output, img_mode\n\n\nclass PrefetchReader(threading.Thread):\n    \"\"\"Prefetch images.\n\n    Args:\n        img_list (list[str]): A image list of image paths to be read.\n        num_prefetch_queue (int): Number of prefetch queue.\n    \"\"\"\n\n    def __init__(self, img_list, num_prefetch_queue):\n        super().__init__()\n        self.que = queue.Queue(num_prefetch_queue)\n        self.img_list = img_list\n\n    def run(self):\n        for img_path in self.img_list:\n            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)\n            self.que.put(img)\n\n        self.que.put(None)\n\n    def __next__(self):\n        next_item = self.que.get()\n        if next_item is None:\n            raise StopIteration\n        return next_item\n\n    def __iter__(self):\n        return self\n\nfrom basicsr.utils.registry import ARCH_REGISTRY\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\n\n@ARCH_REGISTRY.register()\nclass SRVGGNetCompact(nn.Module):\n    \"\"\"A compact VGG-style network structure for super-resolution.\n\n    It is a compact network structure, which performs upsampling in the last layer and no convolution is\n    conducted on the HR feature space.\n\n    Args:\n        num_in_ch (int): Channel number of inputs. Default: 3.\n        num_out_ch (int): Channel number of outputs. Default: 3.\n        num_feat (int): Channel number of intermediate features. Default: 64.\n        num_conv (int): Number of convolution layers in the body network. Default: 16.\n        upscale (int): Upsampling factor. Default: 4.\n        act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.\n    \"\"\"\n\n    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):\n        super(SRVGGNetCompact, self).__init__()\n        self.num_in_ch = num_in_ch\n        self.num_out_ch = num_out_ch\n        self.num_feat = num_feat\n        self.num_conv = num_conv\n        self.upscale = upscale\n        self.act_type = act_type\n\n        self.body = nn.ModuleList()\n        # the first conv\n        self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))\n        # the first activation\n        if act_type == 'relu':\n            activation = nn.ReLU(inplace=True)\n        elif act_type == 'prelu':\n            activation = nn.PReLU(num_parameters=num_feat)\n        elif act_type == 'leakyrelu':\n            activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)\n        self.body.append(activation)\n\n        # the body structure\n        for _ in range(num_conv):\n            self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))\n            # activation\n            if act_type == 'relu':\n                activation = nn.ReLU(inplace=True)\n            elif act_type == 'prelu':\n                activation = nn.PReLU(num_parameters=num_feat)\n            elif act_type == 'leakyrelu':\n                activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)\n            self.body.append(activation)\n\n        # the last conv\n        self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))\n        # upsample\n        self.upsampler = nn.PixelShuffle(upscale)\n\n    def forward(self, x):\n        out = x\n        for i in range(0, len(self.body)):\n            out = self.body[i](out)\n\n        out = self.upsampler(out)\n        # add the nearest upsampled image, so that the network learns the residual\n        base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')\n        out += base\n        return out\n\n\n\ndef main():\n    \"\"\"Inference demo for Real-ESRGAN.\n    \"\"\"\n    parser = argparse.ArgumentParser()\n    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')\n    parser.add_argument(\n        '-n',\n        '--model_name',\n        type=str,\n        default='RealESRGAN_x4plus',\n        help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus | '\n              'realesr-animevideov3 | realesr-general-x4v3'))\n    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')\n    parser.add_argument(\n        '-dn',\n        '--denoise_strength',\n        type=float,\n        default=0.5,\n        help=('Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. '\n              'Only used for the realesr-general-x4v3 model'))\n    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')\n    parser.add_argument(\n        '--model_path', type=str, default=None, help='[Option] Model path. Usually, you do not need to specify it')\n    parser.add_argument('--suffix', type=str, default='', help='Suffix of the restored image')\n    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')\n    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')\n    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')\n    parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')\n    parser.add_argument(\n        '--fp32', action='store_true', help='Use fp32 precision during inference. Default: fp16 (half precision).')\n    parser.add_argument(\n        '--alpha_upsampler',\n        type=str,\n        default='realesrgan',\n        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')\n    parser.add_argument(\n        '--ext',\n        type=str,\n        default='auto',\n        help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')\n    parser.add_argument(\n        '-g', '--gpu-id', type=int, default=None, help='gpu device to use (default=None) can be 0,1,2 for multi-gpu')\n\n    args = parser.parse_args()\n\n    # determine models according to model names\n    args.model_name = args.model_name.split('.')[0]\n    if args.model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model\n        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)\n        netscale = 4\n        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']\n    elif args.model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model\n        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)\n        netscale = 4\n        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']\n    elif args.model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks\n        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)\n        netscale = 4\n        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']\n    elif args.model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model\n        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)\n        netscale = 2\n        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']\n    elif args.model_name == 'realesr-animevideov3':  # x4 VGG-style model (XS size)\n        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')\n        netscale = 4\n        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth']\n    elif args.model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)\n        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')\n        netscale = 4\n        file_url = [\n            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',\n            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'\n        ]\n\n    # determine model paths\n    if args.model_path is not None:\n        model_path = args.model_path\n    else:\n        model_path = os.path.join('weights', args.model_name + '.pth')\n        if not os.path.isfile(model_path):\n            ROOT_DIR = os.path.dirname(os.path.abspath(__file__))\n            for url in file_url:\n                # model_path will be updated\n                model_path = load_file_from_url(\n                    url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)\n\n    # use dni to control the denoise strength\n    dni_weight = None\n    if args.model_name == 'realesr-general-x4v3' and args.denoise_strength != 1:\n        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')\n        model_path = [model_path, wdn_model_path]\n        dni_weight = [args.denoise_strength, 1 - args.denoise_strength]\n\n    # restorer\n    upsampler = RealESRGANer(\n        scale=netscale,\n        model_path=model_path,\n        dni_weight=dni_weight,\n        model=model,\n        tile=args.tile,\n        tile_pad=args.tile_pad,\n        pre_pad=args.pre_pad,\n        half=not args.fp32,\n        gpu_id=args.gpu_id)\n\n    if args.face_enhance:  # Use GFPGAN for face enhancement\n        from gfpgan import GFPGANer\n        face_enhancer = GFPGANer(\n            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',\n            upscale=args.outscale,\n            arch='clean',\n            channel_multiplier=2,\n            bg_upsampler=upsampler)\n    os.makedirs(args.output, exist_ok=True)\n\n    if os.path.isfile(args.input):\n        paths = [args.input]\n    else:\n        paths = sorted(glob.glob(os.path.join(args.input, '*')))\n\n    for idx, path in enumerate(paths):\n        imgname, extension = os.path.splitext(os.path.basename(path))\n        print('Testing', idx, imgname)\n\n        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n        if len(img.shape) == 3 and img.shape[2] == 4:\n            img_mode = 'RGBA'\n        else:\n            img_mode = None\n\n        try:\n            if args.face_enhance:\n                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)\n            else:\n                output, _ = upsampler.enhance(img, outscale=args.outscale)\n        except RuntimeError as error:\n            print('Error', error)\n            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')\n        else:\n            if args.ext == 'auto':\n                extension = extension[1:]\n            else:\n                extension = args.ext\n            if img_mode == 'RGBA':  # RGBA images should be saved in png format\n                extension = 'png'\n            if args.suffix == '':\n                save_path = os.path.join(args.output, f'{imgname}.{extension}')\n            else:\n                save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')\n            cv2.imwrite(save_path, output)\n\n\nif __name__ == '__main__':\n    main()","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:26:18.295570Z","iopub.execute_input":"2024-08-28T06:26:18.295966Z","iopub.status.idle":"2024-08-28T06:26:18.323168Z","shell.execute_reply.started":"2024-08-28T06:26:18.295925Z","shell.execute_reply":"2024-08-28T06:26:18.322212Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Запуск скрипта инференса","metadata":{}},{"cell_type":"code","source":"%%capture\n\n!python /kaggle/working/infer.py \\\n--model_path /kaggle/input/net-g-20000/net_g_20000.pth \\\n-i /kaggle/input/super-resolution-in-video-games/test/lr \\\n-o /kaggle/working/results","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:28:24.759653Z","iopub.execute_input":"2024-08-28T06:28:24.760199Z","iopub.status.idle":"2024-08-28T06:34:33.288938Z","shell.execute_reply.started":"2024-08-28T06:28:24.760124Z","shell.execute_reply":"2024-08-28T06:34:33.287644Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Скрипт энкодера для оформления результатов","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/encode.py\n\nimport base64\nimport os\nimport zlib\nfrom argparse import ArgumentParser\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\n\ndef encode(img: np.ndarray) -> bytes:\n    \"\"\"\n    Lossless encoding of images for submission on kaggle platform.\n    Kaggle doesn't support submissons not in csv format.\n    In the same way target images should be stored.\n\n    Parameters\n    ----------\n    img : np.ndarray\n        cv2.imread(f) - BGR image in (h, w, c) format, c = 3 (even for png format).\n\n    Returns\n    -------\n    bytes\n        Encoded image as bytes.\n    \"\"\"\n    img_to_encode = img.astype(np.uint8)\n    img_to_encode = img_to_encode.flatten()\n    img_to_encode = np.append(img_to_encode, -1)\n\n    cnt, rle = 1, []\n    for i in range(1, img_to_encode.shape[0]):\n        if img_to_encode[i] == img_to_encode[i - 1]:\n            cnt += 1\n            if cnt > 255:\n                rle += [img_to_encode[i - 1], 255]\n                cnt = 1\n        else:\n            rle += [img_to_encode[i - 1], cnt]\n            cnt = 1\n\n    compressed = zlib.compress(bytes(rle), zlib.Z_BEST_COMPRESSION)\n    base64_bytes = base64.b64encode(compressed)\n    return base64_bytes\n\n\ndef decode(encoded_img: bytes) -> np.ndarray:\n    \"\"\"\n    Reverse operation for encode function to get original images.\n\n    Parameters\n    ----------\n    encoded_img : bytes\n        Encoded image as bytes.\n\n    Returns\n    -------\n    np.ndarray\n        BGR image in (h, w, c) format, c = 3 (even for png format).\n    \"\"\"\n    rle = zlib.decompress(base64.b64decode(encoded_img))\n    decoded_img = []\n    for i in range(0, len(rle), 2):\n        decoded_img += [rle[i]] * rle[i + 1]\n    return np.array(decoded_img, dtype=np.uint8)\n\n\ndef encode_folder(\n    folder: str, save_path: str = \"solution.csv\", public_size: float = 0.3\n) -> None:\n    \"\"\"\n    Encode images from folder (target images from test set) and save into solution.csv.\n\n    Parameters\n    ----------\n    folder : str\n        Path to directory with target images.\n    save_path : str\n        Path to save solution.csv.\n    public_size : float\n        Size of public part of the dataset for the leaderboard.\n\n    Returns\n    -------\n    None\n    \"\"\"\n    dct = {\"filename\": [], \"rle\": []}\n\n    files = os.listdir(folder)\n    public_ind = np.random.choice(\n        [i for i in range(len(files))],\n        size=int(public_size * len(files)),\n        replace=False,\n    )\n    for i in tqdm(range(len(files))):\n        img = cv2.imread(os.path.join(folder, files[i]))\n        h, w, c = img.shape\n        encoded_img = encode(img)\n        decoded_img = decode(encoded_img).reshape(h, w, c)\n        assert (img - decoded_img).sum() == 0, \"encoded != decoded\"\n        dct[\"filename\"] += [files[i]]\n        dct[\"rle\"] += [encode(img)]\n\n    df = pd.DataFrame(dct)\n    df.to_csv(save_path, index=True, index_label=\"id\")\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\n        \"-f\",\n        \"--folder\",\n        type=str,\n        required=True,\n        help=\"path to folder with target images\",\n    )\n    parser.add_argument(\n        \"-s\",\n        \"--save-path\",\n        type=str,\n        default=\"solution.csv\",\n        help=\"path to save solution.csv\",\n    )\n    parser.add_argument(\n        \"-p\",\n        \"--public-size\",\n        type=float,\n        default=0.3,\n        help=\"size of public part of the dataset for the leaderboard\",\n    )\n    parser.add_argument(\n        \"--seed\",\n        type=int,\n        default=42,\n        help=\"random seed value\",\n    )\n    args = parser.parse_args()\n\n    np.random.seed(args.seed)\n    encode_folder(args.folder, save_path=args.save_path, public_size=args.public_size)","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:34:33.291057Z","iopub.execute_input":"2024-08-28T06:34:33.291422Z","iopub.status.idle":"2024-08-28T06:34:33.300095Z","shell.execute_reply.started":"2024-08-28T06:34:33.291382Z","shell.execute_reply":"2024-08-28T06:34:33.299190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Запуск скрипта энкодера","metadata":{}},{"cell_type":"code","source":"!python /kaggle/working/encode.py \\\n-f /kaggle/working/results \\\n-s /kaggle/working/solution.csv","metadata":{"execution":{"iopub.status.busy":"2024-08-28T06:35:33.420824Z","iopub.execute_input":"2024-08-28T06:35:33.421456Z","iopub.status.idle":"2024-08-28T07:12:08.206741Z","shell.execute_reply.started":"2024-08-28T06:35:33.421413Z","shell.execute_reply":"2024-08-28T07:12:08.205747Z"},"trusted":true},"execution_count":null,"outputs":[]}]}