# 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 numpy as np from ....utils.deps import class_requires_deps, is_dep_available from ...common.result import BaseCVResult, JsonMixin if is_dep_available("opencv-contrib-python"): import cv2 @class_requires_deps("opencv-contrib-python") class TableRecResult(BaseCVResult): """SaveTableResults""" def _to_img(self): image = self["input_img"] bbox_res = self["bbox"] if len(bbox_res) > 0 and len(bbox_res[0]) == 4: vis_img = self.draw_rectangle(image, bbox_res) else: vis_img = self.draw_bbox(image, bbox_res) return {"res": vis_img} def draw_rectangle(self, image, boxes): """draw_rectangle""" boxes = np.array(boxes) img_show = image.copy() for box in boxes.astype(int): x1, y1, x2, y2 = box cv2.rectangle(img_show, (x1, y1), (x2, y2), (255, 0, 0), 2) return img_show def draw_bbox(self, image, boxes): """draw_bbox""" for box in boxes: box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64) image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2) return image def _to_str(self, *args, **kwargs): data = copy.deepcopy(self) data.pop("input_img") return JsonMixin._to_str(data, *args, **kwargs) def _to_json(self, *args, **kwargs): data = copy.deepcopy(self) data.pop("input_img") return JsonMixin._to_json(data, *args, **kwargs)