angioPySegmentation.py

import os
import os.path
import matplotlib.pyplot as plt
import numpy
import pandas as pd
import streamlit as st
import SimpleITK as sitk
import pydicom
import glob
import mpld3
import streamlit.components.v1 as components
import plotly.express as px
import plotly.graph_objects as go
import tifffile
from streamlit_plotly_events import plotly_events
from streamlit_drawable_canvas import st_canvas
from PIL import Image
# from streamlit_image_coordinates import streamlit_image_coordinates
import predict
import angioPyFunctions
import scipy
import cv2
import json

import ssl

ssl._create_default_https_context = ssl._create_unverified_context

st.set_page_config(page_title="Apec Segmentation", layout="wide")

if 'stage' not in st.session_state:
    st.session_state.stage = 0



# Make output folder
# os.makedirs(name=outputPath, exist_ok=True)

# arteryDictionary = {
#     'LAD':       {'colour': "#f03b20"},
#     'CX':        {'colour': "#31a354"},
#     'OM':    {'colour' : "#74c476"},
#     'RCA':       {'colour': "#08519c"},
#     'AM':   {'colour' : "#3182bd"},
#     'LM':        {'colour' : "#984ea3"},
# }

# def file_selector(folder_path='.'):
#     fileNames = [file for file in glob.glob(f"{folder_path}/*")]
#     selectedDicom = st.sidebar.selectbox('Select a DICOM file:', fileNames)
#     if selectedDicom is None:
#         return None

#     return selectedDicom

@st.cache_data
def selectSlice(slice_ix, pixelArray, fileName):

    # Save the selected frame 
    tifffile.imwrite(f"{outputPath}/{fileName}", pixelArray[slice_ix, :, :])

    # Set the button as clicked
    st.session_state.btnSelectSlice = True


