{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 🐁 MABe - Fast Iterative LightGBM Baseline 🚀\n\n## Introduction\n\nWelcome! This notebook implements a fast and memory-efficient baseline for the **MABe - Social Action Recognition in Mice** competition.\n\nInstead of attempting to load and process the entire large tracking dataset at once (which often leads to memory issues), this notebook adopts an **iterative, video-by-video approach**.\n\n* **Reference/Inspiration**: This notebook's structure is heavily inspired by the efficient baseline created by [VectorQL\n in this notebook](https://www.kaggle.com/code/lordpatil/fast-eda-and-fast-baseline).\n\n## Core Strategy\n\n1.  **Iterative Processing**: Load tracking data for one video at a time.\n2.  **Lightweight Feature Engineering**: Calculate essential kinematic and geometric features for the current video.\n3.  **Negative Sampling**: Reduce the training data size by sampling non-event frames, making training much faster.\n4.  **Binary LGBM Models**: Train a separate, simple LightGBM classifier for each target behavior.\n5.  **Memory Management**: Explicitly delete processed dataframes within the loop to keep memory usage low.\n\n## Workflow Outline\n\n1.  **Setup & Config**: Import libraries and define key parameters.\n2.  **Load Initial Data**: Load metadata (`train.csv`, `test.csv`) and all annotations.\n3.  **Feature Engineering Function**: Define a function to process tracking data from a single video.\n4.  **Iterative Training Loop**: Loop through training videos, generate features, sample data, and train models per behavior.\n5.  **Iterative Prediction Loop**: Loop through test videos, generate features, predict probabilities, apply post-processing, and collect events.\n6.  **Submission**: Format predictions and save `submission.csv`.\n\nLet's build a solid baseline! If you find this useful, please consider an upvote! 👍","metadata":{}},{"cell_type":"markdown","source":"## 1. Setup & Configuration","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nimport lightgbm as lgb\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.notebook import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\nverbose = True # Set to False to reduce output messages\n\nclass CONFIG:\n    BASE_PATH = Path('/kaggle/input/MABe-mouse-behavior-detection/')\n    OUTPUT_PATH = Path('/kaggle/working/')\n    BEHAVIORS_TO_TRAIN = ['sniff', 'attack', 'rear', 'approach', 'selfgroom']\n    \n    LGB_PARAMS = {\n        'objective': 'binary',\n        'metric': 'binary_logloss',\n        'boosting_type': 'gbdt',\n        'n_estimators': 250,\n        'learning_rate': 0.05,\n        'num_leaves': 20,\n        'max_depth': 5,\n        'seed': 42,\n        'n_jobs': -1,\n        'verbose': -1,\n        'colsample_bytree': 0.7,\n        'subsample': 0.7,\n        'reg_alpha': 0.1,\n        'reg_lambda': 0.1,\n    }\n    \n    PROB_THRESHOLD = 0.15\n    MIN_FRAMES_FOR_EVENT = 1\n    NEGATIVE_SAMPLING_RATIO = 4\n    RANDOM_STATE = 42\n\nCONFIG.OUTPUT_PATH.mkdir(exist_ok=True)\nprint(\"Configuration loaded.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T23:14:37.009401Z","iopub.execute_input":"2025-10-18T23:14:37.009779Z","iopub.status.idle":"2025-10-18T23:14:37.017812Z","shell.execute_reply.started":"2025-10-18T23:14:37.009753Z","shell.execute_reply":"2025-10-18T23:14:37.016849Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Load Initial Data","metadata":{}},{"cell_type":"code","source":"def load_all_annotations(base_path):\n    \"\"\"\n    Loads all training annotation files, adds lab_id and video_id.\n    \"\"\"\n    print(\"Loading all annotations...\")\n    annotation_files = list(base_path.glob('train_annotation/*/*.parquet'))\n    all_annotations_list = []\n    # Use tqdm for progress bar\n    for f in tqdm(annotation_files, desc=\"Loading annotation files\"):\n        try:\n            lab_id = f.parts[-2] # Get lab_id from folder name\n            video_id = int(f.stem) # Get video_id from file name (without extension)\n            ann_df = pd.read_parquet(f)\n            ann_df['lab_id'] = lab_id\n            ann_df['video_id'] = video_id\n            all_annotations_list.append(ann_df)\n        except Exception as e:\n            print(f\"Error loading {f}: {e}\") # Handle potential errors\n            continue # Skip this file if error occurs\n            \n    if not all_annotations_list:\n        raise FileNotFoundError(\"No annotation files loaded. Check paths.\")\n        \n    # Concatenate all loaded dataframes\n    full_annotations_df = pd.concat(all_annotations_list).reset_index(drop=True)\n    print(f\"Loaded {len(full_annotations_df)} annotation events from {len(annotation_files)} files.\")\n    return full_annotations_df\n\n# Load metadata\ntrain_meta_df = pd.read_csv(CONFIG.BASE_PATH / 'train.csv')\ntest_meta_df = pd.read_csv(CONFIG.BASE_PATH / 'test.csv')\n\n# Load annotations using the function\nall_annotations_df = load_all_annotations(CONFIG.BASE_PATH)\n\n# Display some info\nprint(\"\\nTrain metadata shape:\", train_meta_df.shape)\nprint(\"Test metadata shape:\", test_meta_df.shape)\nprint(\"Annotations shape:\", all_annotations_df.shape)\nprint(\"\\nSample annotations:\")\ndisplay(all_annotations_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T22:36:59.559024Z","iopub.execute_input":"2025-10-18T22:36:59.559574Z","iopub.status.idle":"2025-10-18T22:37:05.805422Z","shell.execute_reply.started":"2025-10-18T22:36:59.559550Z","shell.execute_reply":"2025-10-18T22:37:05.804652Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Feature Engineering Function","metadata":{}},{"cell_type":"code","source":"# Define the body parts we consistently want features for\n# This helps handle variations in tracked parts across labs\nCANONICAL_SKELETON = ['nose', 'ear_left', 'ear_right', 'neck', 'body_center', 'tail_base']\n\n# Define some useful distances between parts for social interaction features\nSOCIAL_FEATURE_DEFS = [\n    ('mouse1', 'nose', 'mouse2', 'nose'),\n    ('mouse1', 'nose', 'mouse2', 'body_center'),\n    ('mouse1', 'nose', 'mouse2', 'tail_base'),\n    ('mouse1', 'body_center', 'mouse2', 'body_center'),\n]\n\ndef process_video_data(video_id, lab_id, tracking_path, is_train=True):\n    \"\"\"\n    Loads tracking data for a single video, performs lightweight feature engineering.\n    Returns a dataframe with features for each frame.\n    \"\"\"\n    # Construct the file path\n    folder = 'train_tracking' if is_train else 'test_tracking'\n    file_path = tracking_path / f\"{lab_id}/{video_id}.parquet\"\n\n    try:\n        # Load the raw tracking data\n        tracking_df = pd.read_parquet(file_path)\n    except FileNotFoundError:\n        print(f\"Warning: File not found {file_path}. Skipping video.\")\n        return pd.DataFrame() # Return empty dataframe if file not found\n\n    # Standardize mouse_id format (e.g., to 'mouse1', 'mouse2')\n    tracking_df['mouse_id'] = 'mouse' + tracking_df['mouse_id'].astype(str)\n\n    # --- Simple Imputation: Linear Interpolation ---\n    # Interpolate missing 'x' and 'y' values linearly within each mouse/bodypart group\n    cols_to_interpolate = ['x', 'y']\n    tracking_df[cols_to_interpolate] = tracking_df.groupby(['mouse_id', 'bodypart'])[cols_to_interpolate].transform(\n        lambda s: s.interpolate(method='linear', limit_direction='both', limit_area='inside')\n    )\n    # Fill any remaining NaNs (e.g., at the very start/end if interpolation didn't cover) with 0\n    tracking_df[cols_to_interpolate] = tracking_df[cols_to_interpolate].fillna(0)\n\n\n    # --- Pivot to Wide Format ---\n    # Get x, y for each bodypart/mouse onto a single row per frame\n    try:\n        wide_df = tracking_df.pivot_table(\n            index='video_frame',\n            columns=['mouse_id', 'bodypart'],\n            values=['x', 'y']\n        )\n        # Flatten the multi-level column index (e.g., ('x', 'mouse1', 'nose') -> 'x_mouse1_nose')\n        wide_df.columns = ['_'.join(col) for col in wide_df.columns.values]\n        wide_df = wide_df.reset_index() # Make video_frame a regular column\n    except Exception as e:\n         print(f\"Error pivoting video {video_id} ({lab_id}): {e}. Skipping video.\")\n         return pd.DataFrame() # Return empty if pivot fails\n\n    # --- Feature Creation ---\n    features_list = []\n    \n    # Check if wide_df is empty after pivot before proceeding\n    if wide_df.empty:\n        print(f\"Warning: wide_df is empty after pivot for video {video_id} ({lab_id}). Skipping.\")\n        return pd.DataFrame()\n        \n    df_index = wide_df.index # Get index for creating placeholder series\n\n    # 1. Normalized Body Part Coordinates (relative to body_center)\n    for mouse in ['mouse1', 'mouse2']:\n        center_x_col = f'x_{mouse}_body_center'\n        center_y_col = f'y_{mouse}_body_center'\n\n        # Use body_center if available, otherwise default to 0 (or mean coord)\n        center_x = wide_df[center_x_col] if center_x_col in wide_df.columns else 0\n        center_y = wide_df[center_y_col] if center_y_col in wide_df.columns else 0\n\n    for part in CANONICAL_SKELETON:\n            part_x_col = f'x_{mouse}_{part}'\n            part_y_col = f'y_{mouse}_{part}'\n\n            # Get part coordinate Series, or create a Series of zeros if missing\n            part_x = wide_df[part_x_col] if part_x_col in wide_df.columns else pd.Series(0, index=df_index, dtype=float) # Ensure float type\n            part_y = wide_df[part_y_col] if part_y_col in wide_df.columns else pd.Series(0, index=df_index, dtype=float) # Ensure float type\n\n            # Calculate the difference\n            diff_x = part_x - center_x\n            diff_y = part_y - center_y\n\n            # Explicitly create a named Series FROM the difference result\n            # This handles cases where diff_x/diff_y might be scalar if wide_df has 1 row\n            x_norm = pd.Series(diff_x, index=df_index, name=f'{mouse}_{part}_x_norm', dtype=float)\n            y_norm = pd.Series(diff_y, index=df_index, name=f'{mouse}_{part}_y_norm', dtype=float)\n\n            features_list.extend([x_norm, y_norm])\n\n    for m1, p1, m2, p2 in SOCIAL_FEATURE_DEFS:\n        m1_x_col, m1_y_col = f'x_{m1}_{p1}', f'y_{m1}_{p1}'\n        m2_x_col, m2_y_col = f'x_{m2}_{p2}', f'y_{m2}_{p2}'\n\n        if m1_x_col in wide_df.columns and m2_x_col in wide_df.columns:\n            m1_x, m1_y = wide_df[m1_x_col], wide_df[m1_y_col]\n            m2_x, m2_y = wide_df[m2_x_col], wide_df[m2_y_col]\n            # Ensure calculations handle potential NaNs from pivot before sqrt\n            dx = (m1_x - m2_x).fillna(0)\n            dy = (m1_y - m2_y).fillna(0)\n            dist = np.sqrt(dx**2 + dy**2)\n            dist.name = f'dist_{m1}{p1}_{m2}{p2}'\n            features_list.append(dist.astype(float)) # Ensure float type\n        else:\n             placeholder = pd.Series(0, index=df_index, name=f'dist_{m1}{p1}_{m2}{p2}', dtype=float) # Ensure float type\n             features_list.append(placeholder)\n\n    # Combine all feature series into a single DataFrame\n    if not features_list:\n        return pd.DataFrame() # Return empty if no features could be created\n        \n    features_df = pd.concat(features_list, axis=1)\n    \n    # Add video_frame back for potential merging later if needed (optional)\n    # features_df['video_frame'] = wide_df['video_frame']\n\n    return features_df.reset_index(drop=True) # Ensure clean index\n\n\n# --- Test the function on one sample video ---\ntry:\n    sample_video_id_test = train_meta_df.iloc[0]['video_id']\n    sample_lab_id_test = train_meta_df.iloc[0]['lab_id']\n    print(f\"Testing feature engineering on video: {sample_video_id_test} (Lab: {sample_lab_id_test})\")\n    sample_features = process_video_data(sample_video_id_test, sample_lab_id_test, CONFIG.BASE_PATH / 'train_tracking', is_train=True)\n    print(\"Feature engineering successful. Sample features shape:\", sample_features.shape)\n    display(sample_features.head())\n    del sample_features # Clean up memory\n    gc.collect()\nexcept Exception as e:\n    print(f\"Error during feature engineering test: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T22:37:05.806288Z","iopub.execute_input":"2025-10-18T22:37:05.806518Z","iopub.status.idle":"2025-10-18T22:37:07.774409Z","shell.execute_reply.started":"2025-10-18T22:37:05.806501Z","shell.execute_reply":"2025-10-18T22:37:07.773566Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Training Preparation","metadata":{}},{"cell_type":"code","source":"# Filter train_meta_df to only include videos that actually have annotations\n# This avoids processing videos for which we have no labels during training\nannotated_video_ids = all_annotations_df['video_id'].unique()\ntrain_meta_df_filtered = train_meta_df[train_meta_df['video_id'].isin(annotated_video_ids)].copy()\n\nprint(f\"Original train videos: {len(train_meta_df)}\")\nprint(f\"Using {len(train_meta_df_filtered)} videos with annotations for training.\")\n\n# Dictionary to store the trained models, one for each behavior\nmodels = {}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T22:37:07.775431Z","iopub.execute_input":"2025-10-18T22:37:07.775725Z","iopub.status.idle":"2025-10-18T22:37:07.787924Z","shell.execute_reply.started":"2025-10-18T22:37:07.775695Z","shell.execute_reply":"2025-10-18T22:37:07.787119Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Iterative Training Loop","metadata":{}},{"cell_type":"code","source":"# Set seed for numpy's random operations (like sampling)\nnp.random.seed(CONFIG.RANDOM_STATE)\n\n# Loop through each behavior we want to train a model for\nfor behavior in CONFIG.BEHAVIORS_TO_TRAIN:\n    print(f\"\\n--- Processing & Sampling data for behavior: {behavior} ---\")\n    \n    # Lists to collect sampled features and labels from all videos\n    X_train_sampled_list = []\n    y_train_sampled_list = []\n    \n    # Iterate through each training video (that has annotations)\n    # Use tqdm for progress bar\n    for row in tqdm(train_meta_df_filtered.itertuples(), total=len(train_meta_df_filtered), desc=f\"Videos for {behavior}\"):\n        \n        # Process the video to get features\n        features_df = process_video_data(row.video_id, row.lab_id, CONFIG.BASE_PATH / 'train_tracking', is_train=True)\n        \n        # Skip if feature extraction failed or returned empty\n        if features_df.empty:\n            # print(f\"Skipping video {row.video_id} (Lab: {row.lab_id}) due to empty features.\")\n            continue\n            \n        # --- Create Labels for the current behavior ---\n        # Initialize labels series with all zeros (non-event)\n        labels = pd.Series(np.zeros(len(features_df)), name=behavior, dtype=np.int8)\n        \n        # Get annotations specific to this video and behavior\n        video_annotations = all_annotations_df[\n            (all_annotations_df['video_id'] == row.video_id) &\n            (all_annotations_df['action'] == behavior)\n        ]\n        \n        # Mark frames within annotated event ranges as 1 (event)\n        for _, ann_row in video_annotations.iterrows():\n            start_f = ann_row['start_frame']\n            stop_f = ann_row['stop_frame']\n            # Ensure indices are within the bounds of the features_df\n            if start_f < len(labels):\n                # +1 because iloc is exclusive of the stop index\n                labels.iloc[start_f : min(stop_f + 1, len(labels))] = 1\n\n        # --- Negative Sampling ---\n        positive_indices = labels[labels == 1].index\n        negative_indices = labels[labels == 0].index\n        \n        # Only proceed if there are positive examples for this behavior in this video\n        if len(positive_indices) == 0:\n            # print(f\"No positive examples for {behavior} in video {row.video_id}. Skipping sampling.\")\n            continue\n            \n        # Calculate number of negative samples needed based on ratio\n        num_neg_to_sample = min(int(len(positive_indices) * CONFIG.NEGATIVE_SAMPLING_RATIO), len(negative_indices))\n        \n        # Randomly choose negative indices if there are any negatives to sample from\n        if num_neg_to_sample > 0 and len(negative_indices) > 0:\n            sampled_negative_indices = np.random.choice(negative_indices, size=num_neg_to_sample, replace=False)\n            # Combine positive and sampled negative indices\n            final_indices = np.concatenate([positive_indices, sampled_negative_indices])\n        else:\n            # If no negatives to sample or ratio is 0, just use positives\n            final_indices = positive_indices\n            \n        # Append the sampled data to our lists\n        if len(final_indices) > 0:\n            X_train_sampled_list.append(features_df.iloc[final_indices])\n            y_train_sampled_list.append(labels.iloc[final_indices])\n            \n        # --- Memory Management ---\n        del features_df, labels, video_annotations, positive_indices, negative_indices, final_indices\n        # gc.collect() # Optional: Force garbage collection more frequently if memory is tight\n\n    # --- Train Model for the Behavior ---\n    # Check if we collected any data at all for this behavior\n    if not X_train_sampled_list:\n        print(f\"No data collected for behavior '{behavior}' after sampling across all videos. Skipping training.\")\n        continue\n        \n    # Concatenate all sampled data into final training sets\n    X_train = pd.concat(X_train_sampled_list).reset_index(drop=True)\n    y_train = pd.concat(y_train_sampled_list).reset_index(drop=True)\n    \n    print(f\"--- Training model for: {behavior} ---\")\n    print(f\"Sampled training data shape: X={X_train.shape}, y={y_train.shape}\")\n    \n    # Initialize and train the LightGBM model\n    model = lgb.LGBMClassifier(**CONFIG.LGB_PARAMS)\n    model.fit(X_train, y_train)\n    \n    # Store the trained model\n    models[behavior] = model\n    \n    # --- Memory Management ---\n    print(f\"Finished training for {behavior}. Cleaning up memory.\")\n    del X_train, y_train, X_train_sampled_list, y_train_sampled_list\n    gc.collect() # Force garbage collection after training each model\n\nprint(\"\\n--- Model training complete for selected behaviors ---\")\nprint(f\"Trained models for: {list(models.keys())}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T22:37:07.788727Z","iopub.execute_input":"2025-10-18T22:37:07.788998Z","iopub.status.idle":"2025-10-18T23:07:16.749010Z","shell.execute_reply.started":"2025-10-18T22:37:07.788973Z","shell.execute_reply":"2025-10-18T23:07:16.748115Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Prediction Loop & Post-Processing","metadata":{}},{"cell_type":"code","source":"def probabilities_to_events(probs, threshold, min_frames):\n    \"\"\"\n    Converts frame-wise probabilities into a list of event tuples (start_frame, stop_frame).\n    (Code is omitted here for brevity, assume the correct function is loaded)\n    \"\"\"\n    if probs is None or len(probs) == 0:\n        return []\n\n    # Apply threshold to get binary predictions (0 or 1)\n    binary_preds = (probs > threshold).astype(np.int8)\n\n    # Find where sequences of 1s start and end\n    diffs = np.diff(binary_preds, prepend=0, append=0)\n\n    start_frames = np.where(diffs == 1)[0]\n    stop_frames = np.where(diffs == -1)[0] # stop_frames index is exclusive\n\n    events = []\n    # Pair up start and stop frames\n    for start, stop in zip(start_frames, stop_frames):\n        # Check if the duration meets the minimum requirement\n        if (stop - start) >= min_frames:\n            # Add event (inclusive frame indices)\n            events.append((start, stop - 1)) # -1 because stop_frames was exclusive\n\n    return events\n\n\nprint(\"\\n--- Starting Inference on Test Set ---\")\nall_predictions = [] # List to store prediction dictionaries\n\n# Define individual behaviors for target assignment (MUST be kept updated)\nINDIVIDUAL_BEHAVIORS = ['selfgroom', 'rear'] # Add 'dig', 'investigation' etc., if you train them\n\n# Iterate through each test video\nfor row in tqdm(test_meta_df.itertuples(), total=len(test_meta_df), desc=\"Inferring on test videos\"):\n\n    # Process the test video to get features\n    features_df = process_video_data(row.video_id, row.lab_id, CONFIG.BASE_PATH / 'test_tracking', is_train=False)\n\n    # Skip if feature extraction failed\n    if features_df.empty:\n        continue\n\n    # Predict for each behavior using the corresponding trained model\n    for behavior, model in models.items():\n        # Only predict if the model was successfully trained\n        if behavior in models:\n            \n            # Predict probabilities (get probability of the positive class, index 1)\n            probs = model.predict_proba(features_df)[:, 1]\n\n            # Convert probabilities to events using post-processing\n            events = probabilities_to_events(probs, CONFIG.PROB_THRESHOLD, CONFIG.MIN_FRAMES_FOR_EVENT)\n\n            # Format events for submission\n            for start, stop in events:\n                \n                # --- CORRECTED AGENT/TARGET LOGIC ---\n                agent_id = 'mouse1' \n                \n                if behavior in INDIVIDUAL_BEHAVIORS:\n                    target_id = agent_id # Target is self for individual behaviors\n                else:\n                    target_id = 'mouse2' # Target is the other mouse for social behaviors\n                # --- END CORRECTION ---\n\n                all_predictions.append({\n                    'video_id': row.video_id,\n                    'agent_id': agent_id,\n                    'target_id': target_id,\n                    'action': behavior,\n                    'start_frame': start,\n                    'stop_frame': stop\n                })\n        # else:\n            # print(f\"Warning: Model for behavior '{behavior}' not found. Skipping prediction.\")\n\n\n    # --- Memory Management ---\n    del features_df\n    gc.collect()\n\nprint(\"\\n--- Inference complete ---\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T23:07:16.750030Z","iopub.execute_input":"2025-10-18T23:07:16.750876Z","iopub.status.idle":"2025-10-18T23:07:18.804612Z","shell.execute_reply.started":"2025-10-18T23:07:16.750850Z","shell.execute_reply":"2025-10-18T23:07:18.803575Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Create Submission File","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\nprint(\"Creating submission dataframe...\")\n\n# Convert list of prediction dictionaries to a DataFrame\nsubmission_df = pd.DataFrame(all_predictions)\n\n# Ensure required columns exist even if no predictions were made initially\nrequired_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\nfor col in required_cols:\n    if col not in submission_df.columns:\n        submission_df[col] = []\n\n# --- CLEANING FUNCTION ---\ndef clean_overlaps(df):\n    \"\"\"Remove overlapping predictions for same agent/target pair\"\"\"\n    if df.empty:\n        return df\n    \n    cleaned = []\n    for (vid, agent, target), group in df.groupby(['video_id', 'agent_id', 'target_id']):\n        group = group.sort_values('start_frame').reset_index(drop=True)\n        \n        if len(group) == 0:\n            continue\n            \n        valid_mask = [True]\n        last_stop = group.iloc[0]['stop_frame']\n        \n        for i in range(1, len(group)):\n            current_start = group.iloc[i]['start_frame']\n            if current_start < last_stop:\n                # Overlap detected - skip this prediction\n                valid_mask.append(False)\n                if verbose:\n                    print(f\"  Removing overlap: video={vid}, agent={agent}, target={target}, frame={current_start}\")\n            else:\n                valid_mask.append(True)\n                last_stop = group.iloc[i]['stop_frame']\n        \n        cleaned.append(group[valid_mask])\n    \n    return pd.concat(cleaned, ignore_index=True) if cleaned else pd.DataFrame(columns=df.columns)\n\n# --- MERGE OVERLAPPING INTERVALS FUNCTION ---\ndef merge_overlapping_intervals(df_group):\n    \"\"\"Merges overlapping or adjacent intervals within a group.\"\"\"\n    if df_group.empty:\n        return df_group\n\n    # Sort by start_frame is crucial\n    df_group = df_group.sort_values('start_frame')\n\n    merged = []\n    current_start = -1\n    current_stop = -1\n\n    for _, row in df_group.iterrows():\n        if current_start == -1:\n            # Initialize with the first interval\n            current_start = row['start_frame']\n            current_stop = row['stop_frame']\n        else:\n            # Check for overlap or adjacency (stop >= start - 1)\n            if row['start_frame'] <= current_stop + 1:\n                # Merge by extending the current stop frame if necessary\n                current_stop = max(current_stop, row['stop_frame'])\n            else:\n                # No overlap, finalize the previous interval and start a new one\n                merged.append({'start_frame': current_start, 'stop_frame': current_stop})\n                current_start = row['start_frame']\n                current_stop = row['stop_frame']\n\n    # Add the last processed interval\n    if current_start != -1:\n        merged.append({'start_frame': current_start, 'stop_frame': current_stop})\n\n    # Return a new dataframe with merged intervals\n    if merged:\n        merged_df = pd.DataFrame(merged)\n        # Get the keys from the first row of the original group\n        keys = df_group.iloc[0][['video_id', 'agent_id', 'target_id', 'action']]\n        for col, value in keys.items():\n            merged_df[col] = value\n        return merged_df\n    else:\n        return pd.DataFrame(columns=df_group.columns)\n\n# --- APPLY CLEANING AND MERGING ---\nif not submission_df.empty:\n    print(f\"Shape before cleaning overlaps: {submission_df.shape}\")\n    \n    # Step 1: Clean hard overlaps (same start frame or overlapping ranges)\n    submission_df = clean_overlaps(submission_df)\n    print(f\"Shape after cleaning overlaps: {submission_df.shape}\")\n    \n    # Step 2: Merge adjacent/overlapping intervals for same action\n    print(f\"Shape before merging adjacent intervals: {submission_df.shape}\")\n    submission_merged = submission_df.groupby(\n        ['video_id', 'agent_id', 'target_id', 'action'], \n        group_keys=False\n    ).apply(merge_overlapping_intervals).reset_index(drop=True)\n    print(f\"Shape after merging adjacent intervals: {submission_merged.shape}\")\n    submission_df = submission_merged\n    \n    # Step 3: Filter out very short events (likely noise)\n    duration = submission_df['stop_frame'] - submission_df['start_frame']\n    before_filter = len(submission_df)\n    submission_df = submission_df[duration >= 3].reset_index(drop=True)\n    after_filter = len(submission_df)\n    if before_filter > after_filter:\n        print(f\"Filtered out {before_filter - after_filter} short events (duration < 3 frames)\")\nelse:\n    print(\"Submission DataFrame is empty, skipping cleaning and merging.\")\n\n# --- ADD row_id COLUMN ---\nif not submission_df.empty:\n    # Sort for final consistency before adding row_id\n    submission_df = submission_df.sort_values(\n        by=['video_id', 'start_frame', 'action']\n    ).reset_index(drop=True)\n    submission_df['row_id'] = submission_df.index\nelse:\n    submission_df['row_id'] = []\n\n# Reorder columns\nfinal_cols = ['row_id', 'video_id', 'agent_id', 'target_id', 'action', 'start_frame', 'stop_frame']\nif 'row_id' not in submission_df.columns:\n    submission_df['row_id'] = []\nsubmission_df = submission_df.reindex(columns=final_cols, fill_value=np.nan)\n\n# --- VALIDATION BLOCK ---\nprint(\"\\n--- Submission DataFrame Validation ---\")\nprint(\"Shape:\", submission_df.shape)\nprint(\"Columns:\", submission_df.columns.tolist())\nprint(\"Data Types:\\n\", submission_df.dtypes)\nprint(\"Any NaNs:\", submission_df.isnull().values.any())\nprint(\"NaNs per column:\\n\", submission_df.isnull().sum())\n\n# Cast types safely ONLY IF not empty AND no NaNs in integer columns\nif not submission_df.empty and not submission_df[['row_id', 'start_frame', 'stop_frame']].isnull().values.any():\n    try:\n        submission_df['row_id'] = submission_df['row_id'].astype(int)\n        submission_df['video_id'] = submission_df['video_id'].astype(int)\n        submission_df['start_frame'] = submission_df['start_frame'].astype(int)\n        submission_df['stop_frame'] = submission_df['stop_frame'].astype(int)\n        print(\"\\n✓ Type casting successful\")\n        \n        # Additional validations\n        assert (submission_df['start_frame'] < submission_df['stop_frame']).all(), \"ERROR: Some start_frame >= stop_frame\"\n        print(\"✓ All start_frame < stop_frame\")\n        \n        # Check for duplicate rows\n        dup_cols = ['video_id', 'agent_id', 'target_id', 'action', 'start_frame']\n        duplicates = submission_df.duplicated(subset=dup_cols, keep=False)\n        if duplicates.any():\n            print(f\"WARNING: Found {duplicates.sum()} duplicate rows\")\n            print(submission_df[duplicates][dup_cols])\n        else:\n            print(\"✓ No duplicate rows found\")\n        \n        # Final overlap check\n        print(\"\\nChecking for remaining overlaps...\")\n        overlap_found = False\n        for (vid, agent, target), group in submission_df.groupby(['video_id', 'agent_id', 'target_id']):\n            group = group.sort_values('start_frame')\n            frames_used = set()\n            for _, row in group.iterrows():\n                new_frames = set(range(row['start_frame'], row['stop_frame']))\n                if frames_used.intersection(new_frames):\n                    print(f\"ERROR: Overlap still exists for video={vid}, agent={agent}, target={target}\")\n                    overlap_found = True\n                    break\n                frames_used.update(new_frames)\n            if overlap_found:\n                break\n        \n        if not overlap_found:\n            print(\"✓ No overlaps detected\")\n        \n    except Exception as e:\n        print(f\"\\nERROR during validation: {e}\")\nelse:\n    print(\"\\nWARNING: Submission DataFrame is empty or contains NaNs, skipping casting/validation\")\n\nprint(\"------------------------------------\\n\")\n\n# --- HANDLE EMPTY SUBMISSION ---\nif submission_df.empty or len(submission_df) == 0:\n    print(\"WARNING: Empty submission detected. Creating dummy submission.\")\n    submission_df = pd.DataFrame({\n        'row_id': [0],\n        'video_id': [test['video_id'].iloc[0] if len(test) > 0 else 438887472],\n        'agent_id': ['mouse1'],\n        'target_id': ['self'],\n        'action': ['rear'],\n        'start_frame': [0],\n        'stop_frame': [100]\n    })\n\n# Save to csv file\nsubmission_path = CONFIG.OUTPUT_PATH / 'submission.csv'\nsubmission_df.to_csv(submission_path, index=False)\n\nprint(f\"Submission file saved to: {submission_path}\")\nprint(f\"Total predictions: {len(submission_df)}\")\n\nif not submission_df.empty:\n    print(\"\\nSample of final submission file:\")\n    display(submission_df.head(10))\n    \n    # Show statistics\n    print(\"\\nSubmission Statistics:\")\n    print(f\"  Unique videos: {submission_df['video_id'].nunique()}\")\n    print(f\"  Unique actions: {submission_df['action'].nunique()}\")\n    print(\"  Actions distribution:\")\n    print(submission_df['action'].value_counts())\n    print(f\"\\n  Average event duration: {(submission_df['stop_frame'] - submission_df['start_frame']).mean():.1f} frames\")\nelse:\n    print(\"\\nFinal submission file is empty.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-18T23:14:47.016244Z","iopub.execute_input":"2025-10-18T23:14:47.016556Z","iopub.status.idle":"2025-10-18T23:14:47.392252Z","shell.execute_reply.started":"2025-10-18T23:14:47.016534Z","shell.execute_reply":"2025-10-18T23:14:47.391464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}