97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
import cv2
|
|
import numpy as np
|
|
import os
|
|
from rembg import remove
|
|
from PIL import Image
|
|
|
|
def process_video(input_path, output_path):
|
|
# Check if input exists
|
|
if not os.path.exists(input_path):
|
|
print(f"Error: Input file '{input_path}' not found.")
|
|
return
|
|
|
|
print(f"Processing video: {input_path}")
|
|
|
|
cap = cv2.VideoCapture(input_path)
|
|
|
|
# Get video properties
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
print(f"Resolution: {width}x{height}, FPS: {fps}, Total Frames: {total_frames}")
|
|
|
|
# Initialize video writer
|
|
# Using 'mp4v' codec for MP4. Note that standard MP4 does not support alpha channel easily.
|
|
# For web transparency, we usually need WebM with VP9 or a specific MOV codec (ProRes 4444).
|
|
# Here we will try to create a WebM file (VP9) which supports alpha.
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'VP90')
|
|
output_ext = os.path.splitext(output_path)[1].lower()
|
|
|
|
if output_ext == '.mp4':
|
|
print("Warning: MP4 container often doesn't support alpha transparency widely. Switching codec might be needed.")
|
|
# Try mp4v just in case, but alpha might be lost or black
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
elif output_ext == '.webm':
|
|
fourcc = cv2.VideoWriter_fourcc(*'VP90')
|
|
|
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
|
|
|
|
frame_count = 0
|
|
|
|
try:
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Convert BGR (OpenCV) to RGB (PIL/rembg)
|
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
pil_im = Image.fromarray(frame_rgb)
|
|
|
|
# Remove background using rembg
|
|
output_pil = remove(pil_im)
|
|
|
|
# Convert back to numpy
|
|
output_np = np.array(output_pil)
|
|
|
|
# Convert RGB to BGR for OpenCV handling (if we were saving a normal video)
|
|
# But wait, OpenCV VideoWriter expects BGR.
|
|
# If we want transparency, we need a 4-channel write.
|
|
# Standard cv2.VideoWriter might struggle with 4 channels depending on backend.
|
|
|
|
# Let's check if the output has alpha
|
|
if output_np.shape[2] == 4:
|
|
# If we are writing to a format that supports alpha (like VP9 WebM), we should pass the alpha.
|
|
# However, basic cv2 VideoWriter might not support 4 channels.
|
|
# A safer bet for a simple script is to save as a sequence of PNGs or find a writer that supports it.
|
|
# For this PoC, let's try writing the frame.
|
|
|
|
# If VideoWriter fails with 4 channels, we fallback to black background.
|
|
frame_bgr_alpha = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA)
|
|
out.write(frame_bgr_alpha)
|
|
else:
|
|
frame_bgr = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGR)
|
|
out.write(frame_bgr)
|
|
|
|
frame_count += 1
|
|
if frame_count % 10 == 0:
|
|
print(f"Processed {frame_count}/{total_frames} frames...")
|
|
|
|
except Exception as e:
|
|
print(f"Error processing frames: {e}")
|
|
|
|
finally:
|
|
cap.release()
|
|
out.release()
|
|
print("Done.")
|
|
|
|
if __name__ == "__main__":
|
|
# Ensure dependencies: pip install rembg opencv-python pillow numpy
|
|
input_file = "cup_spin.mp4"
|
|
output_file = "cup_spin_no_bg.webm" # using webm for transparency support
|
|
|
|
process_video(input_file, output_file)
|