polar map visualization shows segments and perfusion scores

mu_map/polar_map/visualize.py

@@ -6,6 +6,8 @@ from typing import Tuple
 import cv2 as cv
 import numpy as np
 
+from mu_map.polar_map.get_perfusion import SEGMENTS
+
 
 def get_circular_mask(shape: Tuple[int, int], channels: int = 1) -> np.ndarray:
     """
@@ -38,6 +40,80 @@ def get_circular_mask(shape: Tuple[int, int], channels: int = 1) -> np.ndarray:
     return mask[:, :, 0] if channels == 1 else mask
 
 
+def get_radial_line(
+    center: np.ndarray, r1: float, r2: float, angle: float, degree=True
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Get the start and end point of a line drawn from a central point at an angle.
+
+    Parameters
+    ----------
+    center: np.ndarray
+        the center coordinate as a numpy array
+    r1: float
+        the radius at which the line should start
+    r2: float
+        the radius at which the line should end
+    angle: float
+        the angle of the line
+    degree: bool, optional
+        true if the angle is given in degree of false if in radians
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray]
+        the start and the end point of the line
+    """
+    angle_rad = angle * np.pi / 180.0 if degree else angle
+    direction = np.array([np.cos(angle_rad), np.sin(angle_rad)])
+
+    start = (center + r1 * direction).round().astype(int)
+    end = (center + r2 * direction).round().astype(int)
+
+    return start, end
+
+
+def draw_segments_grid(
+    img: np.ndarray, color: Tuple[int] = (0, 0, 0), line_thickness: int = 1
+) -> np.ndarray:
+    """
+    Draw the 17-segment perfusion score grid on a polar map image.
+
+    Parameters
+    ----------
+    img: np.ndarray
+        the polar map image
+    color: Tuple[int]
+        the color with which the lines are drawn
+    line_thickness: int
+        the line thickness
+    """
+    center = np.array(img.shape[:2]) // 2
+    cy, cx = center
+
+    base_radius = img.shape[0] // 8
+    radii = np.array([1, 2, 3, 4]) * base_radius + np.array([0, 1, 1, 1])
+
+    for radius in radii:
+        img = cv.circle(
+            img, (cx, cy), radius, color, line_thickness, lineType=cv.LINE_AA
+        )
+
+    for angle in np.arange(4) * 90 + 50:
+        line_start, line_end = get_radial_line(center, radii[0], radii[1], angle)
+        img = cv.line(
+            img, line_start, line_end, color, line_thickness, lineType=cv.LINE_AA
+        )
+
+    for angle in np.arange(6) * 60:
+        line_start, line_end = get_radial_line(center, radii[1], radii[-1], angle)
+        img = cv.line(
+            img, line_start, line_end, color, line_thickness, lineType=cv.LINE_AA
+        )
+
+    return img
+
+
 if __name__ == "__main__":
     import argparse
     import os
@@ -49,6 +125,7 @@ if __name__ == "__main__":
     plt.rcParams.update(
         {
             "text.usetex": True,
+            "font.size": 15,
         }
     )
 
@@ -76,6 +153,12 @@ if __name__ == "__main__":
         default="polar_maps.csv",
         help="file under <polar_map_dir> containing meta information about the polar maps",
     )
+    parser.add_argument(
+        "--perfusion_csv",
+        type=str,
+        default="perfusion.csv",
+        help="file under <polar_map_dir> containing perfusion scores",
+    )
     parser.add_argument(
         "--baseline",
         choices=["symbia", "stir"],
@@ -100,9 +183,11 @@ if __name__ == "__main__":
 
     args.images_dir = os.path.join(args.polar_map_dir, args.images_dir)
     args.csv = os.path.join(args.polar_map_dir, args.csv)
+    args.perfusion_csv = os.path.join(args.polar_map_dir, args.perfusion_csv)
 
     meta = pd.read_csv(args.csv)
     ids = meta[headers.id].unique()
+    perfusion = pd.read_csv(args.perfusion_csv)
 
     if args.id:
         assert args.id in ids, f"Id {args.id} is not available. Chose one of {ids}."
@@ -118,9 +203,13 @@ if __name__ == "__main__":
         print(f"Show id {_id:03d}")
         _meta = meta[(meta[headers.id] == _id) & ~(meta[headers.segments])]
 
-        file_recon_ctac = _meta[(_meta[headers.type] == "symbia") & _meta[headers.ac]][
-            headers.file
-        ].values[0]
+        file_recon_ctac = (
+            _meta[(_meta[headers.type] == "symbia") & _meta[headers.ac]][
+                headers.file
+            ].values[0]
+            if args.baseline == "symbia"
+            else _meta[_meta[headers.type] == "ct"][headers.file].values[0]
+        )
         file_recon_noac = _meta[~_meta[headers.ac]][headers.file].values[0]
         file_recon_dlac = _meta[_meta[headers.type] == "dl"][headers.file].values[0]
 
@@ -134,8 +223,18 @@ if __name__ == "__main__":
             os.path.join(args.images_dir, file_recon_dlac), cv.IMREAD_GRAYSCALE
         )
 
+        _perfusion = perfusion[perfusion[headers.id] == _id]
+        perfusion_ctac = (
+            _perfusion[(_perfusion[headers.type] == "symbia") & _perfusion[headers.ac]]
+            if args.baseline == "symbia"
+            else _perfusion[_perfusion[headers.type] == "ct"]
+        ).squeeze()
+        perfusion_noac = _perfusion[~_perfusion[headers.ac]].squeeze()
+        perfusion_dlac = _perfusion[_perfusion[headers.type] == "dl"].squeeze()
+
         recons = [recon_ctac, recon_dlac, recon_noac]
         labels = ["CTAC", "DLAC", "No AC"]
+        perfusions = [perfusion_ctac, perfusion_dlac, perfusion_noac]
 
         fig, axs = plt.subplots(1, 3, figsize=(15, 5))
         for ax in axs.flatten():
@@ -144,11 +243,16 @@ if __name__ == "__main__":
         mask = get_circular_mask(recon_ctac.shape, channels=4)
         black = np.zeros(mask.shape, np.uint8)
 
-        for ax, recon, label in zip(axs, recons, labels):
+        for ax, recon, label, _perfusion in zip(axs, recons, labels, perfusions):
             polar_map = color_map(recon)
             polar_map = np.where(mask, polar_map, black)
+            polar_map = draw_segments_grid(polar_map, color=(0, 0, 0, 1))
 
             ax.imshow(polar_map)
+
+            for segment in SEGMENTS:
+                x, y = segment.rect[1:3]
+                ax.text(x, y, _perfusion[f"segment_{segment.id}"])
             ax.set_title(label)
 
         plt.tight_layout()