Refactor metadata and config

ChartExtractor/extraction/extraction.py

@@ -2,6 +2,7 @@
 
 # Built-in imports
 from functools import partial, reduce
+import json
 from operator import concat
 import os
 from pathlib import Path
@@ -55,7 +56,6 @@
     non_maximum_suppression,
 )
 from ..utilities.image_conversion import pil_to_cv2
-from ..utilities.read_config import read_config
 from ..utilities.tiling import tile_image
 
 # External Imports
@@ -65,40 +65,40 @@
 PATH_TO_DATA: Path = (Path(os.path.dirname(__file__)) / ".." / ".." / "data").resolve()
 PATH_TO_MODELS: Path = PATH_TO_DATA / "models"
 PATH_TO_MODEL_METADATA = PATH_TO_DATA / "model_metadata"
-MODEL_CONFIG: Dict = read_config()
+MODEL_CONFIG: Dict = json.loads(open(str(PATH_TO_DATA/"config.json"), 'r').read())
 INTRAOP_DOC_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["intraoperative_document_landmarks"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["intraoperative_document_landmarks"]["name"].replace(".onnx", ".yaml")
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["intraoperative_document_landmarks"]["name"].replace(".onnx", ".json")
 )
 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", ".yaml"),
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["preop_postop_document_landmarks"]["name"].replace(".onnx", ".json"),
 )
 NUMBERS_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["numbers"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["numbers"]["name"].replace(".onnx", ".yaml")
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["numbers"]["name"].replace(".onnx", ".json")
 )
 SYSTOLIC_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["systolic"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["systolic"]["name"].replace(".onnx", ".yaml"),
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["systolic"]["name"].replace(".onnx", ".json"),
     608,
     608,
 )
 DIASTOLIC_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["diastolic"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["diastolic"]["name"].replace(".onnx", ".yaml"),
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["diastolic"]["name"].replace(".onnx", ".json"),
     608,
     608,
 )
 HEART_RATE_MODEL = OnnxYolov11PoseSingle(
     PATH_TO_MODELS / MODEL_CONFIG["heart_rate"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["heart_rate"]["name"].replace(".onnx", ".yaml"),
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["heart_rate"]["name"].replace(".onnx", ".json"),
     608,
     608,
 )
 CHECKBOXES_MODEL = OnnxYolov11Detection(
     PATH_TO_MODELS / MODEL_CONFIG["checkboxes"]["name"],
-    PATH_TO_MODEL_METADATA / MODEL_CONFIG["checkboxes"]["name"].replace(".onnx", ".yaml"),
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["checkboxes"]["name"].replace(".onnx", ".json"),
 )
 
 

ChartExtractor/object_detection_models/onnx_yolov11_detection.py

@@ -15,6 +15,7 @@
 """
 
 # Built-in imports
+import json
 from pathlib import Path
 from typing import Dict, List, Literal, Tuple
 
@@ -28,7 +29,6 @@
 from ..utilities.annotations import BoundingBox
 from ..utilities.detections import Detection
 from ..utilities.detection_reassembly import non_maximum_suppression
-from ..utilities.read_config import read_yaml_file
 
 
 class OnnxYolov11Detection(ObjectDetectionModel):
@@ -45,7 +45,7 @@ class OnnxYolov11Detection(ObjectDetectionModel):
     def __init__(
         self,
         model_weights_filepath: Path,
-        model_metadata_filepath: Path,
+        model_classes_filepath: Path,
         input_im_width: int = 640,
         input_im_height: int = 640,
     ):
@@ -54,8 +54,8 @@ def __init__(
         Args:
             model_weights_filepath (Path):
                 The filepath to the model's weights.
-            model_metadata_filepath (Path):
-                The filepath to the metadata (for class names).
+            model_classes_filepath (Path):
+                The filepath to a json file with all the classes.
             input_im_width (int):
                 The image width that the model accepts.
                 Defaults to 640.
@@ -66,7 +66,7 @@ def __init__(
         self.model = ort.InferenceSession(model_weights_filepath)
         self.input_im_width = input_im_width
         self.input_im_height = input_im_height
-        self.classes = self.load_classes(model_metadata_filepath)
+        self.classes = self.load_classes(model_classes_filepath)
 
     @staticmethod
     def load_classes(model_metadata_filepath: Path) -> Dict:
@@ -86,7 +86,11 @@ def load_classes(model_metadata_filepath: Path) -> Dict:
         potential_err_msg = "An exception has occured while loading the classes "
         potential_err_msg += "yaml file. Ensure the model metadata filepath is "
         potential_err_msg += "correct and the model's yaml file is correctly formatted."
-        classes: Dict = read_yaml_file(model_metadata_filepath, potential_err_msg)
+        try:
+            classes: Dict[str, str] = json.loads(open(model_metadata_filepath, 'r').read())
+        except FileNotFoundError as e:
+            print(potential_err_msg)
+            print(e)
         return classes
 
     def __call__(

ChartExtractor/object_detection_models/onnx_yolov11_pose_single.py

@@ -19,6 +19,7 @@
 """
 
 # Built-in imports
+import json
 from pathlib import Path
 from typing import Dict, List, Literal, Tuple
 
@@ -32,7 +33,6 @@
 from ..utilities.annotations import BoundingBox, Keypoint, Point
 from ..utilities.detections import Detection
 from ..utilities.detection_reassembly import non_maximum_suppression
-from ..utilities.read_config import read_yaml_file
 
 
 class OnnxYolov11PoseSingle(ObjectDetectionModel):
@@ -49,7 +49,7 @@ class OnnxYolov11PoseSingle(ObjectDetectionModel):
     def __init__(
         self,
         model_weights_filepath: Path,
-        model_metadata_filepath: Path,
+        model_classes_filepath: Path,
         input_im_width: int = 640,
         input_im_height: int = 640,
     ):
@@ -58,8 +58,8 @@ def __init__(
         Args:
             model_weights_filepath (Path):
                 The filepath to the model's weights.
-            model_metadata_filepath (Path):
-                The filepath to the metadata (for class names).
+            model_classes_filepath (Path):
+                The filepath to a json file with all the classes.
             input_im_width (int):
                 The image width that the model accepts.
                 Defaults to 640.
@@ -70,7 +70,7 @@ def __init__(
         self.model = ort.InferenceSession(model_weights_filepath)
         self.input_im_width = input_im_width
         self.input_im_height = input_im_height
-        self.classes = self.load_classes(model_metadata_filepath)
+        self.classes = self.load_classes(model_classes_filepath)
 
     @staticmethod
     def load_classes(model_metadata_filepath: Path) -> Dict:
@@ -90,7 +90,11 @@ def load_classes(model_metadata_filepath: Path) -> Dict:
         potential_err_msg = "An exception has occured while loading the classes "
         potential_err_msg += "yaml file. Ensure the model metadata filepath is "
         potential_err_msg += "correct and the model's yaml file is correctly formatted."
-        classes: Dict = read_yaml_file(model_metadata_filepath, potential_err_msg)
+        try:
+            classes: Dict[str, str] = json.loads(open(model_metadata_filepath, 'r').read())
+        except FileNotFoundError as e:
+            print(potential_err_msg)
+            print(e)
         return classes
 
     def __call__(

pyproject.toml

@@ -19,7 +19,6 @@ opencv-python-headless = "^4.10.0.84"
 numpy = "^2.1.2"
 pillow = "^10.4.0"
 scikit-learn = "^1.5.2"
-pyyaml = "^6.0.2"
 onnxruntime = "^1.20.1"
 
 [tool.poetry.group.dev.dependencies]