Upload visualizer.py with huggingface_hub

visualizer.py

@@ -0,0 +1,196 @@
+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