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("Apec Segmentation
", unsafe_allow_html=True)
st.markdown(" Welcome to Apec Segmentation, an AI-driven, coronary angiography segmentation tool.", 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 = '''
'''
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 = "
Beginning with the desired start point and finishing at the desired end point, click along the artery aiming for ~5-10 points. Segmentation is automatic.
"
st.markdown(f"Selected frame
", 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"■ {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"Predicted mask", unsafe_allow_html=True)
st.markdown(f"
If the predicted mask has errors, restart and select more points to help the segmentation model.
", 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"
Artery profile
", 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"
Contours
", 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)