VASR_ / src /video_utils.py
sophiemaw's picture
Update src/video_utils.py
dd68cd0 verified
import cv2
import os
import glob
import shutil
import subprocess
def extract_frames(video_path, output_folder, overwrite=True):
# Clear output folder if needed
if overwrite and os.path.exists(output_folder):
shutil.rmtree(output_folder)
os.makedirs(output_folder, exist_ok=True)
# Run ffmpeg to extract frames
output_pattern = os.path.join(output_folder, "%d.jpg")
command = [
"ffmpeg", "-i", video_path,
"-q:v", "2", # Quality for JPEG (lower is better)
output_pattern
]
subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
print(f"βœ… Extracted frames to: {output_folder}")
def extract_frames_old(video_path, output_folder):
# Clear output folder if it exists
if os.path.exists(output_folder):
for f in glob.glob(os.path.join(output_folder, "*")):
os.remove(f)
else:
os.makedirs(output_folder)
cap = cv2.VideoCapture(video_path)
frame_id = 0
while True:
ret, frame = cap.read()
if not ret:
break
frame_path = os.path.join(output_folder, f"{frame_id}.jpg")
cv2.imwrite(frame_path, frame)
frame_id += 1
cap.release()
print(f"βœ… Extracted {frame_id} frames to: {output_folder}")
def recombine_frames_(video_path, frames_folder, output_video=None):
# Default output path if not provided
if output_video is None:
base, ext = os.path.splitext(video_path)
output_video = f"{base}_anonymised.mp4"
# Get frame size and FPS from original video
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
cap.release()
# Load all frame paths and sort by frame number
frame_paths = sorted(glob.glob(os.path.join(frames_folder, "*.jpg")), key=lambda x: int(os.path.basename(x).split('.')[0]))
if not frame_paths:
print("❌ No frames found to recombine.")
return
# Get frame size from first frame
frame_sample = cv2.imread(frame_paths[0])
height, width = frame_sample.shape[:2]
# Define the video writer
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video, fourcc, fps, (width, height))
for frame_path in frame_paths:
frame = cv2.imread(frame_path)
out.write(frame)
out.release()
print(f"βœ… Video saved to: {output_video}")
def recombine_frames(video_path, frames_folder, output_video=None, temp_video_path="temp_video_no_audio.mp4"):
"""
Recombines frames into a video and adds back the original audio using ffmpeg.
"""
if output_video is None:
base, ext = os.path.splitext(video_path)
output_video = f"{base}_anonymised.mp4"
# Get frame size and FPS from original video
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
cap.release()
# Collect and sort frame paths
frame_paths = sorted(
glob.glob(os.path.join(frames_folder, "*.jpg")),
key=lambda x: int(os.path.splitext(os.path.basename(x))[0])
)
if not frame_paths:
print("❌ No frames found to recombine.")
return
# Get frame size
sample_frame = cv2.imread(frame_paths[0])
height, width = sample_frame.shape[:2]
# Write video without audio
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
for frame_path in frame_paths:
frame = cv2.imread(frame_path)
out.write(frame)
out.release()
# Use ffmpeg to add original audio
ffmpeg_command = [
"ffmpeg",
"-y", # Overwrite without asking
"-i", temp_video_path,
"-i", video_path,
"-c:v", "copy",
"-map", "0:v:0",
"-map", "1:a:0?",
"-shortest",
output_video
]
subprocess.run(ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
# Cleanup
#os.remove(temp_video_path)
print(f"βœ… Final video with audio saved to: {output_video}")
if __name__ == "__main__":
input_video = "/Users/sophiemaw/Downloads/CONFIDENTIAL DO NOT SHARE Edna & Paul 29.10.10 Part 2 00.12.46.531.mov"
output_folder = 'frames'
extract_frames(input_video,output_folder)