{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":" # Setup & Imports","metadata":{}},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom IPython.display import Video, display\n\ndataset_path = Path('/kaggle/input/MABe-mouse-behavior-detection')\ntrain = pd.read_csv(dataset_path / 'train.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T09:58:50.290228Z","iopub.execute_input":"2025-09-28T09:58:50.290536Z","iopub.status.idle":"2025-09-28T09:58:50.370054Z","shell.execute_reply.started":"2025-09-28T09:58:50.290515Z","shell.execute_reply":"2025-09-28T09:58:50.368935Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load Dataset Metadata","metadata":{}},{"cell_type":"code","source":"# Load train metadata\ntrain_csv_path = \"/kaggle/input/MABe-mouse-behavior-detection/train.csv\"\ntrain_df = pd.read_csv(train_csv_path)\n\n# Filter out MABe22 if needed (as per baseline)\ntrain_df = train_df[~train_df['lab_id'].str.startswith('MABe22')].reset_index(drop=True)\n\nprint(f\"Loaded {len(train_df)} training episodes.\")\ntrain_df.head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T09:58:53.830543Z","iopub.execute_input":"2025-09-28T09:58:53.830903Z","iopub.status.idle":"2025-09-28T09:58:53.933053Z","shell.execute_reply.started":"2025-09-28T09:58:53.830879Z","shell.execute_reply":"2025-09-28T09:58:53.931902Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Function: Render Episode Video with Tracking Overlay","metadata":{}},{"cell_type":"code","source":"import os, json, math, tempfile, subprocess, base64\nfrom functools import lru_cache\nfrom typing import Optional, List, Tuple, Dict\nimport numpy as np, pandas as pd, cv2\nfrom IPython.display import Video\nfrom tqdm.auto import tqdm\n\ndef distinct_colors_bgr(n, sat=200, val=235, hue_offset=0):\n    if n <= 0: return []\n    hues = (np.linspace(0,179,n,endpoint=False)+hue_offset)%180\n    hsv  = np.stack([hues, np.full(n,sat), np.full(n,val)],1).astype(np.uint8)[None,...]\n    return [tuple(map(int,c)) for c in cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0]]\n\ndef _extract_bps(meta: Optional[pd.DataFrame]):\n    if meta is None or meta.empty: return None\n    row = meta.iloc[0]\n    for col in [\"body_parts_tracked\",\"bodyparts_tracked\",\"body_parts\",\"tracked_body_parts\"]:\n        if col in meta.columns:\n            v = row[col]\n            if isinstance(v,list): return v\n            if isinstance(v,str):\n                try:\n                    p = json.loads(v)\n                    if isinstance(p,list): return p\n                except: pass\n    return None\n\ndef _edges(available: List[str]):\n    cand=[(\"nose\",\"neck\"),(\"neck\",\"ear_left\"),(\"neck\",\"ear_right\"),\n          (\"neck\",\"hip_left\"),(\"neck\",\"hip_right\"),\n          (\"hip_left\",\"tail_base\"),(\"hip_right\",\"tail_base\"),\n          (\"body_center\",\"neck\"),(\"body_center\",\"nose\"),(\"body_center\",\"tail_base\")]\n    seen=set(); out=[]\n    for a,b in cand:\n        if a in available and b in available:\n            k=tuple(sorted((a,b)))\n            if k not in seen: seen.add(k); out.append((a,b))\n    return out\n\ndef export_square_nocrop_mp4_embed(\n    df_episode: pd.DataFrame,\n    df_episode_meta: Optional[pd.DataFrame]=None,\n    df_annotation: Optional[pd.DataFrame]=None,   # <— NEW\n    *,\n    size_px:int=640, out_fps:Optional[int]=None,\n    frame_start:Optional[int]=None, frame_end:Optional[int]=None, frame_stride:int=1,\n    show_skeleton:bool=True, show_trails:bool=False, trail_len:int=25,\n    bodyparts:Optional[List[str]]=None, marker_radius:int=3, border_thickness:int=2, line_thickness:int=2,\n    # action caption options (NEW)\n    show_actions: bool=True,\n    action_anchor_priority: Tuple[str,...]=(\"nose\",\"neck\",\"body_center\"),  # where to place label\n    action_font_scale: float=0.5,\n    action_thickness: int=1,\n    action_text_color: Tuple[int,int,int]=(0,0,0),        # BGR\n    action_box_bg: Tuple[int,int,int]=(255,255,255),      # BGR\n    action_box_pad: int=3,\n    action_box_border: int=1,\n    action_box_max_width_ratio: float=0.9,                # clamp box within canvas\n    action_joiner: str=\" | \",                             # if multiple concurrent actions\n    # output / encoding\n    out_path:str=\"episode_square.mp4\",\n    ffmpeg_preset:str=\"veryfast\", ffmpeg_crf:int=23,\n    return_embed:bool=True, max_embed_mb:int=120, display_width:int=720\n):\n    # basic validation\n    need={\"video_frame\",\"mouse_id\",\"bodypart\",\"x\",\"y\"}\n    miss=need-set(df_episode.columns)\n    if miss: raise ValueError(f\"df_episode missing: {sorted(miss)}\")\n    df=df_episode.sort_values([\"video_frame\",\"mouse_id\",\"bodypart\"]).reset_index(drop=True)\n\n    # source plane (no crop)\n    if df_episode_meta is not None and not df_episode_meta.empty:\n        r=df_episode_meta.iloc[0]\n        W=r.get(\"video_width_pix\",np.nan); H=r.get(\"video_height_pix\",np.nan)\n        if np.isfinite(W) and np.isfinite(H):\n            xmin=ymin=0.0; src_w=float(W); src_h=float(H)\n        else:\n            xmin, xmax = float(df.x.min()), float(df.x.max())\n            ymin, ymax = float(df.y.min()), float(df.y.max())\n            src_w=max(1.0,xmax-xmin); src_h=max(1.0,ymax-ymin)\n    else:\n        xmin, xmax = float(df.x.min()), float(df.x.max())\n        ymin, ymax = float(df.y.min()), float(df.y.max())\n        src_w=max(1.0,xmax-xmin); src_h=max(1.0,ymax-ymin); xmin=ymin=0.0\n\n    N=int(size_px); s=min(N/src_w, N/src_h); pad_x=0.5*(N-s*src_w); pad_y=0.5*(N-s*src_h)\n    def map_xy(x,y):\n        X=pad_x + s*(float(x)-xmin); Y=pad_y + s*(float(y)-ymin)\n        return max(0,min(N-1,int(round(X)))), max(0,min(N-1,int(round(Y))))\n\n    # fps\n    if out_fps is None:\n        v = df_episode_meta.iloc[0][\"frames_per_second\"] if (df_episode_meta is not None and not df_episode_meta.empty and \"frames_per_second\" in df_episode_meta.columns) else None\n        out_fps = int(v) if (v is not None and pd.notna(v)) else 12\n\n    # frames\n    all_frames=np.sort(df.video_frame.dropna().unique())\n    if all_frames.size==0: raise ValueError(\"No frames.\")\n    if frame_start is None: frame_start=int(all_frames[0])\n    if frame_end   is None: frame_end  =int(all_frames[-1])\n    frames=[int(f) for f in all_frames if frame_start<=f<=frame_end][::max(1,int(frame_stride))]\n    if not frames: raise ValueError(\"No frames after filtering.\")\n    frames_np = np.array(frames, dtype=int)\n\n    # bodyparts/colors\n    data_bps=sorted(df.bodypart.dropna().unique().tolist())\n    meta_bps=_extract_bps(df_episode_meta)\n    available_bps=meta_bps if meta_bps else data_bps\n    if bodyparts is None: bodyparts=available_bps\n    else: bodyparts=[bp for bp in bodyparts if bp in available_bps]\n    edges=_edges(bodyparts) if show_skeleton else []\n    bp_cols=distinct_colors_bgr(len(available_bps), sat=200, val=235, hue_offset=0)\n    bp_to_bgr={bp: bp_cols[i] for i,bp in enumerate(available_bps)}\n    mice=sorted(df.mouse_id.dropna().unique().tolist())\n    mouse_cols=distinct_colors_bgr(len(mice), sat=240, val=220, hue_offset=17)\n    mouse_to_bgr={mid: mouse_cols[i] for i,mid in enumerate(mice)}\n\n    # --- Build fast lookup for annotations: (frame_idx, agent_id) -> [actions]\n    ann_map: Dict[Tuple[int,int], List[str]] = {}\n    \n    if show_actions and df_annotation is not None and not df_annotation.empty:\n        req_cols = {\"agent_id\",\"action\",\"start_frame\",\"stop_frame\"}\n        miss = req_cols - set(df_annotation.columns)\n        if miss:\n            raise ValueError(f\"df_annotation missing: {sorted(miss)}\")\n        # keep only agents in this episode\n        ann = df_annotation.copy()\n        ann = ann[ann[\"agent_id\"].isin(mice)]\n        # normalize numeric\n        for c in (\"start_frame\",\"stop_frame\"):\n            ann[c] = pd.to_numeric(ann[c], errors=\"coerce\").astype(\"Int64\")\n        ann = ann.dropna(subset=[\"start_frame\",\"stop_frame\",\"action\",\"agent_id\"])\n        if not ann.empty:\n            ann[\"start_frame\"] = ann[\"start_frame\"].astype(int)\n            ann[\"stop_frame\"]  = ann[\"stop_frame\"].astype(int)\n            if ann[\"start_frame\"].gt(ann[\"stop_frame\"]).any():\n                # swap where needed\n                bad = ann[\"start_frame\"] > ann[\"stop_frame\"]\n                tmp = ann.loc[bad, \"start_frame\"].values\n                ann.loc[bad, \"start_frame\"] = ann.loc[bad, \"stop_frame\"].values\n                ann.loc[bad, \"stop_frame\"]  = tmp\n            # For each annotation row, add entries for intersecting frames (use binary search into frames)\n            for row in ann.itertuples(index=False):\n                a0, a1 = int(row.start_frame), int(row.stop_frame)\n                # overlap with rendered frames\n                i0 = int(np.searchsorted(frames_np, a0, side=\"left\"))\n                i1 = int(np.searchsorted(frames_np, a1, side=\"right\"))\n                if i0 >= i1: \n                    continue\n                agent = int(row.agent_id) if pd.api.types.is_numeric_dtype(type(row.agent_id)) or isinstance(row.agent_id, (int,np.integer)) else row.agent_id\n                act   = str(row.action)\n                for f in frames_np[i0:i1]:\n                    ann_map.setdefault((int(f), agent), []).append(act)\n\n    @lru_cache(None)\n    def fm(frame, mid):\n        sub=df[(df.video_frame==frame)&(df.mouse_id==mid)]\n        return sub[sub.bodypart.isin(bodyparts)]\n\n    trails={mid:[] for mid in mice}\n\n    # write tmp with OpenCV, then ffmpeg → H.264\n    tmp_output_path = \"tmp_\" + out_path\n    vw=cv2.VideoWriter(tmp_output_path, cv2.VideoWriter_fourcc(*\"mp4v\"), out_fps, (N,N))\n    if not vw.isOpened(): raise RuntimeError(\"OpenCV VideoWriter failed.\")\n\n    try:\n        for f in tqdm(frames, total=len(frames), desc=\"Rendering frames\", unit=\"f\",\n              smoothing=0.1, mininterval=0.1, leave=False):\n            canvas=np.full((N,N,3),(255,255,255),dtype=np.uint8)  # white background\n            for mid in mice:\n                sub=fm(f,mid)\n                if sub.empty: \n                    # even if no body part this frame, we might still want to write the action near last trail point (optional).\n                    continue\n                bp_map={}\n                for r in sub.itertuples(index=False):\n                    x,y=map_xy(r.x,r.y); bp_map[r.bodypart]=(x,y)\n                # skeleton\n                if show_skeleton:\n                    for p1,p2 in edges:\n                        if p1 in bp_map and p2 in bp_map:\n                            x1,y1=bp_map[p1]; x2,y2=bp_map[p2]\n                            cv2.line(canvas,(x1,y1),(x2,y2), bp_to_bgr.get(p1,(0,0,0)), thickness=line_thickness, lineType=cv2.LINE_AA)\n                # joints\n                for r in sub.itertuples(index=False):\n                    x,y=bp_map[r.bodypart]\n                    cv2.circle(canvas,(x,y),marker_radius, bp_to_bgr.get(r.bodypart,(0,0,0)), -1, cv2.LINE_AA)\n                    cv2.circle(canvas,(x,y),marker_radius, mouse_to_bgr.get(mid,(0,0,0)), border_thickness, cv2.LINE_AA)\n\n                # trails\n                ref = None\n                for k in action_anchor_priority:\n                    if k in bp_map: \n                        ref = bp_map[k]; break\n                if ref is None and bp_map:\n                    xs,ys=zip(*bp_map.values()); ref=(int(round(np.mean(xs))), int(round(np.mean(ys))))\n\n                if show_trails and ref is not None:\n                    t=trails[mid]; t.append(ref); \n                    if len(t)>trail_len: trails[mid]=t[-trail_len:]\n                    pts=np.array(trails[mid],np.int32).reshape(-1,1,2)\n                    cv2.polylines(canvas,[pts],False, mouse_to_bgr[mid], 1, cv2.LINE_AA)\n\n                # action captions\n                if show_actions and ref is not None:\n                    acts = ann_map.get((int(f), mid))\n                    if acts:\n                        # make unique but stable\n                        uniq = list(dict.fromkeys(a.strip() for a in acts if str(a).strip()))\n                        if uniq:\n                            label = action_joiner.join(uniq)\n                            # measure text\n                            font = cv2.FONT_HERSHEY_SIMPLEX\n                            (tw, th), baseline = cv2.getTextSize(label, font, action_font_scale, action_thickness)\n                            pad = int(action_box_pad)\n                            box_w = min(tw + 2*pad, int(action_box_max_width_ratio * N))\n                            # position the box above the anchor, clamp within canvas\n                            x0 = int(ref[0] - box_w//2)\n                            y0 = int(ref[1] - marker_radius - 6 - th - baseline - 2*pad)\n                            x0 = max(0, min(N - box_w, x0))\n                            y0 = max(0, y0)\n                            # background box\n                            x1 = x0 + box_w\n                            y1 = y0 + th + baseline + 2*pad\n                            cv2.rectangle(canvas, (x0,y0), (x1,y1), action_box_bg, thickness=-1)\n                            if action_box_border > 0:\n                                cv2.rectangle(canvas, (x0,y0), (x1,y1), mouse_to_bgr.get(mid,(0,0,0)), thickness=action_box_border, lineType=cv2.LINE_AA)\n                            # text (left padded)\n                            tx = x0 + pad\n                            ty = y1 - baseline - pad\n                            cv2.putText(canvas, label, (tx,ty), font, action_font_scale, action_text_color, action_thickness, cv2.LINE_AA)\n\n            # Draw frame number (top-left corner)\n            font = cv2.FONT_HERSHEY_SIMPLEX\n            text = f\"Frame {f}\"\n            scale = 0.5\n            thickness = 1\n            color = (0, 0, 0)   # black\n            margin = 50\n            \n            (text_w, text_h), baseline = cv2.getTextSize(text, font, scale, thickness)\n            x = margin\n            y = margin\n            \n            cv2.putText(canvas, text, (x, y), font, scale, color, thickness, cv2.LINE_AA)\n\n            vw.write(canvas)\n    finally:\n        vw.release()\n\n    # Re-encode with libx264 baseline/yuv420p-ish quality settings\n    subprocess.run(\n        [\"ffmpeg\", \"-y\", \"-i\", tmp_output_path, \"-crf\", str(ffmpeg_crf), \"-preset\", ffmpeg_preset, \"-vcodec\", \"libx264\", \"-pix_fmt\", \"yuv420p\", out_path],\n        check=False,\n        stdout=subprocess.DEVNULL,\n        stderr=subprocess.DEVNULL\n    )\n    try:\n        os.remove(tmp_output_path)\n    except OSError:\n        pass\n\n    return out_path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T09:59:01.920145Z","iopub.execute_input":"2025-09-28T09:59:01.920467Z","iopub.status.idle":"2025-09-28T09:59:02.094885Z","shell.execute_reply.started":"2025-09-28T09:59:01.920444Z","shell.execute_reply":"2025-09-28T09:59:02.093821Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pick Episode & Render!","metadata":{}},{"cell_type":"code","source":"from IPython.display import Video, display\n\ndef create_video(video_idx):\n    if video_idx is None or not (0 <= int(video_idx) < len(train)):\n        raise IndexError(f\"video_idx must be in [0, {len(train)-1}]\")\n        \n    # Example: iterate over *all* episodes described in the meta dataframe\n    row = train.iloc[int(video_idx)]\n            \n    lab = str(row[\"lab_id\"])\n    vid = str(int(row[\"video_id\"]))  # cast to int then str for clean filename\n\n    # Build the parquet file path\n    pathto_tracking_data = os.path.join(\n        dataset_path,\n        \"train_tracking\",\n        lab,\n        f\"{vid}.parquet\"\n    )\n\n    print(f\"Episode Tracking {video_idx} → {pathto_tracking_data}\")\n\n    # Load the tracking dataframe\n    df_tracking = pd.read_parquet(pathto_tracking_data)\n\n    # Build the parquet file path\n    pathto_annotation_data = os.path.join(\n        dataset_path,\n        \"train_annotation\",\n        lab,\n        f\"{vid}.parquet\"\n    )\n\n    print(f\"Episode Annotation {video_idx} → {pathto_annotation_data}\")\n\n    # Load the tracking dataframe\n    df_annotation = pd.read_parquet(pathto_annotation_data)    \n\n    vid = export_square_nocrop_mp4_embed(\n        df_tracking, train.iloc[[video_idx]], df_annotation,\n        size_px=640, frame_stride=2, show_trails=False,\n        out_path=f\"episode-{video_idx}.mp4\",  # you also keep a file\n        return_embed=True, max_embed_mb=120, display_width=720\n    )\n\n    display(Video(data=vid,\n              embed=True,\n              height=int(640),\n              width=int(640))\n       )\n\n    return df_annotation # For debugging purposes","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T09:59:09.202517Z","iopub.execute_input":"2025-09-28T09:59:09.202886Z","iopub.status.idle":"2025-09-28T09:59:09.213623Z","shell.execute_reply.started":"2025-09-28T09:59:09.202860Z","shell.execute_reply":"2025-09-28T09:59:09.212157Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"index = 1\nannot = create_video(index)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T10:56:59.057507Z","iopub.status.idle":"2025-09-28T10:56:59.057864Z","shell.execute_reply.started":"2025-09-28T10:56:59.057684Z","shell.execute_reply":"2025-09-28T10:56:59.057724Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"index = 0\nannot = create_video(index)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-28T09:59:12.663510Z","iopub.execute_input":"2025-09-28T09:59:12.663864Z","iopub.status.idle":"2025-09-28T10:03:19.365019Z","shell.execute_reply.started":"2025-09-28T09:59:12.663841Z","shell.execute_reply":"2025-09-28T10:03:19.363654Z"}},"outputs":[],"execution_count":null}]}