73 lines
3.5 KiB
Python
73 lines
3.5 KiB
Python
import os, cv2
|
|
import torch
|
|
from pathlib import Path
|
|
from multiprocessing import freeze_support
|
|
|
|
def extract_frames(input_video_path, output_imgs_path):
|
|
# Open the video file
|
|
vidcap = cv2.VideoCapture(input_video_path)
|
|
|
|
# Get the total number of frames in the video
|
|
frame_count = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
# Create the output directory if it does not exist
|
|
if not os.path.exists(output_imgs_path):
|
|
os.makedirs(output_imgs_path)
|
|
|
|
# Extract the frames
|
|
for i in range(frame_count):
|
|
success, image = vidcap.read()
|
|
if success:
|
|
cv2.imwrite(os.path.join(output_imgs_path, f"frame{i}.png"), image)
|
|
print(f"{frame_count} frames extracted and saved to {output_imgs_path}")
|
|
|
|
|
|
def video2humanmasks(input_frames_path, output_folder_path, output_type, fps):
|
|
# freeze support is needed for video outputting
|
|
freeze_support()
|
|
|
|
# check if input path exists and is a directory
|
|
if not os.path.exists(input_frames_path) or not os.path.isdir(input_frames_path):
|
|
raise ValueError("Invalid input path: {}".format(input_frames_path))
|
|
|
|
# check if output path exists and is a directory
|
|
if not os.path.exists(output_folder_path) or not os.path.isdir(output_folder_path):
|
|
raise ValueError("Invalid output path: {}".format(output_folder_path))
|
|
|
|
# check if output_type is valid
|
|
valid_output_types = ["video", "pngs", "both"]
|
|
if output_type.lower() not in valid_output_types:
|
|
raise ValueError("Invalid output type: {}. Must be one of {}".format(output_type, valid_output_types))
|
|
|
|
# try to predict where torch cache lives, so we can try and fetch models from cache in the next step
|
|
predicted_torch_model_cache_path = os.path.join(Path.home(), ".cache", "torch", "hub", "hithereai_RobustVideoMatting_master")
|
|
predicted_rvm_cache_testilfe = os.path.join(predicted_torch_model_cache_path, "hubconf.py")
|
|
|
|
# try to fetch the models from cache, and only if it can't be find, download from the internet (to enable offline usage)
|
|
try:
|
|
# Try to fetch the models from cache
|
|
convert_video = torch.hub.load(predicted_torch_model_cache_path, "converter", source='local')
|
|
model = torch.hub.load(predicted_torch_model_cache_path, "mobilenetv3", source='local').cuda()
|
|
except:
|
|
# Download from the internet if not found in cache
|
|
convert_video = torch.hub.load("hithereai/RobustVideoMatting", "converter")
|
|
model = torch.hub.load("hithereai/RobustVideoMatting", "mobilenetv3").cuda()
|
|
|
|
output_alpha_vid_path = os.path.join(output_folder_path, "human_masked_video.mp4")
|
|
# extract humans masks from the input folder' imgs.
|
|
# in this step PNGs will be extracted only if output_type is set to PNGs. Otherwise a video will be made, and in the case of Both, the video will be extracted in the next step to PNGs
|
|
convert_video(
|
|
model,
|
|
input_source=input_frames_path, # full path of the folder that contains all of the extracted input imgs
|
|
output_type='video' if output_type.upper() in ("VIDEO", "BOTH") else 'png_sequence',
|
|
output_alpha=output_alpha_vid_path if output_type.upper() in ("VIDEO", "BOTH") else output_folder_path,
|
|
output_video_mbps=4,
|
|
output_video_fps=fps,
|
|
downsample_ratio=None, # None for auto
|
|
seq_chunk=12, # Process n frames at once for better parallelism
|
|
progress=True # show extraction progress
|
|
)
|
|
|
|
if output_type.lower() == "both":
|
|
extract_frames(output_alpha_vid_path, output_folder_path)
|