data_processing/select_subset_gemini.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0",
   "metadata": {},
   "source": [
    "# Select Subset: Videos Containing Humans (Gemini-based)\n",
    "\n",
    "Filter VidGen-1M captions to retain only videos that contain humans,\n",
    "using Gemini to classify captions rather than keyword matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xinyuy/anaconda3/envs/flexvideo/lib/python3.10/site-packages/google/api_core/_python_version_support.py:255: FutureWarning: You are using a Python version (3.10.20) which Google will stop supporting in new releases of google.api_core once it reaches its end of life (2026-10-04). Please upgrade to the latest Python version, or at least Python 3.11, to continue receiving updates for google.api_core past that date.\n",
      "  warnings.warn(message, FutureWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total videos: 3000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/xinyuy/anaconda3/envs/flexvideo/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/tmp/ipykernel_2704548/1473380499.py:4: FutureWarning: \n",
      "\n",
      "All support for the `google.generativeai` package has ended. It will no longer be receiving \n",
      "updates or bug fixes. Please switch to the `google.genai` package as soon as possible.\n",
      "See README for more details:\n",
      "\n",
      "https://github.com/google-gemini/deprecated-generative-ai-python/blob/main/README.md\n",
      "\n",
      "  import google.generativeai as genai\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import time\n",
    "from pathlib import Path\n",
    "import google.generativeai as genai\n",
    "\n",
    "CAPTION_FILE = \"/mnt/data/xinyuy/datasets/VIDGEN-1M/VidGen_1M_sample_3000.json\"\n",
    "OUTPUT_FILE  = \"vidgen_humans_subset_gemini.json\"\n",
    "GEMINI_API_KEY = \"AQ.Ab8RN6Kur2zYgr--GaMVHH-tDHFB9ZLlY0ShPEIdgFQAm02D6w\"  # or use os.environ[\"GEMINI_API_KEY\"]\n",
    "\n",
    "genai.configure(api_key=GEMINI_API_KEY)\n",
    "model = genai.GenerativeModel(\"gemini-2.5-pro\")  # fast + cheap; swap to gemini-1.5-pro if needed\n",
    "\n",
    "with open(CAPTION_FILE) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "print(f\"Total videos: {len(data)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt and classify_batch() defined.\n"
     ]
    }
   ],
   "source": [
    "# --------------------------------------------------------------------------\n",
    "# Prompt\n",
    "# --------------------------------------------------------------------------\n",
    "# Each call sends a batch of captions so we minimise API round-trips.\n",
    "# Gemini returns a JSON list of booleans in the same order as the input.\n",
    "\n",
    "SYSTEM_PROMPT = \"\"\"\\\n",
    "You are a video-content classifier. Your task is to decide whether each\n",
    "video caption describes a scene that contains at least one human being\n",
    "(only real, not animated or cartoon).\n",
    "\n",
    "Rules:\n",
    "- Count any depiction of a person, regardless of age, gender.\n",
    "- Animated or cartoon human characters doesn't count.\n",
    "- Anthropomorphised animals (e.g. a cartoon pig wearing clothes and acting\n",
    "  human) do NOT count unless explicitly described as human.\n",
    "- If the caption is ambiguous, lean toward False.\n",
    "- If they are shown only partially, it doesn't count (e.g. just hands doesn't count as a human).\n",
    "\n",
    "You will receive a JSON array of caption strings.\n",
    "Respond with ONLY a JSON array of booleans (true/false), one per caption,\n",
    "in the same order. No explanation, no markdown fences — raw JSON only.\n",
    "\n",
    "Example input:\n",
    "[\"A woman jogs through the park.\", \"A dog chases a ball on the grass.\"]\n",
    "\n",
    "Example output:\n",
    "[true, false]\n",
    "\"\"\"\n",
    "\n",
    "def classify_batch(captions: list[str]) -> list[bool]:\n",
    "    \"\"\"Send a batch of captions to Gemini; return a bool list.\"\"\"\n",
    "    prompt = SYSTEM_PROMPT + \"\\n\" + json.dumps(captions, ensure_ascii=False)\n",
    "    response = model.generate_content(prompt)\n",
    "    text = response.text.strip()\n",
    "    # Strip accidental markdown fences if the model adds them\n",
    "    if text.startswith(\"```\"):\n",
    "        text = text.split(\"\\n\", 1)[1].rsplit(\"```\", 1)[0].strip()\n",
    "    results = json.loads(text)\n",
    "    if len(results) != len(captions):\n",
    "        raise ValueError(f\"Got {len(results)} results for {len(captions)} captions\")\n",
    "    return [bool(r) for r in results]\n",
    "\n",
    "print(\"Prompt and classify_batch() defined.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  YES  A woman sits in the driver's seat of a car, smiling at the camera.\n",
      "  NO   A close-up of two glasses of water on a car seat.\n",
      "  NO   A cartoon shark lifts weights in a gym.\n",
      "  YES  A chef chops vegetables on a wooden cutting board.\n"
     ]
    }
   ],
   "source": [
    "# Quick smoke test on a few known examples before the full run\n",
    "test_captions = [\n",
    "    \"A woman sits in the driver's seat of a car, smiling at the camera.\",\n",
    "    \"A close-up of two glasses of water on a car seat.\",\n",
    "    \"A cartoon shark lifts weights in a gym.\",\n",
    "    \"A chef chops vegetables on a wooden cutting board.\",\n",
    "]\n",
    "\n",
    "test_results = classify_batch(test_captions)\n",
    "for cap, res in zip(test_captions, test_results):\n",
    "    print(f\"  {'YES' if res else 'NO ':3s}  {cap[:80]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# --------------------------------------------------------------------------\n",
    "# Full classification run — saves output after every batch\n",
    "# --------------------------------------------------------------------------\n",
    "BATCH_SIZE  = 50   # captions per API call — tune down if you hit token limits\n",
    "SLEEP_SEC   = 1.0  # pause between batches to stay within rate limits\n",
    "\n",
    "labels: list[bool] = []\n",
    "human_videos: list[dict] = []\n",
    "\n",
    "for start in range(0, len(data), BATCH_SIZE):\n",
    "    batch = data[start : start + BATCH_SIZE]\n",
    "    captions = [item[\"caption\"] for item in batch]\n",
    "\n",
    "    # Retry once on transient errors\n",
    "    for attempt in range(2):\n",
    "        try:\n",
    "            batch_labels = classify_batch(captions)\n",
    "            labels.extend(batch_labels)\n",
    "            break\n",
    "        except Exception as e:\n",
    "            if attempt == 0:\n",
    "                print(f\"  Retrying batch {start}–{start+len(batch)} after error: {e}\")\n",
    "                time.sleep(3)\n",
    "            else:\n",
    "                print(f\"  Batch {start}–{start+len(batch)} FAILED twice; marking all False\")\n",
    "                batch_labels = [False] * len(batch)\n",
    "                labels.extend(batch_labels)\n",
    "\n",
    "    # Append human videos from this batch and save immediately\n",
    "    batch_humans = [item for item, has_human in zip(batch, batch_labels) if has_human]\n",
    "    human_videos.extend(batch_humans)\n",
    "    with open(OUTPUT_FILE, \"w\") as f:\n",
    "        json.dump(human_videos, f, indent=2)\n",
    "\n",
    "    end = start + len(batch)\n",
    "    print(f\"  Processed {end}/{len(data)}  \"\n",
    "          f\"(+{len(batch_humans)} human in this batch, {len(human_videos)} total saved)\", flush=True)\n",
    "    time.sleep(SLEEP_SEC)\n",
    "\n",
    "print(f\"\\nDone. Videos with humans: {len(human_videos)} / {len(data)} \"\n",
    "      f\"({100 * len(human_videos) / len(data):.1f}%)\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spot-check a few positives and negatives\n",
    "non_human = [item for item, has_human in zip(data, labels) if not has_human]\n",
    "\n",
    "print(\"=== HUMAN (first 5) ===\")\n",
    "for item in human_videos[:5]:\n",
    "    print(f\"  {item['vid']}\")\n",
    "    print(f\"  {item['caption'][:150]}\")\n",
    "    print()\n",
    "\n",
    "print(\"=== NON-HUMAN (first 5) ===\")\n",
    "for item in non_human[:5]:\n",
    "    print(f\"  {item['vid']}\")\n",
    "    print(f\"  {item['caption'][:150]}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18da51cf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "742a1a9b",
   "metadata": {},
   "source": [
    "## Load saved results and inspect random samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f5bdaf55",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total human videos loaded: 1705\n",
      "\n",
      "=== 10 randomly sampled captions ===\n",
      "[ 1] Y5KW2DKBNyE-Scene-0064\n",
      "     The video shows a football player wearing a red jersey with the number 45 on it. He is seen running on a football field and catching a ball thrown to him. The player is wearing a helmet and gloves, an\n",
      "\n",
      "[ 2] 6wYSDPGJ9B4-Scene-0020\n",
      "     In this video, a woman is seen speaking to the camera. She is wearing a brown shirt and has her hair styled in a bun. The woman is holding a green bracelet on her wrist and is seen pointing to her lef\n",
      "\n",
      "[ 3] ACYXDsw__3A-Scene-0188\n",
      "     The video shows a person's hand holding a white plastic bag while riding a bicycle. The person is seen adjusting the handlebars of the bike and looking at the camera. The background is blurry, but it \n",
      "\n",
      "[ 4] G2HYBR-KvpY-Scene-0004\n",
      "     The video shows a bald man in a pink shirt and tie standing in a news studio. He is holding a piece of paper in his hand and appears to be speaking to the camera. Behind him, there is a news desk with\n",
      "\n",
      "[ 5] g-BqDFLHHec-Scene-0016\n",
      "     In the video, a man is standing in a boxing ring, holding a microphone and speaking into it. He is wearing a black shirt and has a championship belt around his waist. Behind him, there are several peo\n",
      "\n",
      "[ 6] dDbXzGA1-6o-Scene-0002\n",
      "     The video shows a man riding a bicycle on a road. He is wearing a blue shirt and a white helmet. The bicycle is black and white and has a water bottle attached to it. The man is pedaling and appears t\n",
      "\n",
      "[ 7] fqjooS7TYKc-Scene-0073\n",
      "     In the video, a woman is seen wearing an apron and oven mitts while removing a pie from an oven. She carefully opens the oven door and pulls out a tray with the pie on it. The woman then proceeds to c\n",
      "\n",
      "[ 8] swum4mMzsNI-Scene-0020\n",
      "     In the video, a person is seen grilling chicken wings on a barbecue grill. The person uses tongs to turn the wings over to ensure they are cooked evenly on both sides. After turning the wings, the per\n",
      "\n",
      "[ 9] 54tH9emw5PM-Scene-0014\n",
      "     In the video, a man is sitting at a table and eating a sandwich. He takes a bite out of the sandwich and then shows the inside of it to the camera. The man continues to eat the sandwich while looking \n",
      "\n",
      "[10] j9cEoRuRdag-Scene-0014\n",
      "     In the video, a woman is seen standing in front of an arcade game. She is wearing a blue jacket and has blonde hair. The arcade game has a purple and green color scheme, and there are other arcade gam\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import random\n",
    "\n",
    "OUTPUT_FILE = \"vidgen_humans_subset_gemini.json\"\n",
    "PREVIEW_N   = 10  # number of random captions to print\n",
    "\n",
    "with open(OUTPUT_FILE) as f:\n",
    "    human_videos = json.load(f)\n",
    "\n",
    "print(f\"Total human videos loaded: {len(human_videos)}\\n\")\n",
    "\n",
    "sample_preview = random.sample(human_videos, min(PREVIEW_N, len(human_videos)))\n",
    "print(f\"=== {PREVIEW_N} randomly sampled captions ===\")\n",
    "for i, item in enumerate(sample_preview, 1):\n",
    "    print(f\"[{i:2d}] {item['vid']}\")\n",
    "    print(f\"     {item['caption'][:200]}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65c03292",
   "metadata": {},
   "source": [
    "## Select 50 samples and copy videos to pose directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c5129e04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copied   : 50 videos  →  /mnt/data/xinyuy/datasets/flexcombine-bench/single/pose\n",
      "Missing  : 0 videos\n",
      "JSON saved: /mnt/data/xinyuy/datasets/flexcombine-bench/single/pose/pose_subset_50.json\n"
     ]
    }
   ],
   "source": [
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "VIDEO_SRC_DIR = Path(\"/mnt/data/xinyuy/datasets/VIDGEN-1M/videos\")\n",
    "DEST_DIR      = Path(\"/mnt/data/xinyuy/datasets/flexcombine-bench/single/pose\")\n",
    "SUBSET_JSON   = DEST_DIR / \"pose_subset_50.json\"\n",
    "N_SAMPLES     = 50\n",
    "\n",
    "random.seed(42)  # reproducible selection\n",
    "subset = random.sample(human_videos, N_SAMPLES)\n",
    "\n",
    "DEST_DIR.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "missing, copied = [], []\n",
    "for item in subset:\n",
    "    src = VIDEO_SRC_DIR / f\"{item['vid']}.mp4\"\n",
    "    dst = DEST_DIR / f\"{item['vid']}.mp4\"\n",
    "    if src.exists():\n",
    "        shutil.copy2(src, dst)\n",
    "        copied.append(item)\n",
    "    else:\n",
    "        print(f\"  WARNING: video not found — {src}\")\n",
    "        missing.append(item)\n",
    "\n",
    "with open(SUBSET_JSON, \"w\") as f:\n",
    "    json.dump(subset, f, indent=2)\n",
    "\n",
    "print(f\"Copied   : {len(copied)} videos  →  {DEST_DIR}\")\n",
    "print(f\"Missing  : {len(missing)} videos\")\n",
    "print(f\"JSON saved: {SUBSET_JSON}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "792eebb0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flexvideo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}