{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# install mediapipe\n!pip install mediapipe\n!pip install ultralytics","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:05:51.681778Z","iopub.execute_input":"2023-04-13T21:05:51.682042Z","iopub.status.idle":"2023-04-13T21:06:20.547507Z","shell.execute_reply.started":"2023-04-13T21:05:51.682013Z","shell.execute_reply":"2023-04-13T21:06:20.546276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import ultralytics\nultralytics.checks()","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:06:20.550034Z","iopub.execute_input":"2023-04-13T21:06:20.550355Z","iopub.status.idle":"2023-04-13T21:06:23.280856Z","shell.execute_reply.started":"2023-04-13T21:06:20.550318Z","shell.execute_reply":"2023-04-13T21:06:23.279663Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import dependencies\nimport os\nimport subprocess\nimport IPython\nfrom IPython.display import Video, display\n\nimport numpy as np\nimport pandas as pd\n\nimport cv2\nimport mediapipe as  mp","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:06:23.282989Z","iopub.execute_input":"2023-04-13T21:06:23.283896Z","iopub.status.idle":"2023-04-13T21:06:23.376775Z","shell.execute_reply.started":"2023-04-13T21:06:23.283852Z","shell.execute_reply":"2023-04-13T21:06:23.375790Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def play_video(video_path: str):\n    frac = 0.65 # scaling factor for display \n    display(\n        Video(data=video_path, embed=True, height=int(720*frac), width=int(1280*frac))\n    )","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:06:27.723039Z","iopub.execute_input":"2023-04-13T21:06:27.723544Z","iopub.status.idle":"2023-04-13T21:06:27.729775Z","shell.execute_reply.started":"2023-04-13T21:06:27.723505Z","shell.execute_reply":"2023-04-13T21:06:27.727906Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def detect_player_bbox(video_path:str) -> str:\n    \"\"\"\n    Annotates video with object detection bounding boxes using YOLOv8 model.\n    \"\"\"\n    \n    video_name = video_path.split('/')[-1]\n    folder_name = 'predict'\n    \n    # YOLOv8 Object Detection\n    # you can edit the model if you want to use a different yolo8 model\n    os.system(f\"yolo predict model=yolov8x.pt source={video_path} save_txt=True\")\n    \n    # path to latest output \n    max = 0\n    for dir_name in os.listdir('/kaggle/working/runs/detect'):\n        if dir_name!='predict':\n            if int(dir_name[-1])>max:\n                max = int(dir_name[-1])\n    if max>0:\n        folder_name += str(max)\n    output_path = f'/kaggle/working/runs/detect/{folder_name}/labels'\n    \n    return output_path","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:06:32.107678Z","iopub.execute_input":"2023-04-13T21:06:32.108307Z","iopub.status.idle":"2023-04-13T21:06:32.115468Z","shell.execute_reply.started":"2023-04-13T21:06:32.108245Z","shell.execute_reply":"2023-04-13T21:06:32.114356Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def detect_player_bbox(video_path:str) -> str:\n    \"\"\"\n    Annotates video with object detection bounding boxes using YOLOv8 model.\n    \"\"\"\n    \n    video_name = video_path.split('/')[-1]\n    folder_name = 'predict'\n    \n    # YOLOv8 Object Detection\n    # you can edit the model if you want to use a different yolo8 model\n    os.system(f\"yolo predict model=yolov8x.pt source={video_path} save_txt=True\")\n    \n    # path to latest output \n    max = 0\n    for dir_name in os.listdir('/kaggle/working/runs/detect'):\n        if dir_name!='predict':\n            if int(dir_name[-1])>max:\n                max = int(dir_name[-1])\n    if max>0:\n        folder_name += str(max)\n    output_path = f'/kaggle/working/runs/detect/{folder_name}/labels'\n    \n    return output_path","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:08:13.964173Z","iopub.execute_input":"2023-04-13T21:08:13.965026Z","iopub.status.idle":"2023-04-13T21:08:13.972036Z","shell.execute_reply.started":"2023-04-13T21:08:13.964984Z","shell.execute_reply":"2023-04-13T21:08:13.970866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bbox_labels_to_dataframe(bbox_labels_path: str) -> pd.DataFrame:\n\n    bbox_labels = {\n        'video_name':[],\n        'frame':[],\n        'class_id':[],\n        'center_x':[],\n        'center_y':[],\n        'width':[],\n        'height':[]\n    }\n\n    for filename in os.listdir(bbox_labels_path):\n        video_name = \"_\".join(filename.split('_')[0:3]) + '.mp4'\n        frame = filename.split('_')[-1]\n        frame = int(frame.split('.')[0])\n\n        with open(bbox_labels_path + '/' + filename, 'r') as f:\n            for line in f:\n                line = line.split(\" \")\n                class_id = int(line[0])\n                center_x = float(line[1])\n                center_y = float(line[2])\n                width = float(line[3])\n                height = float(line[4])\n\n                if class_id == 0: # if person\n                    # append to dict\n                    bbox_labels['video_name'].append(video_name)\n                    bbox_labels['frame'].append(frame)\n                    bbox_labels['class_id'].append(class_id)\n                    bbox_labels['center_x'].append(center_x)\n                    bbox_labels['center_y'].append(center_y)\n                    bbox_labels['width'].append(width)\n                    bbox_labels['height'].append(height)\n                    \n    \n    return pd.DataFrame(bbox_labels)","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:13:20.975594Z","iopub.execute_input":"2023-04-13T21:13:20.976487Z","iopub.status.idle":"2023-04-13T21:13:20.987394Z","shell.execute_reply.started":"2023-04-13T21:13:20.976447Z","shell.execute_reply":"2023-04-13T21:13:20.986203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def multi_pose_estimation(video_path:str, bbox_labels: pd.DataFrame, verbose=True) -> str:\n    \"\"\"\n    Performs multi-shot multi-pose estimation by obtianing person bbox from YOLOv8\n    and performing single pose estimation on the bbox crop.\n    \"\"\"\n    \n    # intializing mediapipe utils\n    mp_drawing = mp.solutions.drawing_utils\n    mp_drawing_styles = mp.solutions.drawing_styles\n    mp_pose = mp.solutions.pose\n    \n    # video name\n    video_name = video_path.split('/')[-1]\n\n    # VideoCapture Object\n    cap = cv2.VideoCapture(video_path)\n    \n    # video variables\n    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n    fps = cap.get(cv2.CAP_PROP_FPS)\n    total_frames = bbox_labels['frame'].max()\n    \n    # VideoWriter Object\n    output_path  = \"labeled_\" + video_name\n    tmp_output_path = 'tmp_' + output_path\n    out = cv2.VideoWriter(tmp_output_path, cv2.VideoWriter_fourcc(*'MP4V'), fps, (width, height))\n\n    # check if camera opened successfully\n    if (cap.isOpened()==False):\n        print('Error opening video file!')\n\n    # multipose estimation\n    with mp_pose.Pose(\n        min_detection_confidence = 0.8,\n        min_tracking_confidence = 0.8) as pose:\n        frame = 1\n        while (cap.isOpened()):\n            success, image = cap.read()\n            if success:\n                # selecting the frame\n                bbox_set = bbox_labels.query('frame==@frame')\n                # iterating through bboxs in the frame\n                for idx, annot in bbox_set.iterrows():\n                    bbox_center_x = annot['center_x'] * width\n                    bbox_center_y = annot['center_y'] * height\n                    bbox_width = annot['width'] * width\n                    bbox_height = annot['height'] * height\n\n                    # finding top-left and bottom-right bbox cooridnates\n                    bbox_top_left_x = int(bbox_center_x - (bbox_width/2))\n                    bbox_top_left_y = int(bbox_center_y - (bbox_height/2))\n                    bbox_bottom_right_x = int(bbox_center_x + (bbox_width/2))\n                    bbox_bottom_right_y = int(bbox_center_y + (bbox_height/2))\n\n                    # cropping image to bbox\n                    image_crop = image[bbox_top_left_y:bbox_bottom_right_y, bbox_top_left_x:bbox_bottom_right_x]\n\n                    # pose estimation\n                    # set image as not writeable to improve perfromance\n                    image_crop.flags.writeable = False\n                    image_crop = cv2.cvtColor(image_crop, cv2.COLOR_BGR2RGB)\n                    results = pose.process(image_crop)\n\n                    # transposing results to be drawn on the original image\n                    if results.pose_landmarks != None:\n                        for landmark in results.pose_landmarks.landmark:\n                            landmark.x = ((abs(bbox_bottom_right_x - bbox_top_left_x) / width) * landmark.x) + (bbox_top_left_x/width)\n                            landmark.y = ((abs(bbox_bottom_right_y - bbox_top_left_y) / height) * landmark.y) + (bbox_top_left_y/height)\n\n                        # draw the pose annotations on the image\n                        # set image as writeable\n                        image.flags.writeable = True\n                        mp_drawing.draw_landmarks(\n                            image,\n                            results.pose_landmarks,\n                            mp_pose.POSE_CONNECTIONS,\n                            landmark_drawing_spec = mp_drawing_styles.get_default_pose_landmarks_style())\n\n\n\n                # save video\n                out.write(image)\n                if verbose:\n                    print(f'Frame: {frame}/{total_frames}')\n                frame += 1\n            else:\n                break\n\n\n        cap.release()\n        out.release()\n        \n\n    # Not all browsers support the codec, we will re-load the file at tmp_output_path\n    # and convert to a codec that is more broadly readable using ffmpeg\n    if os.path.exists(output_path):\n        os.remove(output_path)\n    subprocess.run(\n            [\n                \"ffmpeg\",\n                \"-i\",\n                tmp_output_path,\n                \"-crf\",\n                \"18\",\n                \"-preset\",\n                \"veryfast\",\n                \"-hide_banner\",\n                \"-loglevel\",\n                \"error\",\n                \"-vcodec\",\n                \"libx264\",\n                output_path,\n            ]\n        )\n    os.remove(tmp_output_path)\n    \n    return output_path","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:08:51.846856Z","iopub.execute_input":"2023-04-13T21:08:51.847271Z","iopub.status.idle":"2023-04-13T21:08:51.863897Z","shell.execute_reply.started":"2023-04-13T21:08:51.847235Z","shell.execute_reply":"2023-04-13T21:08:51.862515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"video_path = '/kaggle/input/nfl-player-contact-detection/train/58168_003392_Endzone.mp4'\nbbox_labels_path = detect_player_bbox(video_path)\nbbox_labels = bbox_labels_to_dataframe(bbox_labels_path)","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:13:27.779693Z","iopub.execute_input":"2023-04-13T21:13:27.780089Z","iopub.status.idle":"2023-04-13T21:14:13.665294Z","shell.execute_reply.started":"2023-04-13T21:13:27.780053Z","shell.execute_reply":"2023-04-13T21:14:13.664192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output_path = multi_pose_estimation(video_path, bbox_labels)","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:14:13.667254Z","iopub.execute_input":"2023-04-13T21:14:13.667660Z","iopub.status.idle":"2023-04-13T21:18:46.446564Z","shell.execute_reply.started":"2023-04-13T21:14:13.667624Z","shell.execute_reply":"2023-04-13T21:18:46.445130Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"play_video(output_path)","metadata":{"execution":{"iopub.status.busy":"2023-04-13T21:18:46.449401Z","iopub.execute_input":"2023-04-13T21:18:46.449814Z","iopub.status.idle":"2023-04-13T21:18:46.855003Z","shell.execute_reply.started":"2023-04-13T21:18:46.449764Z","shell.execute_reply":"2023-04-13T21:18:46.853563Z"},"trusted":true},"execution_count":null,"outputs":[]}]}