def parse_uploaded_annotations(annotation_data, image_label=None, image_path=None, exclude_categories=None):
    """
    Return a list of dicts with `top` and `left` keys from various annotation schemas.
    Supports Streamlit canvas JSON and COCO polygon segmentations.
    """
    if exclude_categories is None:
        exclude_categories = set()
    else:
        exclude_categories = {str(name).lower() for name in exclude_categories}

    if not isinstance(annotation_data, dict):
        return [], []

    streamlit_objects = annotation_data.get("objects")
    if isinstance(streamlit_objects, list) and streamlit_objects:
        return streamlit_objects, []

    annotations = annotation_data.get("annotations")
    images = annotation_data.get("images", [])
    categories = annotation_data.get("categories", [])

    if not isinstance(annotations, list) or not annotations:
        return [], []

    category_lookup = {}
    for category in categories:
        category_id = category.get("id")
        category_name = category.get("name", "")
        if category_id is not None:
            category_lookup[category_id] = category_name

    target_filename = None
    if image_path:
        target_filename = os.path.basename(str(image_path))
    elif image_label:
        target_filename = os.path.basename(str(image_label))

    image_id = None
    if target_filename:
        for image_entry in images:
            file_name = image_entry.get("file_name")
            if file_name and os.path.basename(str(file_name)) == target_filename:
                image_id = image_entry.get("id")
                break

    if image_id is None and images:
        image_id = images[0].get("id")

    if image_id is None:
        return [], []

    matching_annotations = [
        ann for ann in annotations
        if ann.get("image_id") == image_id and ann.get("segmentation")
    ]

    collected_polygons = []
    primary_points = []

    for ann in matching_annotations:
        segmentation = ann.get("segmentation")
        polygon = None
        category_name = category_lookup.get(ann.get("category_id"), "")
        if category_name and category_name.lower() in exclude_categories:
            continue

        if isinstance(segmentation, list):
            if segmentation and isinstance(segmentation[0], (list, tuple)):
                polygon = segmentation[0]
            else:
                polygon = segmentation

        if not polygon or not isinstance(polygon, (list, tuple)):
            continue

        coords = [float(coord) for coord in polygon if isinstance(coord, (int, float))]
        if len(coords) < 4:
            continue

        even_length = (len(coords) // 2) * 2
        coords = coords[:even_length]

        polygon_points = []
        for idx in range(0, even_length, 2):
            x = coords[idx]
            y = coords[idx + 1]
            polygon_points.append([x, y])

        if not polygon_points:
            continue

        collected_polygons.append(
            {
                "points": numpy.array(polygon_points, dtype=numpy.float32),
                "category": category_name or f"category_{ann.get('category_id')}",
            }
        )

        if not primary_points:
            primary_points = [
                {"top": point[1], "left": point[0] - 3.5, "source": "coco"}
                for point in polygon_points
            ]

    return primary_points, collected_polygons


DicomFolder = "Dicoms/"
# exampleDicoms = {
#     'RCA2' : 'Dicoms/RCA1',
#     'RCA1' : 'Dicoms/RCA4',
#     # 'RCA2' : 'Dicoms/RCA2',
#     # 'RCA3' : 'Dicoms/RCA3',
#     # 'LCA1' : 'Dicoms/LCA1',
#     # 'LCA2' : 'Dicoms/LCA2',
# 
# }
exampleDicoms = {}
files = sorted(glob.glob(DicomFolder+"/*"))
for file in files:
    exampleDicoms[os.path.basename(file)] = file

# Main text
st.markdown("<h1 style='text-align: center;'>Apec Segmentation</h1>", unsafe_allow_html=True)
st.markdown("<h5 style='text-align: center;'> Welcome to <b>Apec Segmentation</b>, an AI-driven, coronary angiography segmentation tool.</h1>", unsafe_allow_html=True)
st.markdown("")

# Build the sidebar
# Select DICOM file: here eventually we will use the file_uploader widget, but for the demo this is deactivate. Instead we will have a choice of 3 anonymised DICOMs to pick from
# selectedDicom = st.sidebar.file_uploader("Upload DICOM file:",type=["dcm"], accept_multiple_files=False)

# def changeSessionState():

#     # value += 1

#     print("CHANGED!")

input_mode = st.sidebar.radio(
    "Input source",
    ("Example DICOM", "Upload Image"),
    key="input_mode_selector",
)

pixelArray = None
selected_label = None
selected_path = None

if input_mode == "Example DICOM":
    if exampleDicoms:
        DropDownDicom = st.sidebar.selectbox(
            "Select example DICOM file:",
            options=list(exampleDicoms.keys()),
            key="dicomDropDown",
        )

        selected_label = DropDownDicom
        selected_path = exampleDicoms[DropDownDicom]

        try:
            print(f"Trying to load {selected_path}")
            dcm = pydicom.dcmread(selected_path, force=True)

            pixelArray = dcm.pixel_array

            # Just take first channel if it's RGB?
            if len(pixelArray.shape) == 4:
                pixelArray = pixelArray[:, :, :, 0]
        except Exception as err:
            st.sidebar.error(f"Unable to read DICOM '{selected_label}': {err}")
            pixelArray = None
    else:
        st.sidebar.info("Add DICOM files to the `Dicoms/` folder or switch to image upload.")
else:
    uploaded_file = st.sidebar.file_uploader(
        "Upload angiography frame (PNG or JPG)",
        type=["png", "jpg", "jpeg"],
        key="uploaded_frame",
    )

    if uploaded_file is not None:
        selected_label = uploaded_file.name
        selected_path = uploaded_file.name

        try:
            uploaded_image = Image.open(uploaded_file)
            if uploaded_image.mode != "L":
                uploaded_image = uploaded_image.convert("L")

            image_array = numpy.array(uploaded_image)
            pixelArray = numpy.expand_dims(image_array, axis=0)
        except Exception as err:
            st.sidebar.error(f"Could not read uploaded image: {err}")
            pixelArray = None

stepOne = st.sidebar.expander("STEP ONE", True)
stepTwo = st.sidebar.expander("STEP TWO", True)

# Create tabs 
tab1, tab2 = st.tabs(["Segmentation", "Analysis"])

# Increase tab font size
css = '''
<style>
    .stTabs [data-baseweb="tab-list"] button [data-testid="stMarkdownContainer"] p {
    font-size:16px;
    }
</style>
'''

st.markdown(css, unsafe_allow_html=True)

# while True:
# Once a file is uploaded, the following annotation sequence is initiated 
if pixelArray is None:
        st.info("Select an example DICOM or upload a PNG/JPG frame to start the segmentation workflow.")
else:
        if pixelArray.ndim == 4:
            pixelArray = pixelArray[:, :, :, 0]
        if pixelArray.ndim == 2:
            pixelArray = numpy.expand_dims(pixelArray, axis=0)

        n_slices = pixelArray.shape[0]
        slice_ix = 0

        with tab1:

            with stepOne:
                st.write("Select frame for annotation. Aim for an end-diastolic frame with good visualisation of the artery of interest.")

                if n_slices > 1:
                    slice_ix = st.slider('Frame', 0, n_slices-1, int(n_slices/2), key='sliceSlider')


                predictedMask = numpy.zeros_like(pixelArray[slice_ix, :, :])


            with stepTwo:

                artery_display_options = {
                    "LAD - Left Anterior Descending": "LAD",
                    "CX - Left Circumflex": "CX",
                    "RCA - Right Coronary Artery": "RCA",
                    "LM - Left Main (LMCA)": "LM",
                    "OM - Obtuse Marginal branch (of the CX)": "OM",
                    "AM - Acute Marginal branch (of the RCA)": "AM",
                    "D - Diagonal branch (of the LAD)": "D",
                }

                selected_display = st.selectbox(
                    "Select artery for annotation:",
                    list(artery_display_options.keys()),
                    key="arteryDropMenu",
                )

                selectedArtery = artery_display_options[selected_display]

                st.write("Beginning with the desired start point and finishing at the desired end point, click along the artery aiming for ~5-10 points.")

                uploaded_annotation_points = []
                uploaded_annotation_polygons = []
                annotation_upload = st.file_uploader(
                    "Optional: Load annotation JSON",
                    type=["json"],
                    key="annotation_json_upload",
                    help="Upload previously saved canvas annotations to reuse the same points.",
                )

                if annotation_upload is not None:
                    try:
                        uploaded_json_raw = annotation_upload.getvalue()
                        uploaded_annotation_data = json.loads(uploaded_json_raw.decode("utf-8"))
                        (
                            uploaded_annotation_points,
                            uploaded_annotation_polygons,
                        ) = parse_uploaded_annotations(
                            uploaded_annotation_data,
                            image_label=selected_label,
                            image_path=selected_path,
                            exclude_categories={"Stenosis_region"},
                        )
                        if not uploaded_annotation_points and not uploaded_annotation_polygons:
                            st.warning("Uploaded JSON did not contain any usable annotation points for this view (Stenosis regions are ignored).")
                    except (json.JSONDecodeError, UnicodeDecodeError) as err:
                        st.error(f"Could not read annotation JSON: {err}")
                        uploaded_annotation_points = []
                        uploaded_annotation_polygons = []

                stroke_color = angioPyFunctions.colourTableList[selectedArtery]

                catheter_mode = st.checkbox(
                    "Use catheter calibration (requires known catheter diameter)",
                    value=False,
                    key="catheter_calibration_toggle",
                    help="Enable this if the selected region corresponds to a catheter and you know its diameter in millimetres.",
                )

                catheter_diameter_mm = None
                if catheter_mode:
                    catheter_diameter_input = st.number_input(
                        "Known catheter diameter (mm)",
                        min_value=0.1,
                        max_value=20.0,
                        value=2.0,
                        step=0.1,
                        key="catheter_diameter_mm_input",
                    )
                    if catheter_diameter_input and catheter_diameter_input > 0:
                        catheter_diameter_mm = float(catheter_diameter_input)
                    else:
                        st.warning("Enter a positive catheter diameter to enable millimetre scaling.")


            col1, col2 = st.columns((15,15))

            with col1:
                col1a, col1b, col1c = st.columns((1,10,1))

                with col1b:

                    leftImageText = "<p style='text-align: center; color: white;'>Beginning with the desired <u><b>start point</b></u> and finishing at the desired <u><b>end point</b></u>, click along the artery aiming for ~5-10 points. Segmentation is automatic.</p>"

                    st.markdown(f"<h5 style='text-align: center; color: white;'>Selected frame</h5>", unsafe_allow_html=True)

                    st.markdown(leftImageText, unsafe_allow_html=True)

                    original_frame = pixelArray[slice_ix, :, :]
                    original_height, original_width = original_frame.shape

                    selectedFrame = cv2.resize(original_frame, (512,512))

                    if selectedFrame.dtype != numpy.uint8:
                        selectedFrameDisplay = numpy.clip(selectedFrame, 0, 255).astype(numpy.uint8)
                    else:
                        selectedFrameDisplay = selectedFrame.copy()

                    canvas_background = selectedFrameDisplay
                    legend_html = ""
                    if uploaded_annotation_polygons:
                        scale_y = 512.0 / original_height
                        scale_x = 512.0 / original_width
                        if canvas_background.ndim == 2:
                            canvas_background_color = numpy.stack([canvas_background] * 3, axis=-1)
                        else:
                            canvas_background_color = canvas_background.copy()

                        category_palette_rgb = [
                            (228, 26, 28),
                            (55, 126, 184),
                            (77, 175, 74),
                            (152, 78, 163),
                            (255, 127, 0),
                            (166, 86, 40),
                            (247, 129, 191),
                            (153, 153, 153),
                            (102, 194, 165),
                            (141, 160, 203),
                        ]
                        category_colours = {}

                        for polygon_entry in uploaded_annotation_polygons:
                            polygon = polygon_entry.get("points")
                            if polygon is None or polygon.size < 4:
                                continue

                            category_name = polygon_entry.get("category") or "annotation"
                            if category_name not in category_colours:
                                rgb = category_palette_rgb[len(category_colours) % len(category_palette_rgb)]
                                category_colours[category_name] = (rgb[2], rgb[1], rgb[0])

                            overlay_colour = category_colours[category_name]

                            scaled_polygon = polygon.copy()
                            scaled_polygon[:, 0] = scaled_polygon[:, 0] * scale_x
                            scaled_polygon[:, 1] = scaled_polygon[:, 1] * scale_y

                            polygon_path = scaled_polygon.reshape((-1, 1, 2)).astype(numpy.int32)
                            cv2.polylines(
                                canvas_background_color,
                                [polygon_path],
                                isClosed=True,
                                color=overlay_colour,
                                thickness=2,
                            )
                        canvas_background = canvas_background_color
                        if category_colours:
                            legend_items = []
                            for category_name, colour_bgr in category_colours.items():
                                colour_rgb = (colour_bgr[2], colour_bgr[1], colour_bgr[0])
                                legend_items.append(
                                    f"<span style='color: rgb({colour_rgb[0]}, {colour_rgb[1]}, {colour_rgb[2]});'>&#9632;</span> {category_name}"
                                )
                            legend_html = " ".join(legend_items)

                    if canvas_background.ndim == 3:
                        background_np = cv2.cvtColor(canvas_background, cv2.COLOR_BGR2RGB)
                    else:
                        background_np = canvas_background

                    canvas_key = f"canvas-{selected_label}" if selected_label else "canvas-default"

                    # Create a canvas component
                    annotationCanvas = st_canvas(
                        fill_color="red",  # Fixed fill color with some opacity
                        stroke_width=1,
                        stroke_color="red",
                        background_color='black',
                        background_image= Image.fromarray(background_np),
                        update_streamlit=True,
                        height=512,
                        width=512,
                        drawing_mode="point",
                        point_display_radius=2,
                        key=canvas_key,
                    )

                    if legend_html:
                        st.markdown(legend_html, unsafe_allow_html=True)

                    # Do something interesting with the image data and paths
                    objects = pd.DataFrame()
                    raw_annotation_objects = []

                    if annotationCanvas.json_data:
                        raw_annotation_objects = annotationCanvas.json_data.get("objects", [])

                    if not raw_annotation_objects and uploaded_annotation_points:
                        raw_annotation_objects = uploaded_annotation_points
                        st.caption(f"Loaded {len(raw_annotation_objects)} annotation points from uploaded JSON.")

                    if raw_annotation_objects:
                        objects = pd.json_normalize(raw_annotation_objects) # need to convert obj to str because PyArrow

                        if len(objects) != 0:

                            for col in objects.select_dtypes(include=['object']).columns:
                                objects[col] = objects[col].astype("str")

                            groundTruthPoints = numpy.vstack(
                                (
                                    numpy.array(objects['top']),
                                    numpy.array(objects['left']+3.5) # compensate for some streamlit offset or something
                                )
                            ).T

                            mask = angioPyFunctions.arterySegmentation(
                                pixelArray[slice_ix],
                                groundTruthPoints,
                            )
                            predictedMask = predict.CoronaryDataset.mask2image(mask)
                            # predictedMask = predictedMask.crop((0, 0, imageSize[0], imageSize[1]))
                            predictedMask = numpy.asarray(predictedMask)
            with col2:
                col2a, col2b, col2c = st.columns((1,10,1))

                with col2b:
                    st.markdown(f"<h5 style='text-align: center; color: white;'>Predicted mask</h1>", unsafe_allow_html=True)
                    st.markdown(f"<p style='text-align: center; color: white;'>If the predicted mask has errors, restart and select more points to help the segmentation model. </p>", unsafe_allow_html=True)

                    stroke_color = "rgba(255, 255, 255, 255)"

                    maskCanvas = st_canvas(
                        fill_color=angioPyFunctions.colourTableList[selectedArtery],  # Fixed fill color with some opacity
                        stroke_width=0,
                        stroke_color=stroke_color,
                        background_color='black',
                        background_image= Image.fromarray(predictedMask),
                        update_streamlit=True,
                        height=512,
                        width=512,
                        drawing_mode="freedraw",
                        point_display_radius=3,
                        key="maskCanvas",
                    )


                    # Check that the mask array is not blank
                    if numpy.sum(predictedMask) > 0 and len(objects)>4:
                        # add alpha channel to predict mask in order to merge
                        b_channel, g_channel, r_channel = cv2.split(predictedMask)
                        a_channel = numpy.full_like(predictedMask[:,:,0], fill_value=255)

                        predictedMaskRGBA = cv2.merge((predictedMask, a_channel))


                        with tab2:
                            # combinedMask = cv2.cvtColor(predictedMaskRGBA, cv2.COLOR_RGBA2RGB)

                            # print(combinedMask.shape)
                            # tifffile.imwrite(f"{outputPath}/test.tif", combinedMask)


                            # tab2Col1, tab2Col2, tab2Col3 = st.columns([1,15,1])
                            tab2Col1, tab2Col2 = st.columns([20,10])

                            with tab2Col1:
                                st.markdown(f"<h5 style='text-align: center; color: white;'><br>Artery profile</h5>", unsafe_allow_html=True)

                                # Extract thickness information from mask
                                EDT = scipy.ndimage.distance_transform_edt(cv2.cvtColor(predictedMaskRGBA, cv2.COLOR_RGBA2GRAY))

                                # Skeletonise, get a list of ordered centreline points, and spline them
                                skel = angioPyFunctions.skeletonise(predictedMaskRGBA)
                                tck = angioPyFunctions.skelSplinerWithThickness(skel=skel, EDT=EDT)

                                # Interogate the spline function over 1000 points
                                splinePointsY, splinePointsX, splineThicknesses = scipy.interpolate.splev(
                                numpy.linspace(
                                    0.0,
                                    1.0,
                                    1000), 
                                    tck)

                                clippingLength = 20

                                vesselThicknesses = splineThicknesses[clippingLength:-clippingLength]*2
                                thickness_unit = "pixels"
                                vesselThicknessesDisplay = vesselThicknesses
                                calibration_message = None
                                if catheter_mode and catheter_diameter_mm:
                                    catheter_diameter_pixels = numpy.median(vesselThicknesses)
                                    if catheter_diameter_pixels > 0:
                                        mm_per_pixel = catheter_diameter_mm / catheter_diameter_pixels
                                        vesselThicknessesDisplay = vesselThicknesses * mm_per_pixel
                                        thickness_unit = "mm"
                                        calibration_message = (
                                            f"Calibrated using catheter diameter {catheter_diameter_mm:.2f} mm "
                                            f"(median profile thickness {catheter_diameter_pixels:.2f} pixels → {mm_per_pixel:.4f} mm/pixel)."
                                        )
                                    else:
                                        st.warning("Unable to calibrate: catheter profile thickness is zero pixels.")

                                fig = px.line(
                                    x=numpy.arange(1,len(vesselThicknessesDisplay)+1),
                                    y=vesselThicknessesDisplay,
                                    labels=dict(x="Centreline point", y=f"Thickness ({thickness_unit})"),
                                    width=800,
                                )
                                # fig.update_layout(showlegend=False, xaxis={'showgrid': False, 'zeroline': True})
                                fig.update_traces(line_color='rgb(31, 119, 180)', textfont_color="white", line={'width':4})
                                fig.update_xaxes(showline=True, linewidth=2, linecolor='white', showgrid=False,gridcolor='white')
                                fig.update_yaxes(showline=True, linewidth=2, linecolor='white', gridcolor='white')

                                fig.update_layout(yaxis_range=[0,numpy.max(vesselThicknessesDisplay)*1.2])
                                fig.update_layout(font_color="white",title_font_color="white")
                                fig.update_layout({'plot_bgcolor': 'rgba(0, 0, 0, 0)','paper_bgcolor': 'rgba(0, 0, 0, 0)'})

                                if calibration_message:
                                    st.caption(calibration_message)


                                selected_points = plotly_events(
                                    fig,
                                    hover_event=True,
                                    click_event=True,
                                )

                                if selected_points:
                                    # Persist the latest hover/click event so the highlight remains visible
                                    st.session_state["artery_profile_hover"] = selected_points[0]

                                hover_event_data = st.session_state.get("artery_profile_hover")



                            with tab2Col2:

                                st.markdown(f"<h5 style='text-align: center; color: white;'><br>Contours</h5>", unsafe_allow_html=True)


                                selectedFrameRGBA = cv2.cvtColor(selectedFrame, cv2.COLOR_GRAY2RGBA)

                                contour = angioPyFunctions.maskOutliner(labelledArtery=predictedMaskRGBA[:,:,0], outlineThickness=1)

                                selectedFrameRGBA[contour, :] =    [angioPyFunctions.colourTableList[selectedArtery][2],
                                                                    angioPyFunctions.colourTableList[selectedArtery][1],
                                                                    angioPyFunctions.colourTableList[selectedArtery][0],
                                                                    255]

                                highlight_center = None
                                highlight_radius = None

                                if hover_event_data and "pointNumber" in hover_event_data:
                                    hover_index = hover_event_data.get("pointNumber")
                                    if isinstance(hover_index, (int, numpy.integer)) and 0 <= hover_index < len(vesselThicknesses):
                                        spline_index = clippingLength + hover_index
                                        highlight_center = (
                                            float(splinePointsX[spline_index]),
                                            float(splinePointsY[spline_index]),
                                        )
                                        highlight_radius = float(vesselThicknesses[hover_index] / 2.0)
                                    else:
                                        # Clear stale hover data if it no longer matches the current profile length
                                        st.session_state.pop("artery_profile_hover", None)
                                        hover_event_data = None

                                fig2 = px.imshow(selectedFrameRGBA)


                                fig2.update_xaxes(visible=False)
                                fig2.update_yaxes(visible=False)
                                fig2.update_layout(margin={"t": 0, "b": 0, "r": 0, "l": 0, "pad": 0},) #remove margins
                                # fig2.coloraxis(visible=False)

                                fig2.update_traces(dict(
                                    showscale=False, 
                                    coloraxis=None, 
                                    colorscale='gray'), selector={'type':'heatmap'})

                                fig2.add_trace(go.Scatter(x=splinePointsX[clippingLength:-clippingLength], y=splinePointsY[clippingLength:-clippingLength], line=dict(width=1)))

                                if highlight_center:
                                    fig2.add_trace(
                                        go.Scatter(
                                            x=[highlight_center[0]],
                                            y=[highlight_center[1]],
                                            mode="markers",
                                            marker=dict(size=12, color="yellow", symbol="circle"),
                                            name="Selected location",
                                            showlegend=False,
                                            hoverinfo="skip",
                                        )
                                    )

                                    if highlight_radius and highlight_radius > 0:
                                        fig2.add_shape(
                                            type="circle",
                                            xref="x",
                                            yref="y",
                                            x0=highlight_center[0] - highlight_radius,
                                            x1=highlight_center[0] + highlight_radius,
                                            y0=highlight_center[1] - highlight_radius,
                                            y1=highlight_center[1] + highlight_radius,
                                            line=dict(color="yellow", width=2),
                                        )

                                st.plotly_chart(fig2, use_container_width=True)