[add] added louvain clustering

requirements.txt

@@ -1,7 +1,8 @@
-streamlit
+streamlit==1.24.0
 numpy
 plotly
 scipy
 mne
 h5py
-networkx
+networkx
+python-louvain

src/loader.py

@@ -42,8 +42,103 @@ def same_stem(a_name: str, b_name: str) -> bool:
     return a_stem == b_stem
 
 
+def extract_electrode_positions_2d(set_path: str) -> np.ndarray:
+    """
+    EEGLABファイルから電極位置（2D）を抽出。
+    
+    Returns:
+        pos: (C, 2) 電極の2D座標、取得できない場合はNone
+    """
+    try:
+        # MNEで読み込み
+        raw = mne.io.read_raw_eeglab(set_path, preload=False, verbose=False)
+        montage = raw.get_montage()
+        
+        if montage is None:
+            return None
+        
+        # 3D座標を取得
+        pos_3d = montage.get_positions()['ch_pos']
+        
+        if not pos_3d:
+            return None
+        
+        # チャンネル名順に並べ替え
+        ch_names = raw.ch_names
+        positions = []
+        for ch_name in ch_names:
+            if ch_name in pos_3d:
+                positions.append(pos_3d[ch_name])
+            else:
+                # 座標がないチャンネルは原点に配置
+                positions.append([0, 0, 0])
+        
+        positions = np.array(positions)
+        
+        # 3D -> 2D 投影（上から見た図）
+        # x, y座標を使用し、正規化
+        pos_2d = positions[:, :2]
+        
+        # 正規化: 最大距離が1になるようにスケーリング
+        max_dist = np.max(np.sqrt(np.sum(pos_2d**2, axis=1)))
+        if max_dist > 0:
+            pos_2d = pos_2d / max_dist * 0.85  # 0.85倍で頭の輪郭内に収める
+        
+        return pos_2d.astype(np.float32)
+        
+    except Exception as e:
+        print(f"電極位置の抽出に失敗: {e}")
+        return None
+
+
 def _load_eeglab_hdf5(set_path: str, fdt_path: Optional[str] = None, debug: bool = False) -> Tuple[np.ndarray, float]:
     """
+    EEGLABファイルから電極位置（2D）を抽出。
+    
+    Returns:
+        pos: (C, 2) 電極の2D座標、取得できない場合はNone
+    """
+    try:
+        # MNEで読み込み
+        raw = mne.io.read_raw_eeglab(set_path, preload=False, verbose=False)
+        montage = raw.get_montage()
+        
+        if montage is None:
+            return None
+        
+        # 3D座標を取得
+        pos_3d = montage.get_positions()['ch_pos']
+        
+        if not pos_3d:
+            return None
+        
+        # チャンネル名順に並べ替え
+        ch_names = raw.ch_names
+        positions = []
+        for ch_name in ch_names:
+            if ch_name in pos_3d:
+                positions.append(pos_3d[ch_name])
+            else:
+                # 座標がないチャンネルは原点に配置
+                positions.append([0, 0, 0])
+        
+        positions = np.array(positions)
+        
+        # 3D -> 2D 投影（上から見た図）
+        # x, y座標を使用し、正規化
+        pos_2d = positions[:, :2]
+        
+        # 正規化: 最大距離が1になるようにスケーリング
+        max_dist = np.max(np.sqrt(np.sum(pos_2d**2, axis=1)))
+        if max_dist > 0:
+            pos_2d = pos_2d / max_dist * 0.85  # 0.85倍で頭の輪郭内に収める
+        
+        return pos_2d.astype(np.float32)
+        
+    except Exception as e:
+        print(f"電極位置の抽出に失敗: {e}")
+        return None
+    """
     Load EEGLAB .set file saved in MATLAB v7.3 (HDF5) format using h5py.
     Returns: (x_tc, fs) where x_tc is (T, C)
     """
