mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
Fix: terminate onnx inference session manually (#10076)
### What problem does this PR solve? terminate onnx inference session and release memory manually. Issue #5050 Issue #9992 Issue #8805 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
@ -13,7 +13,7 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
import gc
|
||||||
import logging
|
import logging
|
||||||
import copy
|
import copy
|
||||||
import time
|
import time
|
||||||
@ -348,6 +348,12 @@ class TextRecognizer:
|
|||||||
|
|
||||||
return img
|
return img
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
# close session and release manually
|
||||||
|
logging.info('Close TextRecognizer.')
|
||||||
|
del self.predictor
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
def __call__(self, img_list):
|
def __call__(self, img_list):
|
||||||
img_num = len(img_list)
|
img_num = len(img_list)
|
||||||
# Calculate the aspect ratio of all text bars
|
# Calculate the aspect ratio of all text bars
|
||||||
@ -395,6 +401,9 @@ class TextRecognizer:
|
|||||||
|
|
||||||
return rec_res, time.time() - st
|
return rec_res, time.time() - st
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self.close()
|
||||||
|
|
||||||
|
|
||||||
class TextDetector:
|
class TextDetector:
|
||||||
def __init__(self, model_dir, device_id: int | None = None):
|
def __init__(self, model_dir, device_id: int | None = None):
|
||||||
@ -479,6 +488,11 @@ class TextDetector:
|
|||||||
dt_boxes = np.array(dt_boxes_new)
|
dt_boxes = np.array(dt_boxes_new)
|
||||||
return dt_boxes
|
return dt_boxes
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
logging.info("Close TextDetector.")
|
||||||
|
del self.predictor
|
||||||
|
gc.collect()
|
||||||
|
|
||||||
def __call__(self, img):
|
def __call__(self, img):
|
||||||
ori_im = img.copy()
|
ori_im = img.copy()
|
||||||
data = {'image': img}
|
data = {'image': img}
|
||||||
@ -508,6 +522,9 @@ class TextDetector:
|
|||||||
|
|
||||||
return dt_boxes, time.time() - st
|
return dt_boxes, time.time() - st
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self.close()
|
||||||
|
|
||||||
|
|
||||||
class OCR:
|
class OCR:
|
||||||
def __init__(self, model_dir=None):
|
def __init__(self, model_dir=None):
|
||||||
|
|||||||
@ -13,7 +13,7 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
import gc
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import math
|
import math
|
||||||
@ -406,6 +406,11 @@ class Recognizer:
|
|||||||
"score": float(scores[i])
|
"score": float(scores[i])
|
||||||
} for i in indices]
|
} 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):
|
def __call__(self, image_list, thr=0.7, batch_size=16):
|
||||||
res = []
|
res = []
|
||||||
images = []
|
images = []
|
||||||
@ -430,5 +435,7 @@ class Recognizer:
|
|||||||
|
|
||||||
return res
|
return res
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self.close()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user