pipeline/process_videos.py

import cv2
import pandas as pd
from PIL import Image
from tqdm import tqdm

from pipeline.clip_wrapper import MODEL_DIM, ClipWrapper
from pipeline.download_videos import DATA_DIR, REPO_ROOT, VIDEO_DIR

FRAME_EXTRACT_RATE_SECONDS = 5 # Extract a frame every 5 seconds
IMAGES_DIR = DATA_DIR / "images"
DATAFRAME_PATH = DATA_DIR / "dataset.parquet"


def process_videos() -> None:
 "Runs clip on video frames, saves results to a parquet file"
 clip_wrapper = ClipWrapper()
 results = []
 for video_path in tqdm(list(VIDEO_DIR.glob("*.mp4")), desc="Processing videos"):
 video_id = video_path.stem
 extracted_images_dir = IMAGES_DIR / video_id
 extracted_images_dir.mkdir(exist_ok=True, parents=True)
 complete_file = extracted_images_dir / "complete"
 if complete_file.exists():
 continue
 for clip_vector, image, timestamp_secs, frame_idx in get_clip_vectors(
 video_path, clip_wrapper
 ):
 image_path = extracted_images_dir / f"{frame_idx}.jpg"
 image.save(image_path)
 results.append(
 [
 video_id,
 frame_idx,
 timestamp_secs,
 str(image_path.relative_to(REPO_ROOT)),
 *clip_vector,
 ]
 )
 complete_file.touch()
 df = pd.DataFrame(
 results,
 columns=["video_id", "frame_idx", "timestamp", "image_path"]
 + [f"dim_{i}" for i in range(MODEL_DIM)],
 )
 print(f"Saving data to {DATAFRAME_PATH}")
 df.to_parquet(DATAFRAME_PATH, index=False)


def get_clip_vectors(video_path, clip_wrapper):
 cap = cv2.VideoCapture(str(video_path))
 num_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 fps = int(cap.get(cv2.CAP_PROP_FPS))
 extract_every_n_frames = FRAME_EXTRACT_RATE_SECONDS * fps
 for frame_idx in tqdm(range(num_video_frames), desc="Running CLIP on video"):
 ret, frame = cap.read()
 if frame_idx % extract_every_n_frames != 0:
 continue
 image = Image.fromarray(frame[..., ::-1])
 clip_vector = clip_wrapper.images2vec([image]).squeeze().numpy()
 timestamp_secs = frame_idx / fps
 yield clip_vector, image, timestamp_secs, frame_idx
 cap.release()


if __name__ == "__main__":
 process_videos()