Add lazy loading option for models.

Author: RyanDoesMath

src/ChartExtractor/extraction/extraction.py

@@ -64,46 +64,53 @@
     ),
     MODEL_CONFIG["intraoperative_document_landmarks"]["imgsz"],
     MODEL_CONFIG["intraoperative_document_landmarks"]["imgsz"],
+    lazy_loading=True,
 )
 PREOP_POSTOP_DOC_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["preop_postop_document_landmarks"]["name"],
     PATH_TO_MODEL_METADATA
     / MODEL_CONFIG["preop_postop_document_landmarks"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["preop_postop_document_landmarks"]["imgsz"],
     MODEL_CONFIG["preop_postop_document_landmarks"]["imgsz"],
+    lazy_loading=True,
 )
 NUMBERS_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["numbers"]["name"],
     PATH_TO_MODEL_METADATA / MODEL_CONFIG["numbers"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["numbers"]["imgsz"],
     MODEL_CONFIG["numbers"]["imgsz"],
+    lazy_loading=True,
 )
 SYSTOLIC_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["systolic"]["name"],
     PATH_TO_MODEL_METADATA / MODEL_CONFIG["systolic"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["systolic"]["imgsz"],
     MODEL_CONFIG["systolic"]["imgsz"],
+    lazy_loading=True,
 )
 DIASTOLIC_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["diastolic"]["name"],
     PATH_TO_MODEL_METADATA
     / MODEL_CONFIG["diastolic"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["diastolic"]["imgsz"],
     MODEL_CONFIG["diastolic"]["imgsz"],
+    lazy_loading=True,
 )
 HEART_RATE_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["heart_rate"]["name"],
     PATH_TO_MODEL_METADATA
     / MODEL_CONFIG["heart_rate"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["heart_rate"]["imgsz"],
     MODEL_CONFIG["heart_rate"]["imgsz"],
+    lazy_loading=True,
 )
 CHECKBOXES_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["checkboxes"]["name"],
     PATH_TO_MODEL_METADATA
     / MODEL_CONFIG["checkboxes"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["checkboxes"]["imgsz"],
     MODEL_CONFIG["checkboxes"]["imgsz"],
+    lazy_loading=True
 )
 
 

src/ChartExtractor/object_detection_models/onnx_yolov11_detection.py

@@ -48,6 +48,7 @@ def __init__(
         model_classes_filepath: Path,
         input_im_width: int = 640,
         input_im_height: int = 640,
+        lazy_loading: bool = False
     ):
         """Initializes the onnx model.
 
@@ -62,11 +63,17 @@ def __init__(
             input_im_height (int):
                 The image height that the model accepts.
                 Defaults to 640.
+            lazy_loading (bool):
+                Whether or not to load the model only when it is called for detection.
+                Defaults to False.
         """
-        self.model = ort.InferenceSession(model_weights_filepath)
+        self.model_weights_filepath = model_weights_filepath
         self.input_im_width = input_im_width
         self.input_im_height = input_im_height
         self.classes = self.load_classes(model_classes_filepath)
+        self.model_is_loaded = False
+        if not lazy_loading:
+            self.load_model()
 
     @staticmethod
     def load_classes(model_metadata_filepath: Path) -> Dict:
@@ -94,6 +101,11 @@ def load_classes(model_metadata_filepath: Path) -> Dict:
             print(potential_err_msg)
             print(e)
         return classes
+    
+    def load_model(self):
+        """Loads the model."""
+        self.model = ort.InferenceSession(self.model_weights_filepath)
+        self.model_is_loaded = True
 
     def __call__(
         self,
@@ -116,6 +128,8 @@ def __call__(
         Returns:
             A list of detections for each image.
         """
+        if not model_is_loaded:
+            self.load_model()
         if not isinstance(images, list):
             images = [images]
         detections: List[List[Detection]] = [

src/ChartExtractor/object_detection_models/onnx_yolov11_pose_single.py

@@ -52,6 +52,7 @@ def __init__(
         model_classes_filepath: Path,
         input_im_width: int = 640,
         input_im_height: int = 640,
+        lazy_loading: bool = False
     ):
         """Initializes the onnx model.
 
@@ -66,11 +67,18 @@ def __init__(
             input_im_height (int):
                 The image height that the model accepts.
                 Defaults to 640.
+            lazy_loading (bool):
+                Whether or not to load the model only when it is called for detection.
+                Defaults to False.
         """
-        self.model = ort.InferenceSession(model_weights_filepath)
+        self.model_weights_filepath = model_weights_filepath
         self.input_im_width = input_im_width
         self.input_im_height = input_im_height
         self.classes = self.load_classes(model_classes_filepath)
+        self.model_is_loaded = False
+        if not lazy_loading:
+            self.load_model()
+
 
     @staticmethod
     def load_classes(model_metadata_filepath: Path) -> Dict:
@@ -98,6 +106,11 @@ def load_classes(model_metadata_filepath: Path) -> Dict:
             print(potential_err_msg)
             print(e)
         return classes
+    
+    def load_model(self):
+        """Loads the model."""
+        self.model = ort.InferenceSession(self.model_weights_filepath)
+        self.model_is_loaded = True
 
     def __call__(
         self,
@@ -120,6 +133,8 @@ def __call__(
         Returns:
             A list of detections for each image.
         """
+        if not model_is_loaded:
+            self.load_model()
         if not isinstance(images, list):
             images = [images]
         detections: List[List[Detection]] = [