retuve.hip_us.multiframe.models

  1# Copyright 2024 Adam McArthur
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     http://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""
 16All the code for building the different 3D Models and Visualizations is in this file.
 17"""
 18
 19import logging
 20from typing import List, Optional, Tuple
 21
 22import numpy as np
 23import open3d as o3d
 24import plotly.graph_objects as go
 25import trimesh
 26from numpy.typing import NDArray
 27from scipy.spatial import Delaunay, _qhull
 28
 29from retuve.classes.seg import SegFrameObjects, SegObject
 30from retuve.hip_us.classes.enums import HipLabelsUS, Side
 31from retuve.hip_us.classes.general import HipDatasUS
 32from retuve.hip_us.metrics.aca import ACA_COLORS, SideZ, Triangle
 33from retuve.hip_us.typing import CoordinatesArray3D, FemoralHeadSphere
 34from retuve.keyphrases.config import Config
 35
 36# Configure logging
 37logging.basicConfig(
 38    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 39)
 40
 41
 42def ms(
 43    x: float, y: float, z: float, radius: float, resolution: int = 20
 44) -> FemoralHeadSphere:
 45    """
 46    Return the coordinates for plotting a sphere centered at (x,y,z)
 47
 48    :param x: The x-coordinate of the center of the sphere.
 49    :param y: The y-coordinate of the center of the sphere.
 50    :param z: The z-coordinate of the center of the sphere.
 51    :param radius: The radius of the sphere.
 52    :param resolution: The resolution of the sphere.
 53
 54    :return: The coordinates for plotting a sphere.
 55    """
 56    u, v = np.mgrid[
 57        0 : 2 * np.pi : resolution * 2j, 0 : np.pi : resolution * 1j
 58    ]
 59    X = radius * np.cos(u) * np.sin(v) + x
 60    Y = radius * np.sin(u) * np.sin(v) + y
 61    Z = radius * np.cos(v) + z
 62    return (X, Y, Z)
 63
 64
 65def build_3d_visual(
 66    illium_mesh: o3d.geometry.TriangleMesh,
 67    femoral_sphere: FemoralHeadSphere,
 68    avg_normals_data: CoordinatesArray3D,
 69    normals_data: List[Triangle],
 70    cr_points: List[CoordinatesArray3D],
 71    hip_datas: HipDatasUS,
 72    config: Config,
 73) -> go.Figure:
 74    """
 75    Build the 3D visual for the hip US module.
 76
 77    :param illium_mesh: The illium mesh.
 78    :param femoral_sphere: The femoral sphere.
 79    :param avg_normals_data: The average normals data of the ACA Vectors.
 80    :param normals_data: The normals data of the ACA Vectors.
 81    :param cr_points: The CR points.
 82    :param hip_datas: The hip data.
 83    :param config: The configuration object.
 84
 85    :return: The 3D visual.
 86    """
 87
 88    # Get vertices and faces
 89    vertices = np.asarray(illium_mesh.vertices)
 90    faces = np.asarray(illium_mesh.triangles)
 91
 92    com = np.mean(vertices, axis=0)
 93
 94    z_gap = config.hip.z_gap * (200 / len(hip_datas))
 95
 96    # Move the points to the center of mass
 97    vertices = vertices - com
 98
 99    # Create figure
100    fig = go.Figure(
101        data=[
102            go.Mesh3d(
103                x=vertices[:, 0],
104                y=vertices[:, 1],
105                z=vertices[:, 2],
106                i=faces[:, 0],
107                j=faces[:, 1],
108                k=faces[:, 2],
109                color="grey",
110                opacity=0.95,
111            )
112        ]
113    )
114
115    if femoral_sphere is not None:
116        fig.add_trace(
117            go.Surface(
118                x=femoral_sphere[0] - com[0],
119                y=femoral_sphere[1] - com[1],
120                z=femoral_sphere[2] - com[2],
121                opacity=0.2,
122                colorscale=[[0, "blue"], [1, "blue"]],
123                showscale=False,
124            )
125        )
126
127    if cr_points is not None:
128        fig.add_trace(
129            go.Scatter3d(
130                x=[point[0] - com[0] for point in cr_points],
131                y=[point[1] - com[1] for point in cr_points],
132                z=[point[2] - com[2] for point in cr_points],
133                mode="markers",
134                marker=dict(color="black", size=7),
135                showlegend=False,
136            )
137        )
138
139    # Red, Orange, Yellow, Blue, Cyan, Green
140    # [0, 2, 3, 2, 3, 3]
141    # https://github.com/plotly/plotly.js/issues/3613#issuecomment-1750709712
142    scaling_factors = {
143        "red": 0,
144        "orange": 2,
145        "yellow": 3,
146        "blue": 2,
147        "cyan": 3,
148        "lightgreen": 3,
149    }
150
151    for i, color in enumerate(ACA_COLORS.values()):
152        triangles = [
153            triangle for triangle in normals_data if triangle.color == color
154        ]
155        fig.add_trace(
156            go.Cone(
157                x=[triangle.centroid[0] - com[0] for triangle in triangles],
158                y=[triangle.centroid[1] - com[1] for triangle in triangles],
159                z=[triangle.centroid[2] - com[2] for triangle in triangles],
160                u=[triangle.normal[0] for triangle in triangles],
161                v=[triangle.normal[1] for triangle in triangles],
162                w=[triangle.normal[2] for triangle in triangles],
163                colorscale=[
164                    [0, color],
165                    [1, color],
166                ],
167                showscale=False,
168                sizemode="absolute",
169                sizeref=5 - scaling_factors[color],
170                anchor="tail",
171            )
172        )
173
174    for avg_normal, avg_centroid, color in avg_normals_data:
175        fig.add_trace(
176            go.Cone(
177                x=[avg_centroid[0] - com[0]],
178                y=[avg_centroid[1] - com[1]],
179                z=[avg_centroid[2] - com[2]],
180                u=[avg_normal[0]],
181                v=[avg_normal[1]],
182                w=[avg_normal[2]],
183                colorscale=[[0, color], [1, "black"]],
184                showscale=False,
185                sizemode="absolute",
186                sizeref=32,
187                anchor="tail",
188            )
189        )
190
191    illium_landmarks = []
192    fem_landmarks = []
193    for hip in hip_datas:
194        landmark = hip.landmarks
195        if landmark is None:
196            continue
197
198        z_pos = z_gap * hip.frame_no
199        for point in ["left", "right", "apex"]:
200            illium_landmarks.append(list(getattr(landmark, point)) + [z_pos])
201
202        if landmark.point_D is not None:
203            fem_landmarks.extend(
204                [
205                    list(getattr(landmark, attr)) + [z_pos]
206                    for attr in ["point_D", "point_d"]
207                ]
208            )
209
210    for landmarks, color in [
211        (illium_landmarks, "black"),
212        (fem_landmarks, "blue"),
213    ]:
214        fig.add_trace(
215            go.Scatter3d(
216                x=[point[0] - com[0] for point in landmarks],
217                y=[point[1] - com[1] for point in landmarks],
218                z=[point[2] - com[2] for point in landmarks],
219                mode="markers",
220                marker=dict(color=color, size=2),
221                showlegend=False,
222            )
223        )
224
225    # https://plotly.com/python/3d-camera-controls/
226    camera = dict(
227        up=dict(x=0, y=-1, z=0),  # Pointing in the negative y direction
228        eye=dict(
229            x=1.5, y=-1.5, z=-1.5
230        ),  # Position in the positive x, negative y, negative z region
231    )
232
233    fig.update_layout(scene_camera=camera)
234
235    fig.update_layout(scene_dragmode="orbit")
236
237    # add legend for each color
238    for apex_side in [SideZ.LEFT, SideZ.RIGHT]:
239        for third in Side.ALL():
240            fig.add_trace(
241                go.Scatter3d(
242                    x=[(0, 0, 0)],
243                    y=[(0, 0, 0)],
244                    z=[(0, 0, 0)],
245                    mode="markers",
246                    marker=dict(color=ACA_COLORS[apex_side, third]),
247                    name=f"{apex_side} {Side.get_name(third)}",
248                )
249            )
250
251    return fig
252
253
254def get_femoral_sphere(
255    hip_datas: HipDatasUS, config: Config
256) -> Optional[Tuple[FemoralHeadSphere, float]]:
257    """
258    Get the femoral sphere.
259
260    :param hip_datas: The hip data.
261    :param config: The configuration object.
262
263    :return: The femoral sphere.
264    """
265    diameters = []
266
267    z_gap = config.hip.z_gap * (200 / len(hip_datas))
268
269    hips_of_interest = [
270        hip
271        for hip in hip_datas
272        if hip.landmarks and hip.landmarks.point_d is not None
273    ]
274
275    if len(hips_of_interest) != 0:
276        middle_percent = 0.95
277        # Calculate how many elements to skip from both ends
278        skip_count = int((1 - middle_percent) / 2 * len(diameters))
279
280        # Extract the middle 'middle_percent' percentage of elements
281        hips_of_interest = hips_of_interest[
282            skip_count : -skip_count if skip_count != 0 else None
283        ]
284
285        middle_frame = (
286            hips_of_interest[0].frame_no + hips_of_interest[-1].frame_no
287        ) // 2
288        middle_hip = min(
289            hips_of_interest,
290            key=lambda hip: abs(hip.frame_no - middle_frame),
291        )
292
293        diameter = np.linalg.norm(
294            np.array(middle_hip.landmarks.point_D)
295            - np.array(middle_hip.landmarks.point_d)
296        )
297
298        fem_center = (
299            middle_hip.landmarks.point_D[0]
300            + (
301                middle_hip.landmarks.point_d[0]
302                - middle_hip.landmarks.point_D[0]
303            )
304            / 2,
305            middle_hip.landmarks.point_D[1]
306            + (
307                middle_hip.landmarks.point_d[1]
308                - middle_hip.landmarks.point_D[1]
309            )
310            / 2,
311            middle_hip.frame_no * z_gap,
312        )
313
314        radius = (diameter / 2) + (diameter * 0.05)
315
316        femoral_head_sphere = ms(
317            fem_center[0],
318            fem_center[1],
319            fem_center[2],
320            radius,
321            resolution=20,
322        )
323
324        return femoral_head_sphere, radius
325
326    return None, None
327
328
329def get_illium_mesh(
330    hip_datas: HipDatasUS, results: List[SegFrameObjects], config: Config
331) -> Optional[Tuple[o3d.geometry.TriangleMesh, NDArray[np.float64]]]:
332    """
333    Get the illium mesh.
334
335    :param hip_datas: The hip data.
336    :param results: The segmentation results.
337    :param config: The configuration object.
338
339    :return: The illium mesh and the apex points.
340    """
341    illium_pc = []
342    apex_points = []
343
344    even_count = 0
345
346    z_gap = config.hip.z_gap * (200 / len(hip_datas))
347
348    for hip_data, seg_frame_objs in zip(hip_datas, results):
349        illium = [
350            seg_obj
351            for seg_obj in seg_frame_objs
352            if seg_obj.cls == HipLabelsUS.IlliumAndAcetabulum
353        ]
354
355        if (
356            len(illium) != 0
357            and illium[0].midline is not None
358            and hip_data.landmarks
359            and hip_data.landmarks.apex is not None
360        ):
361            illium: SegObject = illium[0]
362
363            # Convert white_points to a convenient shape for vector operations
364            midline_moved = np.array(illium.midline_moved)[
365                :, ::-1
366            ]  # Reverse each point only once
367
368            if even_count % 2 == 0:
369                # Pick points a 10 pixel intervals on x axis
370                chosen_indexs = np.arange(0, midline_moved.shape[0], 10)
371            else:
372                # choose every 10 on a 5 point offset
373                chosen_indexs = np.arange(5, midline_moved.shape[0], 10)
374
375            for points in midline_moved[chosen_indexs]:
376                points = np.append(points, [hip_data.frame_no * z_gap], axis=0)
377                illium_pc.append(points)
378                apex = hip_data.landmarks.apex
379
380                apex_points.append(
381                    [apex[0], apex[1], hip_data.frame_no * z_gap]
382                )
383
384        even_count += 1
385
386    if len(illium_pc) == 0:
387        return None, None
388
389    illium_pc = np.array(illium_pc)
390
391    # use these points to create a surface
392    # Create Open3D mesh
393    illium_pc_o3d = o3d.geometry.PointCloud()
394    illium_pc_o3d.points = o3d.utility.Vector3dVector(illium_pc)
395
396    illium_points_xz = illium_pc[:, [0, 2]]  # Selecting X and Z coordinates
397
398    try:
399        # Apply Delaunay triangulation on the XZ plane
400        tri = Delaunay(illium_points_xz)
401    except _qhull.QhullError:
402        return None, None
403
404    # Reconstruct 3D triangles using original Y-coordinates
405    vertices = illium_pc  # Using original 3D points as vertices
406    triangles = tri.simplices
407
408    # convert to trimesh
409    illium_mesh = trimesh.Trimesh(
410        vertices=np.asarray(o3d.utility.Vector3dVector(vertices)),
411        faces=np.asarray(o3d.utility.Vector3iVector(triangles)),
412    )
413
414    # apply humphrey smoothing
415    illium_mesh = trimesh.smoothing.filter_humphrey(
416        illium_mesh, iterations=20, alpha=0.1, beta=1
417    )
418
419    # convert back to open3d
420    illium_mesh = o3d.geometry.TriangleMesh(
421        vertices=o3d.utility.Vector3dVector(illium_mesh.vertices),
422        triangles=o3d.utility.Vector3iVector(illium_mesh.faces),
423    )
424
425    return illium_mesh, apex_points
426
427
428def circle_radius_at_z(
429    sphere_radius: float, z_center: float, z_input: float
430) -> float:
431    """
432    Calculate the radius of a circle on a sphere at a given Z-coordinate.
433
434    :param sphere_radius: The radius of the sphere.
435    :param z_center: The Z-coordinate of the center of the sphere.
436    :param z_input: The Z-coordinate of the circle.
437
438    :return: The radius of the circle.
439    """
440    delta_z = abs(z_input - z_center)
441    if delta_z > sphere_radius:
442        return 0  # The input Z-coordinate is outside the sphere.
443    return np.sqrt(sphere_radius**2 - delta_z**2)