models/TextEnhancement.py

# -*- coding: utf-8 -*-
import cv2
import os.path as osp
import torch
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import logging
logging.getLogger('modelscope').disabled = True

from cnstd import CnStd
from utils.utils_transocr import get_alphabet
from utils.yolo_ocr_xloc import get_yolo_ocr_xloc
from ultralytics import YOLO

from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from networks import *
import warnings
warnings.filterwarnings('ignore')

from modelscope import snapshot_download


##########################################################################################
###############Text Restoration Model revised by xiaoming li
##########################################################################################

alphabet_path = './models/benchmark_cvpr23.txt'
CommonWordsForOCR = get_alphabet(alphabet_path)
CommonWords = CommonWordsForOCR[2:-1]



def str2idx(text):
    idx = []
    for t in text:
        idx.append(CommonWords.index(t) if t in CommonWords else 3484) #3955
    return idx

def get_parameter_details(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    return num_params / 1e6

def tensor2numpy(tensor):
    tensor = tensor * 0.5 + 0.5
    tensor = tensor.squeeze(0).permute(1, 2, 0).flip(2)
    return np.clip(tensor.float().cpu().numpy(), 0, 1) * 255.0


class MARCONetPlus(object):
    def __init__(self, WEncoderPath=None, PriorModelPath=None, SRModelPath=None, YoloPath=None, device='cuda'):
        self.device = device

        modelscope_dir = snapshot_download('damo/cv_convnextTiny_ocr-recognition-general_damo', cache_dir='./checkpoints/modelscope_ocr')
        self.modelscope_ocr_recognition = pipeline(Tasks.ocr_recognition, model=modelscope_dir)
        self.yolo_character = YOLO(YoloPath)

        self.modelWEncoder = PSPEncoder() # WEncoder()
        self.modelWEncoder.load_state_dict(torch.load(WEncoderPath)['params'], strict=True)
        self.modelWEncoder.eval()
        self.modelWEncoder.to(device)

        self.modelPrior = TextPriorModel()
        self.modelPrior.load_state_dict(torch.load(PriorModelPath)['params'], strict=True)
        self.modelPrior.eval()
        self.modelPrior.to(device)
    
        self.modelSR = SRNet()
        self.modelSR.load_state_dict(torch.load(SRModelPath)['params'], strict=True)
        self.modelSR.eval()
        self.modelSR.to(device)


        print('='*128)
        print('{:>25s} : {:.2f} M Parameters'.format('modelWEncoder', get_parameter_details(self.modelWEncoder)))
        print('{:>25s} : {:.2f} M Parameters'.format('modelPrior', get_parameter_details(self.modelPrior)))
        print('{:>25s} : {:.2f} M Parameters'.format('modelSR', get_parameter_details(self.modelSR)))
        print('='*128)

        torch.cuda.empty_cache()
        self.cnstd = CnStd(model_name='db_resnet34',rotated_bbox=True, model_backend='pytorch', box_score_thresh=0.3, min_box_size=10, context=device)
        self.insize = 32


    def handle_texts(self, img, bg=None, sf=4, is_aligned=False, lq_label=None):
        '''
        Parameters:
            img: RGB 0~255.
        '''

        height, width = img.shape[:2]
        bg_height, bg_width = bg.shape[:2]
        print(' ' * 25 + f' ... The input->output image size is {bg_height//sf}*{bg_width//sf}->{bg_height}*{bg_width}')
        
        full_mask_blur = np.zeros(bg.shape, dtype=np.float32)
        full_mask_noblur = np.zeros(bg.shape, dtype=np.float32)
        full_text_img = np.zeros(bg.shape, dtype=np.float32) #+255

        orig_texts, enhanced_texts, debug_texts, pred_texts = [], [], [], []
        ocr_scores = []

        if not is_aligned:
            box_infos = self.cnstd.detect(img)
            for iix, box_info in enumerate(box_infos['detected_texts']):
                box = box_info['box'].astype(int)# left top, right top, right bottom, left bottom, [width, height]
                score = box_info['score']
                if score < 0.5:
                    continue

                extend_box = box.copy()
                w = int(np.linalg.norm(box[0] - box[1]))
                h = int(np.linalg.norm(box[0] - box[3]))

                # extend the bounding box
                extend_lr = 0.15 * h  
                extend_tb = 0.05 * h 
                vec_w = (box[1] - box[0]) / w
                vec_h = (box[3] - box[0]) / h

                extend_box[0] = box[0] - vec_w * extend_lr - vec_h * extend_tb
                extend_box[1] = box[1] + vec_w * extend_lr - vec_h * extend_tb
                extend_box[2] = box[2] + vec_w * extend_lr + vec_h * extend_tb
                extend_box[3] = box[3] - vec_w * extend_lr + vec_h * extend_tb
                extend_box = extend_box.astype(int)

                w = int(np.linalg.norm(extend_box[0] - extend_box[1]))
                h = int(np.linalg.norm(extend_box[0] - extend_box[3]))
                
                if w > h:
                    ref_h = self.insize
                    ref_w = int(ref_h * w / h)
                else:
                    print(' ' * 25 + ' ... Can not handle vertical text temporarily')
                    continue

                ref_point = np.float32([[0,0], [ref_w, 0], [ref_w, ref_h], [0, ref_h]])
                det_point = np.float32(extend_box)

                matrix = cv2.getPerspectiveTransform(det_point, ref_point)
                inv_matrix = cv2.getPerspectiveTransform(ref_point*sf, det_point*sf)

                cropped_img = cv2.warpPerspective(img, matrix, (ref_w, ref_h), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_LINEAR)


                in_img, SQ, save_debug, pred_text, preds_locs_txt = self._process_text_line(cropped_img)
                if in_img is None:
                    continue
                h_crop, w_crop = cropped_img.shape[:2]
                SQ = cv2.resize(SQ, (w_crop * sf, h_crop * sf), interpolation=cv2.INTER_LINEAR)
                
                debug_texts.append(save_debug)
                orig_texts.append(in_img)
                enhanced_texts.append(SQ)
                pred_texts.append(''.join(pred_text))

                tmp_mask = np.ones(SQ.shape).astype(float)
                warp_mask = cv2.warpPerspective(tmp_mask, inv_matrix, (bg_width, bg_height), flags=3)
                warp_img = cv2.warpPerspective(SQ, inv_matrix, (bg_width, bg_height), flags=3)
 

                # erode and blur based on the height of text region
                blur_pad = int(h // 6)

                if blur_pad % 2 == 0:
                    blur_pad += 1
                blur_radius = (blur_pad - 1) // 2
                erode_radius = blur_radius + 1   
                erode_pad = 2 * erode_radius + 1

                kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erode_pad, erode_pad))
                warp_mask_erode = cv2.erode(warp_mask, kernel_erode, iterations=1)

                # warp_mask_blur = cv2.GaussianBlur(warp_mask_erode, (blur_pad, blur_pad), 0)
                warp_mask_blur = cv2.blur(warp_mask_erode, (blur_pad, blur_pad))

                full_text_img = full_text_img + warp_img
                full_mask_blur = full_mask_blur + warp_mask_blur
                full_mask_noblur = full_mask_noblur + warp_mask
                
                ocr_scores.append(score)


            index = full_mask_noblur > 0
            full_text_img[index] = full_text_img[index]/full_mask_noblur[index]

            full_mask_blur = np.clip(full_mask_blur, 0, 1)
            # fuse the text region back to the background
            final_img = full_text_img * full_mask_blur + bg * (1 - full_mask_blur)


            return final_img, orig_texts, enhanced_texts, debug_texts, pred_texts #, ocr_scores
        
        else: #aligned
            
            in_img, SQ, save_debug, pred_text, preds_locs_txt = self._process_text_line(img)
            if in_img is not None:
                debug_texts.append(save_debug)
                orig_texts.append(in_img)
                enhanced_texts.append(SQ)
                pred_texts.append(''.join(pred_text))
        
        return img, orig_texts, enhanced_texts, debug_texts, pred_texts #, preds_locs_txt

    def _process_text_line(self, img):
        """
        Process a single text line region for text enhancement.
        
        Args:
            img: Input text image
        
        """


        height, width = img.shape[:2]
        if height > width:
            print(' ' * 25 + ' ... Can not handle vertical text temporarily')
            return (None,) * 5
        
        w_norm = int(self.insize * width / height) // 4 * 4
        h_norm = self.insize

        img = cv2.resize(img, (w_norm*4, h_norm*4), interpolation=cv2.INTER_LINEAR)
        in_img = cv2.resize(img, (w_norm, h_norm), interpolation=cv2.INTER_LINEAR)
        ShowLQ = img[:,:,::-1]

        LQ_HeightNorm = transforms.ToTensor()(in_img)
        LQ_HeightNorm = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(LQ_HeightNorm).unsqueeze(0).to(self.device)


        '''
        Step 1: Predicting the character labels, bounding boxes.
        '''

        recognized_boxes, pred_text, char_x_centers = get_yolo_ocr_xloc(
            img,                        # input image, RGB 0~255
            yolo_model=self.yolo_character,       # YOLO model instance for character detection
            ocr_pipeline=self.modelscope_ocr_recognition,  # OCR pipeline/model for character recognition
            num_cropped_boxes=5,             # Number of adjacent character boxes to include in each cropped segment (window size)
            expand_px=1,                     # Number of pixels to expand each crop region on all sides (except first/last)
            expand_px_for_first_last_cha=12, # Number of pixels to expand the crop region for the first and last character (left/right respectively)
            yolo_iou=0.1,                    # IOU threshold for YOLO non-max suppression (NMS)
            yolo_conf=0.07                   # Confidence threshold for YOLO detection
        )

        print('{:>25s} ... Recognized chars: {}'.format(' ', ''.join(pred_text)))
        loc_sr = torch.tensor(char_x_centers, device=self.device).unsqueeze(0)

        
        # show character location 
        pad = 1
        ShowPredLoc = ShowLQ.copy()
        for l in range(len(pred_text)):
            center_pred_w = int(loc_sr[0][l].item())
            if center_pred_w > 0: 
                ShowPredLoc[:, max(0, center_pred_w-pad):min(center_pred_w+pad, ShowPredLoc.shape[1]), :] = 0
                ShowPredLoc[:, max(0, center_pred_w-pad):min(center_pred_w+pad, ShowPredLoc.shape[1]), 1] = 255


        '''
        Step 2: Character Prior Generation
        '''

        with torch.no_grad():
            w = self.modelWEncoder(LQ_HeightNorm, loc_sr)
        
        predict_characters128 = []
        predict_characters64 = []
        predict_characters32 = []

        for b in range(w.size(0)):
            w0 = w[b,...].clone() #16*512
            pred_label = str2idx(pred_text)
            pred_label = torch.Tensor(pred_label).type(torch.LongTensor).view(-1, 1)#.to(device)

            with torch.no_grad():
                prior_cha, prior_fea64, prior_fea32 = self.modelPrior(styles=w0[:len(pred_text),:], labels=pred_label, noise=None) #b *n * w * h

            predict_characters128.append(prior_cha)
            predict_characters64.append(prior_fea64)
            predict_characters32.append(prior_fea32)
            
        
        '''
        Step 3: Character SR
        '''

        with torch.no_grad():
            extend_right_width = extend_left_width = h_norm // 2
            LQ_HeightNorm_WidthExtend = F.pad(LQ_HeightNorm, (extend_left_width, extend_right_width, 0, 0), mode='replicate')
            
            preds_locs_txt = ''
            loc_for_extend_sr = loc_sr.clone()
            for i in range(len(pred_text)):
                preds_locs_txt += str(int(loc_for_extend_sr[0][i].cpu().item()))+'_'
                loc_for_extend_sr[0][i] = loc_for_extend_sr[0][i] + extend_left_width * 4
                
            SR = self.modelSR(LQ_HeightNorm_WidthExtend, predict_characters64, predict_characters32, loc_for_extend_sr)
    
        SR = tensor2numpy(SR)[:, extend_left_width * 4:extend_left_width * 4 + w_norm*4, ::-1]

        
        # reduce color inconsistency，use ab channel from in_img
        # sr_lab = cv2.cvtColor(SR.astype(np.uint8), cv2.COLOR_BGR2LAB)
        # target_size = (SR.shape[1], SR.shape[0])
        # in_img_resize = cv2.resize(in_img, target_size, interpolation=cv2.INTER_LINEAR)
        # in_img_lab = cv2.cvtColor(in_img_resize.astype(np.uint8), cv2.COLOR_BGR2LAB)
        # sr_lab[:,:,1:] = in_img_lab[:,:,1:]
        # SR = cv2.cvtColor(sr_lab, cv2.COLOR_LAB2BGR)

        
        prior128 = []
        pad = 2
        for prior in predict_characters128:
            for ii, p in enumerate(prior):
                prior128.append(p)
        prior128 = torch.cat(prior128, dim=2)
        prior128 = prior128 * 0.5 + 0.5
        prior128 = prior128.permute(1, 2, 0).flip(2)
        prior128 = np.clip(prior128.float().cpu().numpy(), 0, 1) * 255.0
        prior128 = np.repeat(prior128, 3, axis=2)

        ShowPrior = cv2.resize(prior128, (SR.shape[1], int(128 * SR.shape[1] / prior128.shape[1])), interpolation=cv2.INTER_LINEAR)
        

        #--------Fuse the structure prior to the LR input to show the details of alignment--------------
        fusion_bg = np.zeros_like(SR, dtype=np.float32)
        w4 = w_norm * 4

        for iii, c in enumerate(loc_sr[0].int()):
            current_prior = prior128[:, iii*128:(iii+1)*128, :]
            center_loc = c.item()

            x1 = max(center_loc - 64, 0)
            x2 = min(center_loc + 64, w4)
            y1 = max(64 - center_loc, 0)
            y2 = y1 + (x2 - x1)
            try:
                fusion_bg[:, x1:x2, :] += current_prior[:, y1:y2, :]
            except:
                return (None,) * 5

        
        mask = fusion_bg / 255.0
        fusion_bg[:,:,0] = 0
        fusion_bg[:,:,2] = 0


        ShowLQ = ShowLQ[:,:,::-1]
        fusion_bg = fusion_bg.astype(ShowLQ.dtype)
        fusion_bg = fusion_bg * 0.3 * mask + ShowLQ * 0.7 * mask + (1-mask) * ShowLQ

        ShowPrior = cv2.normalize(ShowPrior, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

        save_debug = np.vstack((ShowLQ, ShowPredLoc[:,:,::-1], SR, ShowPrior, fusion_bg))

        return in_img, SR, save_debug, pred_text, preds_locs_txt



if __name__ == '__main__':
    print('Test')