examples/torchvision_pose_boxmot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import cv2\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from boxmot import BotSort\n",
    "\n",
    "# Load a pre-trained Keypoint R-CNN model from torchvision\n",
    "device = torch.device('cpu')  # Change to 'cuda' if you have a GPU available\n",
    "pose_model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)\n",
    "pose_model.eval().to(device)\n",
    "\n",
    "tracker = BotSort(\n",
    "    reid_weights=Path('osnet_x0_25_msmt17.pt'),  # ReID model to use\n",
    "    device=device,\n",
    "    half=False,\n",
    ")\n",
    "\n",
    "# Open the video file\n",
    "vid = cv2.VideoCapture(0)\n",
    "\n",
    "# Function to generate a unique color for each track ID\n",
    "def get_color(track_id):\n",
    "    np.random.seed(int(track_id))\n",
    "    return tuple(np.random.randint(0, 255, 3).tolist())\n",
    "\n",
    "while True:\n",
    "    ret, im = vid.read()\n",
    "    if not ret:\n",
    "        break\n",
    "\n",
    "    # Convert frame to tensor and move to device\n",
    "    frame_tensor = torchvision.transforms.functional.to_tensor(im).unsqueeze(0).to(device)\n",
    "\n",
    "    # Run the Keypoint R-CNN model to detect keypoints and bounding boxes\n",
    "    with torch.no_grad():\n",
    "        results = pose_model(frame_tensor)[0]\n",
    "\n",
    "    # Extract detections (bounding boxes and keypoints)\n",
    "    dets = []\n",
    "    keypoints = []\n",
    "\n",
    "    confidence_threshold = 0.5\n",
    "    for i, score in enumerate(results['scores']):\n",
    "        if score >= confidence_threshold:\n",
    "            # Extract bounding box and score\n",
    "            x1, y1, x2, y2 = results['boxes'][i].cpu().numpy()\n",
    "            conf = score.item()\n",
    "            cls = results['labels'][i].item()  # Assuming that 'labels' would be person class\n",
    "            dets.append([x1, y1, x2, y2, conf, cls])\n",
    "\n",
    "            # Extract keypoints\n",
    "            keypoint = results['keypoints'][i].cpu().numpy().tolist()\n",
    "            keypoints.append(keypoint)\n",
    "\n",
    "    # Convert detections to a numpy array (N x (x, y, x, y, conf, cls))\n",
    "    dets = np.array(dets)\n",
    "\n",
    "    # Update tracker with detections and image\n",
    "    tracks = tracker.update(dets, im)  # M x (x, y, x, y, id, conf, cls, ind)\n",
    "\n",
    "    if len(tracks) > 0:\n",
    "        inds = tracks[:, 7].astype('int')  # Get track indices as int\n",
    "\n",
    "        # Use the indices to match tracks with keypoints\n",
    "        keypoints = [keypoints[i] for i in inds if i < len(keypoints)]  # Reorder keypoints to match the tracks\n",
    "\n",
    "        # Draw bounding boxes and keypoints in the same loop\n",
    "        for i, track in enumerate(tracks):\n",
    "            x1, y1, x2, y2, track_id, conf, cls = track[:7].astype('int')\n",
    "            color = get_color(track_id)\n",
    "\n",
    "            # Draw bounding box with unique color\n",
    "            cv2.rectangle(im, (x1, y1), (x2, y2), color, 2)\n",
    "\n",
    "            # Add text with ID, confidence, and class\n",
    "            cv2.putText(im, f'ID: {track_id}, Conf: {conf:.2f}, Class: {cls}', \n",
    "                        (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
    "\n",
    "            # Draw keypoints for the corresponding track\n",
    "            if i < len(keypoints):\n",
    "                kp = keypoints[i]\n",
    "                for point in kp:\n",
    "                    x, y, confidence = int(point[0]), int(point[1]), point[2]\n",
    "                    if confidence > 0.5:  # Only draw keypoints with confidence > 0.5\n",
    "                        cv2.circle(im, (x, y), 3, color, -1)  # Draw keypoints in the color of the corresponding track\n",
    "\n",
    "    # Display the image\n",
    "    cv2.imshow('Pose Tracking', im)\n",
    "\n",
    "    # Break on pressing q or space\n",
    "    key = cv2.waitKey(1) & 0xFF\n",
    "    if key == ord(' ') or key == ord('q'):\n",
    "        break\n",
    "\n",
    "vid.release()\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "boxmot-YDNZdsaB-py3.11",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}