@@ -227,12 +322,13 @@ def load_eeglab_tc_from_bytes(
     set_name: str,
     fdt_bytes: Optional[bytes] = None,
     fdt_name: Optional[str] = None,
-) -> Tuple[np.ndarray, float]:
+) -> Tuple[np.ndarray, float, Optional[np.ndarray]]:
     """
     Load EEGLAB .set (and optional .fdt) from bytes using MNE or h5py.
     Returns:
       x_tc: (T, C) float32
       fs: sampling rate (Hz)
+      electrode_pos: (C, 2) float32 or None - 電極の2D座標
 
     Notes:
       - 多くのEEGLABは .set が .fdt を参照するため、同じディレクトリに同名で置く必要があります。
@@ -261,7 +357,11 @@ def load_eeglab_tc_from_bytes(
             raw = mne.io.read_raw_eeglab(set_path, preload=True, verbose=False)
             fs = float(raw.info["sfreq"])
             x_tc = raw.get_data().T  # (T,C)
-            return x_tc.astype(np.float32), fs
+            
+            # 電極位置を取得
+            electrode_pos = extract_electrode_positions_2d(set_path)
+            
+            return x_tc.astype(np.float32), fs, electrode_pos
 
         except Exception as e_raw:
             # 2) Epochsとして読む（エポックデータ用）
@@ -273,7 +373,11 @@ def load_eeglab_tc_from_bytes(
                 # ここは方針を選ぶ：平均 or 連結
                 x_mean = x.mean(axis=0)   # (C,T)
                 x_tc = x_mean.T           # (T,C)
-                return x_tc.astype(np.float32), fs
+                
+                # 電極位置を取得（epochsからも取得可能）
+                electrode_pos = extract_electrode_positions_2d(set_path)
+                
+                return x_tc.astype(np.float32), fs, electrode_pos
 
             except Exception as e_ep:
                 # 3) HDF5形式として読む（MATLAB v7.3）
@@ -285,7 +389,11 @@ def load_eeglab_tc_from_bytes(
                     if 'streamlit' in sys.modules:
                         debug = True
                     x_tc, fs = _load_eeglab_hdf5(set_path, fdt_path=fdt_path, debug=debug)
-                    return x_tc, fs
+                    
+                    # HDF5の場合は電極位置を取得できない（参照形式のため）
+                    electrode_pos = None
+                    
+                    return x_tc, fs, electrode_pos
                 
                 except Exception as e_hdf5:
                     # すべて失敗した場合

src/streamlit_app.py

@@ -107,14 +107,20 @@ def preprocess_all_eeglab(
     EEGLAB bytes -> load -> auto preprocess (bandpass + hilbert).
     fsは読み込んだデータのものを使う。
     """
