Method _process_text_line

models/TextEnhancement.py:208–360

Process a single text line region for text enhancement. Args: img: Input text image

(self, img)

        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
        '''