demo.py

import os.path
import math
import argparse
import time
import random
import cv2
import numpy as np
from collections import OrderedDict
import logging
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch
from tqdm import tqdm
from glob import glob
from utils import utils_logger
from utils import utils_image as util
from utils import utils_option as option
from utils.utils_dist import get_dist_info, init_dist
from utils.metrics_collector import MetricsCollector
import pandas as pd
from data.select_dataset import define_Dataset
from models.select_model import define_Model
from torchvision.io import write_png
from models.loss_crnn2 import CRNNLoss
from models.loss_moran import MORANLoss
from prettytable import PrettyTable
from pathlib import Path
'''f
# --------------------------------------------
# training code for MSRResNet
# --------------------------------------------
# Kai Zhang (cskaizhang@gmail.com)
# github: https://github.com/cszn/KAIR
# --------------------------------------------
# https://github.com/xinntao/BasicSR
# --------------------------------------------
'''



def main(json_path='options/train_msrresnet_psnr.json'):

    '''
    # ----------------------------------------
    # Step--1 (prepare opt)
    # ----------------------------------------
    '''


    parser = argparse.ArgumentParser()
    parser.add_argument('--opt', type=str, default=json_path, help='Path to option JSON file.')
    parser.add_argument('--launcher', default='pytorch', help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--dist', default=False)
    parser.add_argument('--cluster', type=int, default=0)
    opt = option.parse(parser.parse_args().opt, is_train=True)
    output_dir = Path(opt["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    opt['task'] = opt['task'] + str(parser.parse_args().cluster)
    # opt['datasets']['train']['dataroot_H'] = opt['datasets']['train']['dataroot_H'].replace('cluster_id', str(parser.parse_args().cluster)) 
    # opt['datasets']['train']['dataroot_L'] = opt['datasets']['train']['dataroot_L'].replace('cluster_id', str(parser.parse_args().cluster)) 

    opt['dist'] = parser.parse_args().dist

    # ----------------------------------------
    # distributed settings
    # ----------------------------------------
    if opt['dist']:
        init_dist('pytorch')
    opt['rank'], opt['world_size'] = get_dist_info()

    if opt['rank'] == 0:
        util.mkdirs((path for key, path in opt['path'].items() if 'pretrained' not in key))

    border = opt['scale']
    opt = option.dict_to_nonedict(opt)


    # ----------------------------------------
    # configure logger
    # ----------------------------------------
    if opt['rank'] == 0:
        logger_name = 'train'
        utils_logger.logger_info(logger_name, os.path.join(opt['path']['log'], logger_name+'.log'))
        logger = logging.getLogger(logger_name)
        logger.info(option.dict2str(opt))

    # ----------------------------------------
    # seed
    # ----------------------------------------
    # seed = opt['train']['manual_seed']

    seed = random.randint(1, 10000)
    print('Random seed: {}'.format(seed))
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # ----------------------------------------
    # Creat dataloader
    # ----------------------------------------
    metrics_collector = MetricsCollector()
    # ----------------------------------------
    # 1) create_dataset
    # 2) creat_dataloader for train and test
    # ----------------------------------------
    for phase, dataset_opt in opt['datasets'].items():
        if phase == 'test':
            test_set = define_Dataset(dataset_opt)
            test_loader = DataLoader(test_set, batch_size=1,
                                     shuffle=False, num_workers=1,
                                     drop_last=False, pin_memory=True)
        else:
            raise NotImplementedError("Phase [%s] is not recognized." % phase)

    # ----------------------------------------
    # INIT
    # ----------------------------------------

    model = define_Model(opt)
    model.init_train()
    if opt['rank'] == 0:
        logger.info(model.info_network())
        # logger.info(model.info_params())

    scorer_crnn = CRNNLoss(model.device).to(model.device)
    scorer_crnn.eval()
    alphabet = '0:1:2:3:4:5:6:7:8:9:a:b:c:d:e:f:g:h:i:j:k:l:m:n:o:p:q:r:s:t:u:v:w:x:y:z:$'.split(':')
    scorer_moran = MORANLoss(
        nc=1,
        nclass=len(alphabet),
        nh=256,
        targetH=32,
        targetW=100,
        BidirDecoder=True,
        CUDA='cuda' in str(model.device)
    ).to(model.device)
    scorer_moran.eval()

    # -------------------------------
    # Testing
    # -------------------------------
    model.OCR_lossfn.eval()
    avg_psnr = 0.0
    avg_psnr_y = 0.0
    idx = 0
    correct_aster, correct_crnn, correct_moran, total = 0,0,0,0

    ocr_results = {'MORAN': {}, 'CRNN': {}, 'ASTER': {}}
    totals = {}
    with torch.no_grad():
        for test_data in test_loader:
            idx += 1
            image_name_ext = os.path.basename(test_data['L_path'][0])
            img_name, ext = os.path.splitext(image_name_ext)

            difficulty = img_name.split('_',2)[1]

            window_size = 8
            scale = 2
            _, _, h_old, w_old = test_data['L'].size()
            img_lq = test_data['L']
            h_pad = (h_old // window_size + 1) * window_size - h_old
            w_pad = (w_old // window_size + 1) * window_size - w_old

            img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
            img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]

            test_data['L'] = img_lq

            model.feed_data(test_data)
            model.test()

            visuals = model.current_visuals()
            
            L_img = util.tensor2uint(visuals['L'])
            E_img = util.tensor2uint(visuals['E'][..., :h_old * scale, :w_old * scale])

            visuals['E'] = visuals['E'][..., :h_old * scale, :w_old * scale]
            visuals['H'] = visuals['H'][..., :h_old * scale, :w_old * scale]

            H_img = util.tensor2uint(visuals['H'])
            cv2.imwrite(f'{output_dir}/{img_name}_SwinIR.png', E_img)
            
            # -----------------------
            # Calculate Accuracy
            # -----------------------
            visuals['E'] = visuals['E'].to(model.device)
            c, t = model.OCR_lossfn.score(visuals['E'], img_name.split('_')[-1])
            correct_aster += c
            total += t

            c_crnn, t_crnn = scorer_crnn.score(visuals['E'], img_name.split('_')[-1])
            c_moran, t_moran = scorer_moran.score(visuals['E'], img_name.split('_')[-1])

            ocr_results['ASTER'][difficulty] = ocr_results['ASTER'].get(difficulty, 0) + c
            ocr_results['MORAN'][difficulty] = ocr_results['MORAN'].get(difficulty, 0) + c_moran
            ocr_results['CRNN'][difficulty] = ocr_results['CRNN'].get(difficulty, 0) + c_crnn
            totals[difficulty] = totals.get(difficulty, 0) + t

            correct_crnn += c_crnn
            correct_moran += c_moran

            # -----------------------
            # calculate PSNR
            # -----------------------
            current_psnr = util.calculate_psnr(E_img, H_img, border=border)
            current_psnr_y = util.calculate_psnr_y(E_img, H_img, border=border)
            metrics_collector.save_mean('Test PSNR', current_psnr)
            metrics_collector.save_mean('Test PSNR Y', current_psnr_y)
            avg_psnr += current_psnr
            avg_psnr_y += current_psnr_y

            logger.info('{:->4d}--> {:>10s} | {:<4.2f}dB | {:<4.2f}dBy | {:<4.2f} | {:<4.2f} | {:<4.2f}'.format(idx, img_name, current_psnr, current_psnr_y, c/max(t,1), c_crnn/max(t_crnn,1), c_moran/max(t_moran, 1)))


    table = PrettyTable()
    table.field_names = ['Model', 'Easy Word Acc', 'Medium Word Acc', 'Hard Word Acc', 'Total Word Acc']
    for model in ['ASTER', 'MORAN', 'CRNN']:
        row = [model]
        for difficulty in ['easy', 'medium', 'hard']:
            row.append(f'{round(ocr_results[model][difficulty] / totals[difficulty] * 100, 1)}%')
        row.append(f'{round(sum(list(ocr_results[model].values())) / sum(list(totals.values())) * 100, 1)}%')
        table.add_row(row)
    logger.info(table)
    metrics_collector.save('Test OCR Accuracy', correct_aster / total)
    metrics_collector.save('Test CRNN Accuracy', correct_crnn / total)
    metrics_collector.save('Test MORAN Accuracy', correct_moran / total)
    logger.info(metrics_collector.get_metric_string())

if __name__ == '__main__':
    main()