ComfyUI-ShotSplitter/shot_splitter_node.py

"""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 = ["<no videos in input/>"]
        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,
            ),
        }