-    x_tc, fs = load_eeglab_tc_from_bytes(
+    x_tc, fs, electrode_pos = load_eeglab_tc_from_bytes(
         set_bytes=set_bytes,
         set_name=set_name,
         fdt_bytes=fdt_bytes,
         fdt_name=fdt_name,
     )
     cfg = PreprocessConfig(fs=float(fs), f_low=float(f_low), f_high=float(f_high))
-    return preprocess_tc(x_tc, cfg)
+    result = preprocess_tc(x_tc, cfg)
+    
+    # 電極位置を追加
+    if electrode_pos is not None:
+        result["electrode_pos"] = electrode_pos
+    
+    return result
 
 
 @st.cache_data(show_spinner=False)
@@ -351,11 +357,76 @@ def get_cluster_colors(clusters: np.ndarray) -> List[str]:
     return colors
 
 
-def make_network_figure(W: np.ndarray, thr: float, use_louvain: bool = True) -> tuple[go.Figure, int]:
-    C = W.shape[0]
+def get_electrode_positions(prep: dict) -> np.ndarray:
+    """
+    電極位置を取得する。
+    
+    Returns:
+        pos: (C, 2) 電極の2D座標 (x, y)
+             取得できない場合は円形配置を返す
+    """
+    # prepに電極位置が保存されているかチェック
+    if "electrode_pos" in prep:
+        return prep["electrode_pos"]
+    
+    # デフォルト: 円形配置
+    C = prep["raw"].shape[1]
     angles = np.linspace(0, 2 * np.pi, C, endpoint=False)
     xs = np.cos(angles)
     ys = np.sin(angles)
+    return np.column_stack([xs, ys])
+
+
+def get_head_outline() -> dict:
+    """
+    脳の輪郭（頭のアウトライン）を生成。
+    
+    Returns:
+        outline: {'head': (x, y), 'nose': (x, y), 'ears': [(x_left, y_left), (x_right, y_right)]}
+    """
+    # 頭の円
+    theta = np.linspace(0, 2*np.pi, 100)
+    head_x = np.cos(theta)
+    head_y = np.sin(theta)
+    
+    # 鼻（上部の三角形）
+    nose_x = np.array([0, -0.1, 0.1, 0])
+    nose_y = np.array([1.0, 1.15, 1.15, 1.0])
+    
+    # 耳（左右の突起）
+    ear_theta = np.linspace(-np.pi/4, np.pi/4, 20)
+    ear_left_x = -1.0 + 0.08 * np.cos(ear_theta)
+    ear_left_y = 0.08 * np.sin(ear_theta)
+    
+    ear_right_x = 1.0 - 0.08 * np.cos(ear_theta)
+    ear_right_y = 0.08 * np.sin(ear_theta)
+    
+    return {
+        'head': (head_x, head_y),
+        'nose': (nose_x, nose_y),
+        'ear_left': (ear_left_x, ear_left_y),
+        'ear_right': (ear_right_x, ear_right_y),
+    }
+
+
+def make_network_figure(
+    W: np.ndarray, 
+    thr: float, 
+    use_louvain: bool = True,
+    electrode_pos: np.ndarray = None,
+    show_head: bool = True,
+) -> tuple[go.Figure, int]:
+    C = W.shape[0]
+    
+    # 電極位置を取得
+    if electrode_pos is None or electrode_pos.shape[0] != C:
+        # デフォルト: 円形配置
+        angles = np.linspace(0, 2 * np.pi, C, endpoint=False)
+        xs = np.cos(angles)
+        ys = np.sin(angles)
+    else:
+        xs = electrode_pos[:, 0]
+        ys = electrode_pos[:, 1]
 
     edges = threshold_edges(W, thr)
     fig = go.Figure()
@@ -380,6 +451,46 @@ def make_network_figure(W: np.ndarray, thr: float, use_louvain: bool = True) ->
         r, g, b = colorsys.hsv_to_rgb(hue, 0.9, 0.95)
         return f'rgba({int(255*r)}, {int(255*g)}, {int(255*b)}, 0.7)'
 
+    # 脳の輪郭を描画
+    if show_head:
+        outline = get_head_outline()
+        
+        # 頭の円
+        fig.add_trace(go.Scatter(
+            x=outline['head'][0], y=outline['head'][1],
+            mode='lines',
+            line=dict(color='rgba(150,150,150,0.5)', width=2),
+            showlegend=False,
+            hoverinfo='skip',
+        ))
+        
+        # 鼻
+        fig.add_trace(go.Scatter(
+            x=outline['nose'][0], y=outline['nose'][1],
+            mode='lines',
+            line=dict(color='rgba(150,150,150,0.5)', width=2),
+            showlegend=False,
+            hoverinfo='skip',
+        ))
+        
+        # 左耳
+        fig.add_trace(go.Scatter(
+            x=outline['ear_left'][0], y=outline['ear_left'][1],
+            mode='lines',
+            line=dict(color='rgba(150,150,150,0.5)', width=2),
+            showlegend=False,
+            hoverinfo='skip',
+        ))
+        
+        # 右耳
+        fig.add_trace(go.Scatter(
+            x=outline['ear_right'][0], y=outline['ear_right'][1],
+            mode='lines',
+            line=dict(color='rgba(150,150,150,0.5)', width=2),
+            showlegend=False,
+            hoverinfo='skip',
+        ))
+
     # エッジを描画（重みに応じて色と太さを変える）
     for (i, j, w) in edges:
         # 正規化された重み (0-1)
@@ -422,6 +533,7 @@ def make_network_figure(W: np.ndarray, thr: float, use_louvain: bool = True) ->
             mode="markers+text",
             text=[f"{k}" for k in range(C)],
             textposition="bottom center",
+            textfont=dict(size=8),
             marker=dict(
                 size=14, 
                 color=node_colors,
@@ -435,10 +547,10 @@ def make_network_figure(W: np.ndarray, thr: float, use_louvain: bool = True) ->
 
     fig.update_layout(
         title=f"Estimated Network (thr={thr:.3f})  edges={len(edges)}{title_suffix}",
-        height=500,
+        height=600,
         xaxis=dict(visible=False),
         yaxis=dict(visible=False),
-        margin=dict(l=10, r=10, t=50, b=10),
+        margin=dict(l=10, r=10, t=50, b=50),
         paper_bgcolor='rgba(0,0,0,0.9)',
         plot_bgcolor='rgba(0,0,0,0.9)',
     )
@@ -829,7 +941,7 @@ else:
 # 閾値スライダーとネットワーク図の表示
 wmax = float(np.max(W)) if np.isfinite(np.max(W)) else 1.0
 
-col_thr1, col_thr2 = st.columns([3, 1])
+col_thr1, col_thr2, col_thr3 = st.columns([2, 1, 1])
 with col_thr1:
     thr = st.slider(
         "閾値 (threshold)  ※下げるほどエッジが増えます",
@@ -840,15 +952,35 @@ with col_thr1:
     )
 with col_thr2:
     use_louvain = st.checkbox(
-        "Louvainクラスタリング", 
+        "Louvainクラスタ", 
         value=True, 
         disabled=not LOUVAIN_AVAILABLE,
         help="ノードの色をコミュニティ検出結果で塗り分けます"
     )
+with col_thr3:
+    show_head = st.checkbox(
+        "脳の輪郭を表示",
+        value=True,
+        help="頭部のアウトラインを表示します"
+    )
+
+# 電極位置を取得
+electrode_pos = prep.get("electrode_pos", None)
+
+if electrode_pos is not None:
+    st.info(f"✓ 電極位置を使用してネットワークを配置 ({electrode_pos.shape[0]} channels)")
+else:
+    st.info("ℹ️ 電極位置が取得できなかったため、円形配置を使用します")
 
 net_col1, net_col2 = st.columns([2, 1])
 with net_col1:
-    fig_net, edge_n = make_network_figure(W, float(thr), use_louvain=use_louvain)
+    fig_net, edge_n = make_network_figure(
+        W, 
+        float(thr), 
+        use_louvain=use_louvain, 
+        electrode_pos=electrode_pos,
+        show_head=show_head,
+    )
     st.plotly_chart(fig_net)
 
 with net_col2: