getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with Data Generation SDK\n",
    "\n",
    "We are going to generate a dataset of squat videos with instructions how to perform them, so that we can train an AI pesonal trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from datagen import DatagenConfig\n",
    "\n",
    "config_params = {\n",
    "    'openai': {\n",
    "        'type': 'azure', # openai/azure\n",
    "        'temperature': '1',\n",
    "        'deployment': 'gpt4o' # model for openai / deployment for azure\n",
    "    },\n",
    "    'data_dir': './tmp/squats'\n",
    "}\n",
    "\n",
    "!mkdir -p {config_params['data_dir']}\n",
    "\n",
    "# this config handles all the bookeeping so you need to pass it everywhere.\n",
    "config = DatagenConfig(**config_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get a list of search queries to search for videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['how to do squats', 'squat exercise tutorial']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datagen import get_queries\n",
    "queries = get_queries(\n",
    "    config=config,\n",
    "    prompt='I want to find instructional videos about how to do squats.',\n",
    "    num_queries=2\n",
    ")\n",
    "queries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download video information for each query.\n",
    "\n",
    "We'll get 2 videos for each query.<br>\n",
    "One video might be found with multiple queries, so we might get less than `n_queries*videos_per_query` videos.<br>\n",
    "If you want to get all youtube videos for a query, don't pass `videos_per_query` parameter.\n",
    "\n",
    "You can limit the search to only videos licensed with Creative Commons (as indicated by youtube).<br>\n",
    "As this search isn't directly implemented in searching libraries yet, we search for all videos and filter for license afterwards.<br>\n",
    "Unfortunately, this way you will likely get very few results, so use with caution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:01<00:00,  1.50it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['EbOPpWi4L8s', 'gcNh17Ckjgg', 'xqvCmoLULNY', 'IB_icWRzi4E']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datagen import get_video_ids\n",
    "ids = get_video_ids(queries, config=config, videos_per_query=2,only_creative_commons=False)\n",
    "ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open(config.data_dir / 'video_ids.json') as f:\n",
    "    ids = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download videos and autogenerated subtitles\n",
    "\n",
    "You can change sub languages, formats etc with `yt_dlp_opts` dictionary (refer to https://github.com/yt-dlp/yt-dlp).<br>\n",
    "The SDK is expecting `.mp4` video files (for now), so don't change that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[youtube] Extracting URL: https://www.youtube.com/watch?v=gcNh17Ckjgg\n",
      "[youtube] gcNh17Ckjgg: Downloading webpage\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[youtube] gcNh17Ckjgg: Downloading ios player API JSON\n",
      "[youtube] gcNh17Ckjgg: Downloading web creator player API JSON\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: [youtube] gcNh17Ckjgg: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:23:14.371408 Error at video gcNh17Ckjgg, skipping\n",
      "2024-08-09 11:23:14.371447 ERROR: [youtube] gcNh17Ckjgg: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n",
      "[youtube] Extracting URL: https://www.youtube.com/watch?v=IB_icWRzi4E\n",
      "[youtube] IB_icWRzi4E: Downloading webpage\n",
      "[youtube] IB_icWRzi4E: Downloading ios player API JSON\n",
      "[youtube] IB_icWRzi4E: Downloading web creator player API JSON\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: [youtube] IB_icWRzi4E: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:23:14.974700 Error at video IB_icWRzi4E, skipping\n",
      "2024-08-09 11:23:14.974742 ERROR: [youtube] IB_icWRzi4E: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n",
      "[youtube] Extracting URL: https://www.youtube.com/watch?v=xqvCmoLULNY\n",
      "[youtube] xqvCmoLULNY: Downloading webpage\n",
      "[youtube] xqvCmoLULNY: Downloading ios player API JSON\n",
      "[youtube] xqvCmoLULNY: Downloading web creator player API JSON\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: [youtube] xqvCmoLULNY: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:23:15.656008 Error at video xqvCmoLULNY, skipping\n",
      "2024-08-09 11:23:15.656044 ERROR: [youtube] xqvCmoLULNY: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n",
      "[youtube] Extracting URL: https://www.youtube.com/watch?v=EbOPpWi4L8s\n",
      "[youtube] EbOPpWi4L8s: Downloading webpage\n",
      "[youtube] EbOPpWi4L8s: Downloading ios player API JSON\n",
      "[youtube] EbOPpWi4L8s: Downloading web creator player API JSON\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR: [youtube] EbOPpWi4L8s: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:23:16.219252 Error at video EbOPpWi4L8s, skipping\n",
      "2024-08-09 11:23:16.219294 ERROR: [youtube] EbOPpWi4L8s: Sign in to confirm you’re not a bot. This helps protect our community. Learn more\n"
     ]
    }
   ],
   "source": [
    "from datagen import download_videos\n",
    "# a proxy is necessary if getting \"Sign in to confirm you’re not a bot.\"\n",
    "download_videos(ids, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Detect segments from video\n",
    "\n",
    "We will use the clip version because it's much faster than gpt4o, but we'll need a GPU.\n",
    "You can also try using CPU for debugging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoProcessor, AutoModel\n",
    "\n",
    "# remove .cuda() for cpu\n",
    "# SIGLIP outputs independent probs as opposed to CLIP that outputs multiclass probs\n",
    "device = 'cuda' # or 'cuda:0'\n",
    "model = AutoModel.from_pretrained(\"google/siglip-so400m-patch14-384\").to(device)\n",
    "processor = AutoProcessor.from_pretrained(\"google/siglip-so400m-patch14-384\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/4 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:29:42.322734 grabbing video EbOPpWi4L8s: 193 frames\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running clip on batch [(0, 100, 'EbOPpWi4L8s')]...\n",
      "2024-08-09 11:30:03.028306 grabbing video xqvCmoLULNY: 97 frames\n",
      "running clip on batch [(0, 93, 'EbOPpWi4L8s'), (93, 100, 'xqvCmoLULNY')]...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 1/4 [00:39<01:58, 39.37s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video EbOPpWi4L8s completed - 4 segments detected\n",
      "[Segment(start_timestamp='00:00:04.250', end_timestamp='00:00:13.250', fps=23.976023976023978, segment_info=None, video_id='EbOPpWi4L8s'), Segment(start_timestamp='00:00:35.750', end_timestamp='00:00:52.750', fps=23.976023976023978, segment_info=None, video_id='EbOPpWi4L8s'), Segment(start_timestamp='00:01:01.250', end_timestamp='00:01:06.750', fps=23.976023976023978, segment_info=None, video_id='EbOPpWi4L8s'), Segment(start_timestamp='00:01:08.250', end_timestamp='00:01:17.250', fps=23.976023976023978, segment_info=None, video_id='EbOPpWi4L8s')]\n",
      "2024-08-09 11:30:22.492963 grabbing video gcNh17Ckjgg: 870 frames\n",
      "running clip on batch [(0, 90, 'xqvCmoLULNY'), (90, 100, 'gcNh17Ckjgg')]...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 3/4 [00:59<00:17, 17.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video xqvCmoLULNY completed - 2 segments detected\n",
      "[Segment(start_timestamp='00:00:19.750', end_timestamp='00:00:27.250', fps=23.976023976023978, segment_info=None, video_id='xqvCmoLULNY'), Segment(start_timestamp='00:00:28.750', end_timestamp='00:00:41.750', fps=23.976023976023978, segment_info=None, video_id='xqvCmoLULNY')]\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "running clip on batch [(0, 100, 'gcNh17Ckjgg')]...\n",
      "2024-08-09 11:33:23.247812 grabbing video bEv6CCg2BC8: 1318 frames\n",
      "running clip on batch [(0, 60, 'gcNh17Ckjgg'), (60, 100, 'bEv6CCg2BC8')]...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6it [04:00, 43.86s/it]                       "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video gcNh17Ckjgg completed - 14 segments detected\n",
      "[Segment(start_timestamp='00:00:00.750', end_timestamp='00:00:33.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:00:40.750', end_timestamp='00:00:52.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:01:41.750', end_timestamp='00:01:48.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:01:56.750', end_timestamp='00:02:04.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:02:16.250', end_timestamp='00:02:20.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:02:27.750', end_timestamp='00:02:38.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:02:40.250', end_timestamp='00:02:54.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:03:10.250', end_timestamp='00:03:16.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:03:20.250', end_timestamp='00:03:28.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:03:31.750', end_timestamp='00:03:54.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:04:04.250', end_timestamp='00:04:35.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:04:38.750', end_timestamp='00:05:46.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:06:09.250', end_timestamp='00:06:51.250', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg'), Segment(start_timestamp='00:06:55.750', end_timestamp='00:07:07.750', fps=29.97002997002997, segment_info=None, video_id='gcNh17Ckjgg')]\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 100, 'bEv6CCg2BC8')]...\n",
      "running clip on batch [(0, 78, 'bEv6CCg2BC8')]...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "10it [08:15, 49.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video bEv6CCg2BC8 completed - 13 segments detected\n",
      "[Segment(start_timestamp='00:00:10.250', end_timestamp='00:00:30.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:00:58.250', end_timestamp='00:01:15.250', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:01:25.250', end_timestamp='00:01:55.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:02:03.750', end_timestamp='00:02:24.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:02:29.750', end_timestamp='00:02:43.250', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:03:03.750', end_timestamp='00:03:43.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:03:54.750', end_timestamp='00:04:46.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:04:53.750', end_timestamp='00:06:08.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:06:14.250', end_timestamp='00:06:32.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:06:35.250', end_timestamp='00:08:25.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:08:33.750', end_timestamp='00:08:40.250', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:08:46.750', end_timestamp='00:09:07.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8'), Segment(start_timestamp='00:09:27.250', end_timestamp='00:09:33.750', fps=24.0, segment_info=None, video_id='bEv6CCg2BC8')]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from datagen import detect_segments_clip\n",
    "\n",
    "from typing import Optional\n",
    "from langchain.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "detect_segments_clip(\n",
    "    # video_ids=['KvRK5Owqzgw'],\n",
    "    text_prompts='a person doing squats', # that's the text for CLIP to compare to images. You can provide a list of texts to use average distance.\n",
    "    model=model,\n",
    "    processor=processor,\n",
    "    fps_sampling=2, # the more fps, the more granular segment borders and more precise segments, at the cost of speed.\n",
    "    device='cuda', # 'cpu' for local\n",
    "    frames_per_batch=100, # 100 frames use about 10GB GPU RAM, so batch to fill your GPU RAM.\n",
    "    config=config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each video we get a list of segments:\n",
    "```\n",
    "[\n",
    "    ...\n",
    "    {\n",
    "        \"start_timestamp\": \"00:00:32.500\",\n",
    "        \"end_timestamp\": \"00:00:41.500\",\n",
    "        \"fps\": 29.97002997002997,\n",
    "        \"segment_info\": null, # not used with clip, but could be used with gpt4o\n",
    "        \"video_id\": \"KvRK5Owqzgw\"\n",
    "    },\n",
    "    ...\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annotaion step 1: extract information (clues) from transcript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:40:24.266342 bEv6CCg2BC8 - started\n",
      "2024-08-09 11:40:24.266955 bEv6CCg2BC8 part 0 - started\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:40:46.369636 bEv6CCg2BC8 part 1 - started\n",
      "2024-08-09 11:41:27.258525 bEv6CCg2BC8 part 2 - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:41:29.519412 bEv6CCg2BC8 - done\n",
      "2024-08-09 11:41:29.521122 gcNh17Ckjgg - started\n",
      "2024-08-09 11:41:29.522229 gcNh17Ckjgg part 0 - started\n",
      "2024-08-09 11:41:47.972894 gcNh17Ckjgg part 1 - started\n",
      "2024-08-09 11:42:19.042026 gcNh17Ckjgg part 2 - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:42:48.376612 gcNh17Ckjgg - done\n",
      "2024-08-09 11:42:48.378026 xqvCmoLULNY - started\n",
      "2024-08-09 11:42:48.379028 xqvCmoLULNY part 0 - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:43:08.990347 xqvCmoLULNY - done\n",
      "2024-08-09 11:43:08.992046 EbOPpWi4L8s - started\n",
      "2024-08-09 11:43:08.993245 EbOPpWi4L8s part 0 - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [03:04<00:00, 46.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:43:28.475395 EbOPpWi4L8s - done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from datagen import generate_clues\n",
    "\n",
    "human_prompt = \"\"\"\n",
    "The provided video is a tutorial about how to perform squats. \n",
    "\n",
    "I need to understand HOW THE PERSON SHOWN IN EACH SEGMENT PERFORMS SQUATS IN THIS SEGMENT.\n",
    "What is done correctly.\n",
    "What mistakes they make. Why these mistakes happen.\n",
    "How these mistakes could be improved.\n",
    "\n",
    "It is very improtant that the information that you provide would describe how the person shown in the segment is doing squats, and not some generic advice that is unrelated to the visual information.\n",
    "\"\"\"\n",
    "\n",
    "from time import sleep\n",
    "clues = generate_clues(\n",
    "    # video_ids=['byxWus7BwfQ'],\n",
    "    config=config,\n",
    "    human_prompt=human_prompt,\n",
    "    segments_per_call=5, # the output might be quite long, so need to limit number of segments per gpt call to respect max output legnth\n",
    "    raise_on_error=True, # interrupt when encountering an error. Useful for debugging.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annotaion step 2: extract information from transcript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:43:36.608390 bEv6CCg2BC8 - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:43:51.283055 bEv6CCg2BC8 - done\n",
      "2024-08-09 11:43:51.284728 gcNh17Ckjgg - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:44:03.635522 gcNh17Ckjgg - done\n",
      "2024-08-09 11:44:03.637116 xqvCmoLULNY - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:44:07.351109 xqvCmoLULNY - done\n",
      "2024-08-09 11:44:07.352673 EbOPpWi4L8s - started\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:40<00:00, 10.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-08-09 11:44:17.226695 EbOPpWi4L8s - done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from datagen import generate_annotations\n",
    "from langchain.pydantic_v1 import BaseModel, Field\n",
    "from typing import Optional\n",
    "\n",
    "# This information that will be extracted for each segment from the transcript and data from the previous step.\n",
    "# This is the most important part for the annotation, and getting good results requires a lot of experimenting.\n",
    "\n",
    "\n",
    "human_prompt = '''\n",
    "You are given a JSON object that contains clues about segments of a video with timecodes.\n",
    "!!!! For each segment provided in a JSON object you need to answer on the following questions:\n",
    "1. Given the data found in the JSON object, what is a probability that this part contains a footage of a person doing squats? [the answer could be only \"high\", \"medium\", \"low\", or null (if impossible to infer from the provided data)]\n",
    "2. Given the data found in the JSON object and even if the answer on the previous question is \"low\", does this person do squats right, wrong, or mixed? [the answer could be only \"right\", \"wrong\", \"mixed\", or null (if impossible to infer from the provided data)]\n",
    "3. Given the data found in the JSON object, what exactly does thing person do right and/or wrong regarding their squats technique? [the answer should be clear and focused on body parts]\n",
    "4. If the answer on the previous question contains description of wrong technique, explain how to fix these mistakes using your \"own knowledge\" like you are a sports coach.\n",
    "'''\n",
    "\n",
    "class SegmentFeedback(BaseModel):\n",
    "    '''\n",
    "—> GOOD EXAMPLES:\n",
    "    \"wrong\":\"Knees caving in: This can stress the knees and reduce effectiveness\"\n",
    "    \"correction\":\"Focus on keeping knees aligned with your toes.\"\n",
    "    \"wrong\":\"Rounding the back: This increases the risk of back injuries\"\n",
    "    \"correction\":\"Keep your chest up and maintain a neutral spine throughout the movement.\"\n",
    "    \"wrong\":\"Heels are lifting off the ground: this shifts the weight forward, reducing stability\"\n",
    "    \"correction\":\" Keep your weight on your heels and press through them as you rise.\"\n",
    "    \"right\":\"Chest and shoulders: The chest is up, and the shoulders are back, maintaining an upright torso.\"\n",
    "    \"correction\":null\n",
    "—> BAD EXAMPLES:\n",
    "    \"wrong\":\"knees\"\n",
    "    \"correction\":\"fix knees\"\n",
    "    \"wrong\":\"back looks funny\"\n",
    "    \"correction\":\"make back better\"\n",
    "    \"wrong\":\"feet are doing something\"\n",
    "    \"correction\":\"feet should be different\"\n",
    "    \"right\":\"arms\"\n",
    "    \"correction\":\"arms are fine i think\"\n",
    "—> BAD EXAMPLES END HERE\n",
    "    '''\n",
    "    right: Optional[str] = Field(description='what was right in the performance')\n",
    "    wrong: Optional[str] = Field(description='what was wrong in the performance')\n",
    "    correction: Optional[str] = Field(description='how and in what ways it the performance could be improved')\n",
    "\n",
    "# The segment timestamps are taken from the provided information.\n",
    "class SegmentAnnotation(BaseModel):\n",
    "    squats_probability: Optional[str] = Field(description='how high is the probability that the person is doing squats in the segment: low, medium, high, unknown(null)')\n",
    "    squats_technique_correctness: Optional[str] = Field(description='correctness of the squat technique.')\n",
    "    squats_feedback: Optional[SegmentFeedback] = Field(description='what was right and wrong in the squat perfomance in the segment. When the technique is incorrect, provide instructions how to correct them.')\n",
    "\n",
    "# we will only take the segments where the \"doing_squats\" field is positive.\n",
    "annotations = generate_annotations(\n",
    "    human_prompt=human_prompt,\n",
    "    config=config,\n",
    "    segments_per_call=5,\n",
    "    annotation_schema=SegmentAnnotation,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00, 7225.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total segments: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'start_timestamp': '00:03:31.750',\n",
       " 'end_timestamp': '00:03:54.250',\n",
       " 'segment_annotation': {'squats_probability': 'high',\n",
       "  'squats_technique_correctness': 'right',\n",
       "  'squats_feedback': {'right': 'Controlled descent, proper knee and butt positioning, upper back tightness, deep and comfortable squat.',\n",
       "   'wrong': None,\n",
       "   'correction': None}},\n",
       " 'video_id': 'gcNh17Ckjgg',\n",
       " 'id': 'gcNh17Ckjgg_0',\n",
       " 'video_path': 'gcNh17Ckjgg_0.mp4'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datagen import aggregate_annotations\n",
    "\n",
    "def filter_annotations(ann):\n",
    "    if ann['squats_probability'] in [None, 'low', 'None', 'null']:\n",
    "        # if we're not able to infer probability or prob is low, we don't need it\n",
    "        return False\n",
    "    if ann['squats_technique_correctness'] in [None, 'null', 'None']:\n",
    "        # if we couldnt establish correctness at all, the feedback is probably useless\n",
    "        return False\n",
    "    if ann['squats_technique_correctness'] in ['mixed']:\n",
    "        # discard empty segment if correctness isn't clear since there isn't any information to use for training\n",
    "        if ann['squats_feedback'] is None:\n",
    "            return False\n",
    "        if set(ann['squats_feedback'].values()) == set([None]):\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "annotations = aggregate_annotations(config, filter_func=filter_annotations, annotation_file='annotations.json')\n",
    "print('Total segments:', len(annotations))\n",
    "annotations[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The last step is to cut video clips for annotated segments from original videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:28<00:00,  2.86s/it]\n"
     ]
    }
   ],
   "source": [
    "from datagen import cut_videos\n",
    "cut_videos(config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So as a result we generated:\n",
    "- `<data_dir>/clips/` with video clips that you can use for training\n",
    "- `<data_dir>/annotations.json` with list of items with fields:\n",
    "    - video_id: 11-char youtube video id (youtube.com/watch?v=<id>)\n",
    "    - start_timestamp/end_timestamp of the clip relative to the youtube video it's taken from\n",
    "    - video_path of the clip relative to `<data_dir>/clips/`\n",
    "    - segment_annotation that you can use for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}