Spaces:
Runtime error
Runtime error
File size: 3,597 Bytes
6cf4883 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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)
|