# -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com import math import random from functools import lru_cache from pathlib import Path from typing import List, Optional, Sequence, Tuple, Union import cv2 import numpy as np from omegaconf import OmegaConf from PIL import Image, ImageDraw, ImageFont from .download_file import DownloadFile, DownloadFileInput from .load_image import LoadImage from .log import logger from .model_resolver import normalize_lang from .typings import LangRec root_dir = Path(__file__).resolve().parent.parent InputType = Union[str, np.ndarray, bytes, Path, Image.Image] DEFAULT_FONT_DIR = root_dir / "models" DEFAULT_FONT_PATH = DEFAULT_FONT_DIR / "FZYTK.TTF" FONT_YAML_PATH = root_dir / "default_models.yaml" @lru_cache() def font_cfg(): return OmegaConf.load(FONT_YAML_PATH).fonts @lru_cache() def get_font_path( font_path: Optional[Union[str, Path]] = None, lang_type: Union[LangRec, None, str] = None, ) -> str: default_info = font_cfg()["ch"] default_input_params = DownloadFileInput( file_url=default_info["path"], sha256=default_info["SHA256"], save_path=DEFAULT_FONT_PATH, logger=logger, ) if lang_type is None: # 没有指定语种，用默认字体文件 DownloadFile.run(default_input_params) return str(DEFAULT_FONT_PATH) lang_type = normalize_lang(lang_type) if font_path is None: # 指定了语种，但是没有指定字体文件，根据语种选择字体文件 font_info = font_cfg().get(lang_type, None) if font_info is None: logger.warning( "Font file for %s is not found in the supported font list. Default font file will be used.", lang_type, ) DownloadFile.run(default_input_params) return str(DEFAULT_FONT_PATH) font_url, font_sha256 = font_info["path"], font_info["SHA256"] save_font_path = DEFAULT_FONT_DIR / f"{Path(font_url).name}" input_param = DownloadFileInput( file_url=font_url, sha256=font_sha256, save_path=save_font_path, logger=logger, verbose=False, ) DownloadFile.run(input_param) return str(save_font_path) return str(font_path) class VisRes: load_img = LoadImage() def __init__( self, text_score: float = 0.5, lang_type: Union[LangRec, None, str] = None, font_path: Optional[str] = None, ): self.text_score = text_score self._font_path = font_path self._lang_type = lang_type @property def font_path(self): return get_font_path(self._font_path, self._lang_type) def __call__( self, img_content: InputType, dt_boxes: np.ndarray, txts: Optional[Union[List[str], Tuple[str]]] = None, scores: Optional[List[float]] = None, ) -> np.ndarray: logger.info(f"Using {self.font_path} to visualize results.") if txts is None: return self.draw_dt_boxes(img_content, dt_boxes, scores) return self.draw_ocr_box_txt(img_content, dt_boxes, txts, scores) def draw_dt_boxes( self, img_content: InputType, dt_boxes: np.ndarray, scores: Optional[List[float]] = None, ) -> np.ndarray: img = self.load_img(img_content) if scores is None: scores = [1.0] * len(dt_boxes) for idx, (box, score) in enumerate(zip(dt_boxes, scores)): color = self.get_random_color() points = np.array(box) cv2.polylines(img, np.int32([points]), 1, color=color, thickness=1) start_point = round(points[0][0]), round(points[0][1]) cv2.putText( img, f"{idx}:{score:.3f}", start_point, cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, ) return img def draw_rec_res( self, imgs: Sequence[InputType], txts: Union[List[str], Tuple[str]], scores: List[float], ) -> np.ndarray: result_imgs = [] for img, txt, score in zip(imgs, txts, scores): vis_img = self.draw_one_rec_res(img, txt, score) result_imgs.append(vis_img) return self.concat_imgs(result_imgs, direction="vertical") def draw_one_rec_res( self, img_content: InputType, txt: str, score: float ) -> np.ndarray: image = Image.fromarray(self.load_img(img_content)) h, w = image.height, image.width if image.mode == "L": image = image.convert("RGB") img_left = image.copy() img_right = Image.new("RGB", (w, h), (255, 255, 255)) draw_right = ImageDraw.Draw(img_right) box = [[0, 0], [w, 0], [w, h], [0, h]] box_height = self.get_box_height(box) box_width = self.get_box_width(box) if box_height > 2 * box_width: font_size = max(int(box_width * 0.9), 10) font = ImageFont.truetype(self.font_path, font_size, encoding="utf-8") cur_y = box[0][1] for c in txt: draw_right.text((box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font) cur_y += self.get_char_size(font, c) else: font_size = max(int(box_height * 0.8), 10) font = ImageFont.truetype(self.font_path, font_size, encoding="utf-8") draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) img_left = Image.blend(image, img_left, 0.5) img_show = Image.new("RGB", (w * 2, h), 255) img_show.paste(img_left, (0, 0, w, h)) img_show.paste(img_right, (w, 0, w * 2, h)) return np.array(img_show) def draw_ocr_box_txt( self, img_content: InputType, dt_boxes: np.ndarray, txts: Union[List[str], Tuple[str]], scores: Optional[List[float]] = None, ) -> np.ndarray: image = Image.fromarray(self.load_img(img_content)) h, w = image.height, image.width if image.mode == "L": image = image.convert("RGB") img_left = image.copy() img_right = Image.new("RGB", (w, h), (255, 255, 255)) random.seed(0) draw_left = ImageDraw.Draw(img_left) draw_right = ImageDraw.Draw(img_right) for idx, (box, txt) in enumerate(zip(dt_boxes, txts)): if scores is not None and float(scores[idx]) < self.text_score: continue color = self.get_random_color() box_list = np.array(box).reshape(8).tolist() draw_left.polygon(box_list, fill=color) draw_right.polygon(box_list, outline=color) box_height = self.get_box_height(box) box_width = self.get_box_width(box) if box_height > 2 * box_width: font_size = max(int(box_width * 0.9), 10) font = ImageFont.truetype(self.font_path, font_size, encoding="utf-8") cur_y = box[0][1] for c in txt: draw_right.text( (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font ) cur_y += self.get_char_size(font, c) else: font_size = max(int(box_height * 0.8), 10) font = ImageFont.truetype(self.font_path, font_size, encoding="utf-8") draw_right.text([box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font) img_left = Image.blend(image, img_left, 0.5) img_show = Image.new("RGB", (w * 2, h), (255, 255, 255)) img_show.paste(img_left, (0, 0, w, h)) img_show.paste(img_right, (w, 0, w * 2, h)) return np.array(img_show) @staticmethod def get_random_color() -> Tuple[int, int, int]: return ( random.randint(0, 255), random.randint(0, 255), random.randint(0, 255), ) @staticmethod def get_box_height(box: List[List[Union[float, int]]]) -> float: return math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2) @staticmethod def get_box_width(box: List[List[Union[float, int]]]) -> float: return math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2) @staticmethod def get_char_size(font, char_str: str) -> float: # compatible with Pillow v9 and v10. if hasattr(font, "getsize"): get_size_func = getattr(font, "getsize") return get_size_func(char_str)[1] if hasattr(font, "getlength"): get_size_func = getattr(font, "getlength") return get_size_func(char_str) raise ValueError( "The Pillow ImageFont instance has not getsize or getlength func." ) @staticmethod def concat_imgs( imgs: List[np.ndarray], direction: str = "horizontal" ) -> np.ndarray: img_list = [Image.fromarray(img_path) for img_path in imgs] img_sizes = np.array([v.size for v in img_list]) if direction == "horizontal": width = np.sum(img_sizes[:, 0]) height = np.max(img_sizes[:, 1]) elif direction == "vertical": width = np.max(img_sizes[:, 0]) height = np.sum(img_sizes[:, 1]) else: raise ValueError(f"{direction} is not supported.") new_image = Image.new("RGB", (width, height), color="white") for i, img in enumerate(img_list): if i == 0: new_image.paste(img, (0, 0)) continue if direction == "horizontal": x = np.sum(img_sizes[:, 0][:i]) new_image.paste(img, (x, 0)) elif direction == "vertical": y = np.sum(img_sizes[:, 1][:i]) new_image.paste(img, (0, y)) return np.array(new_image)