Spaces:
Runtime error
Runtime error
import argparse | |
import os | |
from copy import deepcopy | |
from pathlib import Path | |
from multiprocessing import Pool | |
import pandas as pd | |
from scenedetect import open_video, SceneManager | |
from scenedetect.detectors import ContentDetector | |
from tqdm import tqdm | |
from utils.logger import logger | |
def cutscene_detection_star(args): | |
return cutscene_detection(*args) | |
def cutscene_detection(video_path, saved_path, cutscene_threshold=27, min_scene_len=15): | |
try: | |
if os.path.exists(saved_path): | |
logger.info(f"{video_path} has been processed.") | |
return | |
# Use PyAV as the backend to avoid (to some exent) containing the last frame of the previous scene. | |
# https://github.com/Breakthrough/PySceneDetect/issues/279#issuecomment-2152596761. | |
video = open_video(video_path, backend="pyav") | |
frame_rate, frame_size = video.frame_rate, video.frame_size | |
duration = deepcopy(video.duration) | |
frame_points, frame_timecode = [], {} | |
scene_manager = SceneManager() | |
scene_manager.add_detector( | |
# [ContentDetector, ThresholdDetector, AdaptiveDetector] | |
ContentDetector(threshold=cutscene_threshold, min_scene_len=min_scene_len) | |
) | |
scene_manager.detect_scenes(video, show_progress=False) | |
scene_list = scene_manager.get_scene_list() | |
for scene in scene_list: | |
for frame_time_code in scene: | |
frame_index = frame_time_code.get_frames() | |
if frame_index not in frame_points: | |
frame_points.append(frame_index) | |
frame_timecode[frame_index] = frame_time_code | |
del video, scene_manager | |
frame_points = sorted(frame_points) | |
output_scene_list = [] | |
for idx in range(len(frame_points) - 1): | |
output_scene_list.append((frame_timecode[frame_points[idx]], frame_timecode[frame_points[idx+1]])) | |
timecode_list = [(frame_timecode_tuple[0].get_timecode(), frame_timecode_tuple[1].get_timecode()) for frame_timecode_tuple in output_scene_list] | |
meta_scene = [{ | |
"video_path": Path(video_path).name, | |
"timecode_list": timecode_list, | |
"fram_rate": frame_rate, | |
"frame_size": frame_size, | |
"duration": str(duration) # __repr__ | |
}] | |
pd.DataFrame(meta_scene).to_json(saved_path, orient="records", lines=True) | |
except Exception as e: | |
logger.warning(f"Cutscene detection with {video_path} failed. Error is: {e}.") | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser(description="Cutscene Detection") | |
parser.add_argument( | |
"--video_metadata_path", type=str, required=True, help="The path to the video dataset metadata (csv/jsonl)." | |
) | |
parser.add_argument( | |
"--video_path_column", | |
type=str, | |
default="video_path", | |
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).", | |
) | |
parser.add_argument("--video_folder", type=str, default="", help="The video folder.") | |
parser.add_argument("--saved_folder", type=str, required=True, help="The save path to the output results (csv/jsonl).") | |
parser.add_argument("--n_jobs", type=int, default=1, help="The number of processes.") | |
args = parser.parse_args() | |
metadata_df = pd.read_json(args.video_metadata_path, lines=True) | |
video_path_list = metadata_df[args.video_path_column].tolist() | |
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list] | |
if not os.path.exists(args.saved_folder): | |
os.makedirs(args.saved_folder, exist_ok=True) | |
# The glob can be slow when there are many small jsonl files. | |
saved_path_list = [os.path.join(args.saved_folder, Path(video_path).stem + ".jsonl") for video_path in video_path_list] | |
args_list = [ | |
(video_path, saved_path) | |
for video_path, saved_path in zip(video_path_list, saved_path_list) | |
] | |
# Since the length of the video is not uniform, the gather operation is not performed. | |
# We need to run easyanimate/video_caption/utils/gather_jsonl.py after the program finised. | |
with Pool(args.n_jobs) as pool: | |
results = list(tqdm(pool.imap(cutscene_detection_star, args_list), total=len(video_path_list))) | |