Update models/TextEnhancement.py

models/TextEnhancement.py

@@ -1,367 +1,367 @@
-# -*- 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_CUBIC)
-                
-                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_CUBIC)
-        in_img = cv2.resize(img, (w_norm, h_norm), interpolation=cv2.INTER_CUBIC)
-        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_CUBIC)
-        
-
-        #--------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')
-
-
+# -*- 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')
+
+