Files
pikasTech-HWLAB/skills/arm2d-skill/python/video_rembg_tool.py
T
Codex Agent 712db9c597 feat: arm2d-skill directory seed for CaseRun agent workspace
Add B1 directory-type seed mechanism to AGENT_WORKSPACE_SEED_FILES:
- Extended agentWorkspaceFilesForRun with collectDirectorySeedFiles helper
- arm2d-skill (SKILL.md + references/ + python/) auto-injected to .agents/skills/arm2d-skill/
- Enables Code Agent to follow ARM-2D constraints and call asset scripts in case

Skill source: https://github.com/notLabyet/HWLabOA
2026-06-07 14:40:48 +08:00

154 lines
6.3 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
视频抠图转序列帧工具 - 专为ARM2D动效/B站素材设计
========== 所有参数都在下面配置区改,不用动下面逻辑 ==========
"""
import warnings
warnings.filterwarnings("ignore") # 屏蔽无关性能警告
import cv2
from rembg import remove
from PIL import Image
import numpy as np
import os
# ====================== 配置区(你直接改这里就行) ======================
INPUT_VIDEO_PATH = r"C:\Users\11791\OneDrive\Desktop\images\girl.mp4" # 输入视频路径
OUTPUT_DIR = r"C:\Users\11791\OneDrive\Desktop\images\output_frames" # 输出帧保存路径
# 帧配置
TARGET_FRAME_COUNT = 100 # 要输出的总帧数,均匀从视频里抽取
TARGET_SIZE = (240, 240) # 输出图片尺寸(宽,高),和你LCD一致就行
AUTO_CROP_CENTER = True # 自动居中裁剪,避免拉伸变形,建议开
# 画质配置
CONTRAST = 1.0 # 对比度倍数,>1提高,视频暗就往大调
BRIGHTNESS = 10 # 亮度增加值,0~255,视频暗往大调
ENABLE_ALPHA_MATTING = False # 开边缘优化,抠出来无锯齿
ONLY_OUTPUT_MASK = False # 只输出黑白Alpha遮罩,做局部刷新时开
ENABLE_REMOVE_BG = True # 是否开启抠图,关闭后直接抽帧转GIF(速度提升10倍以上)
USE_GPU = False # 要是CUDA装好了就改成True,速度快10倍
# GIF输出配置
ENABLE_GIF_OUTPUT = True # 是否输出合并后的GIF
OUTPUT_GIF_PATH = r"C:\Users\11791\OneDrive\Desktop\images\output.gif" # GIF输出路径
GIF_FPS = 10 # GIF的播放帧率,10帧每秒就很流畅
GIF_LOOP_FOREVER = True # GIF是否无限循环
KEEP_FRAMES_AFTER_GIF = True # 生成GIF后是否保留单独的序列帧
# =======================================================================
# 自动创建输出目录
os.makedirs(OUTPUT_DIR, exist_ok=True)
# 配置校验
if not ENABLE_REMOVE_BG and ONLY_OUTPUT_MASK:
print("⚠️ 提示:不抠图模式下ONLY_OUTPUT_MASK选项无效,已自动关闭")
ONLY_OUTPUT_MASK = False
# 选择运行后端
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if USE_GPU else ['CPUExecutionProvider']
# 读取视频信息
cap = cv2.VideoCapture(INPUT_VIDEO_PATH, cv2.CAP_FFMPEG)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# 计算要抽取的帧索引(均匀分布在整个视频)
extract_indices = [int(i * total_frames / TARGET_FRAME_COUNT) for i in range(TARGET_FRAME_COUNT)]
print(f"=== 开始处理 ===")
print(f"输入视频总帧数:{total_frames}")
print(f"将抽取{len(extract_indices)}帧,输出尺寸:{TARGET_SIZE[0]}×{TARGET_SIZE[1]}")
print(f"输出路径:{OUTPUT_DIR}\n")
for idx, frame_idx in enumerate(extract_indices):
# 跳到指定帧
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
ret, frame = cap.read()
if not ret:
print(f"跳过第{idx+1}帧:读取失败")
continue
# 1. 调整对比度亮度
frame = cv2.convertScaleAbs(frame, alpha=CONTRAST, beta=BRIGHTNESS)
# 2. 裁剪+缩放
if AUTO_CROP_CENTER:
h, w = frame.shape[:2]
crop_size = min(w, h)
x_start = (w - crop_size) // 2
y_start = (h - crop_size) // 2
frame = frame[y_start:y_start+crop_size, x_start:x_start+crop_size]
frame = cv2.resize(frame, TARGET_SIZE)
# 3. 抠图/直接使用原图
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img_pil = Image.fromarray(frame_rgb)
if ENABLE_REMOVE_BG:
img_rembg = remove(
img_pil,
providers=providers,
alpha_matting=ENABLE_ALPHA_MATTING,
alpha_matting_discard_threshold=1e-5, # 优化矩阵计算,消除警告
alpha_matting_shift=1e-2,
only_mask=ONLY_OUTPUT_MASK
)
else:
# 不抠图模式,直接保留原图,添加全不透明alpha通道兼容后续逻辑
img_rembg = img_pil.convert("RGBA")
# 4. 保存
if ONLY_OUTPUT_MASK:
# 遮罩直接存灰度图
img_out = Image.fromarray(np.array(img_rembg) * 255).convert("L")
else:
# 存带透明通道的RGBA图
img_out = Image.fromarray(np.array(img_rembg))
output_path = os.path.join(OUTPUT_DIR, f"frame_{idx+1:02d}.png")
img_out.save(output_path)
print(f"已完成第{idx+1}/{TARGET_FRAME_COUNT}帧:{os.path.basename(output_path)}")
cap.release()
frame_count = len(os.listdir(OUTPUT_DIR))
print(f"\n=== 帧处理完成,共生成{frame_count}帧 ===")
# ========== 生成GIF ==========
if ENABLE_GIF_OUTPUT and frame_count > 0:
print("正在生成GIF...")
# 读取所有帧
frames = []
for i in range(1, frame_count+1):
frame_path = os.path.join(OUTPUT_DIR, f"frame_{i:02d}.png")
if os.path.exists(frame_path):
img = Image.open(frame_path)
# 处理透明通道,避免GIF透明异常
if img.mode == "RGBA":
# 转索引色,保留透明
alpha = img.getchannel("A")
img = img.convert("P", palette=Image.ADAPTIVE, colors=256)
mask = Image.eval(alpha, lambda a: 255 if a <= 128 else 0)
img.paste(255, mask=mask)
img.info["transparency"] = 255
frames.append(img)
if len(frames) > 0:
# 计算每帧延迟时间(毫秒)
frame_delay = int(1000 / GIF_FPS)
loop = 0 if GIF_LOOP_FOREVER else 1
# 保存GIF
frames[0].save(
OUTPUT_GIF_PATH,
save_all=True,
append_images=frames[1:],
duration=frame_delay,
loop=loop,
disposal=2, # 每帧播放完清除,避免透明残留
optimize=False
)
print(f"GIF生成成功:{OUTPUT_GIF_PATH}")
# 删除单独帧(如果配置了)
if not KEEP_FRAMES_AFTER_GIF:
for f in os.listdir(OUTPUT_DIR):
os.remove(os.path.join(OUTPUT_DIR, f))
os.rmdir(OUTPUT_DIR)
print("已清理临时序列帧文件")
print(f"\n=== 全部处理完成 ===")
if KEEP_FRAMES_AFTER_GIF or not ENABLE_GIF_OUTPUT:
print(f"序列帧保存在:{OUTPUT_DIR}")
if ENABLE_GIF_OUTPUT:
print(f"GIF保存在:{OUTPUT_GIF_PATH}")