PaddleX/paddlex/inference/models/text_recognition/result.py

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy

import PIL
from PIL import Image, ImageDraw, ImageFont

from ....utils.fonts import SIMFANG_FONT
from ...common.result import BaseCVResult, JsonMixin


class TextRecResult(BaseCVResult):

    def _to_str(self, *args, **kwargs):
        data = copy.deepcopy(self)
        data.pop("input_img")
        data.pop("vis_font")
        return JsonMixin._to_str(data, *args, **kwargs)

    def _to_json(self, *args, **kwargs):
        data = copy.deepcopy(self)
        data.pop("input_img")
        data.pop("vis_font")
        return JsonMixin._to_json(data, *args, **kwargs)

    def _to_img(self):
        """Draw label on image"""
        image = Image.fromarray(self["input_img"][:, :, ::-1])
        rec_text = self["rec_text"]
        rec_score = self["rec_score"]
        vis_font = self["vis_font"] if self["vis_font"] is not None else SIMFANG_FONT
        image = image.convert("RGB")
        image_width, image_height = image.size
        text = f"{rec_text} ({rec_score})"
        font = self.adjust_font_size(image_width, text, vis_font.path)
        row_height = font.getbbox(text)[3]
        new_image_height = image_height + int(row_height * 1.2)
        new_image = Image.new("RGB", (image_width, new_image_height), (255, 255, 255))
        new_image.paste(image, (0, 0))

        draw = ImageDraw.Draw(new_image)
        draw.text(
            (0, image_height),
            text,
            fill=(0, 0, 0),
            font=font,
        )
        return {"res": new_image}

    def adjust_font_size(self, image_width, text, font_path):
        font_size = int(image_width * 0.06)
        font = ImageFont.truetype(font_path, font_size)

        if int(PIL.__version__.split(".")[0]) < 10:
            text_width, _ = font.getsize(text)
        else:
            text_width, _ = font.getbbox(text)[2:]

        while text_width > image_width:
            font_size -= 1
            font = ImageFont.truetype(font_path, font_size)
            if int(PIL.__version__.split(".")[0]) < 10:
                text_width, _ = font.getsize(text)
            else:
                text_width, _ = font.getbbox(text)[2:]

        return font