Improve legend detection

src/ChartExtractor/extraction/blood_pressure_and_heart_rate.py

@@ -9,100 +9,96 @@
 from ..utilities.detections import Detection
 
 
-def find_timestamp(time_legend: List[Cluster], keypoint_x: float) -> str:
+def find_timestamp(legend: Dict[str, Tuple[float, float]], keypoint_x: float) -> str:
     """Given a keypoint on a blood pressure or heart rate detection, finds the timestamp.
 
     Args:
-        `time_legend` (List[Cluster]):
-            The named clusters which form the timestamp legend that runs horizontally on the top
-            side of the blood pressure and heart rate section.
+        `legend` (Dict[str, Tuple[float, float]]):
+            The dictionary that maps the name of legend entries to their locations on the image.
         `keypoint_x` (float):
             The x value of the keypoint.
 
     Returns:
         The label of the closest timestamp cluster.
     """
-    time_legend_centers: Dict[str, float] = {
-        clust.label: clust.bounding_box.center[0] for clust in time_legend
+    time_legend: Dict[str, Tuple[float, float]] = {
+        k:v for (k, v) in legend.items() if "_mins" in k
     }
     distances: Dict[str, float] = {
-        name: abs(legend_loc - keypoint_x)
-        for (name, legend_loc) in time_legend_centers.items()
+        name: abs(legend_loc[0] - keypoint_x)
+        for (name, legend_loc) in time_legend.items()
     }
     return min(distances, key=distances.get)
 
 
-def find_value(value_legend: List[Cluster], keypoint_y: float) -> int:
+def find_value(legend: Dict[str, Tuple[float, float]], keypoint_y: float) -> int:
     """Given a keypoint on a blood pressure or heart rate detection, finds the in mmhg/bpm value.
 
     Finds the closest two legend values, then uses the distance between the detection and both
     of the closest values to find an approximate value in between.
 
     Args:
-        `value_legend` (List[Cluster]):
-            The named clusters which form the mmhg/bpm legend that runs vertically on the left
-            side of the blood pressure and heart rate section.
+        `legend` (Dict[str, Tuple[float, float]]):
+            The dictionary that maps the name of legend entries to their locations on the image.
         `keypoint_y` (float):
             The y value of the keypoint.
 
     Returns:
         The approximate value that the keypoint encodes in mmhg/bpm.
     """
-    value_legend_centers: Dict[str, float] = {
-        clust.label: clust.bounding_box.center[1] for clust in value_legend
+    value_legend: Dict[str, float] = {
+        k:v for (k, v) in legend.items() if "_mmhg" in k
     }
     distances: Dict[str, float] = {
-        name: abs(legend_loc - keypoint_y)
-        for (name, legend_loc) in value_legend_centers.items()
+        name: abs(legend_loc[1] - keypoint_y)
+        for (name, legend_loc) in value_legend.items()
     }
     first_closest: str = min(distances, key=distances.get)
     distances.pop(first_closest)
     second_closest: str = min(distances, key=distances.get)
     total_dist: float = abs(
-        value_legend_centers[first_closest] - value_legend_centers[second_closest]
+        value_legend[first_closest][1] - value_legend[second_closest][1]
     )
     smaller_of_two_values = min(
         [first_closest, second_closest], key=lambda leg: int(leg.split("_")[0])
     )
     fractional_component = (
-        abs(value_legend_centers[smaller_of_two_values] - keypoint_y) / total_dist
+        abs(value_legend[smaller_of_two_values][1] - keypoint_y) / total_dist
     ) * 10
     return int(smaller_of_two_values.split("_")[0]) + int(fractional_component)
 
 
 def extract_heart_rate_and_blood_pressure(
     detections: List[Detection],
-    time_clusters: List[Cluster],
-    value_clusters: List[Cluster],
+    legend: Dict[str, Tuple[float, float]],
 ) -> Dict[str, Dict[str, str]]:
     """Extracts the heart rate and blood pressure data from the detections.
 
     Args:
         `detections` (List[Detection]):
             The keypoint detections of the systolic, diastolic, and heart rate markings.
-        `time_clusters` (List[Cluster]):
-            The clusters corresponding to the timestamps.
-        `value_clusters` (List[Cluster]):
-            The clusters corresponding to the mmhg and bpm values.
+        `legend` (Dict[str, Tuple[float, float]]):
+            The dictionary that maps the name of legend entries to their locations on the image.
 
     Returns:
         A dictionary mapping each timestamp to the systolic, diastolic, and heart rate reading
         that was recorded at that time.
     """
-
+    
     def filter_detections_outside_bp_and_hr_area(detections):
+        leftmost_point: float = min([point[0] for point in legend.values()])
+        topmost_point: float = min([point[1] for point in legend.values()])
+        rightmost_point: float = max([point[0] for point in legend.values()])
+        bottommost_point: float = max([point[1] for point in legend.values()])
+
         return list(
             filter(
                 lambda d: all(
                     [
-                        d.annotation.bottom
-                        > min(vc.bounding_box.top for vc in value_clusters),
-                        d.annotation.top
-                        < max(vc.bounding_box.bottom for vc in value_clusters),
-                        d.annotation.left
-                        > min(tc.bounding_box.left for tc in time_clusters),
-                        d.annotation.right
-                        < max(tc.bounding_box.right for tc in time_clusters),
+                        d.annotation.bottom > topmost_point,
+                        d.annotation.top < bottommost_point,
+                        d.annotation.right > leftmost_point,
+                        d.annotation.left < rightmost_point,
                     ]
                 ),
                 detections,
@@ -115,7 +111,7 @@ def filter_detections_outside_bp_and_hr_area(detections):
     for det in detections:
         point: Tuple[float, float] = det.annotation.keypoint
         category: str = det.annotation.category
-        timestamp: str = find_timestamp(time_clusters, point.x)
+        timestamp: str = find_timestamp(legend, point.x)
         if data.get(timestamp) is None:
             data[timestamp] = {category: det}
         elif data[timestamp].get(category) is None:
@@ -129,6 +125,6 @@ def filter_detections_outside_bp_and_hr_area(detections):
         for category in data[timestamp].keys():
             point: Tuple[float, float] = data[timestamp][category].annotation.keypoint
             suffix: str = "bpm" if category == "heart_rate" else "mmhg"
-            value: int = find_value(value_clusters, point.y)
+            value: int = find_value(legend, point.y)
             data[timestamp][category] = f"{value}_{suffix}"
     return data

src/ChartExtractor/extraction/extraction.py

@@ -18,6 +18,7 @@
     detect_objects_using_tiling,
     label_studio_to_bboxes,
 )
+from ..extraction.find_legend import find_legend
 from ..extraction.inhaled_volatile import extract_inhaled_volatile
 from ..extraction.intraoperative_digit_boxes import (
     extract_drug_codes,
@@ -110,7 +111,14 @@
     / MODEL_CONFIG["checkboxes"]["name"].replace(".onnx", ".json"),
     MODEL_CONFIG["checkboxes"]["imgsz"],
     MODEL_CONFIG["checkboxes"]["imgsz"],
-    lazy_loading=True
+    lazy_loading=True,
+)
+LEGEND_MODEL = OnnxYolov11Detection(
+    PATH_TO_MODELS / MODEL_CONFIG["whole_number_legend"]["name"],
+    PATH_TO_MODEL_METADATA / MODEL_CONFIG["whole_number_legend"]["name"].replace(".onnx", ".json"),
+    MODEL_CONFIG["whole_number_legend"]["imgsz"],
+    MODEL_CONFIG["whole_number_legend"]["imgsz"],
+    lazy_loading=True,
 )
 
 
@@ -222,12 +230,12 @@ def run_intraoperative_models(intraop_image: Image.Image) -> Dict[str, List[Dete
     )
 
     # checkboxes
-    tile_size = compute_tile_size(MODEL_CONFIG["checkboxes"], intraop_image.size)
+    ckbx_tile_size = compute_tile_size(MODEL_CONFIG["checkboxes"], intraop_image.size)
     detections_dict["checkboxes"] = detect_objects_using_tiling(
         intraop_image,
         CHECKBOXES_MODEL,
-        tile_size,
-        tile_size,
+        ckbx_tile_size,
+        ckbx_tile_size,
         MODEL_CONFIG["checkboxes"]["horz_overlap_proportion"],
         MODEL_CONFIG["checkboxes"]["vert_overlap_proportion"],
         nms_threshold=0.8,
@@ -270,6 +278,19 @@ def run_intraoperative_models(intraop_image: Image.Image) -> Dict[str, List[Dete
         MODEL_CONFIG["heart_rate"]["vert_overlap_proportion"],
     )
 
+    # legend
+    legend_tile_size: int = compute_tile_size(
+        MODEL_CONFIG["whole_number_legend"], intraop_image.size
+    )
+    detections_dict["legend"] = detect_objects_using_tiling(
+        intraop_image.copy(),
+        LEGEND_MODEL,
+        legend_tile_size,
+        legend_tile_size,
+        MODEL_CONFIG["whole_number_legend"]["horz_overlap_proportion"],
+        MODEL_CONFIG["whole_number_legend"]["vert_overlap_proportion"],
+    )
+
     return detections_dict
 
 
@@ -403,27 +424,21 @@ def assign_meaning_to_intraoperative_detections(
     extracted_data["codes"] = extract_drug_codes(
         corrected_detections_dict["numbers"], *image_size
     )
-    extracted_data["timing"] = extract_surgical_timing(
+    extracted_data["intraoperative_timing"] = extract_surgical_timing(
         corrected_detections_dict["numbers"], *image_size
     )
     extracted_data["ett_size"] = extract_ett_size(
         corrected_detections_dict["numbers"], *image_size
     )
-
-    # extract inhaled volatile drugs
-    time_boxes, mmhg_boxes = isolate_blood_pressure_legend_bounding_boxes(
-        [det.annotation for det in corrected_detections_dict["landmarks"]], *image_size
-    )
-    time_clusters: List[Cluster] = cluster_boxes(
-        time_boxes, cluster_kmeans, "mins", possible_nclusters=[40, 41, 42]
-    )
-    mmhg_clusters: List[Cluster] = cluster_boxes(
-        mmhg_boxes, cluster_kmeans, "mmhg", possible_nclusters=[18, 19, 20]
+
+    # get legend locations
+    legend_locations: Dict[str, Tuple[float, float]] = find_legend(
+        intraop_detections_dict["legend"],
+        image_size[0],
+        image_size[1],
     )
 
-    legend_locations: Dict[str, Tuple[float, float]] = find_legend_locations(
-        time_clusters + mmhg_clusters
-    )
+    # extract inhaled volatile drugs
     extracted_data["inhaled_volatile"] = extract_inhaled_volatile(
         corrected_detections_dict["numbers"],
         legend_locations,
@@ -443,8 +458,7 @@ def assign_meaning_to_intraoperative_detections(
 
     extracted_data["bp_and_hr"] = extract_heart_rate_and_blood_pressure(
         bp_and_hr_dets,
-        time_clusters,
-        mmhg_clusters,
+        legend_locations,
     )
 
     # extract physiological indicators
@@ -565,24 +579,27 @@ def digitize_intraop_record(image: Image.Image) -> Dict:
 
     # extract drug code and surgical timing
     codes: Dict = {"codes": extract_drug_codes(digit_detections, *image.size)}
-    times: Dict = {"timing": extract_surgical_timing(digit_detections, *image.size)}
+    times: Dict = {"intraoperative_timing": extract_surgical_timing(digit_detections, *image.size)}
     ett_size: Dict = {"ett_size": extract_ett_size(digit_detections, *image.size)}
-
-    # extract inhaled volatile drugs
-    time_boxes, mmhg_boxes = isolate_blood_pressure_legend_bounding_boxes(
-        [det.annotation for det in document_landmark_detections], *image.size
+    
+    # get legend locations
+    legend_tile_size: int = compute_tile_size(
+        MODEL_CONFIG["whole_number_legend"], image.size
     )
-    time_clusters: List[Cluster] = cluster_boxes(
-        time_boxes, cluster_kmeans, "mins", possible_nclusters=[40, 41, 42]
-    )
-    mmhg_clusters: List[Cluster] = cluster_boxes(
-        mmhg_boxes, cluster_kmeans, "mmhg", possible_nclusters=[18, 19, 20]
+    legend_detections = detect_objects_using_tiling(
+        image,
+        LEGEND_MODEL,
+        legend_tile_size,
+        legend_tile_size,
+        MODEL_CONFIG["whole_number_legend"]["horz_overlap_proportion"],
+        MODEL_CONFIG["whole_number_legend"]["vert_overlap_proportion"],
     )
-
-    legend_locations: Dict[str, Tuple[float, float]] = find_legend_locations(
-        time_clusters + mmhg_clusters
+    legend_locations: Dict[str, Tuple[float, float]] = find_legend(
+        legend_detections,
+        *image.size,
     )
-
+
+    # extract inhaled volatile drugs
     inhaled_volatile: Dict = {
         "inhaled_volatile": extract_inhaled_volatile(
             digit_detections, legend_locations, document_landmark_detections
@@ -591,7 +608,7 @@ def digitize_intraop_record(image: Image.Image) -> Dict:
 
     # extract bp and hr
     bp_and_hr: Dict = {
-        "bp_and_hr": make_bp_and_hr_detections(image, time_clusters, mmhg_clusters)
+        "bp_and_hr": make_bp_and_hr_detections(image, legend_locations)
     }
 
     # extract physiological indicators
@@ -888,18 +905,15 @@ def compute_tile_size(model_config: Dict, image_size: Tuple[int, int]) -> int:
 
 def make_bp_and_hr_detections(
     image: Image.Image,
-    time_clusters: List[Cluster],
-    mmhg_clusters: List[Cluster],
+    legend: Dict[str, Tuple[float, float]]
 ) -> Dict:
     """Finds blood pressure symbols and associates a value and timestamp to them.
 
     Args:
         `image` (Image.Image):
             The image to detect on.
-        `time_clusters` (List[Cluster]):
-            A list of Cluster objects encoding the location of the time legend.
-        `mmhg_clusters` (List[Cluster]):
-            A list of Cluster objects encoding the location of the mmhg/bpm legend.
+        `legend` (Dict[str, Tuple[float, float]]):
+            The dictionary that maps the name of legend entries to their locations on the image.
 
     Returns:
         A dictionary mapping timestamps to values for systolic, diastolic, and heart rate.
@@ -934,9 +948,7 @@ def make_bp_and_hr_detections(
     )
 
     dets: List[Detection] = sys_dets + dia_dets + hr_dets
-    bp_and_hr = extract_heart_rate_and_blood_pressure(
-        dets, time_clusters, mmhg_clusters
-    )
+    bp_and_hr = extract_heart_rate_and_blood_pressure(dets, legend)
     return bp_and_hr