# %% [code] # %% [code] #!/usr/bin/env python # coding: utf-8 # ========================================================= # MABe Challenge - Ensemble Pipeline (XGBoost + CatBoost) # Ensemble Logic: # 1. Average EMA-smoothed probabilities. # 2. Candidate selection: Avg Prob >= Avg Threshold. # 3. Tie-breaking: Average of Z-scores ((P-T)/Sigma). # 4. Post-processing: Max(gap_A, gap_B), Max(dur_A, dur_B). # ========================================================= import json import os import glob import gc import itertools from functools import lru_cache import warnings import joblib from pathlib import Path from collections import defaultdict import pandas as pd import polars as pl import numpy as np from sklearn.base import clone from xgboost import XGBClassifier from catboost import CatBoostClassifier from tqdm import tqdm warnings.filterwarnings('ignore') # Unified training dtype TRAIN_DTYPE = np.float32 DEFAULT_FPS = 30.0 # 共通の特徴量スキップリスト SKIP_SINGLE_FEATURES = [ "cam_body_center+hip_left", "cam_body_center+hip_right", "cam_ear_left+hip_left", "cam_ear_left+tail_tip", "cam_hip_left+lateral_left", "cam_hip_left+lateral_right", "cam_hip_left+tail_tip", "cam_hip_right+lateral_left", "cam_hip_right+lateral_right", "cam_s0_40", "cam_s2_20", "cam_trans_20", "ego_body_center+hip_left", "ego_body_center+hip_right", "ego_cx_ac_lag5", "ego_ego_lat_sp", "ego_hip_left+lateral_left", "ego_hip_left+tail_tip", "ego_hip_right+lateral_left", "ego_hip_right+lateral_right", "ego_hip_right+tail_tip", "ego_lateral_left+tail_tip", "ego_lateral_right+tail_tip", "ego_nose+tail_tip", ] SKIP_PAIR_FEATURES = [ "cam_12+body_center+hip_left", "cam_12+body_center+hip_right", "cam_12+hip_left+body_center", "cam_12+hip_left+lateral_right", "cam_12+hip_left+tail_tip", "cam_12+hip_right+body_center", "cam_12+lateral_left+hip_left", "cam_12+lateral_left+hip_right", "cam_12+lateral_right+hip_left", "cam_12+tail_tip+hip_right", "ego_12+body_center+hip_right", "ego_12+hip_left+body_center", "ego_12+hip_left+lateral_left", "ego_12+hip_left+lateral_right", "ego_12+hip_right+body_center", "ego_12+hip_right+lateral_left", "ego_12+hip_right+tail_tip", "ego_12+lateral_left+hip_right", "ego_12+tail_tip+hip_right", "ego_A_ld30", "ego_A_ld60", "ego_co_m30", "ego_co_m5", "ego_co_s15", "ego_co_s30", "ego_co_s5", "ego_co_s60", "ego_far", "ego_sp_cor120", "ego_sp_cor60", "ego_v_cls", "ego_va_0", ] def _skip_single_feature(name: str) -> bool: return name in SKIP_SINGLE_FEATURES def _skip_pair_feature(name: str) -> bool: return name in SKIP_PAIR_FEATURES # ========================================================= # Config # ========================================================= class CFG: # Paths base_path = "/kaggle/input/MABe-mouse-behavior-detection" # Model A: XGBoost (exp132) model_a_name = "lab_xgb" model_a_base = "/kaggle/input/mabe-exp132" model_a_path = f"{model_a_base}/lab" # Model B: CatBoost (exp128) model_b_name = "lab_cat" model_b_base = "/kaggle/input/mabe-exp128" model_b_path = f"{model_b_base}/lab" test_path = f"{base_path}/test.csv" test_tracking_path = f"{base_path}/test_tracking" # Execution mode mode = "submit" # Logit smoothing (EMA) default_logit_ema_seconds = 0.25 smooth_logits_with_ema = True # Lab grouping calms21_group_name = "CalMS21_combined" calms21_prefix = "CalMS21" combine_calms21_groups = False # Inference settings use_fold_ensemble_inference = True prefer_full_model = False # Data augmentation settings use_flip_augmentation = False use_horizontal_flip_augmentation = False use_vertical_flip_augmentation = False try: SCRIPT_DIR = Path(__file__).resolve().parent except NameError: SCRIPT_DIR = Path.cwd() # ========================================================= # Helpers / Normalization / Data Loading # ========================================================= drop_body_parts = [ 'headpiece_bottombackleft', 'headpiece_bottombackright', 'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', 'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft', 'headpiece_topfrontright', 'spine_1', 'spine_2', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint' ] def normalize_body_parts(body_parts_tracked_str: str) -> str: if not isinstance(body_parts_tracked_str, str): return body_parts_tracked_str try: parts = json.loads(body_parts_tracked_str) except (json.JSONDecodeError, TypeError): return body_parts_tracked_str filtered = [p for p in parts if p not in drop_body_parts] filtered = list(dict.fromkeys(filtered)) filtered.sort() return json.dumps(filtered) def map_lab_group(lab_id: str) -> str: if not isinstance(lab_id, str): return lab_id if getattr(CFG, "combine_calms21_groups", True) and lab_id.startswith(CFG.calms21_prefix): return CFG.calms21_group_name return lab_id def _preload_metadata_polars(csv_path: str, *, is_train: bool) -> pl.DataFrame: df = pl.read_csv(csv_path) df = df.with_columns([ pl.col("body_parts_tracked").map_elements(normalize_body_parts, return_dtype=pl.Utf8), ]) if "lab_id" in df.columns: df = df.with_columns( pl.col("lab_id").map_elements(map_lab_group, return_dtype=pl.Utf8).alias("lab_group") ) else: df = df.with_columns(pl.lit(None).alias("lab_group")) if is_train: strain_cols = [c for c in ["mouse1_strain", "mouse2_strain", "mouse3_strain", "mouse4_strain"] if c in df.columns] if strain_cols: null_counts = pl.sum_horizontal([pl.col(c).is_null().cast(pl.Int64) for c in strain_cols]) df = df.with_columns((4 - null_counts).alias("n_mice")) return df def apply_missing_value_heuristics(pvid: pd.DataFrame, *, allow_interpolate: bool = True) -> pd.DataFrame: if pvid.empty: return pvid pvid = pvid.copy() pvid.replace([np.inf, -np.inf], np.nan, inplace=True) numeric_cols = pvid.select_dtypes(include=[np.number]).columns if len(numeric_cols) == 0: return pvid has_obs = pvid[numeric_cols].notna().any(axis=0) cols_to_drop = numeric_cols[~has_obs] if len(cols_to_drop) > 0: pvid.drop(columns=cols_to_drop, inplace=True) numeric_cols = numeric_cols.difference(cols_to_drop) if len(numeric_cols) == 0: return pvid if allow_interpolate: pvid[numeric_cols] = pvid[numeric_cols].interpolate(method='linear', axis=0, limit_direction='both') return pvid def load_tracking_parquet_fast(path: str) -> pd.DataFrame: df = pl.read_parquet(path) if df.is_empty(): return pd.DataFrame() n_bodyparts = df.select(pl.col("bodypart").n_unique()).to_series().item() if n_bodyparts > 5: df = df.filter(~pl.col("bodypart").is_in(drop_body_parts)) melted = df.melt(id_vars=["video_frame", "mouse_id", "bodypart"], value_vars=["x", "y"], variable_name="coord", value_name="value") wide = melted.pivot(values="value", index="video_frame", columns=["mouse_id", "bodypart", "coord"]) if "video_frame" not in wide.columns: vf = df.select("video_frame").unique().sort("video_frame") wide = wide.with_columns(vf["video_frame"]) data_cols = [c for c in wide.columns if c != "video_frame"] col_map = [] for col in data_cols: if isinstance(col, tuple): col_map.append(col); continue if isinstance(col, str) and col.startswith("{"): inner = col[1:-1] parts = [p.strip().strip('"') for p in inner.split(",")] if len(parts) == 3: try: m_id = int(parts[0]) except: m_id = parts[0] col_map.append((m_id, parts[1], parts[2])) continue col_map.append((col, "", "")) pvid = wide.to_pandas() pvid.index = pvid.pop("video_frame") pvid = pvid[data_cols] pvid.columns = pd.MultiIndex.from_tuples(col_map) pvid = pvid.reorder_levels([1, 2, 0], axis=1).T.sort_index().T pvid.columns = pvid.columns.set_names(['bodypart', 'coord', 'mouse_id']) return pvid def _extract_frame_width(meta_row) -> float | None: if meta_row is None: return None for key in ["video_width_pix", "frame_width", "width", "video_width", "tracking_width"]: try: val = meta_row.get(key, np.nan) except: val = np.nan if np.isfinite(val) and val > 0: return float(val) try: arena_w_cm = meta_row.get("arena_width_cm", np.nan) pix_per_cm = meta_row.get("pix_per_cm_approx", np.nan) except: arena_w_cm = pix_per_cm = np.nan if np.isfinite(arena_w_cm) and arena_w_cm > 0 and np.isfinite(pix_per_cm) and pix_per_cm > 0: return float(arena_w_cm * pix_per_cm) return None def _extract_frame_height(meta_row) -> float | None: if meta_row is None: return None for key in ["video_height_pix", "frame_height", "height", "video_height", "tracking_height"]: try: val = meta_row.get(key, np.nan) except: val = np.nan if np.isfinite(val) and val > 0: return float(val) try: arena_h_cm = meta_row.get("arena_height_cm", np.nan) pix_per_cm = meta_row.get("pix_per_cm_approx", np.nan) except: arena_h_cm = pix_per_cm = np.nan if np.isfinite(arena_h_cm) and arena_h_cm > 0 and np.isfinite(pix_per_cm) and pix_per_cm > 0: return float(arena_h_cm * pix_per_cm) return None def _fps_from_meta(meta_df, fallback_lookup, default_fps=DEFAULT_FPS): if 'frames_per_second' in meta_df.columns and pd.notnull(meta_df['frames_per_second']).any(): return float(meta_df['frames_per_second'].iloc[0]) vid = meta_df['video_id'].iloc[0] return float(fallback_lookup.get(vid, default_fps)) def _build_video_fps_lookup(df: pd.DataFrame, default_fps: float = DEFAULT_FPS) -> dict: if df is None or not isinstance(df, pd.DataFrame): return {} if 'video_id' not in df.columns or 'frames_per_second' not in df.columns: return {} fps_df = df[['video_id', 'frames_per_second']].dropna(subset=['video_id']) fps_df = fps_df[pd.notnull(fps_df['frames_per_second'])].drop_duplicates('video_id') fps_map = {} for _, row in fps_df.iterrows(): try: fps_val = float(row['frames_per_second']) except: continue if np.isfinite(fps_val) and fps_val > 0: fps_map[row['video_id']] = fps_val return fps_map # ========================================================= # Feature Engineering (Same logic as XGBoost pipeline) # ========================================================= @lru_cache(maxsize=None) def _get_single_pairwise(body_parts: tuple): combos = list(itertools.combinations(body_parts, 2)) col_suffixes = tuple(f"{p1}+{p2}" for p1, p2 in combos) return combos, col_suffixes @lru_cache(maxsize=None) def _get_pairwise_product(body_parts: tuple): combos = list(itertools.product(body_parts, repeat=2)) col_suffixes = tuple(f"12+{p1}+{p2}" for p1, p2 in combos) return combos, col_suffixes def _scale(n_frames_at_30fps, fps, ref=30.0): return max(1, int(round(n_frames_at_30fps * float(fps) / ref))) def _scale_signed(n_frames_at_30fps, fps, ref=30.0): if n_frames_at_30fps == 0: return 0 s = 1 if n_frames_at_30fps > 0 else -1 mag = max(1, int(round(abs(n_frames_at_30fps) * float(fps) / ref))) return s * mag def _clean_rotation_series(rotation: pd.Series) -> pd.Series: rotation = rotation.replace([np.inf, -np.inf], np.nan) return rotation.ffill().bfill().fillna(0.0) def _maybe_add_body_center(single_mouse: pd.DataFrame): available_parts = single_mouse.columns.get_level_values(0).unique() if 'body_center' in available_parts: return single_mouse, False xs, ys = [], [] for part in available_parts: if part == 'body_center': continue part_df = single_mouse.get(part) if part_df is None or not {'x', 'y'}.issubset(part_df.columns): continue xs.append(part_df['x']); ys.append(part_df['y']) if not xs or not ys: return single_mouse, False updated = single_mouse.copy() updated[('body_center', 'x')] = pd.concat(xs, axis=1).mean(axis=1) updated[('body_center', 'y')] = pd.concat(ys, axis=1).mean(axis=1) return updated.sort_index(axis=1), True def _maybe_add_pair_body_center(mouse_pair: pd.DataFrame) -> pd.DataFrame: blocks, changed = {}, False for who in ['A', 'B']: if who not in mouse_pair.columns.get_level_values(0): continue block = mouse_pair[who] updated, added = _maybe_add_body_center(block) blocks[who] = updated changed = changed or added if not changed: return mouse_pair return pd.concat([blocks[k] for k in blocks], axis=1, keys=list(blocks.keys())) def build_egocentric_single(single_mouse: pd.DataFrame): available = set(single_mouse.columns.get_level_values(0)) front_part = None for cand in ['nose', 'head', 'neck']: if cand in available: front_part = cand; break required_parts = {'body_center', 'tail_base'} if front_part is not None: required_parts.add(front_part) if not required_parts.issubset(available): return None, None, None, None center = single_mouse['body_center'] heading = single_mouse[front_part] - single_mouse['tail_base'] rotation = pd.Series((np.pi / 2) - np.arctan2(heading['y'], heading['x']), index=heading.index) rotation = _clean_rotation_series(rotation) cos_r, sin_r = np.cos(rotation), np.sin(rotation) rotated_cols = {} for part in sorted(available): part_df = single_mouse.get(part) if part_df is None or not {'x', 'y'}.issubset(part_df.columns): continue dx = part_df['x'] - center['x'] dy = part_df['y'] - center['y'] rotated_cols[(part, 'x')] = dx * cos_r - dy * sin_r rotated_cols[(part, 'y')] = dx * sin_r + dy * cos_r if not rotated_cols: return None, rotation, cos_r, sin_r ego_df = pd.DataFrame(rotated_cols) ego_df.columns = pd.MultiIndex.from_tuples(ego_df.columns) return ego_df, rotation, cos_r, sin_r def build_agent_egocentric_pair(mouse_pair: pd.DataFrame): available_A = set(mouse_pair['A'].columns.get_level_values(0)) front_part = None for cand in ['nose', 'head', 'neck']: if cand in available_A: front_part = cand; break required_parts = {'body_center', 'tail_base'} if front_part is not None: required_parts.add(front_part) if not required_parts.issubset(available_A): return None, None, None, None center = mouse_pair['A']['body_center'] heading = mouse_pair['A'][front_part] - mouse_pair['A']['tail_base'] rotation = pd.Series((np.pi / 2) - np.arctan2(heading['y'], heading['x']), index=heading.index) rotation = _clean_rotation_series(rotation) cos_r, sin_r = np.cos(rotation), np.sin(rotation) rotated_cols = {} for who in ['A', 'B']: if who not in mouse_pair.columns.get_level_values(0): continue for part in sorted(mouse_pair[who].columns.get_level_values(0).unique()): part_df = mouse_pair[who].get(part) if part_df is None or not {'x', 'y'}.issubset(part_df.columns): continue dx = part_df['x'] - center['x'] dy = part_df['y'] - center['y'] rotated_cols[(who, part, 'x')] = dx * cos_r - dy * sin_r rotated_cols[(who, part, 'y')] = dx * sin_r + dy * cos_r if not rotated_cols: return None, rotation, cos_r, sin_r ego_df = pd.DataFrame(rotated_cols) ego_df.columns = pd.MultiIndex.from_tuples(ego_df.columns) return ego_df, rotation, cos_r, sin_r # --- Extended Feature Functions (Included as placeholders in truncated response) --- def add_curvature_features(X, center_x, center_y, fps, prefix=""): vel_x = center_x.diff() vel_y = center_y.diff() acc_x = vel_x.diff() acc_y = vel_y.diff() cross_prod = vel_x * acc_y - vel_y * acc_x vel_mag = np.sqrt(vel_x**2 + vel_y**2) curvature = np.abs(cross_prod) / (vel_mag**3 + 1e-6) for w in [25, 50, 75]: ws = _scale(w, fps) X[f'{prefix}curv_mean_{w}'] = curvature.rolling(ws, min_periods=max(1, ws // 5)).mean() angle = np.arctan2(vel_y, vel_x) angle_change = np.abs(angle.diff()) w = 30 ws = _scale(w, fps) X[f'{prefix}turn_rate_{w}'] = angle_change.rolling(ws, min_periods=max(1, ws // 5)).sum() return X def add_multiscale_features(X, center_x, center_y, fps, prefix=""): speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) scales = [20, 40, 60, 80] for scale in scales: ws = _scale(scale, fps) if len(speed) >= ws: X[f'{prefix}sp_m{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).mean() X[f'{prefix}sp_s{scale}'] = speed.rolling(ws, min_periods=max(1, ws // 4)).std() if len(scales) >= 2 and f'{prefix}sp_m{scales[0]}' in X.columns and f'{prefix}sp_m{scales[-1]}' in X.columns: X[f'{prefix}sp_ratio'] = X[f'{prefix}sp_m{scales[0]}'] / (X[f'{prefix}sp_m{scales[-1]}'] + 1e-6) return X def add_state_features(X, center_x, center_y, fps, prefix=""): speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) w_ma = _scale(15, fps) speed_ma = speed.rolling(w_ma, min_periods=max(1, w_ma // 3)).mean() try: bins = [-np.inf, 0.5 * fps, 2.0 * fps, 5.0 * fps, np.inf] speed_states = pd.cut(speed_ma, bins=bins, labels=[0, 1, 2, 3]).astype(float) for window in [20, 40, 60, 80]: ws = _scale(window, fps) if len(speed_states) >= ws: for state in [0, 1, 2, 3]: col_name = f'{prefix}s{state}_{window}' if _skip_single_feature(col_name): continue X[col_name] = ((speed_states == state).astype(float).rolling(ws, min_periods=max(1, ws // 5)).mean()) state_changes = (speed_states != speed_states.shift(1)).astype(float) trans_col = f'{prefix}trans_{window}' if not _skip_single_feature(trans_col): X[trans_col] = state_changes.rolling(ws, min_periods=max(1, ws // 5)).sum() except Exception: pass return X def add_longrange_features(X, center_x, center_y, fps, prefix=""): for window in [30, 60, 120]: ws = _scale(window, fps) if len(center_x) >= ws: X[f'{prefix}x_ml{window}'] = center_x.rolling(ws, min_periods=max(5, ws // 6)).mean() X[f'{prefix}y_ml{window}'] = center_y.rolling(ws, min_periods=max(5, ws // 6)).mean() for span in [30, 60, 120]: s = _scale(span, fps) X[f'{prefix}x_e{span}'] = center_x.ewm(span=s, min_periods=1).mean() X[f'{prefix}y_e{span}'] = center_y.ewm(span=s, min_periods=1).mean() speed = np.sqrt(center_x.diff()**2 + center_y.diff()**2) * float(fps) for window in [30, 60, 120]: ws = _scale(window, fps) if len(speed) >= ws: X[f'{prefix}sp_pct{window}'] = speed.rolling(ws, min_periods=max(5, ws // 6)).rank(pct=True) return X def add_autocorr_features(X, series, fps, prefix="", base_name="ac"): if series is None: return X lags = [5, 15, 30] for lag in lags: l = _scale(lag, fps) if l <= 0 or len(series) <= l: continue shifted = series.shift(l) win = max(l * 2, l + 2) min_p = min(win, max(3, win // 5)) roll = dict(window=win, min_periods=min_p, center=True) mean_s = series.rolling(**roll).mean() mean_sh = shifted.rolling(**roll).mean() std_s = series.rolling(**roll).std() std_sh = shifted.rolling(**roll).std() cov = (series * shifted).rolling(**roll).mean() - (mean_s * mean_sh) corr = cov / (std_s * std_sh + 1e-6) col_name = f'{prefix}{base_name}_lag{lag}' if _skip_single_feature(col_name): continue X[col_name] = corr return X def transform_single(single_mouse, body_parts_tracked, fps): def _build_single_feature_block(source_df, available_parts, center_x, center_y, prefix, extra_motion=None): pairwise, pair_cols = _get_single_pairwise(tuple(body_parts_tracked)) part_arrays = {p: source_df[p].to_numpy(dtype=TRAIN_DTYPE, copy=False) for p in available_parts if p in source_df.columns.get_level_values(0)} base = pd.DataFrame({f"{prefix}{p1}+{p2}": np.square(part_arrays[p1] - part_arrays[p2]).sum(axis=1) for p1, p2 in pairwise if p1 in part_arrays and p2 in part_arrays and not _skip_single_feature(f"{prefix}{p1}+{p2}")}, index=source_df.index) if base.empty: base = pd.DataFrame(index=source_df.index) base = base.reindex(index=source_df.index, columns=[f"{prefix}{col}" for col in pair_cols if not _skip_single_feature(f"{prefix}{col}")], copy=False) if all(p in available_parts for p in ['ear_left', 'ear_right', 'tail_base']): lag = _scale(10, fps) shifted = source_df[['ear_left', 'ear_right', 'tail_base']].shift(lag) speeds = pd.DataFrame({ f'{prefix}sp_lf': np.square(source_df['ear_left'] - shifted['ear_left']).sum(axis=1, skipna=False), f'{prefix}sp_rt': np.square(source_df['ear_right'] - shifted['ear_right']).sum(axis=1, skipna=False), f'{prefix}sp_lf2': np.square(source_df['ear_left'] - shifted['tail_base']).sum(axis=1, skipna=False), f'{prefix}sp_rt2': np.square(source_df['ear_right'] - shifted['tail_base']).sum(axis=1, skipna=False), }) base = pd.concat([base, speeds], axis=1) nt_key = f"{prefix}nose+tail_base" ear_key = f"{prefix}ear_left+ear_right" if nt_key in base.columns and ear_key in base.columns: base[f'{prefix}elong'] = base[nt_key] / (base[ear_key] + 1e-6) if all(p in source_df.columns.get_level_values(0) for p in ['nose', 'body_center', 'tail_base']): v1 = source_df['nose'] - source_df['body_center'] v2 = source_df['tail_base'] - source_df['body_center'] body_ang_col = f'{prefix}body_ang' if not _skip_single_feature(body_ang_col): base[body_ang_col] = (v1['x'] * v2['x'] + v1['y'] * v2['y']) / (np.sqrt(v1['x']**2 + v1['y']**2) * np.sqrt(v2['x']**2 + v2['y']**2) + 1e-6) base = add_autocorr_features(base, base[body_ang_col].diff(), fps, prefix=prefix, base_name='bang_ac') if center_x is not None and center_y is not None: cx_diff = center_x.diff() cy_diff = center_y.diff() speed = np.sqrt(cx_diff**2 + cy_diff**2) * float(fps) for w in [5, 15, 30, 60]: ws = _scale(w, fps) roll = dict(min_periods=1, center=True) for col_name, series in [ (f'{prefix}cx_m{w}', center_x.rolling(ws, **roll).mean()), (f'{prefix}cy_m{w}', center_y.rolling(ws, **roll).mean()), (f'{prefix}cx_s{w}', center_x.rolling(ws, **roll).std()), (f'{prefix}cy_s{w}', center_y.rolling(ws, **roll).std()), (f'{prefix}x_rng{w}', center_x.rolling(ws, **roll).max() - center_x.rolling(ws, **roll).min()), (f'{prefix}y_rng{w}', center_y.rolling(ws, **roll).max() - center_y.rolling(ws, **roll).min()), (f'{prefix}disp{w}', np.sqrt(cx_diff.rolling(ws, min_periods=1).sum()**2 + cy_diff.rolling(ws, min_periods=1).sum()**2)), (f'{prefix}act{w}', np.sqrt(cx_diff.rolling(ws, min_periods=1).var() + cy_diff.rolling(ws, min_periods=1).var())), ]: if _skip_single_feature(col_name): continue base[col_name] = series base = add_autocorr_features(base, cx_diff, fps, prefix=prefix, base_name="cx_ac") base = add_autocorr_features(base, cy_diff, fps, prefix=prefix, base_name="cy_ac") base = add_autocorr_features(base, speed, fps, prefix=prefix, base_name="sp_ac") base = add_curvature_features(base, center_x, center_y, fps, prefix=prefix) base = add_multiscale_features(base, center_x, center_y, fps, prefix=prefix) base = add_state_features(base, center_x, center_y, fps, prefix=prefix) base = add_longrange_features(base, center_x, center_y, fps, prefix=prefix) if all(p in available_parts for p in ['nose', 'tail_base']): nt_dist = np.sqrt((source_df['nose']['x'] - source_df['tail_base']['x'])**2 + (source_df['nose']['y'] - source_df['tail_base']['y'])**2) for lag in [10, 20, 40]: l = _scale(lag, fps) lg_col = f'{prefix}nt_lg{lag}' df_col = f'{prefix}nt_df{lag}' if not _skip_single_feature(lg_col): base[lg_col] = nt_dist.shift(l) if not _skip_single_feature(df_col): base[df_col] = nt_dist - nt_dist.shift(l) if not _skip_single_feature(f'{prefix}ntd_ac_lag5'): base = add_autocorr_features(base, nt_dist.diff(), fps, prefix=prefix, base_name='ntd_ac') if all(p in available_parts for p in ['ear_left', 'ear_right']): ear_d = np.sqrt((source_df['ear_left']['x'] - source_df['ear_right']['x'])**2 + (source_df['ear_left']['y'] - source_df['ear_right']['y'])**2) for off in [-30, -20, -10, 10, 20, 30]: o = _scale_signed(off, fps) col_name = f'{prefix}ear_o{off}' if _skip_single_feature(col_name): continue base[col_name] = ear_d.shift(-o) w = _scale(30, fps) ear_con_col = f'{prefix}ear_con' if not _skip_single_feature(ear_con_col): base[ear_con_col] = ear_d.rolling(w, min_periods=1, center=True).std() / (ear_d.rolling(w, min_periods=1, center=True).mean() + 1e-6) if extra_motion: for k, v in extra_motion.items(): col_name = f'{prefix}{k}' if _skip_single_feature(col_name): continue base[col_name] = v return base features = [] single_mouse, _ = _maybe_add_body_center(single_mouse) orig_available_parts = single_mouse.columns.get_level_values(0) cam_center_x = single_mouse['body_center']['x'] if 'body_center' in orig_available_parts else None cam_center_y = single_mouse['body_center']['y'] if 'body_center' in orig_available_parts else None cam_block = _build_single_feature_block(single_mouse, orig_available_parts, cam_center_x, cam_center_y, prefix="cam_") features.append(cam_block) ego_mouse, ego_rot, ego_cos, ego_sin = build_egocentric_single(single_mouse) if ego_mouse is not None and cam_center_x is not None: cx_diff = cam_center_x.diff() cy_diff = cam_center_y.diff() ego_cx = cx_diff * ego_cos - cy_diff * ego_sin ego_cy = cx_diff * ego_sin + cy_diff * ego_cos cx_ref = ego_cx.fillna(0).cumsum() cy_ref = ego_cy.fillna(0).cumsum() extra_motion = {'ego_fwd_sp': ego_cy * float(fps), 'ego_lat_sp': ego_cx * float(fps)} for w in [5, 15, 30]: ws = _scale(w, fps) roll = dict(min_periods=1, center=True) extra_motion[f'ego_fwd_m{w}'] = ego_cy.rolling(ws, **roll).mean() * float(fps) extra_motion[f'ego_lat_m{w}'] = ego_cx.rolling(ws, **roll).mean() * float(fps) ego_parts = ego_mouse.columns.get_level_values(0) features.append(_build_single_feature_block(ego_mouse, ego_parts, cx_ref, cy_ref, prefix="ego_", extra_motion=extra_motion)) if len(features) == 0: return pd.DataFrame(index=single_mouse.index).astype(TRAIN_DTYPE) X = pd.concat(features, axis=1) drop_cols = [c for c in SKIP_SINGLE_FEATURES if c in X.columns] if drop_cols: X = X.drop(columns=drop_cols) X = X.replace([np.inf, -np.inf], np.nan) return X.astype(TRAIN_DTYPE, copy=False) def add_interaction_features(X, mouse_pair, avail_A, avail_B, fps, prefix=""): if 'body_center' not in avail_A or 'body_center' not in avail_B: return X rel_x = mouse_pair['A']['body_center']['x'] - mouse_pair['B']['body_center']['x'] rel_y = mouse_pair['A']['body_center']['y'] - mouse_pair['B']['body_center']['y'] rel_dist = np.sqrt(rel_x**2 + rel_y**2) A_vx = mouse_pair['A']['body_center']['x'].diff() A_vy = mouse_pair['A']['body_center']['y'].diff() B_vx = mouse_pair['B']['body_center']['x'].diff() B_vy = mouse_pair['B']['body_center']['y'].diff() A_lead = (A_vx * rel_x + A_vy * rel_y) / (np.sqrt(A_vx**2 + A_vy**2) * rel_dist + 1e-6) B_lead = (B_vx * (-rel_x) + B_vy * (-rel_y)) / (np.sqrt(B_vx**2 + B_vy**2) * rel_dist + 1e-6) for window in [30, 60]: ws = _scale(window, fps) col_A = f'{prefix}A_ld{window}' col_B = f'{prefix}B_ld{window}' if not _skip_pair_feature(col_A): X[col_A] = A_lead.rolling(ws, min_periods=max(1, ws // 6)).mean() if not _skip_pair_feature(col_B): X[col_B] = B_lead.rolling(ws, min_periods=max(1, ws // 6)).mean() approach = -rel_dist.diff() chase = approach * B_lead w = 30 ws = _scale(w, fps) chase_col = f'{prefix}chase_{w}' if not _skip_pair_feature(chase_col): X[chase_col] = chase.rolling(ws, min_periods=max(1, ws // 6)).mean() for window in [60, 120]: ws = _scale(window, fps) A_sp = np.sqrt(A_vx**2 + A_vy**2) B_sp = np.sqrt(B_vx**2 + B_vy**2) sp_cor_col = f'{prefix}sp_cor{window}' if _skip_pair_feature(sp_cor_col): continue X[sp_cor_col] = A_sp.rolling(ws, min_periods=max(1, ws // 6)).corr(B_sp) X = add_autocorr_features(X, approach, fps, prefix=prefix, base_name="appr_ac") X = add_autocorr_features(X, chase, fps, prefix=prefix, base_name="chase_ac") X = add_autocorr_features(X, A_lead, fps, prefix=prefix, base_name="Alead_ac") X = add_autocorr_features(X, B_lead, fps, prefix=prefix, base_name="Blead_ac") X = add_autocorr_features(X, rel_x.diff(), fps, prefix=prefix, base_name="relx_ac") X = add_autocorr_features(X, rel_y.diff(), fps, prefix=prefix, base_name="rely_ac") X = add_autocorr_features(X, rel_dist.diff(), fps, prefix=prefix, base_name="reld_ac") return X def transform_pair(mouse_pair, body_parts_tracked, fps): mouse_pair = _maybe_add_pair_body_center(mouse_pair) orig_avail_A = mouse_pair['A'].columns.get_level_values(0) orig_avail_B = mouse_pair['B'].columns.get_level_values(0) ego_pair, ego_rot, ego_cos, ego_sin = build_agent_egocentric_pair(mouse_pair) def _build_pair_feature_block(coord_source, avail_A, avail_B, prefix, extra_motion=None): pairwise_prod, pair_cols = _get_pairwise_product(tuple(body_parts_tracked)) part_arrays_A = {p: coord_source['A'][p].to_numpy(dtype=TRAIN_DTYPE, copy=False) for p in avail_A if ('A', p) in coord_source.columns} part_arrays_B = {p: coord_source['B'][p].to_numpy(dtype=TRAIN_DTYPE, copy=False) for p in avail_B if ('B', p) in coord_source.columns} X_block = pd.DataFrame({f"{prefix}12+{p1}+{p2}": np.square(part_arrays_A[p1] - part_arrays_B[p2]).sum(axis=1) for p1, p2 in pairwise_prod if p1 in part_arrays_A and p2 in part_arrays_B and not _skip_pair_feature(f"{prefix}12+{p1}+{p2}")}, index=coord_source.index) if X_block.empty: X_block = pd.DataFrame(index=coord_source.index) X_block = X_block.reindex(index=coord_source.index, columns=[f"{prefix}{col}" for col in pair_cols if not _skip_pair_feature(f"{prefix}{col}")], copy=False) if ('A', 'ear_left') in coord_source.columns and ('B', 'ear_left') in coord_source.columns: lag = _scale(10, fps) shA = coord_source['A']['ear_left'].shift(lag) shB = coord_source['B']['ear_left'].shift(lag) speeds = pd.DataFrame({ f'{prefix}sp_A': np.square(coord_source['A']['ear_left'] - shA).sum(axis=1, skipna=False), f'{prefix}sp_AB': np.square(coord_source['A']['ear_left'] - shB).sum(axis=1, skipna=False), f'{prefix}sp_B': np.square(coord_source['B']['ear_left'] - shB).sum(axis=1, skipna=False), }) X_block = pd.concat([X_block, speeds], axis=1) nt_key = f"{prefix}nose+tail_base" ear_key = f"{prefix}ear_left+ear_right" if nt_key in X_block.columns and ear_key in X_block.columns: elong_col = f'{prefix}elong' if not _skip_pair_feature(elong_col): X_block[elong_col] = X_block[nt_key] / (X_block[ear_key] + 1e-6) if all(p in avail_A for p in ['nose', 'tail_base']) and all(p in avail_B for p in ['nose', 'tail_base']): dir_A = coord_source['A']['nose'] - coord_source['A']['tail_base'] dir_B = coord_source['B']['nose'] - coord_source['B']['tail_base'] rel_ori_col = f'{prefix}rel_ori' if not _skip_pair_feature(rel_ori_col): X_block[rel_ori_col] = (dir_A['x'] * dir_B['x'] + dir_A['y'] * dir_B['y']) / (np.sqrt(dir_A['x']**2 + dir_A['y']**2) * np.sqrt(dir_B['x']**2 + dir_B['y']**2) + 1e-6) if all(p in avail_A for p in ['nose']) and all(p in avail_B for p in ['nose']): cur = np.square(coord_source['A']['nose'] - coord_source['B']['nose']).sum(axis=1, skipna=False) lag = _scale(10, fps) shA_n = coord_source['A']['nose'].shift(lag) shB_n = coord_source['B']['nose'].shift(lag) past = np.square(shA_n - shB_n).sum(axis=1, skipna=False) appr_col = f'{prefix}appr' if not _skip_pair_feature(appr_col): X_block[appr_col] = cur - past if extra_motion: for k, v in extra_motion.items(): col_name = f'{prefix}{k}' if _skip_pair_feature(col_name): continue X_block[col_name] = v if 'body_center' in avail_A and 'body_center' in avail_B: cd = np.sqrt((coord_source['A']['body_center']['x'] - coord_source['B']['body_center']['x'])**2 + (coord_source['A']['body_center']['y'] - coord_source['B']['body_center']['y'])**2) for col_name, series in [(f'{prefix}v_cls', (cd < 5.0).astype(float)), (f'{prefix}cls', ((cd >= 5.0) & (cd < 15.0)).astype(float)), (f'{prefix}med', ((cd >= 15.0) & (cd < 30.0)).astype(float)), (f'{prefix}far', (cd >= 30.0).astype(float))]: if _skip_pair_feature(col_name): continue X_block[col_name] = series cd_full = np.square(coord_source['A']['body_center'] - coord_source['B']['body_center']).sum(axis=1, skipna=False) for w in [5, 15, 30, 60]: ws = _scale(w, fps) roll = dict(min_periods=1, center=True) for col_name, series in [(f'{prefix}d_m{w}', cd_full.rolling(ws, **roll).mean()), (f'{prefix}d_s{w}', cd_full.rolling(ws, **roll).std()), (f'{prefix}d_mn{w}', cd_full.rolling(ws, **roll).min()), (f'{prefix}d_mx{w}', cd_full.rolling(ws, **roll).max())]: if _skip_pair_feature(col_name): continue X_block[col_name] = series d_var = cd_full.rolling(ws, **roll).var() int_col = f'{prefix}int{w}' if not _skip_pair_feature(int_col): X_block[int_col] = 1 / (1 + d_var) Axd = coord_source['A']['body_center']['x'].diff() Ayd = coord_source['A']['body_center']['y'].diff() Bxd = coord_source['B']['body_center']['x'].diff() Byd = coord_source['B']['body_center']['y'].diff() coord = Axd * Bxd + Ayd * Byd co_m_col = f'{prefix}co_m{w}' co_s_col = f'{prefix}co_s{w}' if not _skip_pair_feature(co_m_col): X_block[co_m_col] = coord.rolling(ws, **roll).mean() if not _skip_pair_feature(co_s_col): X_block[co_s_col] = coord.rolling(ws, **roll).std() if 'nose' in avail_A and 'nose' in avail_B: nn = np.sqrt((coord_source['A']['nose']['x'] - coord_source['B']['nose']['x'])**2 + (coord_source['A']['nose']['y'] - coord_source['B']['nose']['y'])**2) for lag in [10, 20, 40]: l = _scale(lag, fps) nn_lg_col = f'{prefix}nn_lg{lag}' nn_ch_col = f'{prefix}nn_ch{lag}' if not _skip_pair_feature(nn_lg_col): X_block[nn_lg_col] = nn.shift(l) if not _skip_pair_feature(nn_ch_col): X_block[nn_ch_col] = nn - nn.shift(l) is_cl = (nn < 10.0).astype(float) cl_ps_col = f'{prefix}cl_ps{lag}' if not _skip_pair_feature(cl_ps_col): X_block[cl_ps_col] = is_cl.rolling(l, min_periods=1).mean() if 'body_center' in avail_A and 'body_center' in avail_B: Avx = coord_source['A']['body_center']['x'].diff() Avy = coord_source['A']['body_center']['y'].diff() Bvx = coord_source['B']['body_center']['x'].diff() Bvy = coord_source['B']['body_center']['y'].diff() val = (Avx * Bvx + Avy * Bvy) / (np.sqrt(Avx**2 + Avy**2) * np.sqrt(Bvx**2 + Bvy**2) + 1e-6) for off in [-30, -20, -10, 0, 10, 20, 30]: o = _scale_signed(off, fps) va_col = f'{prefix}va_{off}' if _skip_pair_feature(va_col): continue X_block[va_col] = val.shift(-o) w = _scale(30, fps) cd_full = np.square(coord_source['A']['body_center'] - coord_source['B']['body_center']).sum(axis=1, skipna=False) int_con_col = f'{prefix}int_con' if not _skip_pair_feature(int_con_col): X_block[int_con_col] = cd_full.rolling(w, min_periods=1, center=True).std() / (cd_full.rolling(w, min_periods=1, center=True).mean() + 1e-6) X_block = add_interaction_features(X_block, coord_source, avail_A, avail_B, fps, prefix=prefix) return X_block feature_blocks = [] feature_blocks.append(_build_pair_feature_block(mouse_pair, orig_avail_A, orig_avail_B, prefix="cam_")) if ego_pair is not None: try: ego_avail_A = ego_pair['A'].columns.get_level_values(0); ego_avail_B = ego_pair['B'].columns.get_level_values(0) except KeyError: ego_avail_A = orig_avail_A; ego_avail_B = orig_avail_B extra_motion = {} if 'body_center' in orig_avail_A: Avx = mouse_pair['A']['body_center']['x'].diff() Avy = mouse_pair['A']['body_center']['y'].diff() ego_Ax = Avx * ego_cos - Avy * ego_sin ego_Ay = Avx * ego_sin + Avy * ego_cos extra_motion['ego_A_fwd'] = ego_Ay * float(fps) extra_motion['ego_A_lat'] = ego_Ax * float(fps) for w in [10, 30]: ws = _scale(w, fps) roll = dict(min_periods=1, center=True) extra_motion[f'ego_Af_m{w}'] = ego_Ay.rolling(ws, **roll).mean() * float(fps) extra_motion[f'ego_Al_m{w}'] = ego_Ax.rolling(ws, **roll).mean() * float(fps) if 'body_center' in orig_avail_B: Bvx = mouse_pair['B']['body_center']['x'].diff() Bvy = mouse_pair['B']['body_center']['y'].diff() ego_Bx = Bvx * ego_cos - Bvy * ego_sin ego_By = Bvx * ego_sin + Bvy * ego_cos extra_motion['ego_B_fwd'] = ego_By * float(fps) extra_motion['ego_B_lat'] = ego_Bx * float(fps) if 'B' in ego_pair.columns.get_level_values(0) and 'body_center' in ego_pair['B'].columns.get_level_values(0): bx = ego_pair['B']['body_center']['x'] by = ego_pair['B']['body_center']['y'] extra_motion['ego_rel_x'] = bx extra_motion['ego_rel_y'] = by extra_motion['ego_rel_d'] = np.sqrt(bx**2 + by**2) extra_motion['ego_rel_ang'] = np.arctan2(by, bx) for w in [10, 30]: ws = _scale(w, fps) roll = dict(min_periods=1, center=True) extra_motion[f'ego_rel_xm{w}'] = bx.rolling(ws, **roll).mean() extra_motion[f'ego_rel_ym{w}'] = by.rolling(ws, **roll).mean() feature_blocks.append(_build_pair_feature_block(ego_pair, ego_avail_A, ego_avail_B, prefix="ego_", extra_motion=extra_motion)) X = pd.concat(feature_blocks, axis=1) drop_cols = [c for c in SKIP_PAIR_FEATURES if c in X.columns] if drop_cols: X = X.drop(columns=drop_cols) X = X.replace([np.inf, -np.inf], np.nan) return X.astype(TRAIN_DTYPE, copy=False) # ========================================================= # EMA / Probability Utilities # ========================================================= def _prob_to_logit(prob_array, eps=1e-6): prob_array = np.clip(prob_array, eps, 1.0 - eps) return np.log(prob_array / (1.0 - prob_array)) def _logit_to_prob(logit_array): return 1.0 / (1.0 + np.exp(-logit_array)) def _extract_ema_seconds(entry): if isinstance(entry, dict): if 'ema_seconds' in entry: return entry.get('ema_seconds') return None def _build_ema_seconds_map(thresholds): ema_map = {} for action, entry in thresholds.items(): seconds = _extract_ema_seconds(entry) if seconds is None: ema_map[action] = 0.0; continue try: seconds_f = float(seconds) except: continue ema_map[action] = seconds_f return ema_map def _apply_logit_ema_series(prob_series: pd.Series, meta_subset: pd.DataFrame, base_seconds: float) -> pd.Series: if meta_subset is None or meta_subset.empty or base_seconds is None or base_seconds <= 0: return prob_series.astype(TRAIN_DTYPE, copy=False) df = meta_subset.reset_index(drop=True).copy() prob_filled = prob_series.fillna(0.0) df['prob'] = prob_filled.to_numpy(dtype=np.float64, copy=False) df['logit'] = _prob_to_logit(df['prob']) smoothed = df['prob'].to_numpy(dtype=np.float64, copy=True) def _bidirectional_ema(logits: np.ndarray, span: int) -> np.ndarray: forward = pd.Series(logits).ewm(span=span, adjust=False).mean().to_numpy(dtype=np.float64, copy=False) backward = pd.Series(logits[::-1]).ewm(span=span, adjust=False).mean().to_numpy(dtype=np.float64, copy=False)[::-1] return (forward + backward) * 0.5 group_cols = [col for col in ['video_id', 'agent_id', 'target_id'] if col in df.columns] if not group_cols: group_cols = ['video_id'] if 'video_id' in df.columns else [] if not group_cols: fps = 30.0 span = max(1, int(round(fps * base_seconds))) logits = df['logit'].to_numpy(dtype=np.float64, copy=False) smoothed_logits = _bidirectional_ema(logits, span) return pd.Series(_logit_to_prob(smoothed_logits), index=prob_series.index, dtype=TRAIN_DTYPE) for _, idx in df.groupby(group_cols).groups.items(): group = df.loc[list(idx)].sort_values('video_frame') logits = group['logit'].to_numpy(dtype=np.float64, copy=False) if 'frames_per_second' in group.columns and pd.notnull(group['frames_per_second']).any(): fps = float(group['frames_per_second'].iloc[0]) if not np.isfinite(fps) or fps <= 0: fps = 30.0 else: fps = 30.0 span = max(1, int(round(fps * base_seconds))) smoothed_logits = _bidirectional_ema(logits, span) smoothed[group.index] = _logit_to_prob(smoothed_logits) return pd.Series(smoothed, index=prob_series.index, dtype=TRAIN_DTYPE) def _apply_logit_ema_dataframe(prob_df, meta_df, base_seconds, per_column_seconds=None) -> pd.DataFrame: if prob_df.empty or ((base_seconds is None or base_seconds <= 0) and not per_column_seconds): return prob_df meta_reset = meta_df.reset_index(drop=True) prob_reset = prob_df.reset_index(drop=True).copy() for col in prob_reset.columns: series = prob_reset[col] if series.isna().all(): continue seconds = base_seconds if per_column_seconds: seconds = per_column_seconds.get(col, base_seconds) if seconds is None or seconds <= 0: continue prob_reset[col] = _apply_logit_ema_series(series, meta_reset, base_seconds=seconds) prob_reset.index = prob_df.index return prob_reset.astype(TRAIN_DTYPE, copy=False) # ========================================================= # Data Generator # ========================================================= def generate_mouse_data(dataset, traintest, traintest_directory=None, generate_single=True, generate_pair=True): if traintest_directory is None: traintest_directory = f"{CFG.base_path}/{traintest}_tracking" for _, row in dataset.iterrows(): lab_id = row.lab_id if lab_id.startswith('MABe22') or type(row.behaviors_labeled) != str: continue video_id = row.video_id path = f"{traintest_directory}/{lab_id}/{video_id}.parquet" try: vid = load_tracking_parquet_fast(path) except FileNotFoundError: continue fps = row.get('frames_per_second', np.nan) if not np.isfinite(fps) or fps <= 0: fps = 30.0 pix_per_cm = row.get('pix_per_cm_approx', 1.0) if not np.isfinite(pix_per_cm) or pix_per_cm <= 0: pix_per_cm = 1.0 pvid = vid / pix_per_cm allow_interp = lab_id != 'GroovyShrew' pvid = apply_missing_value_heuristics(pvid, allow_interpolate=allow_interp) vid_behaviors = json.loads(row.behaviors_labeled) vid_behaviors = sorted(list({b.replace("'", "") for b in vid_behaviors})) vid_behaviors = [b.split(',') for b in vid_behaviors] vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action']) if generate_single: vid_behaviors_subset = vid_behaviors.query("target == 'self'") for mouse_id_str in np.unique(vid_behaviors_subset.agent): try: mouse_id = int(mouse_id_str[-1]) vid_agent_actions = np.unique(vid_behaviors_subset.query("agent == @mouse_id_str").action) single_mouse = pvid.xs(mouse_id, level='mouse_id', axis=1, drop_level=True) single_mouse_meta = pd.DataFrame({ 'video_id': video_id, 'agent_id': mouse_id_str, 'target_id': 'self', 'video_frame': single_mouse.index, 'frames_per_second': fps, 'is_flipped': False, }) yield 'single', single_mouse, single_mouse_meta, vid_agent_actions except KeyError: pass if generate_pair: vid_behaviors_subset = vid_behaviors.query("target != 'self'") if len(vid_behaviors_subset) > 0: mice_ids = np.unique(pvid.columns.get_level_values('mouse_id')) for agent, target in itertools.permutations(mice_ids, 2): agent_str, target_str = f"mouse{agent}", f"mouse{target}" vid_agent_actions = np.unique(vid_behaviors_subset.query("(agent == @agent_str) & (target == @target_str)").action) mouse_pair = pd.concat([pvid.xs(agent, level='mouse_id', axis=1, drop_level=True), pvid.xs(target, level='mouse_id', axis=1, drop_level=True)], axis=1, keys=['A', 'B']) mouse_pair_meta = pd.DataFrame({ 'video_id': video_id, 'agent_id': agent_str, 'target_id': target_str, 'video_frame': mouse_pair.index, 'frames_per_second': fps, 'is_flipped': False, }) yield 'pair', mouse_pair, mouse_pair_meta, vid_agent_actions # ========================================================= # Ensemble & Prediction Logic # ========================================================= def _resolve_threshold_tuple(entry, default_thresh=0.27, default_sigma=1.0): # Returns (threshold, sigma) t = default_thresh s = default_sigma if isinstance(entry, dict): t = entry.get('threshold', default_thresh) s = entry.get('sigma', default_sigma) if s is None: s = entry.get('std', default_sigma) elif isinstance(entry, (float, int)): t = entry if not np.isfinite(t): t = default_thresh if not np.isfinite(s) or s <= 0: s = default_sigma return t, s def predict_multiclass_ensemble(pred_a, pred_b, meta, thr_a, thr_b): """ Ensemble decision logic: 1. P_avg = (P_a + P_b) / 2 2. T_avg = (T_a + T_b) / 2 3. Z_avg = ((P_a - T_a)/S_a + (P_b - T_b)/S_b) / 2 4. Candidate if P_avg >= T_avg 5. Argmax(Z_avg) among candidates """ actions = pred_a.columns # 1. Average Probs p_avg = (pred_a.values + pred_b.values) / 2.0 # Prepare Arrays t_a_arr = np.zeros(len(actions), dtype=np.float32) s_a_arr = np.zeros(len(actions), dtype=np.float32) t_b_arr = np.zeros(len(actions), dtype=np.float32) s_b_arr = np.zeros(len(actions), dtype=np.float32) for i, act in enumerate(actions): ta, sa = _resolve_threshold_tuple(thr_a.get(act)) tb, sb = _resolve_threshold_tuple(thr_b.get(act)) t_a_arr[i] = ta s_a_arr[i] = sa t_b_arr[i] = tb s_b_arr[i] = sb # 2. Average Threshold t_avg_arr = (t_a_arr + t_b_arr) / 2.0 # 3. Calculate Z-scores # Avoid division by zero s_a_safe = np.where(s_a_arr > 1e-9, s_a_arr, 1.0) s_b_safe = np.where(s_b_arr > 1e-9, s_b_arr, 1.0) z_a = (pred_a.values - t_a_arr) / s_a_safe z_b = (pred_b.values - t_b_arr) / s_b_safe z_avg = (z_a + z_b) / 2.0 # 4. Candidates exceeds = p_avg >= t_avg_arr # 5. Argmax Z-score margin_score = np.where(exceeds, z_avg, -np.inf) has_candidate = exceeds.any(axis=1) ama = np.where(has_candidate, np.argmax(margin_score, axis=1), -1) ama = pd.Series(ama, index=meta.video_frame) # Segment Construction changes_mask = (ama != ama.shift(1)).values ama_changes = ama[changes_mask] meta_changes = meta[changes_mask] mask = ama_changes.values >= 0 if len(mask) > 0: mask[-1] = False submission_part = pd.DataFrame({ 'video_id': meta_changes['video_id'][mask].values, 'agent_id': meta_changes['agent_id'][mask].values, 'target_id': meta_changes['target_id'][mask].values, 'action': actions[ama_changes[mask].values], 'start_frame': ama_changes.index[mask], 'stop_frame': ama_changes.index[1:][mask[:-1]] }) # Cleanup Stop Frames stop_video_id = meta_changes['video_id'][1:][mask[:-1]].values stop_agent_id = meta_changes['agent_id'][1:][mask[:-1]].values stop_target_id = meta_changes['target_id'][1:][mask[:-1]].values # Handle end of video boundaries for i in range(len(submission_part)): video_id = submission_part.video_id.iloc[i] agent_id = submission_part.agent_id.iloc[i] target_id = submission_part.target_id.iloc[i] if (i >= len(stop_video_id) or stop_video_id[i] != video_id or stop_agent_id[i] != agent_id or stop_target_id[i] != target_id): new_stop_frame = meta.query("(video_id == @video_id)").video_frame.max() + 1 submission_part.iat[i, submission_part.columns.get_loc('stop_frame')] = new_stop_frame if not submission_part.empty: submission_part = _merge_close_segments(submission_part, max_gap=0) return submission_part def _merge_close_segments(sub_df: pd.DataFrame, max_gap: int = 0) -> pd.DataFrame: if sub_df.empty: return sub_df cols = ['video_id', 'agent_id', 'target_id', 'action'] sub_sorted = sub_df.sort_values(cols + ['start_frame']).reset_index(drop=True) merged_rows = [] current = sub_sorted.iloc[0].copy() for _, row in sub_sorted.iloc[1:].iterrows(): same_track = all(current[col] == row[col] for col in cols) gap = row['start_frame'] - current['stop_frame'] if same_track and gap <= max_gap: current['stop_frame'] = max(current['stop_frame'], row['stop_frame']) else: merged_rows.append(current.copy()) current = row.copy() merged_rows.append(current) return pd.DataFrame(merged_rows, columns=sub_sorted.columns) def _apply_segment_postprocess_seconds( sub_df: pd.DataFrame, fps_lookup: dict, *, max_gap_seconds: float = 0.0, min_duration_seconds: float = 0.0, default_fps: float = DEFAULT_FPS, block_df: pd.DataFrame | None = None, ) -> pd.DataFrame: if sub_df.empty: return sub_df cols = ['video_id', 'agent_id', 'target_id', 'action'] sub_sorted = sub_df.sort_values(cols + ['start_frame']).reset_index(drop=True) group_cols = ['video_id', 'agent_id', 'target_id'] block_dict = {} if block_df is not None and not block_df.empty: block_use = block_df.copy() if set(group_cols + ['action', 'start_frame', 'stop_frame']).issubset(block_use.columns): block_use = block_use.sort_values(group_cols + ['start_frame']).reset_index(drop=True) for k, g in block_use.groupby(group_cols): block_dict[k] = g merged_rows = [] current = sub_sorted.iloc[0].copy() for _, row in sub_sorted.iloc[1:].iterrows(): same_track = all(current[col] == row[col] for col in cols) gap = int(row['start_frame']) - int(current['stop_frame']) fps_curr = float(fps_lookup.get(current['video_id'], default_fps)) if not np.isfinite(fps_curr) or fps_curr <= 0: fps_curr = default_fps max_gap_frames = max(0, int(round(fps_curr * float(max_gap_seconds)))) if max_gap_seconds > 0 else 0 can_merge = False if same_track and max_gap_frames > 0 and gap <= max_gap_frames: can_merge = True if block_dict: key = (current['video_id'], current['agent_id'], current['target_id']) blockers = block_dict.get(key, None) if blockers is not None and not blockers.empty: act = current['action'] b = blockers[blockers['action'] != act] if not b.empty: s, e = int(current['stop_frame']), int(row['start_frame']) if ((b['start_frame'].astype(int) < e) & (b['stop_frame'].astype(int) > s)).any(): can_merge = False if can_merge: current['stop_frame'] = max(int(current['stop_frame']), int(row['stop_frame'])) else: merged_rows.append(current.copy()) current = row.copy() merged_rows.append(current) merged = pd.DataFrame(merged_rows, columns=sub_sorted.columns) if min_duration_seconds > 0 and not merged.empty: fps_series = merged['video_id'].map(lambda v: fps_lookup.get(v, default_fps)).fillna(default_fps) min_frames = np.rint(fps_series.astype(float) * float(min_duration_seconds)).astype(int) merged = merged[ (merged['stop_frame'].astype(int) - merged['start_frame'].astype(int)) >= min_frames ].reset_index(drop=True) return merged def apply_postprocess_params_seconds(submission, video_to_lab, video_to_fps, params, default_fps=DEFAULT_FPS): if submission.empty or not params: return submission sub = submission.copy() sub['lab_id'] = sub['video_id'].map(video_to_lab) keep_mask = np.ones(len(sub), dtype=bool) tuned_parts = [] for (lab_id, action), cfg in params.items(): if lab_id is None: continue mask = (sub['lab_id'] == lab_id) & (sub['action'] == action) if not mask.any(): continue base = sub.loc[mask, ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']] block_df = sub.loc[(sub['lab_id'] == lab_id), ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']] tuned = _apply_segment_postprocess_seconds( base, fps_lookup=video_to_fps, max_gap_seconds=float(cfg.get("max_gap_seconds", 0)), min_duration_seconds=float(cfg.get("min_duration_seconds", 0)), default_fps=default_fps, block_df=block_df ) tuned_parts.append(tuned) keep_mask &= ~mask.to_numpy() return pd.concat([sub.loc[keep_mask, ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']]] + tuned_parts, axis=0, ignore_index=True) def _drop_overlapping_across_actions(sub_df: pd.DataFrame) -> pd.DataFrame: if sub_df.empty: return sub_df kept_parts = [] group_cols = ['video_id', 'agent_id', 'target_id'] for _, group in sub_df.groupby(group_cols): group_sorted = group.sort_values(['start_frame', 'stop_frame']).reset_index(drop=True) keep_mask = np.ones(len(group_sorted), dtype=bool) last_stop = -1 for i, row in group_sorted.iterrows(): if int(row['start_frame']) < last_stop: keep_mask[i] = False else: last_stop = max(last_stop, int(row['stop_frame'])) kept_parts.append(group_sorted[keep_mask]) return pd.concat(kept_parts, axis=0, ignore_index=True) def robustify(submission, dataset, traintest, traintest_directory=None): if traintest_directory is None: traintest_directory = f"{CFG.base_path}/{traintest}_tracking" submission = submission[submission.start_frame < submission.stop_frame].copy() group_list = [] for _, group in submission.groupby(['video_id', 'agent_id', 'target_id']): group = group.sort_values('start_frame') mask = np.ones(len(group), dtype=bool) last_stop_frame = -1 for i, (_, row) in enumerate(group.iterrows()): if row['start_frame'] < last_stop_frame: mask[i] = False else: last_stop_frame = row['stop_frame'] group_list.append(group[mask]) if group_list: submission = pd.concat(group_list) else: submission = pd.DataFrame(columns=submission.columns) s_list = [] for _, row in dataset.iterrows(): if row.lab_id.startswith('MABe22'): continue if (submission.video_id == row.video_id).any(): continue if type(row.behaviors_labeled) != str: continue try: path = f"{traintest_directory}/{row.lab_id}/{row.video_id}.parquet" vid = pd.read_parquet(path) except: continue vid_behaviors = json.loads(row.behaviors_labeled) vid_behaviors = [b.replace("'", "").split(',') for b in vid_behaviors] vid_behaviors = pd.DataFrame(vid_behaviors, columns=['agent', 'target', 'action']) start_frame, stop_frame = vid.video_frame.min(), vid.video_frame.max() + 1 for (agent, target), actions in vid_behaviors.groupby(['agent', 'target']): batch_length = int(np.ceil((stop_frame - start_frame) / len(actions))) for i, (_, action_row) in enumerate(actions.iterrows()): s = start_frame + i * batch_length e = min(s + batch_length, stop_frame) s_list.append((row.video_id, agent, target, action_row['action'], s, e)) if s_list: submission = pd.concat([submission, pd.DataFrame(s_list, columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame'])]) return submission.reset_index(drop=True) def _postprocess_adaptablesnail_25fps(sub_df: pd.DataFrame, scale: float = 1.2) -> pd.DataFrame: if sub_df.empty: return sub_df filtered = sub_df[~sub_df["action"].isin(["rear", "approach", "submit"])].copy() if filtered.empty: return filtered for col in ("start_frame", "stop_frame"): filtered[col] = np.rint(filtered[col].astype(float) * scale).astype(int) filtered["stop_frame"] = np.maximum(filtered["stop_frame"], filtered["start_frame"] + 1) return filtered def _apply_adaptablesnail_25fps_after_postprocess(submission, video_to_lab, video_to_fps, scale=1.2): if submission.empty: return submission def _is_25fps(vid): try: fps = float(video_to_fps.get(vid, np.nan)) except: return False return np.isfinite(fps) and (24.9 < fps < 25.1) adapt_mask = submission['video_id'].map(lambda v: video_to_lab.get(v) == 'AdaptableSnail' and _is_25fps(v)).fillna(False) if not adapt_mask.any(): return submission adapt_part = _postprocess_adaptablesnail_25fps(submission.loc[adapt_mask].copy(), scale=scale) return pd.concat([submission.loc[~adapt_mask].copy(), adapt_part], axis=0, ignore_index=True) # ========================================================= # Submit Logic (Ensemble) # ========================================================= def submit_ensemble(body_parts_tracked_str, switch_tr, section, thresholds_a, thresholds_b, lab_group): body_parts_tracked = json.loads(body_parts_tracked_str) if len(body_parts_tracked) > 5: body_parts_tracked = [b for b in body_parts_tracked if b not in drop_body_parts] test_pl = _preload_metadata_polars(CFG.test_path, is_train=False) test = test_pl.to_pandas() test_subset = test[(test.body_parts_tracked == body_parts_tracked_str) & (test.lab_group == lab_group)] fps_lookup = test_subset[['video_id', 'frames_per_second']].drop_duplicates('video_id').set_index('video_id')['frames_per_second'].to_dict() generator = generate_mouse_data( test_subset, 'test', traintest_directory=CFG.test_tracking_path, generate_single=(switch_tr == 'single'), generate_pair=(switch_tr == 'pair') ) accumulated_preds_a = [] accumulated_preds_b = [] accumulated_metas = [] ema_map_a = _build_ema_seconds_map(thresholds_a) ema_map_b = _build_ema_seconds_map(thresholds_b) for switch_te, data_te, meta_te, actions_te in generator: try: fps_i = _fps_from_meta(meta_te, fps_lookup, default_fps=30.0) if switch_te == 'single': X_te = transform_single(data_te, body_parts_tracked, fps_i) else: X_te = transform_pair(data_te, body_parts_tracked, fps_i) del data_te; gc.collect() if X_te.shape[1] == 0: continue pred_df_a = pd.DataFrame(index=meta_te.video_frame) pred_df_b = pd.DataFrame(index=meta_te.video_frame) for action in actions_te: # --- Model A (XGB) --- root_a = Path(CFG.model_a_path) / lab_group / str(section) / action model_files_a = sorted(list(glob.glob(str(root_a / "fold_*_model.pkl")))) full_model_a = root_a / "full_model.pkl" preds_fold_a = [] if CFG.prefer_full_model and full_model_a.exists(): try: m = joblib.load(full_model_a) preds_fold_a.append(m.predict_proba(X_te)[:, 1]) except: pass if not preds_fold_a and model_files_a: for mp in model_files_a: try: m = joblib.load(mp) preds_fold_a.append(m.predict_proba(X_te)[:, 1]) except: pass pred_df_a[action] = np.mean(preds_fold_a, axis=0) if preds_fold_a else 0.0 # --- Model B (CatBoost) --- root_b = Path(CFG.model_b_path) / lab_group / str(section) / action model_files_b = sorted(list(glob.glob(str(root_b / "fold_*_model.pkl")))) full_model_b = root_b / "full_model.pkl" preds_fold_b = [] if CFG.prefer_full_model and full_model_b.exists(): try: m = joblib.load(full_model_b) preds_fold_b.append(m.predict_proba(X_te)[:, 1]) except: pass if not preds_fold_b and model_files_b: for mp in model_files_b: try: m = joblib.load(mp) preds_fold_b.append(m.predict_proba(X_te)[:, 1]) except: pass pred_df_b[action] = np.mean(preds_fold_b, axis=0) if preds_fold_b else 0.0 del X_te; gc.collect() if CFG.smooth_logits_with_ema: pred_df_a = _apply_logit_ema_dataframe(pred_df_a, meta_te, CFG.default_logit_ema_seconds, ema_map_a) pred_df_b = _apply_logit_ema_dataframe(pred_df_b, meta_te, CFG.default_logit_ema_seconds, ema_map_b) accumulated_preds_a.append(pred_df_a) accumulated_preds_b.append(pred_df_b) accumulated_metas.append(meta_te) except KeyError: del data_te; gc.collect() submission_list = [] if accumulated_preds_a: full_pred_a = pd.concat(accumulated_preds_a, axis=0, ignore_index=True) full_pred_b = pd.concat(accumulated_preds_b, axis=0, ignore_index=True) full_meta = pd.concat(accumulated_metas, axis=0, ignore_index=True) mask_fixed = np.zeros(len(full_meta), dtype=bool) if lab_group == 'AdaptableSnail' and 'frames_per_second' in full_meta.columns: fps_vals = full_meta['frames_per_second'].values mask_fixed = (fps_vals > 24.9) & (fps_vals < 25.1) mask_rank = ~mask_fixed if np.any(mask_fixed): sub_fixed = predict_multiclass_ensemble( full_pred_a.loc[mask_fixed], full_pred_b.loc[mask_fixed], full_meta.loc[mask_fixed], thresholds_a, thresholds_b ) if lab_group == "AdaptableSnail" and not sub_fixed.empty: sub_fixed = _postprocess_adaptablesnail_25fps(sub_fixed) if not sub_fixed.empty: submission_list.append(sub_fixed) if np.any(mask_rank): sub_rank = predict_multiclass_ensemble( full_pred_a.loc[mask_rank], full_pred_b.loc[mask_rank], full_meta.loc[mask_rank], thresholds_a, thresholds_b ) if not sub_rank.empty: submission_list.append(sub_rank) return submission_list # ========================================================= # Main Execution # ========================================================= if __name__ == "__main__": print(f"--- Starting ENSEMBLE run ---") test_pl = _preload_metadata_polars(CFG.test_path, is_train=False) if 'frames_per_second' in test_pl.columns and 'lab_id' in test_pl.columns: fps_col = pl.col('frames_per_second') test_round_mask = (~(pl.col('lab_id') == 'AdaptableSnail')) & fps_col.is_not_null() test_pl = test_pl.with_columns( pl.when(test_round_mask).then(fps_col.round(0).cast(pl.Float64)).otherwise(fps_col).alias('frames_per_second') ) test = test_pl.to_pandas() default_fps = DEFAULT_FPS test_video_fps = _build_video_fps_lookup(test, default_fps=default_fps) lab_groups = sorted(test.lab_group.dropna().unique()) path_thresh_a = Path(CFG.model_a_path) / "thresholds.pkl" path_thresh_b = Path(CFG.model_b_path) / "thresholds.pkl" path_pp_a = Path(CFG.model_a_path) / "postprocess_params.pkl" path_pp_b = Path(CFG.model_b_path) / "postprocess_params.pkl" thresholds_A = joblib.load(path_thresh_a) thresholds_B = joblib.load(path_thresh_b) try: pp_A = joblib.load(path_pp_a) except: pp_A = {} try: pp_B = joblib.load(path_pp_b) except: pp_B = {} merged_pp = {} all_pp_keys = set(pp_A.keys()) | set(pp_B.keys()) for k in all_pp_keys: va = pp_A.get(k, {}) vb = pp_B.get(k, {}) merged_pp[k] = { 'max_gap_seconds': min(float(va.get('max_gap_seconds', 0)), float(vb.get('max_gap_seconds', 0))), 'min_duration_seconds': min(float(va.get('min_duration_seconds', 0)), float(vb.get('min_duration_seconds', 0))), } submission_list = [] for lab_idx, lab_group in enumerate(lab_groups, start=1): print(f"\n=== Lab {lab_idx}/{len(lab_groups)}: {lab_group} ===") lab_test = test[test.lab_group == lab_group] if lab_test.empty: continue th_a_lab = thresholds_A.get(lab_group, {"single": {}, "pair": {}}) th_b_lab = thresholds_B.get(lab_group, {"single": {}, "pair": {}}) body_parts_tracked_list = sorted(lab_test.body_parts_tracked.unique()) for section_idx, body_parts_tracked_str in enumerate(body_parts_tracked_list, start=1): print(f" Section {section_idx}: {body_parts_tracked_str}") section_key = str(section_idx) th_a_sec = th_a_lab.get("single", {}).get(section_key, {}) th_b_sec = th_b_lab.get("single", {}).get(section_key, {}) submission_list.extend(submit_ensemble(body_parts_tracked_str, 'single', section_idx, th_a_sec, th_b_sec, lab_group)) th_a_sec = th_a_lab.get("pair", {}).get(section_key, {}) th_b_sec = th_b_lab.get("pair", {}).get(section_key, {}) submission_list.extend(submit_ensemble(body_parts_tracked_str, 'pair', section_idx, th_a_sec, th_b_sec, lab_group)) print("\nProcessing submission...") if submission_list: submission = pd.concat(submission_list) else: submission = pd.DataFrame(columns=['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']) submission_robust = robustify(submission, test, 'test', traintest_directory=CFG.test_tracking_path) video_to_lab = test.set_index('video_id')['lab_id'].to_dict() # if merged_pp: # submission_robust = apply_postprocess_params_seconds( # submission_robust, video_to_lab=video_to_lab, video_to_fps=test_video_fps, # params=merged_pp, default_fps=default_fps # ) # submission_robust = _apply_adaptablesnail_25fps_after_postprocess( # submission_robust, video_to_lab=video_to_lab, video_to_fps=test_video_fps, scale=1.2 # ) submission_robust = _drop_overlapping_across_actions(submission_robust) submission_robust.index.name = 'row_id' submission_robust.to_csv('submission.csv') print("submission.csv created successfully.") print(submission_robust.head())