tuar-eeg-dashboard / visualizer.py
Gunin09's picture
Upload visualizer.py with huggingface_hub
ca09663 verified
Raw
History Blame
6.89 kB
import numpy as np
import plotly.graph_objects as go
from data_loader import ARTIFACT_COLORS, ARTIFACT_LABEL_MAP
CHANNEL_SPACING = 1.0
def build_eeg_figure(
signal: np.ndarray,
channel_names: list[str],
sampling_rate_hz: float,
start_time_s: float = 0.0,
annotations: list[dict] | None = None,
title: str = "EEG Signal",
) -> go.Figure:
n_channels, n_samples = signal.shape
time_axis = start_time_s + np.arange(n_samples) / sampling_rate_hz
normalized = np.zeros_like(signal, dtype=np.float32)
for i in range(n_channels):
p2p = np.nanpercentile(signal[i], 97.5) - np.nanpercentile(signal[i], 2.5)
if p2p > 0:
normalized[i] = signal[i] / p2p * CHANNEL_SPACING * 0.8
else:
normalized[i] = 0.0
fig = go.Figure()
if annotations:
_add_artifact_shapes(fig, annotations, start_time_s, time_axis[-1], n_channels)
for i in range(n_channels):
ch_data = normalized[i]
offset = -i * CHANNEL_SPACING
name = channel_names[i] if i < len(channel_names) else f"Ch{i}"
fig.add_trace(go.Scatter(
x=time_axis,
y=ch_data + offset,
mode="lines",
name=name,
line=dict(width=0.8, color=f"hsl({(i * 37) % 360}, 70%, 65%)"),
hovertemplate=(
f"<b>{name}</b><br>"
"Time: %{x:.3f}s<br>"
"Amplitude: %{customdata:.1f} uV<extra></extra>"
),
customdata=signal[i],
))
y_ticks = [-i * CHANNEL_SPACING for i in range(n_channels)]
y_labels = [channel_names[i] if i < len(channel_names) else f"Ch{i}" for i in range(n_channels)]
fig.update_layout(
title=dict(text=title, font=dict(size=14)),
xaxis=dict(
title="Time (s)",
showgrid=True,
gridcolor="rgba(128,128,128,0.2)",
zeroline=False,
dtick=1.0,
),
yaxis=dict(
tickvals=y_ticks,
ticktext=y_labels,
showgrid=False,
zeroline=False,
),
height=max(500, n_channels * 30 + 120),
margin=dict(l=110, r=20, t=50, b=50),
plot_bgcolor="rgb(15, 15, 25)",
paper_bgcolor="rgb(10, 10, 20)",
font=dict(color="rgb(200, 200, 210)", size=11),
showlegend=False,
hovermode="x unified",
dragmode="zoom",
)
return fig
def _add_artifact_shapes(
fig: go.Figure,
annotations: list[dict],
view_start: float,
view_end: float,
n_channels: int,
) -> None:
y_top = CHANNEL_SPACING
y_bottom = -(n_channels - 0.5) * CHANNEL_SPACING
for ann in annotations:
onset = max(ann["onset_s"], view_start)
end = min(ann["end_s"], view_end)
if onset >= end:
continue
label = ann.get("label", ann.get("raw_label", ""))
color = ann.get("color", "rgba(128, 128, 128, 0.3)")
strong_color = color.replace("0.25)", "0.35)")
fig.add_shape(
type="rect",
x0=onset, x1=end,
y0=y_bottom, y1=y_top,
fillcolor=strong_color,
line=dict(width=1, color=strong_color.replace("0.35)", "0.7)")),
layer="below",
)
mid_x = (onset + end) / 2
if (end - onset) > 0.2:
fig.add_annotation(
x=mid_x,
y=y_top + CHANNEL_SPACING * 0.3,
text=f"<b>{label}</b>",
showarrow=False,
font=dict(size=11, color="white"),
bgcolor=strong_color.replace("0.35)", "0.85)"),
borderpad=3,
)
def build_artifact_legend() -> str:
items = []
for label, color in ARTIFACT_COLORS.items():
if label in ("Clean", "Background"):
continue
rgba_solid = color.replace("0.25)", "0.85)")
items.append(
f'<div style="display:flex;align-items:center;margin:4px 0;">'
f'<span style="display:inline-block;width:16px;height:16px;'
f'background:{rgba_solid};border-radius:3px;margin-right:8px;'
f'flex-shrink:0;"></span>'
f'<span style="font-size:13px;">{label}</span></div>'
)
return '<div style="padding:4px;">' + "".join(items) + "</div>"
def build_metadata_html(info: dict) -> str:
rows = [
("Recording ID", info.get("recording_id", "N/A")),
("Dataset", info.get("dataset_id", "N/A")),
("Subject", info.get("subject", "N/A")),
("Session", info.get("session", "N/A")),
("Task", info.get("task", "N/A")),
("Duration", f"{info.get('duration_s', 0):.1f} s"),
("Channels", f"{info.get('n_channels', 0)} ({info.get('n_eeg_channels', 0)} EEG)"),
("Sampling Rate", f"{info.get('sampling_rate_hz', 0):.0f} Hz"),
("Reference", info.get("reference", "N/A")),
("Montage", info.get("montage_name", "N/A")),
("Format", info.get("archival_format", "N/A")),
("Roundtrip", info.get("roundtrip_class", "N/A")),
]
html = '<div style="font-family:monospace; font-size:13px; line-height:1.8;">'
for label, value in rows:
html += (
f'<div><span style="color:#888;min-width:130px;display:inline-block;">'
f'{label}:</span> <span style="color:#e0e0e0;">{value}</span></div>'
)
html += "</div>"
return html
def build_annotation_summary(annotations: list[dict]) -> str:
if not annotations:
return (
'<div style="color:#888; font-size:13px; padding:8px;">'
'No artifact annotations in this window. '
'Try navigating to a different time region.</div>'
)
counts: dict[str, int] = {}
total_dur: dict[str, float] = {}
for ann in annotations:
label = ann.get("label", "Unknown")
counts[label] = counts.get(label, 0) + 1
total_dur[label] = total_dur.get(label, 0) + ann.get("duration_s", 0)
html = '<div style="font-family:monospace; font-size:13px; line-height:1.8; padding:4px;">'
html += f'<div style="color:#fff; margin-bottom:6px;"><b>Artifacts found: {len(annotations)}</b></div>'
for label in sorted(counts.keys()):
color = ARTIFACT_COLORS.get(label, "rgba(128,128,128,0.5)")
rgba_solid = color.replace("0.25)", "0.85)")
html += (
f'<div style="display:flex;align-items:center;margin:2px 0;">'
f'<span style="display:inline-block;width:12px;height:12px;'
f'background:{rgba_solid};border-radius:2px;margin-right:8px;'
f'flex-shrink:0;"></span>'
f'<span style="color:#e0e0e0;">{label}</span>: '
f'<span style="color:#aaa;">{counts[label]}x, {total_dur[label]:.1f}s total</span>'
f'</div>'
)
html += "</div>"
return html