| """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, |
| ), |
| } |
|
|
