Video_inference_new.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model loaded successfully.\n"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "\n",
    "import glob\n",
    "import os\n",
    "import cv2\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torchvision.transforms as transforms\n",
    "from timm import create_model\n",
    "import torch.nn.functional as F\n",
    "from skimage import io, transform\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Paths\n",
    "video_path = \"/Users/annastuckert/Documents/GitHub/facemap/cam1_G7c1_1_10seconds.avi\"\n",
    "output_video_path = \"output_video_with_keypoints_224x224.mp4\"\n",
    "model_path = \"models/best_model.pt\"\n",
    "\n",
    "# Device setup\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Define model loading function\n",
    "def load_model(model_path, num_output_classes=24):\n",
    "    model = create_model(\"vit_base_patch16_224\", pretrained=False, in_chans=1, num_classes=num_output_classes)\n",
    "    model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "# Load model\n",
    "try:\n",
    "    model = load_model(model_path)\n",
    "    print(\"Model loaded successfully.\")\n",
    "except Exception as e:\n",
    "    print(f\"Error loading model: {e}\")\n",
    "    exit()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import cv2\n",
    "# import os\n",
    "\n",
    "# # Define video path and output folder\n",
    "# video_path = '/Users/annastuckert/Documents/GitHub/facemap/cam1_G7c1_1_10seconds.avi'  # Update with the path to your video\n",
    "# output_folder = 'frames'  # Folder to save frames\n",
    "\n",
    "# # Ensure the output folder exists\n",
    "# os.makedirs(output_folder, exist_ok=True)\n",
    "\n",
    "# # Set up video capture\n",
    "# cap = cv2.VideoCapture(video_path)\n",
    "# fps = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "# frame_count = 0\n",
    "\n",
    "# # Process each frame in the video\n",
    "# while cap.isOpened():\n",
    "#     ret, frame = cap.read()\n",
    "#     if not ret:\n",
    "#         break\n",
    "    \n",
    "#     # Resize the frame to 224x224\n",
    "#     #frame = cv2.resize(frame, (224, 224))\n",
    "\n",
    "#     # Save each frame as an image in the output folder\n",
    "#     frame_filename = os.path.join(output_folder, f\"frame_{frame_count:04d}.jpg\")\n",
    "#     cv2.imwrite(frame_filename, frame)\n",
    "\n",
    "#     frame_count += 1\n",
    "\n",
    "# # Release resources\n",
    "# cap.release()\n",
    "# print(f\"Frames saved to '{output_folder}' folder with 224x224 resolution.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Folder exists!\n"
     ]
    }
   ],
   "source": [
    "# load packages\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import pandas as pd\n",
    "from skimage import transform,io\n",
    "import os\n",
    "import pdb\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# where are pictures?\n",
    "IMG_LOC='/Users/annastuckert/Documents/GitHub/facemap/frames/'\n",
    "\n",
    "# make folder called low_res if not exist\n",
    "if os.path.isdir(IMG_LOC+'low_res'):\n",
    "    print('Folder exists!')\n",
    "else:\n",
    "    os.makedirs(IMG_LOC+'low_res')\n",
    "\n",
    "# find all pngs\n",
    "img_files = sorted(glob.glob(IMG_LOC+'*.jpg'))\n",
    "# find labels\n",
    "#labels = pd.read_csv(IMG_LOC+'labels.csv')\n",
    "\n",
    "# specify height and width\n",
    "h = w = 224\n",
    "\n",
    "# read one image\n",
    "img = plt.imread(img_files[0])\n",
    "\n",
    "# find aspects of original image\n",
    "h_org = img.shape[0]\n",
    "w_org = img.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import numpy as np\n",
    "import cv2\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn.functional as F\n",
    "from skimage import io, transform, color\n",
    "\n",
    "# Define the low-res folder and video output path\n",
    "low_res_folder = '/Users/annastuckert/Documents/GitHub/facemap/frames/low_res'\n",
    "os.makedirs(low_res_folder, exist_ok=True)\n",
    "video_output_path = '/Users/annastuckert/Documents/GitHub/facemap/output_video.avi'  # Adjust the output path\n",
    "\n",
    "# Initialize transformation for the video frames\n",
    "vit_transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Grayscale(num_output_channels=1),  # Convert to grayscale\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "target_height = target_width = 224\n",
    "\n",
    "# Get list of image files\n",
    "\n",
    "# Device configuration\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Model should be defined and loaded here\n",
    "# model = YourModelClass().to(device)\n",
    "# model.load_state_dict(torch.load('your_model.pth'))\n",
    "# model.eval()\n",
    "\n",
    "# List to hold frames for video\n",
    "frames = []\n",
    "\n",
    "# Process each image file\n",
    "for img_file in img_files:  # Limit to first 10 images or adjust as necessary\n",
    "    # Read the image\n",
    "    img = io.imread(img_file)\n",
    "\n",
    "    # Get original dimensions\n",
    "    h_org, w_org = img.shape[:2]\n",
    "    \n",
    "    # Determine the center crop region\n",
    "    start_x = (w_org - h_org) // 2  # Horizontal start point\n",
    "    start_y = 0                     # Vertical start point (full height)\n",
    "    \n",
    "    # Crop the center square\n",
    "    cropped_img = img[start_y:start_y + h_org, start_x:start_x + h_org]\n",
    "\n",
    "    # Convert to grayscale\n",
    "    gray_img = color.rgb2gray(cropped_img)\n",
    "\n",
    "    # Resize the image to 224x224 pixels\n",
    "    resized_img = transform.resize(gray_img, (target_height, target_width), anti_aliasing=True)\n",
    "\n",
    "    # Convert grayscale to RGB for drawing (3 channels)\n",
    "    color_img = np.stack([resized_img] * 3, axis=-1)  # Create an RGB version of the grayscale image\n",
    "\n",
    "    # Add key point inference\n",
    "    input_frame = vit_transform(resized_img).unsqueeze(0).to(device)  # Add batch dimension\n",
    "\n",
    "    # Model inference for keypoints\n",
    "    with torch.no_grad():\n",
    "        scores = F.softplus(model(input_frame))\n",
    "        keypoints = scores.squeeze().cpu().numpy()\n",
    "\n",
    "    # Overlay keypoints on the color image\n",
    "    for i in range(0, len(keypoints), 2):\n",
    "        x, y = int(keypoints[i]), int(keypoints[i + 1])\n",
    "        # Draw a red circle on the RGB image\n",
    "        cv2.circle(color_img, (x, y), radius=5, color=(1, 0, 0), thickness=-1)\n",
    "\n",
    "    # Save the RGB image with keypoints to the low_res folder\n",
    "    base_filename = os.path.basename(img_file)\n",
    "    save_path = os.path.join(low_res_folder, base_filename)\n",
    "    io.imsave(save_path, (color_img * 255).astype(np.uint8))  # Convert back to uint8 for saving\n",
    "\n",
    "    # Add processed frame to the frames list for video\n",
    "    frames.append((color_img * 255).astype(np.uint8))  # Convert to uint8 for video\n",
    "\n",
    "# Create a video from the frames\n",
    "if frames:\n",
    "    height, width, _ = frames[0].shape\n",
    "    fourcc = cv2.VideoWriter_fourcc(*'XVID')  # You can change the codec if necessary\n",
    "    video_writer = cv2.VideoWriter(video_output_path, fourcc, 30, (width, height))  # 30 fps\n",
    "\n",
    "    for frame in frames:\n",
    "        video_writer.write(frame)\n",
    "\n",
    "    video_writer.release()  # Finalize the video file\n",
    "    print(f\"Video saved to '{video_output_path}'.\")\n",
    "\n",
    "print(f\"Processed images saved to '{low_res_folder}'.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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
}