# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve. # # 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 os import json from PIL import Image, ImageDraw, ImageFont from pathlib import Path import numpy as np from ....utils.fonts import PINGFANG_FONT_FILE_PATH from ...base import BaseTransform from ...base.predictor.io.writers import ImageWriter from .keys import ClsKeys as K from ....utils import logging __all__ = ["Topk", "NormalizeFeatures", "PrintResult", "SaveClsResults"] def _parse_class_id_map(class_ids): """ parse class id to label map file """ if class_ids is None: return None class_id_map = {id: str(lb) for id, lb in enumerate(class_ids)} return class_id_map class Topk(BaseTransform): """ Topk Transform """ def __init__(self, topk, class_ids=None): super().__init__() assert isinstance(topk, (int, )) self.topk = topk self.class_id_map = _parse_class_id_map(class_ids) def apply(self, data): """ apply """ x = data[K.CLS_PRED] class_id_map = self.class_id_map y = [] index = x.argsort(axis=0)[-self.topk:][::-1].astype("int32") clas_id_list = [] score_list = [] label_name_list = [] for i in index: clas_id_list.append(i.item()) score_list.append(x[i].item()) if class_id_map is not None: label_name_list.append(class_id_map[i.item()]) result = { "class_ids": clas_id_list, "scores": np.around( score_list, decimals=5).tolist() } if label_name_list is not None: result["label_names"] = label_name_list y.append(result) data[K.CLS_RESULT] = y return data @classmethod def get_input_keys(cls): """ get input keys """ return [K.IM_PATH, K.CLS_PRED] @classmethod def get_output_keys(cls): """ get output keys """ return [K.CLS_RESULT] class NormalizeFeatures(BaseTransform): """ Normalize Features Transform """ def apply(self, data): """ apply """ x = data[K.CLS_PRED] feas_norm = np.sqrt(np.sum(np.square(x), axis=0, keepdims=True)) x = np.divide(x, feas_norm) data[K.CLS_RESULT] = x return data @classmethod def get_input_keys(cls): """ get input keys """ return [K.IM_PATH, K.CLS_PRED] @classmethod def get_output_keys(cls): """ get output keys """ return [K.CLS_RESULT] class PrintResult(BaseTransform): """ Print Result Transform """ def apply(self, data): """ apply """ logging.info("The prediction result is:") logging.info(data[K.CLS_RESULT]) return data @classmethod def get_input_keys(cls): """ get input keys """ return [K.CLS_RESULT] @classmethod def get_output_keys(cls): """ get output keys """ return [] class SaveClsResults(BaseTransform): def __init__(self, save_dir, class_ids=None): super().__init__() self.save_dir = save_dir self.class_id_map = _parse_class_id_map(class_ids) self._writer = ImageWriter(backend='pillow') def _get_colormap(self, rgb=False): """ Get colormap """ color_list = np.array([ 0xFF, 0x00, 0x00, 0xCC, 0xFF, 0x00, 0x00, 0xFF, 0x66, 0x00, 0x66, 0xFF, 0xCC, 0x00, 0xFF, 0xFF, 0x4D, 0x00, 0x80, 0xff, 0x00, 0x00, 0xFF, 0xB2, 0x00, 0x1A, 0xFF, 0xFF, 0x00, 0xE5, 0xFF, 0x99, 0x00, 0x33, 0xFF, 0x00, 0x00, 0xFF, 0xFF, 0x33, 0x00, 0xFF, 0xff, 0x00, 0x99, 0xFF, 0xE5, 0x00, 0x00, 0xFF, 0x1A, 0x00, 0xB2, 0xFF, 0x80, 0x00, 0xFF, 0xFF, 0x00, 0x4D ]).astype(np.float32) color_list = (color_list.reshape((-1, 3))) if not rgb: color_list = color_list[:, ::-1] return color_list.astype('int32') def _get_font_colormap(self, color_index): """ Get font colormap """ dark = np.array([0x14, 0x0E, 0x35]) light = np.array([0xFF, 0xFF, 0xFF]) light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19] if color_index in light_indexs: return light.astype('int32') else: return dark.astype('int32') def apply(self, data): """ Draw label on image """ ori_path = data[K.IM_PATH] pred = data[K.CLS_PRED] index = pred.argsort(axis=0)[-1].astype("int32") score = pred[index].item() label = self.class_id_map[int(index)] label_str = f"{label} {score:.2f}" file_name = os.path.basename(ori_path) save_path = os.path.join(self.save_dir, file_name) image = Image.open(ori_path) image = image.convert('RGB') image_size = image.size draw = ImageDraw.Draw(image) min_font_size = int(image_size[0] * 0.02) max_font_size = int(image_size[0] * 0.05) for font_size in range(max_font_size, min_font_size - 1, -1): font = ImageFont.truetype( PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8") text_width_tmp, text_height_tmp = draw.textsize(label_str, font) if text_width_tmp <= image_size[0]: break else: font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, min_font_size) color_list = self._get_colormap(rgb=True) color = tuple(color_list[0]) font_color = tuple(self._get_font_colormap(3)) text_width, text_height = draw.textsize(label_str, font) rect_left = 3 rect_top = 3 rect_right = rect_left + text_width + 3 rect_bottom = rect_top + text_height + 6 draw.rectangle( [(rect_left, rect_top), (rect_right, rect_bottom)], fill=color) text_x = rect_left + 3 text_y = rect_top draw.text((text_x, text_y), label_str, fill=font_color, font=font) self._write_image(save_path, image) return data def _write_image(self, path, image): """ write image """ if os.path.exists(path): logging.warning(f"{path} already exists. Overwriting it.") self._writer.write(path, image) @classmethod def get_input_keys(cls): """ get input keys """ return [K.IM_PATH, K.CLS_PRED] @classmethod def get_output_keys(cls): """ get output keys """ return []