{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"},{"sourceId":13298882,"sourceType":"datasetVersion","datasetId":8429100},{"sourceId":13302786,"sourceType":"datasetVersion","datasetId":8431980},{"sourceId":13325571,"sourceType":"datasetVersion","datasetId":8446658},{"sourceId":13515842,"sourceType":"datasetVersion","datasetId":8581510},{"sourceId":13694923,"sourceType":"datasetVersion","datasetId":8700007},{"sourceId":13695025,"sourceType":"datasetVersion","datasetId":8699423},{"sourceId":13695038,"sourceType":"datasetVersion","datasetId":8699427},{"sourceId":13906574,"sourceType":"datasetVersion","datasetId":8860391},{"sourceId":13906587,"sourceType":"datasetVersion","datasetId":8860402},{"sourceId":13926184,"sourceType":"datasetVersion","datasetId":8874280},{"sourceId":13932126,"sourceType":"datasetVersion","datasetId":8878507},{"sourceId":13939923,"sourceType":"datasetVersion","datasetId":8883983},{"sourceId":13972810,"sourceType":"datasetVersion","datasetId":8834748},{"sourceId":13972858,"sourceType":"datasetVersion","datasetId":8908034},{"sourceId":13972925,"sourceType":"datasetVersion","datasetId":8908084},{"sourceId":14069749,"sourceType":"datasetVersion","datasetId":8955775},{"sourceId":14070283,"sourceType":"datasetVersion","datasetId":8956153},{"sourceId":14073230,"sourceType":"datasetVersion","datasetId":8958342},{"sourceId":14125881,"sourceType":"datasetVersion","datasetId":9000167},{"sourceId":14125884,"sourceType":"datasetVersion","datasetId":9000168},{"sourceId":14126124,"sourceType":"datasetVersion","datasetId":9000349},{"sourceId":14126126,"sourceType":"datasetVersion","datasetId":9000351},{"sourceId":14137701,"sourceType":"datasetVersion","datasetId":9000352},{"sourceId":14151205,"sourceType":"datasetVersion","datasetId":9019204},{"sourceId":14151207,"sourceType":"datasetVersion","datasetId":9019206},{"sourceId":14151443,"sourceType":"datasetVersion","datasetId":9019126},{"sourceId":14151445,"sourceType":"datasetVersion","datasetId":9019129},{"sourceId":14151446,"sourceType":"datasetVersion","datasetId":9009212},{"sourceId":14151512,"sourceType":"datasetVersion","datasetId":9019444},{"sourceId":14164152,"sourceType":"datasetVersion","datasetId":9028257},{"sourceId":14164154,"sourceType":"datasetVersion","datasetId":9028259},{"sourceId":14165291,"sourceType":"datasetVersion","datasetId":9029084},{"sourceId":14165294,"sourceType":"datasetVersion","datasetId":8419449},{"sourceId":14165440,"sourceType":"datasetVersion","datasetId":9029188},{"sourceId":14165611,"sourceType":"datasetVersion","datasetId":9029326},{"sourceId":250296873,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index --find-links /kaggle/input/mdc-dependencies hydra-core","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:18:52.222174Z","iopub.execute_input":"2025-12-15T14:18:52.222409Z","iopub.status.idle":"2025-12-15T14:18:56.821725Z","shell.execute_reply.started":"2025-12-15T14:18:52.222382Z","shell.execute_reply":"2025-12-15T14:18:56.821042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd /kaggle/working\n!rm -r mabe-tenten-code\n!cp -r /kaggle/input/mabe-tenten-code ./\n%cd /kaggle/working/mabe-tenten-code\n!ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:18:56.823553Z","iopub.execute_input":"2025-12-15T14:18:56.823782Z","iopub.status.idle":"2025-12-15T14:19:02.922708Z","shell.execute_reply.started":"2025-12-15T14:18:56.823756Z","shell.execute_reply":"2025-12-15T14:19:02.922092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"competition_name = \"mabe\"\nexp_name = \"exp113\"\nsub_name = \"self\"\n\n!python exps/{exp_name}/create_dataset.py exp.dataset.is_train=False dir=kaggle","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:19:02.923555Z","iopub.execute_input":"2025-12-15T14:19:02.923819Z","iopub.status.idle":"2025-12-15T14:19:22.834438Z","shell.execute_reply.started":"2025-12-15T14:19:02.923783Z","shell.execute_reply":"2025-12-15T14:19:22.833251Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\n\nfeature_statistics = {\n  \"hip_left_x\": {\n    \"mean\": 25.64735788873652,\n    \"std\": 14.188237874465381,\n    \"min\": 0.0,\n    \"max\": 55.73170471191406,\n    \"count\": 26441494\n  },\n  \"tail_midpoint_x\": {\n    \"mean\": 44.33567625119911,\n    \"std\": 17.024571358058733,\n    \"min\": 0.660812497138977,\n    \"max\": 83.12711334228516,\n    \"count\": 2974792\n  },\n  \"hip_right_y\": {\n    \"mean\": 16.14733782796993,\n    \"std\": 7.099667661354833,\n    \"min\": 0.0,\n    \"max\": 44.241207122802734,\n    \"count\": 26441494\n  },\n  \"body_center_y\": {\n    \"mean\": 32.06692215464333,\n    \"std\": 20.28218841429278,\n    \"min\": -0.49807727336883545,\n    \"max\": 127.92010498046875,\n    \"count\": 17240415\n  },\n  \"nose_y\": {\n    \"mean\": 22.19957942893967,\n    \"std\": 15.976416217101793,\n    \"min\": -0.4995335042476654,\n    \"max\": 129.87588500976562,\n    \"count\": 44309408\n  },\n  \"lateral_left_y\": {\n    \"mean\": 29.99313384363584,\n    \"std\": 25.871251701145297,\n    \"min\": -0.14450053870677948,\n    \"max\": 129.43624877929688,\n    \"count\": 7024709\n  },\n  \"ear_left_y\": {\n    \"mean\": 22.13788573163825,\n    \"std\": 15.817787222315642,\n    \"min\": -0.49991700053215027,\n    \"max\": 129.50132751464844,\n    \"count\": 44310521\n  },\n  \"neck_y\": {\n    \"mean\": 17.089995666252424,\n    \"std\": 9.395953409885452,\n    \"min\": 0.0,\n    \"max\": 65.73812866210938,\n    \"count\": 28213393\n  },\n  \"ear_right_x\": {\n    \"mean\": 30.837927364934256,\n    \"std\": 17.43188971084471,\n    \"min\": -0.4988972246646881,\n    \"max\": 128.90309143066406,\n    \"count\": 44310530\n  },\n  \"hip_right_x\": {\n    \"mean\": 25.692085131112908,\n    \"std\": 14.182848258903288,\n    \"min\": 0.0,\n    \"max\": 55.652862548828125,\n    \"count\": 26441494\n  },\n  \"hip_left_y\": {\n    \"mean\": 16.34312315603273,\n    \"std\": 7.003185690698667,\n    \"min\": 0.0,\n    \"max\": 44.06658935546875,\n    \"count\": 26441494\n  },\n  \"ear_right_y\": {\n    \"mean\": 22.059206280847935,\n    \"std\": 15.851551061607308,\n    \"min\": -0.49826285243034363,\n    \"max\": 128.5943603515625,\n    \"count\": 44309989\n  },\n  \"tail_base_x\": {\n    \"mean\": 30.67865107347306,\n    \"std\": 17.213900760973576,\n    \"min\": -0.497320294380188,\n    \"max\": 128.05787658691406,\n    \"count\": 44311369\n  },\n  \"body_center_x\": {\n    \"mean\": 38.986779163396186,\n    \"std\": 18.596210040053382,\n    \"min\": -0.4992671608924866,\n    \"max\": 128.89544677734375,\n    \"count\": 17239912\n  },\n  \"neck_x\": {\n    \"mean\": 27.656770312777198,\n    \"std\": 15.978726168733159,\n    \"min\": 0.0,\n    \"max\": 82.82951354980469,\n    \"count\": 28213393\n  },\n  \"tail_tip_y\": {\n    \"mean\": 32.42177175401586,\n    \"std\": 27.409948203653542,\n    \"min\": -0.3388923108577728,\n    \"max\": 160.06829833984375,\n    \"count\": 5756269\n  },\n  \"tail_base_y\": {\n    \"mean\": 22.176088921562844,\n    \"std\": 15.941853991586138,\n    \"min\": -0.29868677258491516,\n    \"max\": 129.7112579345703,\n    \"count\": 44311508\n  },\n  \"nose_x\": {\n    \"mean\": 30.874007739010192,\n    \"std\": 17.540910162903135,\n    \"min\": -0.49927133321762085,\n    \"max\": 130.93814086914062,\n    \"count\": 44310179\n  },\n  \"lateral_right_y\": {\n    \"mean\": 29.98370120126074,\n    \"std\": 25.925259461069533,\n    \"min\": 0.0,\n    \"max\": 128.92929077148438,\n    \"count\": 7024712\n  },\n  \"tail_midpoint_y\": {\n    \"mean\": 25.94822559616726,\n    \"std\": 19.17516254244732,\n    \"min\": 0.23855669796466827,\n    \"max\": 66.32506561279297,\n    \"count\": 2974792\n  },\n  \"lateral_right_x\": {\n    \"mean\": 37.590458713784656,\n    \"std\": 24.83671324848758,\n    \"min\": 0.06613044440746307,\n    \"max\": 129.25393676757812,\n    \"count\": 7024712\n  },\n  \"tail_tip_x\": {\n    \"mean\": 42.95627311757143,\n    \"std\": 23.662033633289415,\n    \"min\": -0.30277031660079956,\n    \"max\": 127.79247283935547,\n    \"count\": 5756271\n  },\n  \"ear_left_x\": {\n    \"mean\": 30.788103761733325,\n    \"std\": 17.414776628246255,\n    \"min\": -0.4991244077682495,\n    \"max\": 128.8426513671875,\n    \"count\": 44310618\n  },\n  \"lateral_left_x\": {\n    \"mean\": 37.6952454466084,\n    \"std\": 24.917845446359383,\n    \"min\": 0.0,\n    \"max\": 127.85034942626953,\n    \"count\": 7024712\n  }\n}\ndistance_velocity_acceleration_statistics = {\n  \"distance\": {\n    \"mean\": 6.327864085783904,\n    \"std\": 14.118793272336088,\n    \"min\": 0.0,\n    \"max\": 175.05316162109375,\n    \"count\": 7410358644\n  },\n  \"velocity\": {\n    \"mean\": 0.13461063607889667,\n    \"std\": 0.5633761079797619,\n    \"min\": 0.0,\n    \"max\": 59.42070388793945,\n    \"count\": 1140055176\n  },\n  \"acceleration\": {\n    \"mean\": 0.07121916207071484,\n    \"std\": 0.42541316979091875,\n    \"min\": 0.0,\n    \"max\": 50.23505401611328,\n    \"count\": 1140055176\n  }\n}\n\nwith open(\n    f\"/kaggle/working/dataset/mabe-dataset-exp113/feature_statistics.json\", \"w\"\n) as f:\n    json.dump(feature_statistics, f, indent=4)\nwith open(\n    f\"/kaggle/working/dataset/mabe-dataset-exp113/distance_velocity_acceleration_statistics.json\", \"w\"\n) as f:\n    json.dump(distance_velocity_acceleration_statistics, f, indent=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:19:22.836075Z","iopub.execute_input":"2025-12-15T14:19:22.836520Z","iopub.status.idle":"2025-12-15T14:19:22.855594Z","shell.execute_reply.started":"2025-12-15T14:19:22.836487Z","shell.execute_reply":"2025-12-15T14:19:22.854747Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%python\nimport argparse\nimport importlib.util\nimport json\nimport sys\nimport warnings\nfrom pathlib import Path\nfrom typing import List, Tuple\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom omegaconf import OmegaConf\nfrom pytorch_lightning import seed_everything\nfrom pytorch_lightning.loggers import WandbLogger\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\n# Add parent directories to path\nsys.path.append(str(Path(__file__).resolve().parents[2]))\n\nfrom exps.exp078.run import (\n    MyLightningDataModule,\n    SlidingWindowDataset,\n    _accumulate_predictions,\n    _extract_segments_with_lab_action_thresholds,\n    _interpolate_predictions_to_original_fps,\n    _search_optimal_thresholds_per_lab_and_action,\n    _smooth_predictions,\n    load_actions_from_config,\n    load_test_data,\n    prepare_solution_for_scoring,\n)\nfrom src.metrics import score as mabe_score\nfrom utils.memory_time_trace import trace\n\nwarnings.filterwarnings(\"ignore\")\n\n\ndef _adjust_frames_for_adaptable_snail(\n    submission: pd.DataFrame, train_csv_path: Path, test_csv_path: Path = None\n) -> pd.DataFrame:\n    \"\"\"\n    Adjust frame numbers for AdaptableSnail videos with fps ~25.0.\n    Multiply frame numbers by 30/25 = 1.2.\n\n    Args:\n        submission: Submission DataFrame\n        train_csv_path: Path to train CSV to get video metadata\n        test_csv_path: Path to test CSV to get video metadata (optional, for test data)\n\n    Returns:\n        Adjusted submission DataFrame\n    \"\"\"\n    if len(submission) == 0:\n        return submission\n\n    # Load train data to get video metadata\n    train_df = pd.read_csv(\"/kaggle/input/MABe-mouse-behavior-detection/test.csv\")\n\n    # If test_csv_path is provided, load and combine with train data\n    if test_csv_path is not None and test_csv_path.exists():\n        test_df = pd.read_csv(test_csv_path)\n        # Combine train and test metadata\n        combined_df = pd.concat(\n            [\n                train_df[[\"video_id\", \"lab_id\", \"frames_per_second\"]],\n                test_df[[\"video_id\", \"lab_id\", \"frames_per_second\"]],\n            ],\n            ignore_index=True,\n        )\n        # Remove duplicates, keeping first occurrence\n        combined_df = combined_df.drop_duplicates(subset=[\"video_id\"], keep=\"first\")\n        video_metadata = combined_df.set_index(\"video_id\")[\n            [\"lab_id\", \"frames_per_second\"]\n        ].to_dict(\"index\")\n    else:\n        # Create video_id to (lab_id, fps) mapping from train data only\n        video_metadata = train_df.set_index(\"video_id\")[\n            [\"lab_id\", \"frames_per_second\"]\n        ].to_dict(\"index\")\n\n    # Find videos that need adjustment\n    adjustment_needed = []\n    for video_id in submission[\"video_id\"].unique():\n        if video_id in video_metadata:\n            info = video_metadata[video_id]\n            if (\n                info[\"lab_id\"] == \"AdaptableSnail\"\n                and abs(info[\"frames_per_second\"] - 25.0) < 0.1\n            ):\n                adjustment_needed.append(video_id)\n\n    if len(adjustment_needed) == 0:\n        print(\"No AdaptableSnail videos with fps ~25.0 found in submission\")\n        return submission\n\n    print(\n        f\"\\nAdjusting frames for {len(adjustment_needed)} AdaptableSnail videos with fps ~25.0\"\n    )\n    print(f\"Videos to adjust: {adjustment_needed}\")\n\n    # Create a copy to avoid modifying original\n    adjusted_submission = submission.copy()\n\n    # Adjust frames for matching videos\n    mask = adjusted_submission[\"video_id\"].isin(adjustment_needed)\n    num_affected = mask.sum()\n\n    if num_affected > 0:\n        adjusted_submission.loc[mask, \"start_frame\"] = (\n            (adjusted_submission.loc[mask, \"start_frame\"] * 30 / 25).round().astype(int)\n        )\n        adjusted_submission.loc[mask, \"stop_frame\"] = (\n            (adjusted_submission.loc[mask, \"stop_frame\"] * 30 / 25).round().astype(int)\n        )\n\n        print(\n            f\"Adjusted {num_affected} predictions (frame numbers multiplied by 30/25 = 1.2)\"\n        )\n\n    return adjusted_submission\n\n\ndef load_model_class_from_exp(\n    exp_name: str, base_output_dir: str = \"output\", model_dir: Path = None\n):\n    \"\"\"\n    Dynamically load MyModel class from experiment's run.py.\n\n    Args:\n        exp_name: Experiment name (e.g., \"exp081\")\n        base_output_dir: Base directory for outputs\n        model_dir: Specific directory where the model is located (for Kaggle)\n\n    Returns:\n        MyModel class from the experiment\n    \"\"\"\n    # Try standard path first (local environment)\n    run_py_path = Path(base_output_dir).parent / \"exps\" / exp_name / \"run.py\"\n\n    # If not found, try Kaggle-style path (assuming base_output_dir contains the full path)\n    if not run_py_path.exists():\n        # For Kaggle, find run.py in workspace\n        workspace_path = (\n            Path(\"/kaggle/input/mabe-tenten-code/exps\") / exp_name / \"run.py\"\n        )\n        if workspace_path.exists():\n            run_py_path = workspace_path\n        else:\n            raise FileNotFoundError(\n                f\"run.py not found at {run_py_path} or {workspace_path}\"\n            )\n\n    # Load module dynamically\n    spec = importlib.util.spec_from_file_location(f\"{exp_name}.run\", run_py_path)\n    module = importlib.util.module_from_spec(spec)\n    spec.loader.exec_module(module)\n\n    return module.MyModel\n\n\nclass EnsembleModel(nn.Module):\n    \"\"\"Ensemble multiple models from different experiments.\"\"\"\n\n    def __init__(\n        self,\n        exp_configs: List[Tuple[str, str, int]],\n        device: torch.device,\n        base_output_dir: str = \"output\",\n        weights: List[float] = None,\n        ensemble_method: str = \"mean\",\n    ):\n        \"\"\"\n        Initialize ensemble model.\n\n        Args:\n            exp_configs: List of (exp_name, sub_name, fold) tuples\n            device: Device to load models on\n            base_output_dir: Base directory for outputs\n            weights: Optional list of weights for each model (will be normalized to sum to 1)\n            ensemble_method: Ensemble method (\"mean\" or \"max\")\n        \"\"\"\n        super().__init__()\n        self.device = device\n        self.models = []\n        self.configs = []\n        self.ensemble_method = ensemble_method\n\n        if ensemble_method not in [\"mean\", \"max\"]:\n            raise ValueError(\n                f\"ensemble_method must be 'mean' or 'max', got {ensemble_method}\"\n            )\n\n        # Setup weights\n        if weights is None:\n            self.weights = [1.0 / len(exp_configs)] * len(exp_configs)\n        else:\n            if len(weights) != len(exp_configs):\n                raise ValueError(\n                    f\"Number of weights ({len(weights)}) must match number of configs ({len(exp_configs)})\"\n                )\n            # Normalize weights to sum to 1\n            total = sum(weights)\n            self.weights = [w / total for w in weights]\n\n        print(f\"\\nLoading {len(exp_configs)} models for ensemble:\")\n        print(\"=\" * 80)\n        print(f\"Ensemble method: {ensemble_method}\")\n        print(f\"Model weights: {[f'{w:.3f}' for w in self.weights]}\")\n\n        for exp_name, sub_name, fold in exp_configs:\n            # Try standard path first (local environment: output/exp_name/sub_name)\n            exp_path = Path(base_output_dir) / exp_name / sub_name\n\n            # Check if standard path exists\n            standard_config_path = exp_path / \"config.yaml\"\n\n            # If standard path doesn't exist, try Kaggle-style path\n            # (e.g., /kaggle/input/mabe-tenten-models-exp096-base)\n            if not standard_config_path.exists():\n                # Construct Kaggle-style path: /kaggle/input/mabe-tenten-models-{exp_name}-{sub_name}\n                kaggle_exp_path = (\n                    Path(base_output_dir) / f\"mabe-tenten-models-{exp_name}-{sub_name}\"\n                )\n                if (\n                    kaggle_exp_path.exists()\n                    and (kaggle_exp_path / \"config.yaml\").exists()\n                ):\n                    exp_path = kaggle_exp_path\n                    print(f\"  Using Kaggle-style path: {exp_path}\")\n                else:\n                    # Fallback to base_output_dir directly (old behavior)\n                    exp_path = Path(base_output_dir)\n                    print(f\"  Using base output directory: {exp_path}\")\n\n            config_path = exp_path / \"config.yaml\"\n            model_path = exp_path / f\"best_score_model_{fold}.pt\"\n\n            if not config_path.exists():\n                raise FileNotFoundError(f\"Config not found: {config_path}\")\n            if not model_path.exists():\n                raise FileNotFoundError(f\"Model not found: {model_path}\")\n\n            # Load config\n            cfg = OmegaConf.load(config_path)\n            self.configs.append(cfg)\n\n            # Calculate model parameters from config\n            bodyparts = list(cfg.exp.dataset.bodyparts)\n            features_per_mouse = len(bodyparts) * 2\n            input_features = features_per_mouse * 2\n\n            if cfg.exp.mode == \"self\":\n                num_classes = len(list(cfg.exp.dataset.actions_self))\n            else:\n                num_classes = len(list(cfg.exp.dataset.actions_pair))\n\n            num_labs = len(dict(cfg.exp.dataset.lab_id_to_index))\n\n            # Load MyModel class dynamically from experiment's run.py\n            # Pass exp_path so it can find run.py in the same directory as the model (for Kaggle)\n            MyModelClass = load_model_class_from_exp(\n                exp_name, base_output_dir, model_dir=exp_path\n            )\n\n            # Create model with this config\n            model = MyModelClass(\n                cfg,\n                num_classes=num_classes,\n                input_features=input_features,\n                num_labs=num_labs,\n            )\n\n            # Load state dict\n            state_dict = torch.load(model_path, map_location=\"cpu\")\n            # Use strict=False to handle potential architecture mismatches\n            missing_keys, unexpected_keys = model.load_state_dict(\n                state_dict, strict=False\n            )\n            if missing_keys:\n                print(\n                    f\"    Warning: Missing keys: {missing_keys[:3]}{'...' if len(missing_keys) > 3 else ''}\"\n                )\n            if unexpected_keys:\n                print(\n                    f\"    Warning: Unexpected keys: {unexpected_keys[:3]}{'...' if len(unexpected_keys) > 3 else ''}\"\n                )\n\n            # Move to device and set eval mode\n            model = model.to(device)\n            model.eval()\n\n            self.models.append(model)\n\n            print(\n                f\"  Loaded: {exp_name}/{sub_name}/fold_{fold} \"\n                f\"(classes={num_classes}, labs={num_labs})\"\n            )\n\n        # Save first config as representative config for data loading\n        self.cfg = self.configs[0]\n        print(\"=\" * 80)\n        print(\"Using config from first experiment as representative config\\n\")\n\n    def predict(self, x, candidate_labels, lab_id, return_individual=False):\n        \"\"\"\n        Predict using ensemble of models.\n\n        Args:\n            x: Input features\n            candidate_labels: Candidate labels\n            lab_id: Lab ID\n            return_individual: If True, return individual model predictions along with ensemble\n\n        Returns:\n            If return_individual=False: Ensemble probabilities (mean or max)\n            If return_individual=True: (ensemble_probs, list of individual model predictions)\n        \"\"\"\n        ensemble_logits = []\n        for model in self.models:\n            logits = model(x, candidate_labels, lab_id)\n            ensemble_logits.append((torch.sigmoid(logits)))\n\n        # Ensemble predictions based on method\n        if self.ensemble_method == \"mean\":\n            # Weighted average across all models\n            ensemble_result = sum(w * p for w, p in zip(self.weights, ensemble_logits))\n        elif self.ensemble_method == \"max\":\n            # Max across all models\n            ensemble_result = torch.stack(ensemble_logits, dim=0).max(dim=0)[0]\n        else:\n            raise ValueError(f\"Unknown ensemble_method: {self.ensemble_method}\")\n\n        if return_individual:\n            return ensemble_result, ensemble_logits\n        return ensemble_result\n\n\ndef _visualize_model_distributions(\n    model_predictions_list: List[np.ndarray],\n    exp_configs: List[Tuple[str, str, int]],\n    weights: List[float],\n    output_path: Path,\n    fold: int,\n):\n    \"\"\"\n    Visualize prediction probability distributions for each model with overlapping histograms.\n\n    Args:\n        model_predictions_list: List of prediction arrays from each model\n        exp_configs: List of (exp_name, sub_name, fold) tuples\n        weights: List of weights for each model\n        output_path: Output directory\n        fold: Fold number\n    \"\"\"\n    vis_dir = output_path / \"visualizations\"\n    vis_dir.mkdir(parents=True, exist_ok=True)\n\n    # Flatten all predictions\n    flattened_preds = [preds.flatten() for preds in model_predictions_list]\n\n    # Filter predictions >= 0.1\n    filtered_preds = [preds[preds >= 0.1] for preds in flattened_preds]\n\n    # Create figure\n    fig, ax = plt.subplots(figsize=(12, 6))\n\n    # Define colors for each model\n    colors = plt.cm.tab10(np.linspace(0, 1, len(exp_configs)))\n\n    # Plot overlapping histograms with higher transparency\n    for i, (preds, (exp_name, sub_name, _), weight, color) in enumerate(\n        zip(filtered_preds, exp_configs, weights, colors)\n    ):\n        if len(preds) > 0:\n            label = f\"{exp_name}/{sub_name} (weight={weight:.2f}, n={len(preds)})\"\n            ax.hist(\n                preds,\n                bins=100,\n                alpha=0.3,\n                color=color,\n                label=label,\n                density=True,\n                range=(0.1, 1),\n            )\n\n    ax.set_xlabel(\"Prediction Probability\", fontsize=12)\n    ax.set_ylabel(\"Density\", fontsize=12)\n    ax.set_title(\n        f\"Model Prediction Probability Distributions >= 0.1 (Fold {fold})\",\n        fontsize=14,\n    )\n    ax.legend(loc=\"best\")\n    ax.grid(True, alpha=0.3)\n\n    # Save figure\n    save_path = vis_dir / f\"model_distribution_fold_{fold}.png\"\n    plt.tight_layout()\n    plt.savefig(save_path, dpi=150, bbox_inches=\"tight\")\n    plt.close()\n\n\ndef _run_ensemble_inference_single_fold(\n    exp_configs: List[Tuple[str, str, int]],\n    fold: int,\n    output_path: Path,\n    is_test: bool = False,\n    base_output_dir: str = \"output\",\n    weights: List[float] = None,\n    logger: WandbLogger = None,\n    ensemble_method: str = \"mean\",\n    min_frames: int = 0,\n) -> float:\n    \"\"\"\n    Run ensemble inference for a single fold.\n\n    Args:\n        exp_configs: List of (exp_name, sub_name, fold) tuples\n        fold: Fold number\n        output_path: Output directory\n        is_test: Whether to run on test data\n        base_output_dir: Base directory for outputs\n        weights: Optional list of weights for each model\n        logger: WandbLogger for logging metrics\n        ensemble_method: Ensemble method (\"mean\" or \"max\")\n\n    Returns:\n        Validation score (0.0 for test)\n    \"\"\"\n    print(f\"\\n{'=' * 80}\")\n    print(\n        f\"Running {'test' if is_test else 'validation'} ensemble inference for fold {fold}\"\n    )\n    print(f\"{'=' * 80}\\n\")\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    # Create ensemble model\n    ensemble_model = EnsembleModel(\n        exp_configs, device, base_output_dir, weights, ensemble_method\n    )\n    cfg = ensemble_model.cfg\n\n    # Override val_stride for ensemble inference\n    cfg.exp.dataset.val_stride = 128\n\n    # Load data\n    if is_test:\n        df_ann, df_test = load_test_data(cfg)\n        val_video_ids = None\n        df_samples = df_ann.unique(subset=[\"video_id\", \"agent_id\", \"target_id\"])\n    else:\n        dm = MyLightningDataModule(cfg, fold=fold)\n        df_samples = dm.val_dataset.df_ann\n        df_ann = dm.df_ann\n        val_video_ids = dm.val_video_ids\n\n    if len(df_ann) == 0:\n        print(\"No samples found!\")\n        return 0.0\n\n    # Get action and lab mappings\n    unique_actions, action_to_index = load_actions_from_config(cfg)\n    index_to_action = {v: k for k, v in action_to_index.items()}\n    num_classes = len(action_to_index)\n\n    lab_id_to_index = dict(cfg.exp.dataset.lab_id_to_index)\n    index_to_lab_id = {v: k for k, v in lab_id_to_index.items()}\n\n    missing_actions = OmegaConf.to_container(\n        cfg.exp.dataset.missing_actions, resolve=True\n    )\n\n    print(f\"Number of action classes: {num_classes}\")\n    print(f\"Actions: {unique_actions}\")\n    print(f\"Number of samples: {len(df_samples)}\")\n\n    # Load distance/velocity/acceleration statistics\n    dataset_base = Path(cfg.dir.dataset_dir) / cfg.exp.dataset.name\n    stats_path = dataset_base / \"distance_velocity_acceleration_statistics.json\"\n    if stats_path.exists():\n        with open(stats_path, \"r\") as f:\n            distance_velocity_acceleration_stats = json.load(f)\n    else:\n        distance_velocity_acceleration_stats = None\n\n    # Process in chunks\n    chunk_size = cfg.exp.inference.chunk_size\n    num_chunks = (len(df_samples) + chunk_size - 1) // chunk_size\n    sample_chunks = [\n        df_samples.slice(\n            i * chunk_size, min(chunk_size, len(df_samples) - i * chunk_size)\n        )\n        for i in range(num_chunks)\n    ]\n\n    print(f\"Processing {len(sample_chunks)} chunks (chunk_size={chunk_size})\")\n\n    all_sample_predictions = {}\n    all_sample_metadata = {}\n\n    # For visualization: collect predictions from each model\n    model_predictions_for_vis = (\n        [[] for _ in range(len(exp_configs))] if not is_test else None\n    )\n\n    for chunk_idx, chunk_samples in enumerate(\n        tqdm(sample_chunks, desc=\"Processing sample chunks\")\n    ):\n        chunk_dataset = SlidingWindowDataset(\n            cfg=cfg,\n            mode=\"test\",\n            df_ann=chunk_samples,\n            distance_velocity_acceleration_stats=distance_velocity_acceleration_stats,\n        )\n        chunk_loader = DataLoader(\n            chunk_dataset,\n            batch_size=cfg.exp.dataset.val_batch_size,\n            num_workers=cfg.exp.dataset.num_workers,\n            shuffle=False,\n            pin_memory=True,\n            drop_last=False,\n        )\n\n        # Initialize prediction storage for this chunk\n        chunk_predictions = {}\n        chunk_counts = {}\n        chunk_metadata = {}\n\n        for row in chunk_samples.iter_rows(named=True):\n            key = (int(row[\"video_id\"]), int(row[\"agent_id\"]), int(row[\"target_id\"]))\n            chunk_predictions[key] = np.zeros(\n                (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n            )\n            chunk_counts[key] = np.zeros(\n                (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n            )\n            chunk_metadata[key] = {\n                \"num_frames\": int(row[\"num_frames\"]),\n                \"video_frame_min\": int(row.get(\"video_frame_min\", 0)),\n            }\n\n        edge_removal = int(cfg.exp.inference.get(\"edge_frame_removal\", 0))\n\n        # Run inference\n        with trace(\"inference\"):\n            with torch.no_grad():\n                for batch in chunk_loader:\n                    x = batch[\"x\"].to(device)\n                    candidate_labels = batch[\"candidate_labels\"].to(device)\n                    lab_id = batch[\"lab_id\"].to(device)\n\n                    # Ensemble prediction\n                    if not is_test:\n                        # For validation: collect individual model predictions\n                        probs, individual_probs = ensemble_model.predict(\n                            x, candidate_labels, lab_id, return_individual=True\n                        )\n                        probs = probs.cpu().numpy()\n                        individual_probs_np = [\n                            p.cpu().numpy() for p in individual_probs\n                        ]\n\n                        # Collect predictions for visualization\n                        for model_idx, model_preds in enumerate(individual_probs_np):\n                            model_predictions_for_vis[model_idx].append(model_preds)\n                    else:\n                        # For test: just get ensemble predictions\n                        probs = ensemble_model.predict(x, candidate_labels, lab_id)\n                        probs = probs.cpu().numpy()\n\n                    meta_dict = batch[\"meta\"]\n                    batch_size = len(meta_dict[\"video_id\"])\n\n                    for i in range(batch_size):\n                        video_id = meta_dict[\"video_id\"][i].item()\n                        agent_id = meta_dict[\"agent_id\"][i].item()\n                        target_id = meta_dict[\"target_id\"][i].item()\n                        start_frame = meta_dict[\"start_frame\"][i].item()\n                        pred_window = probs[i]\n\n                        key = (video_id, agent_id, target_id)\n                        sample_preds = chunk_predictions[key]\n                        sample_cnts = chunk_counts[key]\n                        metadata = chunk_metadata[key]\n\n                        num_frames = metadata[\"num_frames\"]\n                        video_frame_min = metadata[\"video_frame_min\"]\n\n                        sample_preds, sample_cnts = _accumulate_predictions(\n                            sample_preds,\n                            sample_cnts,\n                            pred_window,\n                            start_frame,\n                            num_frames,\n                            video_frame_min,\n                            edge_removal,\n                        )\n\n        # Process chunk predictions\n        for row in chunk_samples.iter_rows(named=True):\n            key = (int(row[\"video_id\"]), int(row[\"agent_id\"]), int(row[\"target_id\"]))\n            avg_predictions = chunk_predictions[key]\n            counts = chunk_counts[key]\n\n            # Average predictions\n            valid_mask = counts > 0\n            avg_predictions[valid_mask] /= counts[valid_mask]\n\n            # Apply missing actions mask\n            lab_id_idx = int(row.get(\"lab_id_index\", 0))\n            lab_id = index_to_lab_id.get(lab_id_idx)\n            if lab_id in missing_actions:\n                for action in missing_actions[lab_id]:\n                    action_idx = action_to_index.get(action)\n                    if action_idx is not None:\n                        avg_predictions[:, action_idx] = 0.0\n\n            # Interpolate to original FPS\n            video_frame_min = int(row.get(\"video_frame_min\", 0))\n            avg_predictions = _interpolate_predictions_to_original_fps(\n                avg_predictions, float(row[\"fps_scale\"]), video_frame_min\n            )\n\n            # Apply smoothing if configured\n            smooth_window_size = int(cfg.exp.inference.get(\"smooth_window_size\", 0))\n            if smooth_window_size > 1:\n                avg_predictions = _smooth_predictions(\n                    avg_predictions, smooth_window_size\n                )\n\n            candidate_labels = np.array(\n                json.loads(row[\"candidate_labels\"]), dtype=np.float32\n            )\n\n            # Store for threshold search\n            if not is_test:\n                all_sample_predictions[key] = avg_predictions.copy()\n                all_sample_metadata[key] = {\n                    \"video_id\": int(row[\"video_id\"]),\n                    \"agent_id\": int(row[\"agent_id\"]),\n                    \"target_id\": int(row[\"target_id\"]),\n                    \"candidate_labels\": candidate_labels,\n                    \"lab_id_index\": int(row.get(\"lab_id_index\", 0)),\n                }\n\n        del chunk_predictions\n        del chunk_counts\n        del chunk_dataset\n        del chunk_loader\n\n    # For validation, run threshold search and calculate score\n    val_score = 0.0\n    if not is_test:\n        # Visualize model prediction distributions\n        print(\"\\n\" + \"=\" * 80)\n        print(\"Visualizing model prediction distributions\")\n        print(\"=\" * 80)\n\n        # # Concatenate all predictions for each model\n        # model_predictions_concat = [\n        #     np.concatenate(preds, axis=0) for preds in model_predictions_for_vis\n        # ]\n\n        # _visualize_model_distributions(\n        #     model_predictions_list=model_predictions_concat,\n        #     exp_configs=exp_configs,\n        #     weights=weights if weights else [1.0 / len(exp_configs)] * len(exp_configs),\n        #     output_path=output_path,\n        #     fold=fold,\n        # )\n\n        print(\"\\n\" + \"=\" * 80)\n        print(\"Searching for optimal thresholds per lab and action\")\n        print(\"=\" * 80)\n\n        train_csv_path = (\n            Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/train_fixed.csv\"\n        )\n        annotation_path = (\n            Path(cfg.dir.input_dir)\n            / \"MABe-mouse-behavior-detection\"\n            / \"all_annotations_fixed.parquet\"\n        )\n        threshold_search_results = _search_optimal_thresholds_per_lab_and_action(\n            all_sample_predictions=all_sample_predictions,\n            all_sample_metadata=all_sample_metadata,\n            val_video_ids=val_video_ids,\n            train_csv_path=train_csv_path,\n            annotation_path=annotation_path,\n            index_to_action=index_to_action,\n            default_threshold=float(cfg.exp.pred_threshold),\n            threshold_step=0.05,\n        )\n\n        # Save threshold search results\n        threshold_results_path = (\n            output_path / f\"threshold_search_ensemble_fold_{fold}.json\"\n        )\n        with open(threshold_results_path, \"w\") as f:\n            json.dump(threshold_search_results, f, indent=2)\n        print(f\"\\nThreshold search results saved to: {threshold_results_path}\")\n\n        # Extract best thresholds\n        lab_action_thresholds = threshold_search_results[\"lab_action_best_thresholds\"]\n\n        # Create submission with best thresholds\n        submission = _extract_segments_with_lab_action_thresholds(\n            all_sample_predictions=all_sample_predictions,\n            all_sample_metadata=all_sample_metadata,\n            lab_action_thresholds=lab_action_thresholds,\n            default_threshold=float(cfg.exp.pred_threshold),\n            index_to_action=index_to_action,\n            index_to_lab_id=index_to_lab_id,\n        )\n\n        # Filter out predictions with frame length <= min_frames\n        if min_frames > 0:\n            original_len = len(submission)\n            submission[\"frame_length\"] = (\n                submission[\"stop_frame\"] - submission[\"start_frame\"] + 1\n            )\n            submission = submission[submission[\"frame_length\"] > min_frames].drop(\n                columns=[\"frame_length\"]\n            )\n            filtered_len = len(submission)\n            print(\n                f\"Filtered predictions: {original_len} -> {filtered_len} \"\n                f\"(removed {original_len - filtered_len} predictions with <={min_frames} frames)\"\n            )\n\n        # Adjust frames for AdaptableSnail videos with fps ~25.0\n        submission = _adjust_frames_for_adaptable_snail(submission, train_csv_path)\n\n        submission.insert(0, \"row_id\", range(len(submission)))\n\n        # Save submission\n        submission_path = output_path / f\"val_submission_ensemble_fold_{fold}.csv\"\n        submission.to_csv(submission_path, index=False)\n        print(f\"Submission saved to: {submission_path}\")\n\n        # Calculate score\n        solution = prepare_solution_for_scoring(\n            val_video_ids,\n            train_csv_path,\n            annotation_path,\n        )\n\n        val_score = mabe_score(\n            solution,\n            submission,\n            row_id_column_name=\"row_id\",\n            beta=1,\n        )\n        print(f\"\\nValidation Score with best thresholds: {val_score:.6f}\")\n\n        # Per-lab scores\n        print(\"\\nPer-lab validation scores:\")\n        lab_scores = {\"val_inference_score\": float(val_score)}\n\n        for lab_id in solution[\"lab_id\"].unique():\n            lab_solution = solution[solution[\"lab_id\"] == lab_id].copy()\n            lab_submission = submission[\n                submission[\"video_id\"].isin(lab_solution[\"video_id\"].unique())\n            ].copy()\n\n            if len(lab_submission) > 0:\n                lab_submission[\"row_id\"] = range(len(lab_submission))\n                lab_score = mabe_score(\n                    lab_solution,\n                    lab_submission,\n                    row_id_column_name=\"row_id\",\n                    beta=1,\n                )\n            else:\n                lab_score = 0.0\n\n            print(f\"  {lab_id}: {lab_score:.6f}\")\n            lab_scores[f\"val_inference_score/{lab_id}\"] = float(lab_score)\n\n            # Per-action scores within each lab\n            for action in lab_solution[\"action\"].unique():\n                action_solution = lab_solution[lab_solution[\"action\"] == action].copy()\n                action_submission = lab_submission[\n                    lab_submission[\"action\"] == action\n                ].copy()\n\n                if len(action_submission) > 0 and len(action_solution) > 0:\n                    action_submission[\"row_id\"] = range(len(action_submission))\n                    action_score = mabe_score(\n                        action_solution,\n                        action_submission,\n                        row_id_column_name=\"row_id\",\n                        beta=1,\n                    )\n                else:\n                    action_score = 0.0\n\n                print(f\"    {lab_id}/{action}: {action_score:.6f}\")\n                lab_scores[f\"val_inference_score/{lab_id}/{action}\"] = float(\n                    action_score\n                )\n\n        # Log metrics to wandb if logger is provided\n        if logger is not None:\n            logger.log_metrics(lab_scores)\n\n    return val_score\n\n\ndef _search_optimal_thresholds_per_lab_and_action_sorted(\n    all_sample_predictions: dict,\n    all_sample_metadata: dict,\n    val_video_ids: List[int],\n    train_csv_path: Path,\n    annotation_path: Path,\n    index_to_action: dict,\n    default_threshold: float,\n) -> dict:\n    \"\"\"\n    Search optimal thresholds per lab and action using sorted predictions.\n    For each lab-action pair, sort predictions and find the threshold that maximizes F1.\n\n    Args:\n        all_sample_predictions: Dictionary of predictions {(video_id, agent_id, target_id): predictions}\n        all_sample_metadata: Dictionary of metadata {(video_id, agent_id, target_id): metadata}\n        val_video_ids: List of validation video IDs\n        train_csv_path: Path to train CSV\n        annotation_path: Path to annotations\n        index_to_action: Mapping from action index to action name\n        default_threshold: Default threshold value\n\n    Returns:\n        Dictionary containing threshold search results\n    \"\"\"\n    print(\"\\nSearching optimal thresholds using sorted predictions...\")\n\n    # Load ground truth\n    train_df = pd.read_csv(train_csv_path)\n    annotation_df = pd.read_parquet(annotation_path)\n\n    # Convert video_id to int for matching\n    annotation_df[\"video_id\"] = annotation_df[\"video_id\"].astype(int)\n\n    # Filter annotations for validation video IDs only\n    annotation_df = annotation_df[annotation_df[\"video_id\"].isin(val_video_ids)]\n    print(\n        f\"Filtered annotations to {len(annotation_df)} rows for {len(val_video_ids)} validation videos\"\n    )\n\n    # Prepare video_id to lab_id mapping\n    video_id_to_lab_id = train_df.set_index(\"video_id\")[\"lab_id\"].to_dict()\n\n    # Group predictions by lab_id and action\n    lab_action_predictions = {}\n    lab_action_ground_truth = {}\n\n    for key, predictions in tqdm(\n        all_sample_predictions.items(), desc=\"Grouping predictions by lab-action\"\n    ):\n        video_id, agent_id, target_id = key\n\n        # Get lab_id from video_id\n        lab_id = video_id_to_lab_id.get(video_id)\n        if lab_id is None:\n            continue\n\n        # Get ground truth for this sample\n        gt_df = annotation_df[\n            (annotation_df[\"video_id\"] == video_id)\n            & (annotation_df[\"agent_id\"] == agent_id)\n            & (annotation_df[\"target_id\"] == target_id)\n        ]\n\n        num_frames = predictions.shape[0]\n\n        for action_idx, action_name in index_to_action.items():\n            if lab_id not in lab_action_predictions:\n                lab_action_predictions[lab_id] = {}\n                lab_action_ground_truth[lab_id] = {}\n            if action_name not in lab_action_predictions[lab_id]:\n                lab_action_predictions[lab_id][action_name] = []\n                lab_action_ground_truth[lab_id][action_name] = []\n\n            # Get predictions for this action\n            action_preds = predictions[:, action_idx]\n\n            # Get ground truth for this action\n            action_gt = np.zeros(num_frames, dtype=bool)\n            action_gt_df = gt_df[gt_df[\"action\"] == action_name]\n            for _, row in action_gt_df.iterrows():\n                start = int(row[\"start_frame\"])\n                end = int(row[\"stop_frame\"])\n                action_gt[start : end + 1] = True\n\n            # Store predictions and ground truth\n            lab_action_predictions[lab_id][action_name].append(action_preds)\n            lab_action_ground_truth[lab_id][action_name].append(action_gt)\n\n    # Search optimal threshold for each lab-action\n    lab_action_best_thresholds = {}\n\n    print(\"\\nSearching optimal thresholds per lab-action:\")\n    for lab_id in sorted(lab_action_predictions.keys()):\n        lab_action_best_thresholds[lab_id] = {}\n        print(f\"\\nLab {lab_id}:\")\n\n        for action_name in sorted(lab_action_predictions[lab_id].keys()):\n            # Concatenate all predictions and ground truth for this lab-action\n            all_preds = np.concatenate(lab_action_predictions[lab_id][action_name])\n            all_gt = np.concatenate(lab_action_ground_truth[lab_id][action_name])\n\n            # Check if there are any positive ground truth samples\n            num_positive_gt = np.sum(all_gt)\n            if num_positive_gt == 0:\n                # No positive samples for this action, skip without recording\n                print(f\"  {action_name}: No positive GT samples, skipping\")\n                continue\n\n            # Sort predictions to get unique threshold candidates\n            unique_preds = np.unique(all_preds)\n            # Include 0 as a candidate and filter very small values\n            unique_preds = unique_preds[unique_preds >= 1e-6]\n\n            # If no predictions > 1e-6, use a small range around 0\n            if len(unique_preds) == 0:\n                unique_preds = np.array([1e-6, 0.01, 0.05, 0.1])\n\n            # Sample threshold candidates if too many\n            max_candidates = 1000\n            if len(unique_preds) > max_candidates:\n                # Sample uniformly + include min/max\n                indices = np.linspace(\n                    0, len(unique_preds) - 1, max_candidates - 2, dtype=int\n                )\n                sampled_preds = unique_preds[indices]\n                unique_preds = np.concatenate(\n                    [[unique_preds[0]], sampled_preds, [unique_preds[-1]]]\n                )\n                unique_preds = np.unique(unique_preds)\n\n            best_threshold = default_threshold\n            best_f1 = 0.0\n\n            # Test each threshold candidate\n            for threshold in unique_preds:\n                pred_binary = all_preds >= threshold\n\n                tp = np.sum(pred_binary & all_gt)\n                fp = np.sum(pred_binary & ~all_gt)\n                fn = np.sum(~pred_binary & all_gt)\n\n                precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0\n                recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0\n                f1 = (\n                    2 * precision * recall / (precision + recall)\n                    if (precision + recall) > 0\n                    else 0.0\n                )\n\n                if f1 > best_f1:\n                    best_f1 = f1\n                    best_threshold = threshold\n\n            lab_action_best_thresholds[lab_id][action_name] = {\n                \"threshold\": float(best_threshold),\n                \"f1_score\": float(best_f1),\n            }\n\n            # Add debug info\n            max_pred = np.max(all_preds)\n            min_pred = np.min(all_preds)\n            num_candidates = len(unique_preds)\n\n            if best_f1 == 0.0:\n                print(\n                    f\"  {action_name}: threshold={best_threshold:.4f}, F1={best_f1:.4f} \"\n                    f\"[WARNING: pred range=[{min_pred:.4f}, {max_pred:.4f}], candidates={num_candidates}, GT_positive={num_positive_gt}]\"\n                )\n            else:\n                print(\n                    f\"  {action_name}: threshold={best_threshold:.4f}, F1={best_f1:.4f}\"\n                )\n\n    results = {\n        \"default_threshold\": default_threshold,\n        \"lab_action_best_thresholds\": lab_action_best_thresholds,\n    }\n\n    return results\n\n\ndef _run_ensemble_inference_oof(\n    exp_configs: List[Tuple[str, str]],\n    folds: List[int],\n    output_path: Path,\n    base_output_dir: str = \"output\",\n    weights: List[float] = None,\n    logger: WandbLogger = None,\n    ensemble_method: str = \"mean\",\n    min_frames: int = 0,\n) -> float:\n    \"\"\"\n    Run ensemble inference for all folds to create OOF predictions and find optimal thresholds.\n\n    Args:\n        exp_configs: List of (exp_name, sub_name) tuples\n        folds: List of fold numbers\n        output_path: Output directory\n        base_output_dir: Base directory for outputs\n        weights: Optional list of weights for each model\n        logger: WandbLogger for logging metrics\n        ensemble_method: Ensemble method (\"mean\" or \"max\")\n\n    Returns:\n        OOF validation score\n    \"\"\"\n    print(f\"\\n{'=' * 80}\")\n    print(\"Running OOF ensemble inference for all folds\")\n    print(f\"{'=' * 80}\\n\")\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    # Load config from first experiment\n    first_exp, first_sub = exp_configs[0]\n    first_cfg_path = Path(base_output_dir) / first_exp / first_sub / \"config.yaml\"\n\n    if not first_cfg_path.exists():\n        first_cfg_path = Path(base_output_dir) / \"config.yaml\"\n\n    cfg = OmegaConf.load(first_cfg_path)\n\n    # Override val_stride for ensemble inference\n    cfg.exp.dataset.val_stride = 128\n\n    # Get action and lab mappings\n    unique_actions, action_to_index = load_actions_from_config(cfg)\n    index_to_action = {v: k for k, v in action_to_index.items()}\n    num_classes = len(action_to_index)\n\n    lab_id_to_index = dict(cfg.exp.dataset.lab_id_to_index)\n    index_to_lab_id = {v: k for k, v in lab_id_to_index.items()}\n\n    missing_actions = OmegaConf.to_container(\n        cfg.exp.dataset.missing_actions, resolve=True\n    )\n\n    # Load distance/velocity/acceleration statistics\n    dataset_base = Path(cfg.dir.dataset_dir) / cfg.exp.dataset.name\n    stats_path = dataset_base / \"distance_velocity_acceleration_statistics.json\"\n    if stats_path.exists():\n        with open(stats_path, \"r\") as f:\n            distance_velocity_acceleration_stats = json.load(f)\n    else:\n        distance_velocity_acceleration_stats = None\n\n    # Collect OOF predictions from all folds\n    oof_predictions = {}\n    oof_metadata = {}\n    all_val_video_ids = []\n\n    for fold in folds:\n        print(f\"\\n{'=' * 80}\")\n        print(f\"Processing fold {fold}\")\n        print(f\"{'=' * 80}\\n\")\n\n        # Create ensemble model for this fold\n        fold_exp_configs = [(exp, sub, fold) for exp, sub in exp_configs]\n        ensemble_model = EnsembleModel(\n            fold_exp_configs, device, base_output_dir, weights, ensemble_method\n        )\n\n        # Load validation data for this fold\n        dm = MyLightningDataModule(cfg, fold=fold)\n        df_samples = dm.val_dataset.df_ann\n        val_video_ids = dm.val_video_ids\n        all_val_video_ids.extend(val_video_ids)\n\n        print(f\"Number of validation samples in fold {fold}: {len(df_samples)}\")\n\n        # Process in chunks\n        chunk_size = cfg.exp.inference.chunk_size\n        num_chunks = (len(df_samples) + chunk_size - 1) // chunk_size\n        sample_chunks = [\n            df_samples.slice(\n                i * chunk_size, min(chunk_size, len(df_samples) - i * chunk_size)\n            )\n            for i in range(num_chunks)\n        ]\n\n        for chunk_idx, chunk_samples in enumerate(\n            tqdm(sample_chunks, desc=f\"Processing fold {fold} chunks\")\n        ):\n            chunk_dataset = SlidingWindowDataset(\n                cfg=cfg,\n                mode=\"test\",\n                df_ann=chunk_samples,\n                distance_velocity_acceleration_stats=distance_velocity_acceleration_stats,\n            )\n            chunk_loader = DataLoader(\n                chunk_dataset,\n                batch_size=cfg.exp.dataset.val_batch_size,\n                num_workers=cfg.exp.dataset.num_workers,\n                shuffle=False,\n                pin_memory=True,\n                drop_last=False,\n            )\n\n            chunk_predictions = {}\n            chunk_counts = {}\n            chunk_metadata = {}\n\n            for row in chunk_samples.iter_rows(named=True):\n                key = (\n                    int(row[\"video_id\"]),\n                    int(row[\"agent_id\"]),\n                    int(row[\"target_id\"]),\n                )\n                chunk_predictions[key] = np.zeros(\n                    (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                )\n                chunk_counts[key] = np.zeros(\n                    (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                )\n                chunk_metadata[key] = {\n                    \"num_frames\": int(row[\"num_frames\"]),\n                    \"video_frame_min\": int(row.get(\"video_frame_min\", 0)),\n                }\n\n            edge_removal = int(cfg.exp.inference.get(\"edge_frame_removal\", 0))\n\n            # Run inference\n            with torch.no_grad():\n                for batch in chunk_loader:\n                    x = batch[\"x\"].to(device)\n                    candidate_labels = batch[\"candidate_labels\"].to(device)\n                    lab_id = batch[\"lab_id\"].to(device)\n\n                    probs = ensemble_model.predict(x, candidate_labels, lab_id)\n                    probs = probs.cpu().numpy()\n\n                    meta_dict = batch[\"meta\"]\n                    batch_size = len(meta_dict[\"video_id\"])\n\n                    for i in range(batch_size):\n                        video_id = meta_dict[\"video_id\"][i].item()\n                        agent_id = meta_dict[\"agent_id\"][i].item()\n                        target_id = meta_dict[\"target_id\"][i].item()\n                        start_frame = meta_dict[\"start_frame\"][i].item()\n                        pred_window = probs[i]\n\n                        key = (video_id, agent_id, target_id)\n                        sample_preds = chunk_predictions[key]\n                        sample_cnts = chunk_counts[key]\n                        metadata = chunk_metadata[key]\n\n                        num_frames = metadata[\"num_frames\"]\n                        video_frame_min = metadata[\"video_frame_min\"]\n\n                        sample_preds, sample_cnts = _accumulate_predictions(\n                            sample_preds,\n                            sample_cnts,\n                            pred_window,\n                            start_frame,\n                            num_frames,\n                            video_frame_min,\n                            edge_removal,\n                        )\n\n            # Process chunk predictions\n            for row in chunk_samples.iter_rows(named=True):\n                key = (\n                    int(row[\"video_id\"]),\n                    int(row[\"agent_id\"]),\n                    int(row[\"target_id\"]),\n                )\n                avg_predictions = chunk_predictions[key]\n                counts = chunk_counts[key]\n\n                # Average predictions\n                valid_mask = counts > 0\n                avg_predictions[valid_mask] /= counts[valid_mask]\n\n                # Apply missing actions mask\n                lab_id_idx = int(row.get(\"lab_id_index\", 0))\n                lab_id = index_to_lab_id.get(lab_id_idx)\n                if lab_id in missing_actions:\n                    for action in missing_actions[lab_id]:\n                        action_idx = action_to_index.get(action)\n                        if action_idx is not None:\n                            avg_predictions[:, action_idx] = 0.0\n\n                # Interpolate to original FPS\n                video_frame_min = int(row.get(\"video_frame_min\", 0))\n                avg_predictions = _interpolate_predictions_to_original_fps(\n                    avg_predictions, float(row[\"fps_scale\"]), video_frame_min\n                )\n\n                # Apply smoothing if configured\n                smooth_window_size = int(cfg.exp.inference.get(\"smooth_window_size\", 0))\n                if smooth_window_size > 1:\n                    avg_predictions = _smooth_predictions(\n                        avg_predictions, smooth_window_size\n                    )\n\n                candidate_labels = np.array(\n                    json.loads(row[\"candidate_labels\"]), dtype=np.float32\n                )\n\n                # Store OOF predictions\n                oof_predictions[key] = avg_predictions.copy()\n                oof_metadata[key] = {\n                    \"video_id\": int(row[\"video_id\"]),\n                    \"agent_id\": int(row[\"agent_id\"]),\n                    \"target_id\": int(row[\"target_id\"]),\n                    \"candidate_labels\": candidate_labels,\n                    \"lab_id_index\": int(row.get(\"lab_id_index\", 0)),\n                }\n\n            del chunk_predictions\n            del chunk_counts\n            del chunk_dataset\n            del chunk_loader\n\n        del ensemble_model\n\n    # Search optimal thresholds on entire OOF\n    print(\"\\n\" + \"=\" * 80)\n    print(\"Searching optimal thresholds on entire OOF\")\n    print(\"=\" * 80)\n\n    train_csv_path = (\n        Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/train_fixed.csv\"\n    )\n    annotation_path = (\n        Path(cfg.dir.input_dir)\n        / \"MABe-mouse-behavior-detection\"\n        / \"all_annotations_fixed.parquet\"\n    )\n\n    threshold_search_results = _search_optimal_thresholds_per_lab_and_action_sorted(\n        all_sample_predictions=oof_predictions,\n        all_sample_metadata=oof_metadata,\n        val_video_ids=all_val_video_ids,\n        train_csv_path=train_csv_path,\n        annotation_path=annotation_path,\n        index_to_action=index_to_action,\n        default_threshold=float(cfg.exp.pred_threshold),\n    )\n\n    # Save threshold search results\n    threshold_results_path = output_path / \"threshold_search_oof.json\"\n    with open(threshold_results_path, \"w\") as f:\n        json.dump(threshold_search_results, f, indent=2)\n    print(f\"\\nThreshold search results saved to: {threshold_results_path}\")\n\n    # Extract best thresholds\n    lab_action_thresholds = threshold_search_results[\"lab_action_best_thresholds\"]\n\n    # Create submission with best thresholds\n    submission = _extract_segments_with_lab_action_thresholds(\n        all_sample_predictions=oof_predictions,\n        all_sample_metadata=oof_metadata,\n        lab_action_thresholds=lab_action_thresholds,\n        default_threshold=float(cfg.exp.pred_threshold),\n        index_to_action=index_to_action,\n        index_to_lab_id=index_to_lab_id,\n    )\n\n    # Filter out predictions with frame length <= min_frames\n    if min_frames > 0:\n        original_len = len(submission)\n        submission[\"frame_length\"] = (\n            submission[\"stop_frame\"] - submission[\"start_frame\"] + 1\n        )\n        submission = submission[submission[\"frame_length\"] > min_frames].drop(\n            columns=[\"frame_length\"]\n        )\n        filtered_len = len(submission)\n        print(\n            f\"Filtered predictions: {original_len} -> {filtered_len} \"\n            f\"(removed {original_len - filtered_len} predictions with <={min_frames} frames)\"\n        )\n\n    # Adjust frames for AdaptableSnail videos with fps ~25.0\n    submission = _adjust_frames_for_adaptable_snail(submission, train_csv_path)\n\n    submission.insert(0, \"row_id\", range(len(submission)))\n\n    # Save submission\n    submission_path = output_path / \"oof_submission.csv\"\n    submission.to_csv(submission_path, index=False)\n    print(f\"OOF submission saved to: {submission_path}\")\n\n    # Calculate OOF score\n    solution = prepare_solution_for_scoring(\n        all_val_video_ids,\n        train_csv_path,\n        annotation_path,\n    )\n\n    oof_score = mabe_score(\n        solution,\n        submission,\n        row_id_column_name=\"row_id\",\n        beta=1,\n    )\n    print(f\"\\nOOF Score with best thresholds: {oof_score:.6f}\")\n\n    # Per-lab scores\n    print(\"\\nPer-lab OOF scores:\")\n    lab_scores = {\"oof_score\": float(oof_score)}\n\n    for lab_id in solution[\"lab_id\"].unique():\n        lab_solution = solution[solution[\"lab_id\"] == lab_id].copy()\n        lab_submission = submission[\n            submission[\"video_id\"].isin(lab_solution[\"video_id\"].unique())\n        ].copy()\n\n        if len(lab_submission) > 0:\n            lab_submission[\"row_id\"] = range(len(lab_submission))\n            lab_score = mabe_score(\n                lab_solution,\n                lab_submission,\n                row_id_column_name=\"row_id\",\n                beta=1,\n            )\n        else:\n            lab_score = 0.0\n\n        print(f\"  {lab_id}: {lab_score:.6f}\")\n        lab_scores[f\"oof_score/{lab_id}\"] = float(lab_score)\n\n        # Per-action scores within each lab\n        for action in lab_solution[\"action\"].unique():\n            action_solution = lab_solution[lab_solution[\"action\"] == action].copy()\n            action_submission = lab_submission[\n                lab_submission[\"action\"] == action\n            ].copy()\n\n            if len(action_submission) > 0 and len(action_solution) > 0:\n                action_submission[\"row_id\"] = range(len(action_submission))\n                action_score = mabe_score(\n                    action_solution,\n                    action_submission,\n                    row_id_column_name=\"row_id\",\n                    beta=1,\n                )\n            else:\n                action_score = 0.0\n\n            print(f\"    {lab_id}/{action}: {action_score:.6f}\")\n            lab_scores[f\"oof_score/{lab_id}/{action}\"] = float(action_score)\n\n    # Log metrics to wandb if logger is provided\n    if logger is not None:\n        logger.log_metrics(lab_scores)\n\n    return oof_score\n\n\ndef _run_ensemble_inference_test_with_oof_threshold(\n    exp_configs: List[Tuple[str, str]],\n    folds: List[int],\n    output_path: Path,\n    base_output_dir: str = \"output\",\n    weights: List[float] = None,\n    threshold_path: str = None,\n    ensemble_method: str = \"mean\",\n    is_kaggle: bool = False,\n    min_frames: int = 0,\n):\n    print(f\"\\n{'=' * 80}\")\n    print(\"Running test ensemble inference with OOF threshold\")\n    print(f\"{'=' * 80}\\n\")\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    # Load first config for data loading\n    first_exp, first_sub = exp_configs[0]\n    first_cfg_path = (\n        Path(base_output_dir)\n        / f\"mabe-tenten-models-{first_exp}-{first_sub}\"\n        / \"config.yaml\"\n    )\n\n    if not first_cfg_path.exists():\n        raise FileNotFoundError(f\"Config not found at {first_cfg_path}\")\n\n    cfg = OmegaConf.load(first_cfg_path)\n    if is_kaggle:\n        cfg.dir = OmegaConf.load(\"yamls/dir/kaggle.yaml\")\n\n    # Override val_stride for ensemble inference\n    cfg.exp.dataset.val_stride = 128\n\n    # Get action and lab mappings\n    unique_actions, action_to_index = load_actions_from_config(cfg)\n    index_to_action = {v: k for k, v in action_to_index.items()}\n    num_classes = len(action_to_index)\n\n    lab_id_to_index = dict(cfg.exp.dataset.lab_id_to_index)\n    index_to_lab_id = {v: k for k, v in lab_id_to_index.items()}\n\n    missing_actions = OmegaConf.to_container(\n        cfg.exp.dataset.missing_actions, resolve=True\n    )\n\n    # Load distance/velocity/acceleration statistics\n    dataset_base = Path(cfg.dir.dataset_dir) / cfg.exp.dataset.name\n    stats_path = dataset_base / \"distance_velocity_acceleration_statistics.json\"\n    if stats_path.exists():\n        with open(stats_path, \"r\") as f:\n            distance_velocity_acceleration_stats = json.load(f)\n    else:\n        distance_velocity_acceleration_stats = None\n\n    # Load test data (only once)\n    df_ann, df_test = load_test_data(cfg)\n    df_samples = df_ann.unique(subset=[\"video_id\", \"agent_id\", \"target_id\"])\n\n    # Process each fold\n    fold_predictions = {}\n    for fold in folds:\n        print(f\"\\nProcessing fold {fold}...\")\n\n        # Create exp_configs with fold\n        fold_exp_configs = [(exp, sub, fold) for exp, sub in exp_configs]\n\n        # Create ensemble model for this fold\n        ensemble_model = EnsembleModel(\n            fold_exp_configs, device, base_output_dir, weights, ensemble_method\n        )\n\n        # Process in chunks\n        chunk_size = cfg.exp.inference.chunk_size\n        num_chunks = (len(df_samples) + chunk_size - 1) // chunk_size\n        sample_chunks = [\n            df_samples.slice(\n                i * chunk_size, min(chunk_size, len(df_samples) - i * chunk_size)\n            )\n            for i in range(num_chunks)\n        ]\n\n        fold_sample_predictions = {}\n\n        for chunk_idx, chunk_samples in enumerate(\n            tqdm(sample_chunks, desc=f\"Processing fold {fold} chunks\")\n        ):\n            chunk_dataset = SlidingWindowDataset(\n                cfg=cfg,\n                mode=\"test\",\n                df_ann=chunk_samples,\n                distance_velocity_acceleration_stats=distance_velocity_acceleration_stats,\n            )\n            chunk_loader = DataLoader(\n                chunk_dataset,\n                batch_size=cfg.exp.dataset.val_batch_size,\n                num_workers=cfg.exp.dataset.num_workers,\n                shuffle=False,\n                pin_memory=True,\n                drop_last=False,\n            )\n\n            chunk_predictions = {}\n            chunk_counts = {}\n            chunk_metadata = {}\n\n            for row in chunk_samples.iter_rows(named=True):\n                key = (\n                    int(row[\"video_id\"]),\n                    int(row[\"agent_id\"]),\n                    int(row[\"target_id\"]),\n                )\n                chunk_predictions[key] = np.zeros(\n                    (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                )\n                chunk_counts[key] = np.zeros(\n                    (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                )\n                chunk_metadata[key] = {\n                    \"num_frames\": int(row[\"num_frames\"]),\n                    \"video_frame_min\": int(row.get(\"video_frame_min\", 0)),\n                }\n\n            edge_removal = int(cfg.exp.inference.get(\"edge_frame_removal\", 0))\n\n            with torch.no_grad():\n                for batch in chunk_loader:\n                    x = batch[\"x\"].to(device)\n                    candidate_labels = batch[\"candidate_labels\"].to(device)\n                    lab_id = batch[\"lab_id\"].to(device)\n\n                    probs = ensemble_model.predict(x, candidate_labels, lab_id)\n                    probs = probs.cpu().numpy()\n\n                    meta_dict = batch[\"meta\"]\n                    batch_size = len(meta_dict[\"video_id\"])\n\n                    for i in range(batch_size):\n                        video_id = meta_dict[\"video_id\"][i].item()\n                        agent_id = meta_dict[\"agent_id\"][i].item()\n                        target_id = meta_dict[\"target_id\"][i].item()\n                        start_frame = meta_dict[\"start_frame\"][i].item()\n                        pred_window = probs[i]\n\n                        key = (video_id, agent_id, target_id)\n                        sample_preds = chunk_predictions[key]\n                        sample_cnts = chunk_counts[key]\n                        metadata = chunk_metadata[key]\n\n                        num_frames = metadata[\"num_frames\"]\n                        video_frame_min = metadata[\"video_frame_min\"]\n\n                        sample_preds, sample_cnts = _accumulate_predictions(\n                            sample_preds,\n                            sample_cnts,\n                            pred_window,\n                            start_frame,\n                            num_frames,\n                            video_frame_min,\n                            edge_removal,\n                        )\n\n            # Process chunk predictions\n            for row in chunk_samples.iter_rows(named=True):\n                key = (\n                    int(row[\"video_id\"]),\n                    int(row[\"agent_id\"]),\n                    int(row[\"target_id\"]),\n                )\n                avg_predictions = chunk_predictions[key]\n                counts = chunk_counts[key]\n\n                valid_mask = counts > 0\n                avg_predictions[valid_mask] /= counts[valid_mask]\n\n                lab_id_idx = int(row.get(\"lab_id_index\", 0))\n                lab_id = index_to_lab_id.get(lab_id_idx)\n                if lab_id in missing_actions:\n                    for action in missing_actions[lab_id]:\n                        action_idx = action_to_index.get(action)\n                        if action_idx is not None:\n                            avg_predictions[:, action_idx] = 0.0\n\n                video_frame_min = int(row.get(\"video_frame_min\", 0))\n                avg_predictions = _interpolate_predictions_to_original_fps(\n                    avg_predictions, float(row[\"fps_scale\"]), video_frame_min\n                )\n\n                smooth_window_size = int(cfg.exp.inference.get(\"smooth_window_size\", 0))\n                if smooth_window_size > 1:\n                    avg_predictions = _smooth_predictions(\n                        avg_predictions, smooth_window_size\n                    )\n\n                fold_sample_predictions[key] = avg_predictions.copy()\n\n            del chunk_predictions\n            del chunk_counts\n            del chunk_dataset\n            del chunk_loader\n\n        fold_predictions[fold] = fold_sample_predictions\n        del ensemble_model\n\n    # Average predictions across folds\n    print(\"\\nAveraging predictions across folds...\")\n    all_keys = set()\n    for fold_preds in fold_predictions.values():\n        all_keys.update(fold_preds.keys())\n\n    averaged_predictions = {}\n    for key in tqdm(all_keys, desc=\"Averaging predictions\"):\n        fold_preds = [\n            fold_predictions[fold][key]\n            for fold in folds\n            if key in fold_predictions[fold]\n        ]\n        if len(fold_preds) > 0:\n            averaged_predictions[key] = np.mean(fold_preds, axis=0)\n\n    # Create metadata\n    all_sample_metadata = {}\n    for row in df_samples.iter_rows(named=True):\n        key = (int(row[\"video_id\"]), int(row[\"agent_id\"]), int(row[\"target_id\"]))\n        candidate_labels = np.array(\n            json.loads(row[\"candidate_labels\"]), dtype=np.float32\n        )\n        all_sample_metadata[key] = {\n            \"video_id\": int(row[\"video_id\"]),\n            \"agent_id\": int(row[\"agent_id\"]),\n            \"target_id\": int(row[\"target_id\"]),\n            \"candidate_labels\": candidate_labels,\n            \"lab_id_index\": int(row.get(\"lab_id_index\", 0)),\n        }\n\n    # Load OOF threshold\n    print(\"\\n\" + \"=\" * 80)\n    print(\"Loading OOF threshold\")\n    print(\"=\" * 80)\n\n    if threshold_path is None:\n        threshold_path = output_path / \"threshold_search_oof.json\"\n    else:\n        threshold_path = Path(threshold_path)\n\n    print(f\"Threshold path: {threshold_path}\")\n\n    if not threshold_path.exists():\n        raise FileNotFoundError(f\"OOF threshold file not found: {threshold_path}\")\n\n    with open(threshold_path, \"r\") as f:\n        threshold_data = json.load(f)\n        lab_action_thresholds = threshold_data[\"lab_action_best_thresholds\"]\n        default_threshold = threshold_data[\"default_threshold\"]\n\n    print(f\"Loaded OOF thresholds from: {threshold_path}\")\n    print(f\"Default threshold: {default_threshold}\")\n\n    # Display loaded thresholds\n    print(\"\\nOOF thresholds per lab and action:\")\n    for lab_id in sorted(lab_action_thresholds.keys()):\n        if len(lab_action_thresholds[lab_id]) > 0:\n            print(f\"  Lab {lab_id}:\")\n            for action_name in sorted(lab_action_thresholds[lab_id].keys()):\n                threshold = lab_action_thresholds[lab_id][action_name][\"threshold\"]\n                f1_score = lab_action_thresholds[lab_id][action_name][\"f1_score\"]\n                print(\n                    f\"    {action_name}: threshold={threshold:.4f}, F1={f1_score:.4f}\"\n                )\n\n    print(\"=\" * 80 + \"\\n\")\n\n    # Create submission with OOF thresholds\n    print(\"Creating submission with OOF thresholds...\")\n    submission = _extract_segments_with_lab_action_thresholds(\n        all_sample_predictions=averaged_predictions,\n        all_sample_metadata=all_sample_metadata,\n        lab_action_thresholds=lab_action_thresholds,\n        default_threshold=default_threshold,\n        index_to_action=index_to_action,\n        index_to_lab_id=index_to_lab_id,\n    )\n\n    # Filter out predictions with frame length <= min_frames\n    if min_frames > 0:\n        original_len = len(submission)\n        submission[\"frame_length\"] = (\n            submission[\"stop_frame\"] - submission[\"start_frame\"] + 1\n        )\n        submission = submission[submission[\"frame_length\"] > min_frames].drop(\n            columns=[\"frame_length\"]\n        )\n        filtered_len = len(submission)\n        print(\n            f\"Filtered predictions: {original_len} -> {filtered_len} \"\n            f\"(removed {original_len - filtered_len} predictions with <={min_frames} frames)\"\n        )\n\n    # Adjust frames for AdaptableSnail videos with fps ~25.0\n    train_csv_path = (\n        Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/train_fixed.csv\"\n    )\n    test_csv_path = Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/test.csv\"\n    submission = _adjust_frames_for_adaptable_snail(\n        submission, train_csv_path, test_csv_path\n    )\n\n    submission.insert(0, \"row_id\", range(len(submission)))\n\n    submission_path = output_path / \"submission.csv\"\n    submission.to_csv(submission_path, index=False)\n\n    print(f\"\\nTest submission saved to: {submission_path}\")\n    print(f\"Total predictions: {len(submission)}\")\n\n    return 0.0\n\n\ndef run_ensemble_inference(\n    exp_configs: List[Tuple[str, str]],\n    folds: List[int],\n    output_path: Path,\n    is_test: bool = False,\n    base_output_dir: str = \"output\",\n    weights: List[float] = None,\n    threshold_dir: str = None,\n    ensemble_name: str = \"ensemble_001\",\n    ensemble_method: str = \"mean\",\n    is_kaggle: bool = False,\n    min_frames: int = 0,\n) -> float:\n    \"\"\"\n    Run ensemble inference across multiple folds.\n\n    Args:\n        exp_configs: List of (exp_name, sub_name) tuples\n        folds: List of fold numbers\n        output_path: Output directory\n        is_test: Whether to run on test data\n        base_output_dir: Base directory for outputs\n        weights: Optional list of weights for each model\n        threshold_dir: Directory to load thresholds from (for test mode).\n                      If None, uses output_path. Useful for Kaggle environments.\n        ensemble_name: Name for ensemble (used in logger group name)\n        ensemble_method: Ensemble method (\"mean\" or \"max\")\n        is_kaggle: Whether running in Kaggle environment\n\n    Returns:\n        Average validation score (0.0 for test)\n    \"\"\"\n\n    output_path.mkdir(parents=True, exist_ok=True)\n\n    if not is_test:\n        # Validation: run each fold and average scores\n        fold_scores = []\n        for fold in folds:\n            print(f\"\\n{'=' * 80}\")\n            print(f\"Processing ensemble validation fold {fold}\")\n            print(f\"{'=' * 80}\\n\")\n\n            # Create WandbLogger for this fold\n            fold_logger = WandbLogger(\n                project=\"kaggle-mabe\",\n                name=f\"fold_{fold}\",\n                group=ensemble_name,\n            )\n\n            # Create exp_configs with fold\n            fold_exp_configs = [(exp, sub, fold) for exp, sub in exp_configs]\n\n            score = _run_ensemble_inference_single_fold(\n                exp_configs=fold_exp_configs,\n                fold=fold,\n                output_path=output_path,\n                is_test=False,\n                base_output_dir=base_output_dir,\n                weights=weights,\n                logger=fold_logger,\n                ensemble_method=ensemble_method,\n                min_frames=min_frames,\n            )\n            fold_scores.append(score)\n            print(f\"Fold {fold} ensemble validation score: {score:.6f}\")\n\n            fold_logger.experiment.finish()\n\n        avg_score = np.mean(fold_scores) if len(fold_scores) > 0 else 0.0\n        print(f\"\\n{'=' * 80}\")\n        print(\n            f\"Average ensemble validation score across {len(folds)} folds: {avg_score:.6f}\"\n        )\n        print(f\"{'=' * 80}\\n\")\n        return avg_score\n    else:\n        # Test: run each fold, then average thresholds for final submission\n        print(f\"\\n{'=' * 80}\")\n        print(\"Running test ensemble inference\")\n        print(f\"{'=' * 80}\\n\")\n\n        # For test, we need to process all folds together\n        # Load first config for data loading\n        first_exp, first_sub = exp_configs[0]\n\n        # Try standard path first (local environment)\n        first_cfg_path = Path(base_output_dir) / first_exp / first_sub / \"config.yaml\"\n\n        # If not found, try Kaggle-style path\n        if not first_cfg_path.exists():\n            first_cfg_path = (\n                Path(base_output_dir)\n                / f\"mabe-tenten-models-{first_exp}-{first_sub}\"\n                / \"config.yaml\"\n            )\n\n        if not first_cfg_path.exists():\n            raise FileNotFoundError(f\"Config not found at {first_cfg_path}\")\n\n        cfg = OmegaConf.load(first_cfg_path)\n\n        # Override directory config for Kaggle environment\n        if is_kaggle:\n            cfg.dir = OmegaConf.load(\"yamls/dir/kaggle.yaml\")\n\n        # Override val_stride for ensemble inference\n        cfg.exp.dataset.val_stride = 128\n\n        device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n        # Process each fold\n        fold_predictions = {}\n        for fold in folds:\n            print(f\"\\nProcessing fold {fold}...\")\n\n            # Create exp_configs with fold\n            fold_exp_configs = [(exp, sub, fold) for exp, sub in exp_configs]\n\n            # Create ensemble model for this fold\n            ensemble_model = EnsembleModel(\n                fold_exp_configs, device, base_output_dir, weights, ensemble_method\n            )\n\n            # Load test data (only once)\n            if fold == folds[0]:\n                df_ann, df_test = load_test_data(cfg)\n                df_samples = df_ann.unique(subset=[\"video_id\", \"agent_id\", \"target_id\"])\n\n                unique_actions, action_to_index = load_actions_from_config(cfg)\n                index_to_action = {v: k for k, v in action_to_index.items()}\n                num_classes = len(action_to_index)\n\n                lab_id_to_index = dict(cfg.exp.dataset.lab_id_to_index)\n                index_to_lab_id = {v: k for k, v in lab_id_to_index.items()}\n\n                missing_actions = OmegaConf.to_container(\n                    cfg.exp.dataset.missing_actions, resolve=True\n                )\n\n                dataset_base = Path(cfg.dir.dataset_dir) / cfg.exp.dataset.name\n                stats_path = (\n                    dataset_base / \"distance_velocity_acceleration_statistics.json\"\n                )\n                if stats_path.exists():\n                    with open(stats_path, \"r\") as f:\n                        distance_velocity_acceleration_stats = json.load(f)\n                else:\n                    distance_velocity_acceleration_stats = None\n\n            # Process in chunks\n            chunk_size = cfg.exp.inference.chunk_size\n            num_chunks = (len(df_samples) + chunk_size - 1) // chunk_size\n            sample_chunks = [\n                df_samples.slice(\n                    i * chunk_size, min(chunk_size, len(df_samples) - i * chunk_size)\n                )\n                for i in range(num_chunks)\n            ]\n\n            fold_sample_predictions = {}\n\n            for chunk_idx, chunk_samples in enumerate(\n                tqdm(sample_chunks, desc=f\"Processing fold {fold} chunks\")\n            ):\n                chunk_dataset = SlidingWindowDataset(\n                    cfg=cfg,\n                    mode=\"test\",\n                    df_ann=chunk_samples,\n                    distance_velocity_acceleration_stats=distance_velocity_acceleration_stats,\n                )\n                chunk_loader = DataLoader(\n                    chunk_dataset,\n                    batch_size=cfg.exp.dataset.val_batch_size,\n                    num_workers=cfg.exp.dataset.num_workers,\n                    shuffle=False,\n                    pin_memory=True,\n                    drop_last=False,\n                )\n\n                chunk_predictions = {}\n                chunk_counts = {}\n                chunk_metadata = {}\n\n                for row in chunk_samples.iter_rows(named=True):\n                    key = (\n                        int(row[\"video_id\"]),\n                        int(row[\"agent_id\"]),\n                        int(row[\"target_id\"]),\n                    )\n                    chunk_predictions[key] = np.zeros(\n                        (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                    )\n                    chunk_counts[key] = np.zeros(\n                        (int(row[\"num_frames\"]), num_classes), dtype=np.float32\n                    )\n                    chunk_metadata[key] = {\n                        \"num_frames\": int(row[\"num_frames\"]),\n                        \"video_frame_min\": int(row.get(\"video_frame_min\", 0)),\n                    }\n\n                edge_removal = int(cfg.exp.inference.get(\"edge_frame_removal\", 0))\n\n                with torch.no_grad():\n                    for batch in chunk_loader:\n                        x = batch[\"x\"].to(device)\n                        candidate_labels = batch[\"candidate_labels\"].to(device)\n                        lab_id = batch[\"lab_id\"].to(device)\n\n                        probs = ensemble_model.predict(x, candidate_labels, lab_id)\n                        probs = probs.cpu().numpy()\n\n                        meta_dict = batch[\"meta\"]\n                        batch_size = len(meta_dict[\"video_id\"])\n\n                        for i in range(batch_size):\n                            video_id = meta_dict[\"video_id\"][i].item()\n                            agent_id = meta_dict[\"agent_id\"][i].item()\n                            target_id = meta_dict[\"target_id\"][i].item()\n                            start_frame = meta_dict[\"start_frame\"][i].item()\n                            pred_window = probs[i]\n\n                            key = (video_id, agent_id, target_id)\n                            sample_preds = chunk_predictions[key]\n                            sample_cnts = chunk_counts[key]\n                            metadata = chunk_metadata[key]\n\n                            num_frames = metadata[\"num_frames\"]\n                            video_frame_min = metadata[\"video_frame_min\"]\n\n                            sample_preds, sample_cnts = _accumulate_predictions(\n                                sample_preds,\n                                sample_cnts,\n                                pred_window,\n                                start_frame,\n                                num_frames,\n                                video_frame_min,\n                                edge_removal,\n                            )\n\n                # Process chunk predictions\n                for row in chunk_samples.iter_rows(named=True):\n                    key = (\n                        int(row[\"video_id\"]),\n                        int(row[\"agent_id\"]),\n                        int(row[\"target_id\"]),\n                    )\n                    avg_predictions = chunk_predictions[key]\n                    counts = chunk_counts[key]\n\n                    valid_mask = counts > 0\n                    avg_predictions[valid_mask] /= counts[valid_mask]\n\n                    lab_id_idx = int(row.get(\"lab_id_index\", 0))\n                    lab_id = index_to_lab_id.get(lab_id_idx)\n                    if lab_id in missing_actions:\n                        for action in missing_actions[lab_id]:\n                            action_idx = action_to_index.get(action)\n                            if action_idx is not None:\n                                avg_predictions[:, action_idx] = 0.0\n\n                    video_frame_min = int(row.get(\"video_frame_min\", 0))\n                    avg_predictions = _interpolate_predictions_to_original_fps(\n                        avg_predictions, float(row[\"fps_scale\"]), video_frame_min\n                    )\n\n                    smooth_window_size = int(\n                        cfg.exp.inference.get(\"smooth_window_size\", 0)\n                    )\n                    if smooth_window_size > 1:\n                        avg_predictions = _smooth_predictions(\n                            avg_predictions, smooth_window_size\n                        )\n\n                    fold_sample_predictions[key] = avg_predictions.copy()\n\n                del chunk_predictions\n                del chunk_counts\n                del chunk_dataset\n                del chunk_loader\n\n            fold_predictions[fold] = fold_sample_predictions\n            del ensemble_model\n\n        # Average predictions across folds\n        print(\"\\nAveraging predictions across folds...\")\n        all_keys = set()\n        for fold_preds in fold_predictions.values():\n            all_keys.update(fold_preds.keys())\n\n        averaged_predictions = {}\n        for key in tqdm(all_keys, desc=\"Averaging predictions\"):\n            fold_preds = [\n                fold_predictions[fold][key]\n                for fold in folds\n                if key in fold_predictions[fold]\n            ]\n            if len(fold_preds) > 0:\n                averaged_predictions[key] = np.mean(fold_preds, axis=0)\n\n        # Create metadata\n        all_sample_metadata = {}\n        for row in df_samples.iter_rows(named=True):\n            key = (int(row[\"video_id\"]), int(row[\"agent_id\"]), int(row[\"target_id\"]))\n            candidate_labels = np.array(\n                json.loads(row[\"candidate_labels\"]), dtype=np.float32\n            )\n            all_sample_metadata[key] = {\n                \"video_id\": int(row[\"video_id\"]),\n                \"agent_id\": int(row[\"agent_id\"]),\n                \"target_id\": int(row[\"target_id\"]),\n                \"candidate_labels\": candidate_labels,\n                \"lab_id_index\": int(row.get(\"lab_id_index\", 0)),\n            }\n\n        # Load and average thresholds from validation\n        print(\"\\n\" + \"=\" * 80)\n        print(\"Loading and averaging thresholds from validation folds\")\n        print(\"=\" * 80)\n\n        # Determine threshold directory\n        threshold_search_dir = Path(threshold_dir) if threshold_dir else output_path\n        print(f\"Threshold directory: {threshold_search_dir}\")\n\n        fold_lab_action_thresholds = {}\n        for fold in folds:\n            threshold_path = (\n                threshold_search_dir / f\"threshold_search_ensemble_fold_{fold}.json\"\n            )\n            if threshold_path.exists():\n                with open(threshold_path, \"r\") as f:\n                    threshold_data = json.load(f)\n                    fold_lab_action_thresholds[fold] = threshold_data[\n                        \"lab_action_best_thresholds\"\n                    ]\n                    print(f\"  Loaded thresholds from fold {fold}: {threshold_path}\")\n            else:\n                print(\n                    f\"  Warning: Threshold JSON not found for fold {fold} at {threshold_path}\"\n                )\n\n        pred_threshold = float(cfg.exp.pred_threshold)\n\n        # Average thresholds across folds\n        lab_ids = set()\n        for fold_data in fold_lab_action_thresholds.values():\n            lab_ids.update(fold_data.keys())\n\n        average_lab_action_thresholds = {}\n        for lab_id in lab_ids:\n            average_lab_action_thresholds[lab_id] = {}\n\n            all_actions = set()\n            for fold_data in fold_lab_action_thresholds.values():\n                if lab_id in fold_data:\n                    all_actions.update(fold_data[lab_id].keys())\n\n            for action_name in all_actions:\n                thresholds = [\n                    fold_data[lab_id][action_name][\"threshold\"]\n                    for fold_data in fold_lab_action_thresholds.values()\n                    if lab_id in fold_data and action_name in fold_data[lab_id]\n                ]\n                if len(thresholds) > 0:\n                    average_lab_action_thresholds[lab_id][action_name] = float(\n                        np.mean(thresholds)\n                    )\n\n        print(\"\\nAveraged thresholds per lab and action:\")\n        for lab_id in sorted(average_lab_action_thresholds.keys()):\n            print(f\"  Lab {lab_id}:\")\n            for action_name in sorted(average_lab_action_thresholds[lab_id].keys()):\n                threshold = average_lab_action_thresholds[lab_id][action_name]\n                print(f\"    {action_name}: {threshold:.2f}\")\n\n        print(\"=\" * 80 + \"\\n\")\n\n        # Create submission with averaged best thresholds\n        print(\"Creating submission with averaged best thresholds...\")\n        submission = _extract_segments_with_lab_action_thresholds(\n            all_sample_predictions=averaged_predictions,\n            all_sample_metadata=all_sample_metadata,\n            lab_action_thresholds=average_lab_action_thresholds,\n            default_threshold=pred_threshold,\n            index_to_action=index_to_action,\n            index_to_lab_id=index_to_lab_id,\n        )\n\n        # Filter out predictions with frame length <= min_frames\n        if min_frames > 0:\n            original_len = len(submission)\n            submission[\"frame_length\"] = (\n                submission[\"stop_frame\"] - submission[\"start_frame\"] + 1\n            )\n            submission = submission[submission[\"frame_length\"] > min_frames].drop(\n                columns=[\"frame_length\"]\n            )\n            filtered_len = len(submission)\n            print(\n                f\"Filtered predictions: {original_len} -> {filtered_len} \"\n                f\"(removed {original_len - filtered_len} predictions with <={min_frames} frames)\"\n            )\n\n        # Adjust frames for AdaptableSnail videos with fps ~25.0\n        train_csv_path = (\n            Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/train_fixed.csv\"\n        )\n        test_csv_path = (\n            Path(cfg.dir.input_dir) / \"MABe-mouse-behavior-detection/test.csv\"\n        )\n        submission = _adjust_frames_for_adaptable_snail(\n            submission, train_csv_path, test_csv_path\n        )\n\n        submission.insert(0, \"row_id\", range(len(submission)))\n\n        submission_path = output_path / \"submission.csv\"\n        submission.to_csv(submission_path, index=False)\n\n        print(f\"\\nTest submission saved to: {submission_path}\")\n        print(f\"Total predictions: {len(submission)}\")\n\n        return 0.0\n\nif __name__ == \"__main__\":\n    # コマンドライン引数のパース\n    parser = argparse.ArgumentParser(description=\"Ensemble inference\")\n    parser.add_argument(\n        \"--ensemble_name\",\n        type=str,\n        default=\"ensemble_014\",\n        help=\"Ensemble name (default: ensemble_014)\",\n    )\n    parser.add_argument(\n        \"--exp_configs\",\n        type=str,\n        default=\"exp113:self,exp117:self,exp119:self,exp125:self,exp126:self\",\n        help='Experiment configs in format \"exp1:sub1,exp2:sub2,...\" (default: exp113:self,exp117:self,exp119:self,exp125:self,exp126:self)',\n    )\n    parser.add_argument(\n        \"--min_frames\",\n        type=int,\n        default=0,\n        help=\"Minimum frame length for predictions (default: 0, no filtering)\",\n    )\n    args = parser.parse_args()\n\n    seed_everything(55)\n\n    # アンサンブル名の設定\n    ensemble_name = args.ensemble_name\n\n    # exp_configsのパース\n    pair_config = \"exp127:aug,exp118:base,exp126:base,exp130:base,exp132:lr\"\n    pair_configs = []\n    for config_str in pair_config.split(\",\"):\n        exp_name, sub_name = config_str.strip().split(\":\")\n        pair_configs.append((exp_name.strip(), sub_name.strip()))\n\n    self_config = \"exp113:self,exp117:self,exp119:self,exp125:self,exp126:self\"\n    self_configs = []\n    for config_str in self_config.split(\",\"):\n        exp_name, sub_name = config_str.strip().split(\":\")\n        self_configs.append((exp_name.strip(), sub_name.strip()))\n\n    # exp_configs = []\n    # for config_str in args.exp_configs.split(\",\"):\n    #     exp_name, sub_name = config_str.strip().split(\":\")\n    #     exp_configs.append((exp_name.strip(), sub_name.strip()))\n\n    # print(f\"Ensemble name: {ensemble_name}\")\n    # print(f\"Experiment configs: {exp_configs}\")\n    # print(f\"Minimum frames: {args.min_frames}\")\n\n    # モデルの重み設定\n    # 例: [0.25, 0.25, 0.25, 0.25] → 均等\n    weights = None  # 均等な重み\n    # weights = [0.3, 0.3, 0.2, 0.2]  # カスタム重み例\n\n    # アンサンブル方法の設定\n    # \"mean\": 重み付き平均\n    # \"max\": 各位置で全モデルの予測の最大値\n    ensemble_method = \"mean\"\n\n    # 処理するfold\n    folds = [0, 1, 2, 3, 4]\n\n    # 出力ディレクトリ\n    output_path = Path(f\"output/{ensemble_name}\")\n    output_path.mkdir(parents=True, exist_ok=True)\n\n    # ================================================================================\n    # OOF推論（推奨）: すべてのfoldの予測をまとめてOOFを作成し、OOF全体で最適なthresholdを探索\n    # ================================================================================\n    # from pytorch_lightning.loggers import WandbLogger\n\n    # oof_logger = WandbLogger(\n    #     project=\"kaggle-mabe\",\n    #     name=\"oof\",\n    #     group=ensemble_name,\n    # )\n\n    # oof_score = _run_ensemble_inference_oof(\n    #     exp_configs=exp_configs,\n    #     folds=folds,\n    #     output_path=output_path,\n    #     base_output_dir=\"output\",\n    #     weights=weights,\n    #     logger=oof_logger,\n    #     ensemble_method=ensemble_method,\n    #     min_frames=args.min_frames,\n    # )\n\n    # print(f\"\\n{'=' * 80}\")\n    # print(f\"Final OOF Score: {oof_score:.6f}\")\n    # print(f\"{'=' * 80}\\n\")\n\n    # oof_logger.experiment.finish()\n\n    # ================================================================================\n    # 従来の方法（各foldごとに独立してthreshold探索）\n    # ================================================================================\n    # run_ensemble_inference(\n    #     exp_configs=exp_configs,\n    #     folds=folds,\n    #     output_path=output_path,\n    #     is_test=False,  # validation推論\n    #     weights=weights,\n    #     ensemble_name=ensemble_name,\n    #     ensemble_method=ensemble_method,\n    #     min_frames=args.min_frames,\n    # )\n\n    # ================================================================================\n    # Test推論（OOF推論後に実行）\n    # ================================================================================\n    # OOF thresholdを使ったTest推論（推奨）\n    # ローカル環境の場合\n    # _run_ensemble_inference_test_with_oof_threshold(\n    #     exp_configs=exp_configs,\n    #     folds=folds,\n    #     output_path=output_path,\n    #     base_output_dir=\"output\",\n    #     weights=weights,\n    #     threshold_path=None,  # output_path/threshold_search_oof.jsonから読み込む\n    #     ensemble_method=ensemble_method,\n    #     min_frames=args.min_frames,\n    # )\n\n    # Kaggle環境の場合（thresholdを別のdatasetから読み込む）\n    # _run_ensemble_inference_test_with_oof_threshold(\n    #     exp_configs=exp_configs,\n    #     folds=folds,\n    #     output_path=Path(\"/kaggle/working/\"),\n    #     base_output_dir=\"/kaggle/input/\",\n    #     weights=weights,\n    #     threshold_path=\"/kaggle/input/mabe-tenten-models-ensemble-007/threshold_search_oof.json\",\n    #     ensemble_method=ensemble_method,\n    #     is_kaggle=True,\n    #     min_frames=args.min_frames,\n    # )\n\n    run_ensemble_inference(\n        exp_configs=self_configs,\n        folds=folds,\n        output_path=Path(\"/kaggle/working/\"),\n        base_output_dir=\"/kaggle/input/\",\n        is_test=True,\n        weights=weights,\n        threshold_dir=\"/kaggle/input/mabe-tenten-models-ensemble-014\",\n        ensemble_method=ensemble_method,\n        is_kaggle=True,\n        min_frames=args.min_frames,\n    )\n    df_self = pd.read_csv(\"/kaggle/working/submission.csv\")\n    df_self.to_csv(\"/kaggle/working/submission_self.csv\", index=False)\n\n    run_ensemble_inference(\n        exp_configs=pair_configs,\n        folds=folds,\n        output_path=Path(\"/kaggle/working/\"),\n        base_output_dir=\"/kaggle/input/\",\n        is_test=True,\n        weights=weights,\n        threshold_dir=\"/kaggle/input/mabe-tenten-models-ensemble-127-118-126-130-132\",\n        ensemble_method=ensemble_method,\n        is_kaggle=True,\n        min_frames=args.min_frames,\n    )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T15:29:49.044748Z","iopub.execute_input":"2025-12-15T15:29:49.045239Z","iopub.status.idle":"2025-12-15T15:31:47.112636Z","shell.execute_reply.started":"2025-12-15T15:29:49.045216Z","shell.execute_reply":"2025-12-15T15:31:47.111951Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\ndf = pd.read_csv(\"/kaggle/working/submission.csv\")\ndf2 = pd.read_csv(\"/kaggle/working/submission_self.csv\")\n\ndf = pd.concat([df, df2]).sort_values(by=[\"video_id\", \"agent_id\"]).reset_index(drop=True)\ndf[\"row_id\"] = df.index\n\ntest = pd.read_csv(\"/kaggle/input/MABe-mouse-behavior-detection/test.csv\")\nV = test.loc[(test.lab_id==\"AdaptableSnail\")&(test.frames_per_second==25)].video_id.unique()\ndf = df.loc[~df.video_id.isin(V)] \ndf.to_csv(\"/kaggle/working/submission.csv\", index=False)\ndf.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T15:31:47.113692Z","iopub.execute_input":"2025-12-15T15:31:47.113930Z","iopub.status.idle":"2025-12-15T15:31:47.131241Z","shell.execute_reply.started":"2025-12-15T15:31:47.113914Z","shell.execute_reply":"2025-12-15T15:31:47.130705Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%cd ..\n!rm -r /kaggle/working/dataset\n!rm -r /kaggle/working/mabe-tenten-code","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T14:19:33.174928Z","iopub.status.idle":"2025-12-15T14:19:33.175144Z","shell.execute_reply.started":"2025-12-15T14:19:33.175040Z","shell.execute_reply":"2025-12-15T14:19:33.175049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}