{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":686802,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":514546,"modelId":529193}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:15.930331Z","iopub.execute_input":"2026-01-06T01:30:15.930738Z","iopub.status.idle":"2026-01-06T01:30:19.956624Z","shell.execute_reply.started":"2026-01-06T01:30:15.930710Z","shell.execute_reply":"2026-01-06T01:30:19.955215Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import ast\nimport keras\nimport random\nimport numpy as np\nimport tensorflow as tf\nimport pandas as pd\nimport pathlib\nfrom math import ceil, sqrt\nfrom tqdm import tqdm\nimport csv\nimport os\noutput_file = \"submission.csv\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:19.958663Z","iopub.execute_input":"2026-01-06T01:30:19.959088Z","iopub.status.idle":"2026-01-06T01:30:19.965729Z","shell.execute_reply.started":"2026-01-06T01:30:19.959042Z","shell.execute_reply":"2026-01-06T01:30:19.964586Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# These need to be here to load the model\n@keras.saving.register_keras_serializable()\ndef make_pairs(x):\n    # x: (B, M, T, E)\n    M = tf.shape(x)[1]\n    # Expand dims to prepare for broadcasting\n    x1 = tf.expand_dims(x, axis=2)  # (B, M, 1, T, E)\n    x2 = tf.expand_dims(x, axis=1)  # (B, 1, M, T, E)\n    # Tile to get all pairs\n    x1_tiled = tf.tile(x1, [1, 1, M, 1, 1])  # (B, M, M, T, E)\n    x2_tiled = tf.tile(x2, [1, M, 1, 1, 1])  # (B, M, M, T, E)\n    # Stack pair dimension\n    pairs = tf.stack([x1_tiled, x2_tiled], axis=-2)  # (B, M, M, T, 2, E)\n    return pairs\n\n\n@keras.saving.register_keras_serializable()\ndef scale_broadcast(x):\n    scale, embedding = x\n    scale = tf.reshape(scale, (-1, 1, 1, 1))  # (batch, 1, 1, 1)\n    return scale * tf.ones_like(embedding[..., :1])\n\n\n@keras.saving.register_keras_serializable()\ndef time_broadcast(x):\n    scale, embedding = x  # scale: (batch,), embedding: (BATCH, M, T, E)\n\n    BATCH = tf.shape(embedding)[0]\n    N = tf.shape(embedding)[1]\n    X = tf.shape(embedding)[2]\n    DIM = tf.shape(embedding)[3]\n    increments = tf.linspace(0.0, 20.0, X)\n    increments = tf.reshape(increments, (1, 1, X, 1))  # broadcast shape\n    increments = tf.tile(increments, [BATCH, N, 1, 1])  # shape: [BATCH, M, T, 1]\n\n    # multiply by per-batch scale\n    scale = tf.reshape(scale, (-1, 1, 1, 1))  # (BATCH, 1, 1, 1)\n    return scale * increments\n\n\nmodel = keras.models.load_model('/kaggle/input/mabe-resnet/keras/default/15/big_sf_426.keras', safe_mode=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:19.966877Z","iopub.execute_input":"2026-01-06T01:30:19.967222Z","iopub.status.idle":"2026-01-06T01:30:20.234426Z","shell.execute_reply.started":"2026-01-06T01:30:19.967187Z","shell.execute_reply":"2026-01-06T01:30:20.233275Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.summary()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:20.235533Z","iopub.execute_input":"2026-01-06T01:30:20.235792Z","iopub.status.idle":"2026-01-06T01:30:20.280142Z","shell.execute_reply.started":"2026-01-06T01:30:20.235772Z","shell.execute_reply":"2026-01-06T01:30:20.279320Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_SEQ_SIZE = 2000\nseed = 12\nrandom.seed(seed)\nnp.random.seed(seed)\ntf.random.set_seed(seed)\ndata_dir = pathlib.Path('/kaggle/input/MABe-mouse-behavior-detection')\npd.options.mode.chained_assignment = None\npd.set_option('future.no_silent_downcasting', True)\n\ntrain_metadata = pd.read_csv(data_dir / 'train.csv')\ntest_metadata = pd.read_csv(data_dir / 'test.csv')\ntrain_metadata['x_cm_scale'] = train_metadata['video_width_pix']/train_metadata['pix_per_cm_approx']\ntrain_metadata['y_cm_scale'] = train_metadata['video_height_pix']/train_metadata['pix_per_cm_approx']\n\ntest_metadata['x_cm_scale'] = test_metadata['video_width_pix']/test_metadata['pix_per_cm_approx']\ntest_metadata['x_cm_scale'] = test_metadata['x_cm_scale']/(train_metadata['x_cm_scale'].max())\ntest_metadata['y_cm_scale'] = test_metadata['video_height_pix']/test_metadata['pix_per_cm_approx']\ntest_metadata['y_cm_scale'] = test_metadata['y_cm_scale']/(train_metadata['y_cm_scale'].max())\n\nbehaviors_train = (train_metadata['behaviors_labeled']\n             .apply(lambda x: ast.literal_eval(x) if x is not np.nan else x)\n             .explode('behaviors_labeled'))\nbehaviors = (test_metadata['behaviors_labeled']\n             .apply(lambda x: ast.literal_eval(x) if x is not np.nan else x)\n             .explode('behaviors_labeled'))\nif len(behaviors_train) > len(behaviors):\n    behaviors = behaviors_train\nbehaviors = behaviors.dropna().str.split(',').str[2].unique()\nbehaviors = list(set(map(lambda x: x.replace(\"'\", \"\"), behaviors)))\nbehaviors = sorted(behaviors)\nbehaviors = [\"nothing\"] + behaviors\nNUM_BEHAVIORS = len(behaviors)\n# number to behavior\nbehaviors_map = {i: x for i, x in enumerate(behaviors)}\n# behavior to number\nbehaviors_map_rev = {x: i for i, x in enumerate(behaviors)}\n_behaviors = [',' + x for x in behaviors]\n\nlabeled_videos = []\nfor _, row in test_metadata.loc[test_metadata['behaviors_labeled'].notna()].iterrows():\n    behaviors_found = set()\n    for _behavior, behavior in zip(_behaviors, behaviors):\n        if _behavior in str(row['behaviors_labeled']):\n            # get rid of the extra quote pairs\n            behaviors_found.add(behaviors_map_rev[behavior.replace(\"'\", \"\")])\n    behaviors_found.add(behaviors_map_rev[\"nothing\"])\n    labeled_videos.append({\n        'lab': row['lab_id'],\n        'video': row['video_id'],\n        'seconds_per_frame': 1 / row['frames_per_second'],\n        'video_width_pix': row['video_width_pix'],\n        'video_height_pix': row['video_height_pix'],\n        'x_cm_scale': row['x_cm_scale'],\n        'y_cm_scale': row['y_cm_scale'],\n        'behaviors': list(behaviors_found)\n    })\n\nbody_parts_train = (train_metadata['body_parts_tracked']\n              .apply(lambda x: ast.literal_eval(x) if x is not np.nan else x).explode().unique())\nbody_parts = (test_metadata['body_parts_tracked']\n              .apply(lambda x: ast.literal_eval(x) if x is not np.nan else x).explode().unique())\nif len(body_parts_train) > len(body_parts):\n    body_parts = body_parts_train\nbody_parts = sorted(body_parts)\nBODY_PARTS = len(body_parts)\nFEATURES = BODY_PARTS\nbody_parts_map = {x: i for i, x in enumerate(body_parts)}\n\npd.options.mode.chained_assignment = None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:27.409238Z","iopub.execute_input":"2026-01-06T01:30:27.409644Z","iopub.status.idle":"2026-01-06T01:30:27.775083Z","shell.execute_reply.started":"2026-01-06T01:30:27.409616Z","shell.execute_reply":"2026-01-06T01:30:27.773823Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"output_file = \"submission.csv\"\n\nif not os.path.exists(output_file):\n    with open(output_file, \"w\", newline=\"\") as f:\n        writer = csv.writer(f)\n        writer.writerow([\n            \"row_id\",\n            \"video_id\",\n            \"agent_id\",\n            \"target_id\",\n            \"action\",\n            \"start_frame\",\n            \"stop_frame\"\n        ])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:27.776434Z","iopub.execute_input":"2026-01-06T01:30:27.776730Z","iopub.status.idle":"2026-01-06T01:30:27.783314Z","shell.execute_reply.started":"2026-01-06T01:30:27.776706Z","shell.execute_reply":"2026-01-06T01:30:27.782035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_windows(pair, segment_start, video_stop):\n    up_limit = min(len(pair), video_stop - segment_start)\n\n    if up_limit <= 1:\n        return []\n\n    segments = []\n    curr = pair[0]\n    start = 0\n\n    for i in range(1, up_limit):\n        if pair[i] != curr:\n            if behaviors_map[curr] != \"nothing\":\n                global_start = start + segment_start\n                global_end = (i - 1) + segment_start\n\n                if global_start != global_end:\n                    segments.append((behaviors_map[curr], global_start, global_end))\n\n            curr = pair[i]\n            start = i\n\n    if behaviors_map[curr] != \"nothing\":\n        global_start = start + segment_start\n        global_end = (up_limit - 1) + segment_start\n\n        if global_start != global_end:\n            segments.append((behaviors_map[curr], global_start, global_end))\n\n    return segments","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:30:27.784572Z","iopub.execute_input":"2026-01-06T01:30:27.785553Z","iopub.status.idle":"2026-01-06T01:30:27.802676Z","shell.execute_reply.started":"2026-01-06T01:30:27.785528Z","shell.execute_reply":"2026-01-06T01:30:27.801568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_sequence_from_df(track_df, video_width, video_height, start_frame, end_frame):\n    track_frames_filter = (track_df['video_frame'] >= start_frame) & (track_df['video_frame'] < end_frame)\n    track_seq_df = track_df[track_frames_filter]\n    track_seq_df['bodypart'] = track_seq_df['bodypart'].replace(body_parts_map)\n    track_seq_df['x'] = track_seq_df['x'] / video_width\n    track_seq_df['y'] = track_seq_df['y'] / video_height\n    track_seq_df['video_frame'] = track_seq_df['video_frame'] - start_frame\n    # Mice are not always incremented by 1 in order or 0 indexed in a given segment\n    mice_ids = track_df['mouse_id'].unique()\n    num_mice = len(mice_ids)\n    # Literally no mice data\n    if num_mice == 0:\n        return None\n    track_array = np.zeros((num_mice, TRAIN_SEQ_SIZE, FEATURES, 2))\n    for row in track_seq_df.itertuples():\n        track_array[row.mouse_id, row.video_frame, row.bodypart] = [row.x, row.y]\n    return track_array\n\n\ndef ds_generator():\n    for video in labeled_videos:\n        track_df = pd.read_parquet(data_dir/f\"test_tracking/\"\n                                   f\"{video['lab']}/{video['video']}.parquet\")\n        track_df = track_df.fillna(0)\n        mice_ids = track_df['mouse_id'].unique()\n        zero_ind_mice_ids = {x: i for i, x in enumerate(mice_ids)}\n        track_df['mouse_id'] = track_df['mouse_id'].replace(zero_ind_mice_ids)\n        video_start_frame = track_df['video_frame'].min()\n        video_stop_frame = track_df['video_frame'].max()\n        divisions = ceil((video_stop_frame - video_start_frame) / TRAIN_SEQ_SIZE)\n        behaviors_mask = np.zeros((NUM_BEHAVIORS,))\n        behaviors_mask[video['behaviors']] = 1\n        for division in range(divisions):\n            segment_start_frame = division * TRAIN_SEQ_SIZE + video_start_frame\n            segment_stop_frame = (division + 1) * TRAIN_SEQ_SIZE + video_start_frame\n            seq_out = get_sequence_from_df(track_df, video['video_width_pix'],\n                                           video['video_height_pix'], segment_start_frame, segment_stop_frame)\n            if seq_out is None:\n                continue\n            yield (seq_out, video['x_cm_scale'], video['y_cm_scale'],\n                   video['seconds_per_frame'], video['video'],\n                   segment_start_frame, list(zero_ind_mice_ids.keys()),\n                   video_stop_frame, behaviors_mask)\n\n\ntf_ds = tf.data.Dataset.from_generator(ds_generator, output_signature=(\n    tf.TensorSpec((None, TRAIN_SEQ_SIZE, FEATURES, 2)),  # seq\n    tf.TensorSpec(()),  # x_cm\n    tf.TensorSpec(()),  # y_cm\n    tf.TensorSpec(()),  # time,\n    tf.TensorSpec((), dtype=tf.int32),  # video\n    tf.TensorSpec((), dtype=tf.int32),  # segment start\n    tf.TensorSpec((None,), dtype=tf.int32),  # mice ids\n    tf.TensorSpec((), dtype=tf.int32),  # video stop\n    tf.TensorSpec((NUM_BEHAVIORS,)),  # behaviors mask\n))\ntf_ds = tf_ds.batch(1).prefetch(tf.data.AUTOTUNE)\nrow_number = 0\nfor x in tqdm(tf_ds):\n    out = model(x[:4])[0]\n    mask = x[8][0]\n    out = out * mask\n    out = tf.argmax(out, axis=-1)\n    out = out.numpy()\n    segment_start = x[5][0].numpy()\n    video = x[4][0].numpy()\n    id_map = x[6][0].numpy()\n    video_stop = x[7][0].numpy()\n    num_mice = len(id_map)\n    pair_ids = [(id_map[i], id_map[j]) for i in range(num_mice) for j in range(num_mice)]\n    with open('submission.csv', mode='a') as f:\n        writer = csv.writer(f)\n        for pair_id, pair in zip(pair_ids, out):\n            windows = make_windows(pair, segment_start, video_stop)\n            for window in windows:\n                target = f\"mouse{pair_id[1]}\"\n                if pair_id[0] == pair_id[1]:\n                    target = \"self\"\n                writer.writerow([\n                    row_number,\n                    video,\n                    f\"mouse{pair_id[0]}\",\n                    target,\n                    window[0],\n                    window[1],\n                    window[2]\n                ])\n                row_number += 1\n                ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T01:31:57.295604Z","iopub.execute_input":"2026-01-06T01:31:57.295926Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}