"""Oz_ShotSplitter: ComfyUI output node that detects shots and cuts them to mp4.""" import json import logging import os from typing import List import folder_paths from .cutter import cut_shots from .detectors.ensemble import ( apply_min_shot_filter, boundaries_to_intervals, union_boundaries, ) from .first_frame_reader import read_first_frames from .video_probe import probe logger = logging.getLogger(__name__) VIDEO_EXTENSIONS = (".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v") class Oz_ShotSplitter: CATEGORY = "Oz/video" OUTPUT_NODE = True FUNCTION = "run" RETURN_TYPES = ("STRING", "IMAGE", "INT", "STRING") RETURN_NAMES = ("clip_paths", "first_frames", "shot_count", "manifest_json") OUTPUT_IS_LIST = (True, False, False, False) @classmethod def INPUT_TYPES(cls): input_dir = folder_paths.get_input_directory() files: List[str] = [] if os.path.isdir(input_dir): files = sorted( f for f in os.listdir(input_dir) if f.lower().endswith(VIDEO_EXTENSIONS) ) if not files: files = [""] return { "required": { "video": (files,), "detector": (["ensemble", "transnet", "pyscenedetect"], {"default": "ensemble"}), "min_shot_seconds": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 10.0, "step": 0.1}), "merge_window_frames": ("INT", {"default": 3, "min": 0, "max": 30}), "adaptive_threshold": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 10.0, "step": 0.1}), "transnet_threshold": ("FLOAT", {"default": 0.3, "min": 0.1, "max": 0.9, "step": 0.05}), "output_subfolder": ("STRING", {"default": "shots/{source_stem}"}), "filename_prefix": ("STRING", {"default": "shot"}), "crf": ("INT", {"default": 18, "min": 0, "max": 51}), "preset": ( ["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow"], {"default": "slow"}, ), "keep_audio": ("BOOLEAN", {"default": True}), } } def run( self, video: str, detector: str, min_shot_seconds: float, merge_window_frames: int, adaptive_threshold: float, transnet_threshold: float, output_subfolder: str, filename_prefix: str, crf: int, preset: str, keep_audio: bool, ): if video.startswith("<"): raise ValueError("No video file selected") input_dir = folder_paths.get_input_directory() source_path = os.path.join(input_dir, video) if not os.path.isfile(source_path): raise ValueError(f"File not found: {source_path}") info = probe(source_path) logger.info( "ShotSplitter: %dx%d @ %.3f fps, %.2fs, %d frames, audio=%s, vfr=%s", info.width, info.height, info.fps, info.duration_s, info.total_frames, info.has_audio, info.vfr, ) boundaries_all: List[List[int]] = [] if detector in ("ensemble", "transnet"): try: from .detectors.transnet import detect_transnet b = detect_transnet(source_path, threshold=transnet_threshold) logger.info("TransNetV2 boundaries: %d", len(b)) boundaries_all.append(b) except Exception as e: logger.warning("TransNetV2 unavailable/failed (%s), falling back", e) if detector == "transnet": raise if detector in ("ensemble", "pyscenedetect"): from .detectors.pyscenedetect_adapter import detect_pyscenedetect b = detect_pyscenedetect(source_path, adaptive_threshold=adaptive_threshold) logger.info("PySceneDetect boundaries: %d", len(b)) boundaries_all.append(b) merged = union_boundaries(boundaries_all, window=merge_window_frames) intervals = boundaries_to_intervals(merged, total_frames=info.total_frames) intervals = apply_min_shot_filter( intervals, fps=info.fps, min_shot_seconds=min_shot_seconds ) logger.info("Final intervals: %d", len(intervals)) source_stem = os.path.splitext(os.path.basename(source_path))[0] subfolder = output_subfolder.replace("{source_stem}", source_stem) out_base = folder_paths.get_output_directory() out_dir = os.path.join(out_base, subfolder) results = cut_shots( source_path=source_path, intervals=intervals, fps=info.fps, output_dir=out_dir, filename_prefix=filename_prefix, crf=crf, preset=preset, keep_audio=keep_audio, has_audio=info.has_audio, vfr=info.vfr, ) successful = [r for r in results if r.success] clip_paths = [r.path for r in successful] first_frames = read_first_frames(clip_paths) manifest = [ { "index": i + 1, "start_s": round(r.start_s, 4), "end_s": round(r.end_s, 4), "n_frames": int(round((r.end_s - r.start_s) * info.fps)), "path": r.path, } for i, r in enumerate(successful) ] manifest_json = json.dumps(manifest, indent=2) videos_ui = [ { "filename": os.path.basename(r.path), "subfolder": subfolder, "type": "output", "fullpath": r.path, "format": "video/mp4", "frame_rate": info.fps, } for r in successful ] return { "ui": {"videos": videos_ui}, "result": ( clip_paths, first_frames, len(successful), manifest_json, ), }