582 lines
21 KiB
Python
582 lines
21 KiB
Python
'''
|
|
This file is to inference arbitrary video files for Shot Cut
|
|
'''
|
|
import os, sys, shutil
|
|
import argparse
|
|
import numpy as np
|
|
import math
|
|
import subprocess
|
|
import cv2
|
|
import copy
|
|
import ffmpeg
|
|
import json
|
|
import torch
|
|
import torchvision.transforms as T
|
|
from torch.utils.data import DataLoader
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
import warnings
|
|
warnings.filterwarnings("ignore", category=UserWarning)
|
|
|
|
|
|
# Import files from the local folder
|
|
# 使用 __file__ 定位插件根目录，而不是用 '.' (ComfyUI 启动目录)
|
|
root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
# 确保插件根目录始终位于 sys.path 最前面，
|
|
# 防止其他路径中存在同名的 config.py 单文件与插件的 config 包冲突
|
|
if root_path in sys.path:
|
|
sys.path.remove(root_path)
|
|
sys.path.insert(0, root_path)
|
|
|
|
# 清理可能被错误缓存的同名系统包（config、datasets 等），强制从插件根目录重新加载
|
|
for _mod_name in list(sys.modules.keys()):
|
|
if (
|
|
_mod_name == "config" or _mod_name.startswith("config.") or
|
|
_mod_name == "datasets" or _mod_name.startswith("datasets.")
|
|
):
|
|
del sys.modules[_mod_name]
|
|
|
|
from config.argument_setting import get_args_parser
|
|
from architecture.backbone import build_backbone
|
|
from architecture.transformer import build_transformer
|
|
from architecture.model import OmniShotCut
|
|
from datasets.transforms import Video_Augmentation_Transform
|
|
from util.visualization import visualize_concated_frames
|
|
from config.label_correspondence import unique_intra_label_mapping, unique_inter_label_mapping, intra_int2string, inter_int2string
|
|
|
|
|
|
# Video Transform
|
|
video_transform = Video_Augmentation_Transform(set_type = "val")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def load_model(checkpoint_path: str):
|
|
# 1. 检查路径
|
|
checkpoint_path = os.path.abspath(checkpoint_path)
|
|
if not os.path.exists(checkpoint_path):
|
|
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
|
|
|
|
# 2. 加载 state_dict (修复 PyTorch 2.6 安全报错)
|
|
try:
|
|
# 显式设置 weights_only=False 以允许加载 argparse.Namespace 等对象
|
|
state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
|
except TypeError:
|
|
# 兼容不支持 weights_only 参数的旧版本 PyTorch
|
|
state_dict = torch.load(checkpoint_path, map_location="cpu")
|
|
|
|
if "args" not in state_dict or "model" not in state_dict:
|
|
raise ValueError("Checkpoint must contain keys: 'args' and 'model'.")
|
|
|
|
# 3. 初始化模型结构
|
|
model_args = state_dict["args"]
|
|
backbone = build_backbone(model_args)
|
|
transformer = build_transformer(model_args)
|
|
|
|
model = OmniShotCut(
|
|
backbone,
|
|
transformer,
|
|
num_intra_relation_classes = model_args.num_intra_relation_classes,
|
|
num_inter_relation_classes = model_args.num_inter_relation_classes,
|
|
num_frames = model_args.max_process_window_length,
|
|
num_queries = model_args.num_queries,
|
|
aux_loss = model_args.aux_loss,
|
|
)
|
|
|
|
# 4. 加载权重并移动到设备
|
|
model.load_state_dict(state_dict["model"], strict=True)
|
|
|
|
# TA 建议：动态判断设备，增强代码鲁棒性
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
return model, model_args
|
|
|
|
|
|
|
|
def get_video_fps_safe(video_path: str, default_fps: float = 24.0) -> float:
|
|
try:
|
|
cap = cv2.VideoCapture(video_path)
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
cap.release()
|
|
if fps is None or fps <= 1e-6 or math.isnan(fps):
|
|
return default_fps
|
|
return float(fps)
|
|
except Exception:
|
|
return default_fps
|
|
|
|
|
|
|
|
def split_videos(video, chunk_size, num_context_frames):
|
|
|
|
assert video.ndim == 4, "video must be (T, H, W, C)"
|
|
total_num_frames, H, W, C = video.shape
|
|
|
|
|
|
# Padding at the beginning
|
|
black = np.zeros((num_context_frames, H, W, C), dtype=video.dtype)
|
|
video = np.concatenate([black, video], axis=0)
|
|
|
|
|
|
# Split Video to clips
|
|
stride = chunk_size - 2 * num_context_frames
|
|
cur_frame_idx = 0
|
|
return_list = []
|
|
while cur_frame_idx < total_num_frames:
|
|
|
|
# Fetch the range
|
|
cropped_videos = video[cur_frame_idx : cur_frame_idx + chunk_size]
|
|
|
|
# Add padding if needed
|
|
clip_num_adding_frames = chunk_size - len(cropped_videos)
|
|
if clip_num_adding_frames > 0:
|
|
black = np.zeros((clip_num_adding_frames, H, W, C), dtype=video.dtype)
|
|
cropped_videos = np.concatenate([cropped_videos, black], axis=0)
|
|
|
|
# Append all return info: (video_np, clip padding frames, global start frame idx)
|
|
return_list.append([cropped_videos, clip_num_adding_frames])
|
|
|
|
# Update
|
|
cur_frame_idx += stride
|
|
|
|
return return_list
|
|
|
|
|
|
|
|
def prune_non_context_ranges(pred_ranges, pred_intra_labels, pred_inter_labels, inference_window_size, num_context_frames):
|
|
|
|
# Init
|
|
new_pred_ranges, new_pred_intra_labels, new_pred_inter_labels = [], [], []
|
|
|
|
|
|
# Iterate
|
|
for shot_idx in range(len(pred_ranges)):
|
|
|
|
# Fetch
|
|
start_frame_idx, end_frame_idx = pred_ranges[shot_idx]
|
|
|
|
# Check if we should skip
|
|
if end_frame_idx <= num_context_frames: # Beginning
|
|
continue
|
|
if start_frame_idx >= inference_window_size - num_context_frames: # Ending
|
|
break
|
|
|
|
# Re align start & end
|
|
aligned_start_frame_idx = max(start_frame_idx, num_context_frames) - num_context_frames
|
|
aligned_end_frame_idx = min(end_frame_idx, inference_window_size - num_context_frames) - num_context_frames # exclusive on the right range
|
|
|
|
# Append
|
|
new_pred_ranges.append([aligned_start_frame_idx, aligned_end_frame_idx])
|
|
new_pred_intra_labels.append(pred_intra_labels[shot_idx])
|
|
new_pred_inter_labels.append(pred_inter_labels[shot_idx])
|
|
|
|
return new_pred_ranges, new_pred_intra_labels, new_pred_inter_labels
|
|
|
|
|
|
|
|
def merge_ranges(pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full, pred_ranges, pred_intra_labels, pred_inter_labels):
|
|
|
|
# Prepare
|
|
last_frame_idx = pred_ranges_full[-1][-1] if len(pred_intra_labels_full) != 0 else 0
|
|
|
|
|
|
# Merge last one of the list list
|
|
if len(pred_intra_labels_full) != 0 and pred_intra_labels_full[-1] == pred_intra_labels[0] and pred_inter_labels[0] == unique_inter_label_mapping['new_start']:
|
|
pred_ranges_full[-1][-1] = last_frame_idx + pred_ranges[0][-1]
|
|
|
|
# Crop the first one
|
|
pred_ranges = pred_ranges[1:]
|
|
pred_intra_labels = pred_intra_labels[1:]
|
|
pred_inter_labels = pred_inter_labels[1:]
|
|
|
|
|
|
# Extend the following list
|
|
for idx in range(len(pred_ranges)):
|
|
start_frame_idx, end_frame_idx = pred_ranges[idx]
|
|
|
|
pred_ranges_full.append([last_frame_idx + start_frame_idx, last_frame_idx + end_frame_idx])
|
|
pred_intra_labels_full.append(pred_intra_labels[idx])
|
|
pred_inter_labels_full.append(pred_inter_labels[idx])
|
|
|
|
|
|
return pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full
|
|
|
|
|
|
|
|
def _get_ffmpeg_major_version() -> int:
|
|
"""获取系统 FFmpeg 主版本号，失败时返回 0。"""
|
|
try:
|
|
result = subprocess.run(
|
|
["ffmpeg", "-version"],
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE
|
|
)
|
|
first_line = result.stdout.decode("utf-8", errors="ignore").splitlines()[0]
|
|
# 格式: "ffmpeg version 4.4.2-..." 或 "ffmpeg version 6.0-..."
|
|
ver_str = first_line.split("version")[1].strip().split()[0] # e.g. "4.4.2-0ubuntu..."
|
|
major = int(ver_str.split(".")[0])
|
|
return major
|
|
except Exception:
|
|
return 0
|
|
|
|
_FFMPEG_MAJOR = _get_ffmpeg_major_version()
|
|
|
|
|
|
def single_video_inference(video_path, model, model_args, num_context_frames):
|
|
# 1. 强制转换为绝对路径，并且不要手动加引号（ffmpeg-python 会处理）
|
|
video_path = os.path.abspath(video_path)
|
|
|
|
# 调试打印，方便你在控制台看完整路径
|
|
print(f"## [OmniShotCut] FFmpeg 正在打开: {video_path} (FFmpeg major={_FFMPEG_MAJOR})")
|
|
|
|
process_height, process_width = model_args.process_height, model_args.process_width
|
|
max_process_window_length = model_args.max_process_window_length
|
|
|
|
# 2. 调用 FFmpeg
|
|
# fps_mode 是 FFmpeg 5.0+ 的参数；4.x 及以下用 vsync
|
|
try:
|
|
output_kwargs = dict(
|
|
format="rawvideo",
|
|
pix_fmt="rgb24",
|
|
s=f"{process_width}x{process_height}",
|
|
)
|
|
if _FFMPEG_MAJOR >= 5:
|
|
output_kwargs["fps_mode"] = "passthrough"
|
|
else:
|
|
output_kwargs["vsync"] = "passthrough"
|
|
|
|
process = (
|
|
ffmpeg
|
|
.input(filename=video_path)
|
|
.output("pipe:", **output_kwargs)
|
|
.run_async(pipe_stdout=True, pipe_stderr=True)
|
|
)
|
|
video_stream, err = process.communicate()
|
|
|
|
if process.returncode != 0:
|
|
raise ffmpeg.Error(None, video_stream, err)
|
|
|
|
except ffmpeg.Error as e:
|
|
stderr_msg = e.stderr.decode('utf-8') if e.stderr else "FFmpeg 启动失败"
|
|
print(f"## [OmniShotCut] FFmpeg 错误详情:\n{stderr_msg}")
|
|
raise RuntimeError(f"FFmpeg 无法打开视频。详情: {stderr_msg}")
|
|
|
|
# 3. 转换 Buffer
|
|
video_np_full = np.frombuffer(video_stream, np.uint8).reshape(-1, process_height, process_width, 3)
|
|
|
|
# 获取 FPS
|
|
fps = get_video_fps_safe(video_path)
|
|
|
|
# ... 后面的推理逻辑保持不变 ...
|
|
|
|
pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full = [], [], []
|
|
|
|
# 注意：此处确保 video_transform 已经定义
|
|
for clip_idx, (video_np, num_pad_frames) in enumerate(split_videos(video_np_full, max_process_window_length, num_context_frames)):
|
|
|
|
# Transform & Move to GPU
|
|
video_tensor = video_transform(video_np).unsqueeze(0).to("cuda")
|
|
|
|
# Inference
|
|
with torch.inference_mode():
|
|
outputs = model(video_tensor)
|
|
|
|
# 后处理逻辑
|
|
probas_intra = outputs['intra_clip_logits'].softmax(-1)[0, :, :-1]
|
|
probas_inter = outputs['inter_clip_logits'].softmax(-1)[0, :, :-1]
|
|
range_probas = outputs['pred_shot_logits'].softmax(-1)[0, :, :-1]
|
|
|
|
query_intra_idx = probas_intra.argmax(dim=-1)
|
|
query_inter_idx = probas_inter.argmax(dim=-1)
|
|
query_range_idx = range_probas.argmax(dim=-1)
|
|
|
|
pred_ranges, pred_intra_labels, pred_inter_labels = [], [], []
|
|
start_frame_idx = 0
|
|
|
|
for keep_idx in range(len(query_intra_idx)):
|
|
pred_intra_label = int(query_intra_idx[keep_idx].detach().cpu())
|
|
pred_inter_label = int(query_inter_idx[keep_idx].detach().cpu())
|
|
end_frame_idx = int(query_range_idx[keep_idx].detach().cpu())
|
|
|
|
pred_range = [start_frame_idx, end_frame_idx]
|
|
if start_frame_idx >= end_frame_idx:
|
|
continue
|
|
|
|
pred_ranges.append(pred_range)
|
|
pred_intra_labels.append(pred_intra_label)
|
|
pred_inter_labels.append(pred_inter_label)
|
|
start_frame_idx = end_frame_idx
|
|
|
|
if end_frame_idx >= max_process_window_length - num_pad_frames:
|
|
break
|
|
|
|
# 剪枝与合并
|
|
pred_ranges, pred_intra_labels, pred_inter_labels = prune_non_context_ranges(
|
|
pred_ranges, pred_intra_labels, pred_inter_labels, max_process_window_length, num_context_frames
|
|
)
|
|
|
|
pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full = merge_ranges(
|
|
pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full,
|
|
pred_ranges, pred_intra_labels, pred_inter_labels
|
|
)
|
|
|
|
return pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full, video_np_full, fps
|
|
|
|
|
|
|
|
|
|
def dump_list_of_dict(data, save_path, indent=4):
|
|
"""
|
|
Save list[dict] as JSON.
|
|
Convert pred_intra_labels / pred_inter_labels from int to string.
|
|
"""
|
|
|
|
def convert_item(item):
|
|
item = copy.deepcopy(item)
|
|
|
|
if "pred_intra_labels" in item:
|
|
item["pred_intra_labels"] = [
|
|
intra_int2string.get(x, f"Unknown_{x}")
|
|
for x in item["pred_intra_labels"]
|
|
]
|
|
|
|
if "pred_inter_labels" in item:
|
|
item["pred_inter_labels"] = [
|
|
inter_int2string.get(x, f"Unknown_{x}")
|
|
for x in item["pred_inter_labels"]
|
|
]
|
|
|
|
return item
|
|
|
|
def format_dict(d, level):
|
|
indent_str = " " * (indent * level)
|
|
inner_indent = " " * (indent * (level + 1))
|
|
|
|
lines = [indent_str + "{"]
|
|
items = list(d.items())
|
|
|
|
for i, (k, v) in enumerate(items):
|
|
value_str = json.dumps(v, ensure_ascii=False)
|
|
comma = "," if i < len(items) - 1 else ""
|
|
lines.append(f'{inner_indent}"{k}": {value_str}{comma}')
|
|
|
|
lines.append(f"{indent_str}}}")
|
|
return "\n".join(lines)
|
|
|
|
data = [convert_item(item) for item in data]
|
|
|
|
with open(save_path, "w", encoding="utf-8") as f:
|
|
f.write("[\n")
|
|
|
|
for i, item in enumerate(data):
|
|
dict_str = format_dict(item, level=1)
|
|
comma = "," if i < len(data) - 1 else ""
|
|
f.write(dict_str + comma + "\n")
|
|
|
|
f.write("]\n")
|
|
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--checkpoint_path",
|
|
type = str,
|
|
default = "checkpoints/OmniShotCut_ckpt.pth",
|
|
help = "Path to checkpoint file."
|
|
)
|
|
parser.add_argument(
|
|
"--input_video_path",
|
|
type = str,
|
|
default = "__assets__/demo_video1.mp4",
|
|
help = "Path to the input video path."
|
|
)
|
|
parser.add_argument(
|
|
"--result_store_path",
|
|
type = str,
|
|
default = "results.json",
|
|
help="Path to save result json."
|
|
)
|
|
parser.add_argument(
|
|
"--num_context_frames",
|
|
type = int,
|
|
default = 0,
|
|
help = "Number of historical context frames to overlap between each clip window. The default is 100, using 5-10 should be good."
|
|
)
|
|
parser.add_argument(
|
|
"--visual_store_folder_path",
|
|
type = str,
|
|
default = "demo_video_results",
|
|
help = "Path to save the visualization results. Set to None to disable."
|
|
)
|
|
parser.add_argument(
|
|
"--mode",
|
|
type = str,
|
|
default = "default",
|
|
choices = ["default", "clean_shot"],
|
|
help = "Output Mode. 'default' means all Intra and Inter label. 'clean_shot' means only General Shot Cut without transitions. "
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
# Setting
|
|
inference_args = parse_args()
|
|
checkpoint_path = inference_args.checkpoint_path
|
|
input_video_path = inference_args.input_video_path
|
|
assert(os.path.exists(input_video_path))
|
|
result_store_path = inference_args.result_store_path
|
|
visual_store_folder_path = inference_args.visual_store_folder_path
|
|
mode = inference_args.mode
|
|
|
|
|
|
|
|
# Prepare the folder
|
|
if visual_store_folder_path is not None:
|
|
if os.path.exists(visual_store_folder_path):
|
|
shutil.rmtree(visual_store_folder_path)
|
|
os.makedirs(visual_store_folder_path)
|
|
|
|
|
|
# Load Checkpoint & Model Config
|
|
assert(os.path.exists(checkpoint_path))
|
|
state_dict = torch.load(checkpoint_path, map_location='cpu')
|
|
model_args = state_dict['args']
|
|
print("Checkpoint stored args are", model_args)
|
|
|
|
|
|
# Init the Model
|
|
print("Loading OmniShotCut Model!")
|
|
backbone = build_backbone(model_args)
|
|
transformer = build_transformer(model_args)
|
|
model = OmniShotCut(
|
|
backbone,
|
|
transformer,
|
|
num_intra_relation_classes = model_args.num_intra_relation_classes,
|
|
num_inter_relation_classes = model_args.num_inter_relation_classes,
|
|
num_frames = model_args.max_process_window_length,
|
|
num_queries = model_args.num_queries,
|
|
aux_loss = model_args.aux_loss,
|
|
)
|
|
model.load_state_dict(state_dict['model'], strict=True)
|
|
model.to("cuda")
|
|
model.eval()
|
|
|
|
|
|
|
|
|
|
# Do the inference
|
|
print("Do the inference!")
|
|
pred_ranges_full, pred_intra_labels_full, pred_inter_labels_full, video_np_full, fps = single_video_inference(input_video_path, model, model_args, inference_args.num_context_frames)
|
|
|
|
|
|
|
|
# Visualize
|
|
if visual_store_folder_path is not None:
|
|
print("Visualize the results!")
|
|
pred_saved_paths = visualize_concated_frames(video_np_full, visual_store_folder_path, pred_ranges_full, max_frames_per_img=264, end_range_exclusive=True, fps=24, start_index = 0)
|
|
|
|
|
|
|
|
# For Clean Shot mode, we only leave general types
|
|
if mode == "clean_shot":
|
|
|
|
# Clean shot mode (No Transitions in the middle)
|
|
general_type_idx = unique_intra_label_mapping["general"]
|
|
effective_indices = []
|
|
for idx, intra_label in enumerate(pred_intra_labels_full):
|
|
if intra_label == general_type_idx:
|
|
effective_indices.append(idx)
|
|
|
|
# Reassign the shots
|
|
pred_ranges_full = np.array(pred_ranges_full)[effective_indices].tolist()
|
|
pred_intra_labels_full = np.array(pred_intra_labels_full)[effective_indices].tolist()
|
|
pred_inter_labels_full = np.array(pred_inter_labels_full)[effective_indices].tolist()
|
|
|
|
|
|
|
|
# Collect prediction results
|
|
pred_result = {}
|
|
pred_result["video_path"] = input_video_path
|
|
pred_result["pred_ranges"] = pred_ranges_full
|
|
pred_result["pred_intra_labels"] = pred_intra_labels_full
|
|
pred_result["pred_inter_labels"] = pred_inter_labels_full
|
|
|
|
# Dump to json
|
|
dump_list_of_dict([pred_result], result_store_path)
|
|
|
|
|
|
print("Finished!")
|
|
|
|
# 包装类，方便 ComfyUI 调用
|
|
class OmniShotCutDetector:
|
|
def __init__(self, checkpoint_path, device="cuda"):
|
|
self.model, self.model_args = load_model(checkpoint_path)
|
|
self.model.to(device)
|
|
self.device = device
|
|
|
|
def run_inference(self, input_video_path, mode="default", num_context_frames=0):
|
|
# 调用原始推理函数
|
|
pred_ranges, pred_intra_labels, pred_inter_labels, video_np, fps = single_video_inference(
|
|
input_video_path, self.model, self.model_args, num_context_frames
|
|
)
|
|
|
|
# ── 步骤1：先生成预览图（使用完整镜头列表，与原项目 inference.py 第427行一致）
|
|
vis_images = []
|
|
try:
|
|
import tempfile
|
|
tmp_dir = tempfile.mkdtemp(prefix="omni_preview_")
|
|
saved_paths = visualize_concated_frames(
|
|
frames=video_np,
|
|
out_dir=tmp_dir,
|
|
highlight_ranges_closed=pred_ranges,
|
|
max_frames_per_img=264,
|
|
end_range_exclusive=True,
|
|
fps=24, # 固定24fps，与原项目保持一致
|
|
start_index=0,
|
|
)
|
|
for p in saved_paths:
|
|
vis_images.append(Image.open(p).copy()) # copy() 防止文件句柄占用
|
|
# 清理临时文件
|
|
for p in saved_paths:
|
|
try:
|
|
os.remove(p)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
os.rmdir(tmp_dir)
|
|
except Exception:
|
|
pass
|
|
except Exception as e:
|
|
import traceback
|
|
print(f"## [OmniShotCut] 预览图生成失败: {e}")
|
|
traceback.print_exc()
|
|
|
|
# ── 步骤2：clean_shot 模式过滤（仅影响 JSON 输出，不影响预览图）
|
|
if mode == "clean_shot":
|
|
general_type_idx = unique_intra_label_mapping["general"]
|
|
effective_indices = [i for i, lbl in enumerate(pred_intra_labels) if lbl == general_type_idx]
|
|
pred_ranges = [pred_ranges[i] for i in effective_indices]
|
|
pred_intra_labels = [pred_intra_labels[i] for i in effective_indices]
|
|
pred_inter_labels = [pred_inter_labels[i] for i in effective_indices]
|
|
|
|
# ── 步骤3：结果打包，int 标签转字符串，列表包裹，与原项目 output.json 格式一致
|
|
results = [
|
|
{
|
|
"video_path": input_video_path,
|
|
"pred_ranges": pred_ranges,
|
|
"pred_intra_labels": [intra_int2string.get(x, f"Unknown_{x}") for x in pred_intra_labels],
|
|
"pred_inter_labels": [inter_int2string.get(x, f"Unknown_{x}") for x in pred_inter_labels],
|
|
}
|
|
]
|
|
|
|
return results, vis_images # vis_images 为 PIL Image 列表 |