diff --git a/deepdoc/vision/ocr.py b/deepdoc/vision/ocr.py index e9e594274..fb42fba27 100644 --- a/deepdoc/vision/ocr.py +++ b/deepdoc/vision/ocr.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # - +import gc import logging import copy import time @@ -348,6 +348,12 @@ class TextRecognizer: return img + def close(self): + # close session and release manually + logging.info('Close TextRecognizer.') + del self.predictor + gc.collect() + def __call__(self, img_list): img_num = len(img_list) # Calculate the aspect ratio of all text bars @@ -395,6 +401,9 @@ class TextRecognizer: return rec_res, time.time() - st + def __del__(self): + self.close() + class TextDetector: def __init__(self, model_dir, device_id: int | None = None): @@ -479,6 +488,11 @@ class TextDetector: dt_boxes = np.array(dt_boxes_new) return dt_boxes + def close(self): + logging.info("Close TextDetector.") + del self.predictor + gc.collect() + def __call__(self, img): ori_im = img.copy() data = {'image': img} @@ -508,6 +522,9 @@ class TextDetector: return dt_boxes, time.time() - st + def __del__(self): + self.close() + class OCR: def __init__(self, model_dir=None): diff --git a/deepdoc/vision/recognizer.py b/deepdoc/vision/recognizer.py index 9fa82d7f5..a829c1cd3 100644 --- a/deepdoc/vision/recognizer.py +++ b/deepdoc/vision/recognizer.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # - +import gc import logging import os import math @@ -406,6 +406,11 @@ class Recognizer: "score": float(scores[i]) } for i in indices] + def close(self): + logging.info("Close recognizer.") + del self.ort_sess + gc.collect() + def __call__(self, image_list, thr=0.7, batch_size=16): res = [] images = [] @@ -430,5 +435,7 @@ class Recognizer: return res + def __del__(self): + self.close()