Add cache

app.py

@@ -23,11 +23,15 @@ from paddleocr import PaddleOCR
 import postprocess
 
 
-ocr_instance = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=True)
-detection_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/detection_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
-structure_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/structure_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
+@st.cache_resource(ttl=3600)
+def load_models():
+    ocr_instance = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=True)
+    detection_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/detection_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
+    structure_model = torch.hub.load('ultralytics/yolov5', 'custom', 'weights/structure_wts.pt', force_reload=True, skip_validation=True, trust_repo=True)
+    return ocr_instance, detection_model, structure_model
 
-imgsz = 640
+
+ocr_instance, detection_model, structure_model = load_models()
 
 detection_class_names = ['table', 'table rotated', 'no object']
 structure_class_names = [
@@ -62,7 +66,7 @@ def cv_to_PIL(cv_img):
     return PIL.Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
 
 
-def table_detection(pil_img):
+def table_detection(pil_img, imgsz=640):
     image = PIL_to_cv(pil_img)
     pred = detection_model(image, size=imgsz)
     pred = pred.xywhn[0]
@@ -70,7 +74,7 @@ def table_detection(pil_img):
     return result
 
 
-def table_structure(pil_img):
+def table_structure(pil_img, imgsz=640):
     image = PIL_to_cv(pil_img)
     pred = structure_model(image, size=imgsz)
     pred = pred.xywhn[0]