Add GPU support and multi-threaded subtitle detection
Docker Build and Push / check-secrets (push) Successful in 30s
Docker Build and Push / build-and-push (cuda, 12.6) (push) Has been skipped
Docker Build and Push / build-and-push (cuda, 12.8) (push) Has been skipped
Docker Build and Push / build-and-push (directml, latest) (push) Has been skipped
Docker Build and Push / build-and-push (cpu, latest) (push) Has been skipped
Docker Build and Push / build-and-push (cuda, 11.8) (push) Has been skipped
Build Windows CPU / build (push) Has been cancelled
Build Windows CUDA 11.8 / build (push) Has been cancelled
Build Windows CUDA 12.6 / build (push) Has been cancelled
Build Windows CUDA 12.8 / build (push) Has been cancelled
Build Windows DirectML / build (push) Has been cancelled
Docker Build and Push / check-secrets (push) Successful in 30s
Docker Build and Push / build-and-push (cuda, 12.6) (push) Has been skipped
Docker Build and Push / build-and-push (cuda, 12.8) (push) Has been skipped
Docker Build and Push / build-and-push (directml, latest) (push) Has been skipped
Docker Build and Push / build-and-push (cpu, latest) (push) Has been skipped
Docker Build and Push / build-and-push (cuda, 11.8) (push) Has been skipped
Build Windows CPU / build (push) Has been cancelled
Build Windows CUDA 11.8 / build (push) Has been cancelled
Build Windows CUDA 12.6 / build (push) Has been cancelled
Build Windows CUDA 12.8 / build (push) Has been cancelled
Build Windows DirectML / build (push) Has been cancelled
- Add gpu_id parameter to SubtitleDetect for GPU acceleration - Implement multi-threaded frame scanning with thread-safe CUDA context - Add smart_detect module and web application components - Update .gitignore to exclude web cache Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Haiku 4.5
parent
5b780e3186
commit
693d394f55
@@ -375,3 +375,4 @@ test*_no_sub*.mp4
|
||||
/backend/tools/train/release_model/
|
||||
inference.onnx
|
||||
/config/config.json
|
||||
web_cache/
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
"""
|
||||
智能水印检测:通过帧差分找静态区域,用SAM2精确分割。
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
SAM2_PATH = "/opt/sam2"
|
||||
SAM2_CKPT = "/opt/sam2/checkpoints/sam2.1_hiera_tiny.pt"
|
||||
SAM2_CFG = "configs/sam2.1/sam2.1_hiera_t.yaml"
|
||||
|
||||
|
||||
def detect_static_regions(video_path: str, sample_frames: int = 30, threshold: float = 20.0,
|
||||
max_scan_seconds: float = None):
|
||||
"""
|
||||
饱和度签名 + 跨帧对比法检测水印/Logo:
|
||||
1. 采样多帧,通过高饱和度像素比例判断哪些帧含有logo(logo通常色彩鲜艳)
|
||||
2. 在"含logo帧"之间比较,logo像素在多帧中出现在同一位置 → 差值极低
|
||||
3. 能检测周期性出现/动态水印,无需logo持续存在于全部帧
|
||||
"""
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
|
||||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
|
||||
duration_s = total / fps
|
||||
num_samples = max(10, min(24, int(duration_s // 8)))
|
||||
positions = [min(int((i + 0.5) * total / num_samples), total - 1)
|
||||
for i in range(num_samples)]
|
||||
|
||||
bgr_frames = []
|
||||
gray_frames = []
|
||||
for pos in positions:
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
|
||||
ret, frame = cap.read()
|
||||
if ret:
|
||||
bgr_frames.append(frame)
|
||||
gray_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32))
|
||||
cap.release()
|
||||
|
||||
print(f"[SAM2-detect] 采样 {len(bgr_frames)} 帧 (均匀分布于{duration_s:.1f}s视频)", flush=True)
|
||||
|
||||
if len(bgr_frames) < 4:
|
||||
print(f"[SAM2-detect] 帧数不足4帧,返回空", flush=True)
|
||||
return []
|
||||
|
||||
# ── 阶段1:通过高饱和度识别含logo帧 ──────────────────────────────────
|
||||
# 水印/logo通常色彩比视频自然内容更鲜艳(高饱和度)
|
||||
# 在帧的上半部分(logo常出现区域)检测高饱和度像素数
|
||||
top_h = h // 5 # 检查帧顶部 20%
|
||||
logo_sat_thresh = 200 # 饱和度 > 200 认为是logo颜色
|
||||
logo_min_pixels = top_h * w * 0.003 # 至少 0.3% 的上方区域为logo颜色
|
||||
|
||||
logo_frames_gray = []
|
||||
for i, (bgr, gray) in enumerate(zip(bgr_frames, gray_frames)):
|
||||
hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
|
||||
high_sat = int(np.sum(hsv[:top_h, :, 1] > logo_sat_thresh))
|
||||
pos_s = positions[i] / fps
|
||||
if high_sat >= logo_min_pixels:
|
||||
logo_frames_gray.append(gray)
|
||||
print(f"[SAM2-detect] ✓ {pos_s:.1f}s: 高饱和={high_sat}px (含logo)", flush=True)
|
||||
else:
|
||||
print(f"[SAM2-detect] ✗ {pos_s:.1f}s: 高饱和={high_sat}px (无logo)", flush=True)
|
||||
|
||||
print(f"[SAM2-detect] logo含帧: {len(logo_frames_gray)}/{len(bgr_frames)}", flush=True)
|
||||
|
||||
logo_bgr_frames = [bgr_frames[i] for i, g in enumerate(gray_frames)
|
||||
if any(np.array_equal(g, l) for l in logo_frames_gray)]
|
||||
nlogo_bgr_frames = [bgr_frames[i] for i, g in enumerate(gray_frames)
|
||||
if not any(np.array_equal(g, l) for l in logo_frames_gray)]
|
||||
|
||||
if len(logo_frames_gray) < 2:
|
||||
print(f"[SAM2-detect] logo帧不足,回退到时序std法", flush=True)
|
||||
stack = np.stack(gray_frames, axis=0)
|
||||
std_map = np.std(stack, axis=0)
|
||||
static_mask = (std_map < 8.0).astype(np.uint8) * 255
|
||||
else:
|
||||
# 在logo帧中统计每像素平均饱和度
|
||||
sat_acc = np.zeros((h, w), dtype=np.float32)
|
||||
for bgr in logo_bgr_frames:
|
||||
hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
|
||||
sat_acc += (hsv[:, :, 1].astype(np.float32) / 255.0)
|
||||
sat_mean_logo = sat_acc / len(logo_bgr_frames)
|
||||
|
||||
if nlogo_bgr_frames:
|
||||
# ── 有无logo帧对比:准确检测logo区域 ───────────────────────────
|
||||
# Logo像素: logo帧高饱和 + 无logo帧低饱和(差值大)
|
||||
sat_acc_nl = np.zeros((h, w), dtype=np.float32)
|
||||
for bgr in nlogo_bgr_frames[:6]:
|
||||
hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
|
||||
sat_acc_nl += (hsv[:, :, 1].astype(np.float32) / 255.0)
|
||||
sat_mean_nlogo = sat_acc_nl / min(len(nlogo_bgr_frames), 6)
|
||||
sat_diff = sat_mean_logo - sat_mean_nlogo
|
||||
|
||||
sat_thresh = min(0.6, np.percentile(sat_mean_logo, 92))
|
||||
diff_thresh = max(0.25, sat_diff.mean() + sat_diff.std())
|
||||
logo_pixel_mask = ((sat_mean_logo >= sat_thresh) & (sat_diff >= diff_thresh)).astype(np.uint8) * 255
|
||||
|
||||
print(f"[SAM2-detect] logo帧平均饱和: max={sat_mean_logo.max():.3f}", flush=True)
|
||||
print(f"[SAM2-detect] 饱和差值: max={sat_diff.max():.3f} mean={sat_diff.mean():.3f}", flush=True)
|
||||
print(f"[SAM2-detect] sat_thresh={sat_thresh:.3f} diff_thresh={diff_thresh:.3f}, 候选={np.sum(logo_pixel_mask > 0)}", flush=True)
|
||||
else:
|
||||
# ── 无非logo帧(短视频/预览:logo全程可见)──────────────────────
|
||||
# 只在帧的边缘区域(顶部30%)寻找高饱和logo
|
||||
# 避免把视频中央的自然内容误报为logo
|
||||
border_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
border_mask[:h * 3 // 10, :] = 1 # 顶部 30%
|
||||
|
||||
# 取最高饱和的4%像素(极鲜艳 = logo文字/图形)
|
||||
sat_thresh = max(0.6, np.percentile(sat_mean_logo, 96))
|
||||
logo_pixel_mask = ((sat_mean_logo >= sat_thresh).astype(np.uint8) * border_mask * 255)
|
||||
|
||||
print(f"[SAM2-detect] 无非logo帧,限定边缘区域搜索", flush=True)
|
||||
print(f"[SAM2-detect] sat_thresh={sat_thresh:.3f}, 候选={np.sum(logo_pixel_mask > 0)}", flush=True)
|
||||
|
||||
static_mask = logo_pixel_mask
|
||||
|
||||
# 仅排除纯黑区域
|
||||
mean_gray = np.mean(np.stack(gray_frames, axis=0), axis=0)
|
||||
content_mask = (mean_gray > 8).astype(np.uint8) * 255
|
||||
static_mask = cv2.bitwise_and(static_mask, content_mask)
|
||||
print(f"[SAM2-detect] 过滤纯黑后候选像素数={np.sum(static_mask > 0)}", flush=True)
|
||||
|
||||
# 形态学处理:去噪 + 膨胀
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
||||
static_mask = cv2.morphologyEx(static_mask, cv2.MORPH_OPEN, kernel)
|
||||
static_mask = cv2.dilate(static_mask, kernel, iterations=2)
|
||||
|
||||
# 找轮廓
|
||||
contours, _ = cv2.findContours(static_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
regions = []
|
||||
min_area = (w * h) * 0.0005 # 至少占画面 0.05%
|
||||
max_area = (w * h) * 0.25 # 最多占画面 25%
|
||||
|
||||
print(f"[SAM2-detect] 找到 {len(contours)} 个轮廓,min_area={min_area:.0f} max_area={max_area:.0f}", flush=True)
|
||||
for cnt in contours:
|
||||
area = cv2.contourArea(cnt)
|
||||
x, y, bw, bh = cv2.boundingRect(cnt)
|
||||
ratio = bw / bh if bh > 0 else 0
|
||||
if min_area < area < max_area:
|
||||
if 0.1 < ratio < 20:
|
||||
regions.append((x, y, x + bw, y + bh))
|
||||
print(f"[SAM2-detect] ✓ 区域 ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f} 比例={ratio:.2f}", flush=True)
|
||||
else:
|
||||
print(f"[SAM2-detect] ✗ 比例过滤: ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f} 比例={ratio:.2f}", flush=True)
|
||||
else:
|
||||
print(f"[SAM2-detect] ✗ 面积过滤: ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f}", flush=True)
|
||||
|
||||
# 若检测到多个位于帧顶部的区域,合并成一个大框覆盖整个logo
|
||||
# logo通常有高饱和的亮色区域(检测到)和低饱和的暗色区域(未检测到)
|
||||
# 合并+适当扩展可确保完整覆盖
|
||||
if len(regions) >= 2:
|
||||
top_regions = [r for r in regions if r[1] < h * 0.3] # 只合并帧顶30%内的区域
|
||||
if len(top_regions) >= 2:
|
||||
x1 = max(0, min(r[0] for r in top_regions) - 30)
|
||||
y1 = max(0, min(r[1] for r in top_regions) - 15)
|
||||
x2 = min(w, max(r[2] for r in top_regions) + 30)
|
||||
y2 = min(h, max(r[3] for r in top_regions) + 20)
|
||||
other = [r for r in regions if r[1] >= h * 0.3]
|
||||
merged = [(x1, y1, x2, y2)] + other
|
||||
print(f"[SAM2-detect] 合并顶部 {len(top_regions)} 个区域为 ({x1},{y1})-({x2},{y2})", flush=True)
|
||||
regions = merged
|
||||
|
||||
print(f"[SAM2-detect] detect_static_regions 返回 {len(regions)} 个区域", flush=True)
|
||||
return regions
|
||||
|
||||
|
||||
def refine_with_sam2(image: np.ndarray, boxes: list, gpu_id: int = 0):
|
||||
"""
|
||||
用 SAM2 对检测到的区域精确分割,返回合并后的 mask。
|
||||
image: BGR numpy array
|
||||
boxes: [(x1,y1,x2,y2), ...]
|
||||
返回 binary mask (H, W) uint8
|
||||
"""
|
||||
if not boxes:
|
||||
return None
|
||||
|
||||
try:
|
||||
if SAM2_PATH not in sys.path:
|
||||
sys.path.insert(0, SAM2_PATH)
|
||||
|
||||
import torch
|
||||
from sam2.build_sam import build_sam2
|
||||
from sam2.sam2_image_predictor import SAM2ImagePredictor
|
||||
|
||||
device = f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu"
|
||||
model = build_sam2(SAM2_CFG, SAM2_CKPT, device=device)
|
||||
predictor = SAM2ImagePredictor(model)
|
||||
|
||||
rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
predictor.set_image(rgb)
|
||||
|
||||
h, w = image.shape[:2]
|
||||
combined_mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
for (x1, y1, x2, y2) in boxes:
|
||||
box = np.array([x1, y1, x2, y2])
|
||||
masks, scores, _ = predictor.predict(
|
||||
box=box[None],
|
||||
multimask_output=False,
|
||||
)
|
||||
if masks is not None and len(masks) > 0:
|
||||
m = masks[0].astype(np.uint8) * 255
|
||||
combined_mask = cv2.bitwise_or(combined_mask, m)
|
||||
|
||||
return combined_mask
|
||||
|
||||
except Exception as e:
|
||||
print(f"[SAM2] 分割失败: {e}", flush=True)
|
||||
return None
|
||||
|
||||
|
||||
def smart_detect_watermark(video_path: str, sample_frames: int = 30, gpu_id: int = 0):
|
||||
"""
|
||||
完整流程:帧差分(前30秒每秒一帧)→ SAM2 精确分割。
|
||||
限制扫描前30秒:H.264压缩导致长视频里即使静态logo的像素std也会随时间累积超过阈值。
|
||||
返回 [(x1, y1, x2, y2), ...] 像素坐标列表。
|
||||
"""
|
||||
print(f"[SAM2] smart_detect_watermark 开始: {video_path}", flush=True)
|
||||
# 1. 静态区域检测:跨时间段帧对比法(全视频)
|
||||
boxes = detect_static_regions(video_path, sample_frames)
|
||||
print(f"[SAM2] detect_static_regions 返回 {len(boxes)} 个候选区域: {boxes}", flush=True)
|
||||
if not boxes:
|
||||
return []
|
||||
|
||||
# 2. 读取代表帧用于 SAM2
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, total // 2)
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
if not ret:
|
||||
return boxes # SAM2 失败时直接返回粗检测结果
|
||||
|
||||
# 3. SAM2 精确分割
|
||||
mask = refine_with_sam2(frame, boxes, gpu_id=gpu_id)
|
||||
if mask is None:
|
||||
return boxes
|
||||
|
||||
# 4. 从 mask 提取最终 bounding boxes
|
||||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
refined = []
|
||||
for cnt in contours:
|
||||
if cv2.contourArea(cnt) > 100:
|
||||
x, y, bw, bh = cv2.boundingRect(cnt)
|
||||
refined.append((x, y, x + bw, y + bh))
|
||||
|
||||
result = refined if refined else boxes
|
||||
print(f"[SAM2] smart_detect_watermark 最终返回 {len(result)} 个区域: {result}", flush=True)
|
||||
return result
|
||||
@@ -21,9 +21,10 @@ class SubtitleDetect:
|
||||
# 采样间隔,根据视频帧率在 _init_sample_step 中自适应设置
|
||||
SAMPLE_STEP = 3
|
||||
|
||||
def __init__(self, video_path, sub_areas=[]):
|
||||
def __init__(self, video_path, sub_areas=[], gpu_id=0):
|
||||
self.video_path = video_path
|
||||
self.sub_areas = sub_areas
|
||||
self.gpu_id = gpu_id
|
||||
self._init_sample_step()
|
||||
|
||||
def _init_sample_step(self):
|
||||
@@ -46,10 +47,12 @@ class SubtitleDetect:
|
||||
hardware_accelerator = HardwareAccelerator.instance()
|
||||
onnx_providers = hardware_accelerator.onnx_providers
|
||||
model_config = ModelConfig()
|
||||
import torch as _torch
|
||||
_device = f"gpu:{self.gpu_id}" if _torch.cuda.is_available() else "cpu"
|
||||
return TextDetection(
|
||||
model_name=model_config.DET_MODEL_NAME,
|
||||
model_dir=model_config.DET_MODEL_DIR,
|
||||
device="cpu",
|
||||
device=_device,
|
||||
enable_hpi=len(onnx_providers) > 0,
|
||||
)
|
||||
|
||||
@@ -81,34 +84,97 @@ class SubtitleDetect:
|
||||
break
|
||||
return temp_list
|
||||
|
||||
def find_subtitle_frame_no(self, sub_remover=None):
|
||||
video_cap = cv2.VideoCapture(get_readable_path(self.video_path))
|
||||
frame_count = video_cap.get(cv2.CAP_PROP_FRAME_COUNT)
|
||||
tbar = tqdm(total=int(frame_count), unit='frame', position=0, file=sys.__stdout__, desc='Subtitle Finding')
|
||||
current_frame_no = 0
|
||||
# 阶段1:采样检测,仅对每隔 sample_step 帧执行 OCR
|
||||
sampled_results = {} # frame_no -> temp_list
|
||||
if sub_remover:
|
||||
sub_remover.append_output(tr['Main']['ProcessingStartFindingSubtitles'])
|
||||
while video_cap.isOpened():
|
||||
ret, frame = video_cap.read()
|
||||
# 如果读取视频帧失败(视频读到最后一帧)
|
||||
def _scan_range(self, start_frame, end_frame, ab_sections, results, lock,
|
||||
progress_counter, frame_count, sub_remover, gpu_id=None):
|
||||
"""单线程扫描指定帧范围。
|
||||
每个线程在自己内部初始化 OCR 模型(CUDA 上下文归该线程所有),避免跨线程 CUDA 问题。
|
||||
gpu_id: >=0 用 GPU,None 用 CPU。
|
||||
"""
|
||||
# 在本线程内创建并初始化 OCR 模型,CUDA 上下文属于本线程
|
||||
if gpu_id is not None and gpu_id >= 0:
|
||||
import torch as _t
|
||||
if _t.cuda.is_available():
|
||||
_local_det = SubtitleDetect(self.video_path, self.sub_areas, gpu_id=gpu_id)
|
||||
_ = _local_det.text_detector # 在本线程建立 CUDA 上下文
|
||||
else:
|
||||
_local_det = SubtitleDetect(self.video_path, self.sub_areas)
|
||||
else:
|
||||
_local_det = SubtitleDetect(self.video_path, self.sub_areas)
|
||||
|
||||
cap = cv2.VideoCapture(get_readable_path(self.video_path))
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
|
||||
frame_no = start_frame
|
||||
while frame_no < end_frame:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
# 读取视频帧成功
|
||||
current_frame_no += 1
|
||||
if not is_frame_number_in_ab_sections(current_frame_no - 1, sub_remover.ab_sections):
|
||||
tbar.update(1)
|
||||
frame_no += 1
|
||||
if not is_frame_number_in_ab_sections(frame_no - 1, ab_sections):
|
||||
continue
|
||||
# 仅对采样帧执行 OCR 推理
|
||||
if (current_frame_no - 1) % self.SAMPLE_STEP == 0 or self.SAMPLE_STEP <= 1:
|
||||
temp_list = self.detect_subtitle(frame)
|
||||
if len(temp_list) > 0:
|
||||
sampled_results[current_frame_no] = temp_list
|
||||
tbar.update(1)
|
||||
if sub_remover:
|
||||
sub_remover.progress_total = (100 * float(current_frame_no) / float(frame_count)) // 2
|
||||
if (frame_no - 1) % self.SAMPLE_STEP == 0 or self.SAMPLE_STEP <= 1:
|
||||
temp_list = _local_det.detect_subtitle(frame)
|
||||
if temp_list:
|
||||
with lock:
|
||||
results[frame_no] = temp_list
|
||||
with lock:
|
||||
progress_counter[0] += 1
|
||||
if sub_remover:
|
||||
sub_remover.progress_total = (100 * progress_counter[0] / frame_count) // 2
|
||||
sub_remover.notify_progress_listeners()
|
||||
cap.release()
|
||||
|
||||
def find_subtitle_frame_no(self, sub_remover=None, num_workers=4):
|
||||
import threading
|
||||
video_cap = cv2.VideoCapture(get_readable_path(self.video_path))
|
||||
frame_count = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
video_cap.release()
|
||||
tbar = tqdm(total=frame_count, unit='frame', position=0, file=sys.__stdout__, desc='Subtitle Finding')
|
||||
# 阶段1:并行采样检测
|
||||
sampled_results = {}
|
||||
if sub_remover:
|
||||
sub_remover.append_output(tr['Main']['ProcessingStartFindingSubtitles'])
|
||||
ab_sections = sub_remover.ab_sections if sub_remover else []
|
||||
lock = threading.Lock()
|
||||
import torch as _torch
|
||||
cuda_count = _torch.cuda.device_count() if _torch.cuda.is_available() else 0
|
||||
# 每个线程在自己内部初始化模型,CUDA 上下文属于该线程,避免跨线程推理失败
|
||||
# GPU0 用于 OCR,GPU1 留给 STTN/LaMa(若可用);无 GPU 则 CPU
|
||||
_ocr_gpu_id = 0 if cuda_count >= 1 else None
|
||||
print(f"[OCR] {'GPU' + str(_ocr_gpu_id) if _ocr_gpu_id is not None else 'CPU'} 推理,{num_workers} 线程并行(每线程独立模型)", flush=True)
|
||||
progress_counter = [0]
|
||||
chunk = max(1, frame_count // num_workers)
|
||||
ranges = [(i * chunk, min((i + 1) * chunk, frame_count)) for i in range(num_workers)]
|
||||
threads = []
|
||||
for i, (s, e) in enumerate(ranges):
|
||||
t = threading.Thread(
|
||||
target=self._scan_range,
|
||||
args=(s, e, ab_sections, sampled_results, lock, progress_counter, frame_count, sub_remover, _ocr_gpu_id),
|
||||
daemon=True,
|
||||
)
|
||||
threads.append(t)
|
||||
t.start()
|
||||
# tqdm 进度跟踪(单独线程更新)
|
||||
import time as _time
|
||||
last = 0
|
||||
while any(t.is_alive() for t in threads):
|
||||
with lock:
|
||||
cur = progress_counter[0]
|
||||
tbar.update(cur - last)
|
||||
last = cur
|
||||
_time.sleep(0.2)
|
||||
for t in threads:
|
||||
t.join()
|
||||
tbar.update(frame_count - last)
|
||||
tbar.close()
|
||||
# 线程退出后其局部 OCR 模型自动被 GC,触发 GPU 显存释放
|
||||
import gc as _gc
|
||||
_gc.collect()
|
||||
try:
|
||||
import torch as _t2
|
||||
if _t2.cuda.is_available():
|
||||
_t2.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
# 阶段2:插值填充 — 两个采样帧之间都有字幕时,中间帧也标记为有字幕
|
||||
subtitle_frame_no_box_dict = {}
|
||||
detected_nos = sorted(sampled_results.keys())
|
||||
|
||||
+63
-43
@@ -25,7 +25,13 @@ class FramePrefetcher:
|
||||
def _read_loop(self):
|
||||
while not self._stopped:
|
||||
ret, frame = self.cap.read()
|
||||
self._buffer.put((ret, frame))
|
||||
# 用带超时的 put,避免队列满时永久阻塞
|
||||
while not self._stopped:
|
||||
try:
|
||||
self._buffer.put((ret, frame), timeout=0.5)
|
||||
break
|
||||
except queue.Full:
|
||||
pass
|
||||
if not ret:
|
||||
break
|
||||
|
||||
@@ -39,11 +45,12 @@ class FramePrefetcher:
|
||||
def stop(self):
|
||||
"""停止预读取,不释放底层 video_cap。"""
|
||||
self._stopped = True
|
||||
try:
|
||||
while not self._buffer.empty():
|
||||
# 排空队列,解除后台线程可能阻塞的 put()
|
||||
while True:
|
||||
try:
|
||||
self._buffer.get_nowait()
|
||||
except queue.Empty:
|
||||
pass
|
||||
except queue.Empty:
|
||||
break
|
||||
self._thread.join(timeout=5)
|
||||
|
||||
def release(self):
|
||||
@@ -53,52 +60,65 @@ class FramePrefetcher:
|
||||
|
||||
class FFmpegVideoWriter:
|
||||
"""
|
||||
通过 FFmpeg 管道写入帧,使用 libx264 编码。
|
||||
接口兼容 cv2.VideoWriter(write/release)。
|
||||
独立 I/O 线程写帧到 FFmpeg stdin,直接输出 H.264。
|
||||
主线程只往队列里放帧,永不阻塞在 I/O 上。
|
||||
"""
|
||||
|
||||
def __init__(self, output_path, fps, size):
|
||||
w, h = size
|
||||
cmd = [
|
||||
FFmpegCLI.instance().ffmpeg_path,
|
||||
'-y',
|
||||
'-f', 'rawvideo',
|
||||
'-vcodec', 'rawvideo',
|
||||
'-s', f'{w}x{h}',
|
||||
'-pix_fmt', 'bgr24',
|
||||
'-r', str(fps),
|
||||
'-i', '-',
|
||||
'-c:v', 'libx264',
|
||||
'-pix_fmt', 'yuv420p',
|
||||
'-crf', '18',
|
||||
'-preset', 'fast',
|
||||
'-loglevel', 'error',
|
||||
output_path
|
||||
]
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
self._queue = queue.Queue(maxsize=60)
|
||||
self._error = None
|
||||
# 优先用系统 FFmpeg + NVENC(GPU 编码),回退到 libx264
|
||||
import shutil as _shutil
|
||||
_sys_ffmpeg = '/usr/bin/ffmpeg'
|
||||
_nvenc_ok = (
|
||||
_shutil.which(_sys_ffmpeg) is not None and
|
||||
subprocess.run(
|
||||
[_sys_ffmpeg, '-hide_banner', '-encoders'],
|
||||
capture_output=True, text=True
|
||||
).stdout.find('h264_nvenc') >= 0
|
||||
)
|
||||
ffmpeg_bin = _sys_ffmpeg if _nvenc_ok else FFmpegCLI.instance().ffmpeg_path
|
||||
codec_args = (
|
||||
['-vf', 'format=yuv420p', '-vcodec', 'h264_nvenc', '-cq', '18', '-preset', 'p4']
|
||||
if _nvenc_ok else
|
||||
['-vcodec', 'libx264', '-crf', '18', '-preset', 'ultrafast']
|
||||
)
|
||||
cmd = [
|
||||
ffmpeg_bin, '-y',
|
||||
'-f', 'rawvideo', '-vcodec', 'rawvideo',
|
||||
'-s', f'{w}x{h}', '-pix_fmt', 'bgr24',
|
||||
'-r', str(fps), '-i', 'pipe:0',
|
||||
*codec_args,
|
||||
'-loglevel', 'error', output_path,
|
||||
]
|
||||
self._proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
|
||||
self._thread = threading.Thread(target=self._io_loop, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def _io_loop(self):
|
||||
try:
|
||||
while True:
|
||||
frame = self._queue.get()
|
||||
if frame is None: # sentinel
|
||||
break
|
||||
self._proc.stdin.write(frame.tobytes())
|
||||
except Exception as e:
|
||||
self._error = e
|
||||
finally:
|
||||
try:
|
||||
self._proc.stdin.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def write(self, frame):
|
||||
"""写入一帧(numpy BGR 数组)。"""
|
||||
if frame.dtype != np.uint8:
|
||||
frame = np.clip(frame, 0, 255).astype(np.uint8)
|
||||
try:
|
||||
self._process.stdin.write(frame.tobytes())
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
self._queue.put(frame)
|
||||
|
||||
def release(self):
|
||||
"""关闭管道并等待编码完成。"""
|
||||
try:
|
||||
self._process.stdin.close()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
try:
|
||||
self._process.wait(timeout=600)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._process.terminate()
|
||||
self._process.wait(timeout=5)
|
||||
self._queue.put(None) # 通知 I/O 线程退出
|
||||
self._thread.join() # 等待所有帧写完
|
||||
self._proc.wait() # 等待 FFmpeg 完成编码
|
||||
if self._error:
|
||||
raise RuntimeError(f"FFmpegVideoWriter error: {self._error}")
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Minimal qfluentwidgets stub for headless/server use — no Qt required.
|
||||
Provides just enough for backend/config.py to import and function correctly.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
|
||||
class _Signal:
|
||||
def connect(self, fn): pass
|
||||
def disconnect(self, fn=None): pass
|
||||
def emit(self, *args): pass
|
||||
|
||||
|
||||
class ConfigValidator:
|
||||
def validate(self, v): return True
|
||||
def correct(self, v): return v
|
||||
|
||||
|
||||
class BoolValidator(ConfigValidator):
|
||||
def validate(self, v): return isinstance(v, bool)
|
||||
def correct(self, v): return bool(v)
|
||||
|
||||
|
||||
class OptionsValidator(ConfigValidator):
|
||||
def __init__(self, options):
|
||||
try:
|
||||
self._opts = [e for e in options]
|
||||
except TypeError:
|
||||
self._opts = []
|
||||
|
||||
def validate(self, v): return v in self._opts
|
||||
def correct(self, v): return v if v in self._opts else (self._opts[0] if self._opts else v)
|
||||
|
||||
|
||||
class RangeValidator(ConfigValidator):
|
||||
def __init__(self, min_val, max_val):
|
||||
self.min_val = min_val
|
||||
self.max_val = max_val
|
||||
|
||||
def validate(self, v): return self.min_val <= v <= self.max_val
|
||||
def correct(self, v): return max(self.min_val, min(self.max_val, v))
|
||||
|
||||
|
||||
class EnumSerializer:
|
||||
def __init__(self, enum_cls):
|
||||
self.enum_cls = enum_cls
|
||||
|
||||
def serialize(self, v):
|
||||
return v.value if hasattr(v, 'value') else str(v)
|
||||
|
||||
def deserialize(self, s):
|
||||
for e in self.enum_cls:
|
||||
if e.value == s:
|
||||
return e
|
||||
return s
|
||||
|
||||
|
||||
class ConfigItem:
|
||||
def __init__(self, section, key, default, validator=None, restart=False):
|
||||
self.section = section
|
||||
self.key = key
|
||||
self.value = default
|
||||
self.validator = validator or ConfigValidator()
|
||||
self.restart = restart
|
||||
self.valueChanged = _Signal()
|
||||
|
||||
|
||||
class OptionsConfigItem(ConfigItem):
|
||||
def __init__(self, section, key, default, validator=None, serializer=None, restart=False):
|
||||
super().__init__(section, key, default, validator, restart)
|
||||
self.serializer = serializer
|
||||
|
||||
|
||||
class RangeConfigItem(ConfigItem):
|
||||
def __init__(self, section, key, default, validator=None):
|
||||
super().__init__(section, key, default, validator)
|
||||
|
||||
|
||||
class QConfig:
|
||||
appRestartSig = _Signal()
|
||||
|
||||
def set(self, item, value):
|
||||
item.value = value
|
||||
try:
|
||||
item.valueChanged.emit(value)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class _QConfigLoader:
|
||||
def load(self, filepath, config_obj):
|
||||
if not os.path.exists(filepath):
|
||||
return
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
cls = type(config_obj)
|
||||
for attr_name in dir(cls):
|
||||
if attr_name.startswith('_'):
|
||||
continue
|
||||
item = getattr(cls, attr_name, None)
|
||||
if not isinstance(item, ConfigItem):
|
||||
continue
|
||||
section_data = data.get(item.section, {})
|
||||
if item.key not in section_data:
|
||||
continue
|
||||
raw = section_data[item.key]
|
||||
if isinstance(item, OptionsConfigItem) and item.serializer:
|
||||
try:
|
||||
raw = item.serializer.deserialize(raw)
|
||||
except Exception:
|
||||
pass
|
||||
if item.validator.validate(raw):
|
||||
item.value = raw
|
||||
|
||||
|
||||
qconfig = _QConfigLoader()
|
||||
|
||||
# ── additional symbols referenced by gui code (safe no-ops) ──────────────
|
||||
class FluentIcon:
|
||||
HOME = None; SETTING = None; FOLDER = None; PLAY = None
|
||||
|
||||
class Theme:
|
||||
DARK = 'dark'; LIGHT = 'light'
|
||||
|
||||
class NavigationItemPosition:
|
||||
BOTTOM = 'bottom'
|
||||
|
||||
def setTheme(*a, **k): pass
|
||||
def setThemeColor(*a, **k): pass
|
||||
def isDarkTheme(): return True
|
||||
|
||||
# Fluent widget stubs — only imported by GUI code, never used by backend
|
||||
class _Widget:
|
||||
def __init__(self, *a, **k): pass
|
||||
|
||||
CardWidget = _Widget
|
||||
PushButton = _Widget
|
||||
Slider = _Widget
|
||||
ProgressBar = _Widget
|
||||
PlainTextEdit = _Widget
|
||||
TextEdit = _Widget
|
||||
SubtitleLabel = _Widget
|
||||
HollowHandleStyle = _Widget
|
||||
InfoBar = _Widget
|
||||
SettingCardGroup = _Widget
|
||||
ComboBoxSettingCard = _Widget
|
||||
SwitchSettingCard = _Widget
|
||||
FluentWindow = _Widget
|
||||
NavigationWidget = _Widget
|
||||
+1028
File diff suppressed because it is too large
Load Diff
+768
@@ -0,0 +1,768 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Web interface for video-subtitle-remover.
|
||||
|
||||
Usage (from project root):
|
||||
python web_app.py
|
||||
# then open http://localhost:8000
|
||||
|
||||
Or from any directory:
|
||||
VSR_PATH=/path/to/video-subtitle-remover python web_app.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import cv2
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import queue as q_module
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Request
|
||||
from fastapi.responses import FileResponse, HTMLResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
# ── project path ────────────────────────────────────────────────────────────
|
||||
# Default: the directory containing this file (run from project root)
|
||||
VSR_PATH = os.environ.get("VSR_PATH", str(Path(__file__).parent))
|
||||
if VSR_PATH not in sys.path:
|
||||
sys.path.insert(0, VSR_PATH)
|
||||
|
||||
try:
|
||||
from backend.config import config, tr
|
||||
from backend.tools.constant import InpaintMode
|
||||
from backend.main import SubtitleRemover
|
||||
import configparser as _cp
|
||||
# Force Chinese interface for web app
|
||||
_ch_ini = os.path.join(VSR_PATH, 'backend', 'interface', 'ch.ini')
|
||||
tr.read(_ch_ini, encoding='utf-8')
|
||||
VSR_AVAILABLE = True
|
||||
_import_err_msg = ""
|
||||
except ImportError as _import_err:
|
||||
VSR_AVAILABLE = False
|
||||
_import_err_msg = str(_import_err)
|
||||
|
||||
# ── storage dirs ────────────────────────────────────────────────────────────
|
||||
BASE_DIR = Path(__file__).parent
|
||||
CACHE_DIR = BASE_DIR / "web_cache" / "downloads"
|
||||
OUTPUT_DIR = BASE_DIR / "web_cache" / "outputs"
|
||||
DB_PATH = BASE_DIR / "web_cache" / "tasks.db"
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ── database ─────────────────────────────────────────────────────────────
|
||||
def init_db():
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
id TEXT PRIMARY KEY,
|
||||
url TEXT,
|
||||
is_cache_file BOOLEAN,
|
||||
cache_path TEXT,
|
||||
inpaint_mode TEXT,
|
||||
subtitle_areas TEXT,
|
||||
preview_mode BOOLEAN DEFAULT 0,
|
||||
detect_mode TEXT DEFAULT 'text',
|
||||
status TEXT,
|
||||
progress INTEGER,
|
||||
download_progress INTEGER,
|
||||
cached BOOLEAN,
|
||||
output_path TEXT,
|
||||
error TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
# Try to add new columns if they don't exist
|
||||
try:
|
||||
conn.execute("ALTER TABLE tasks ADD COLUMN preview_mode BOOLEAN DEFAULT 0")
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
conn.execute("ALTER TABLE tasks ADD COLUMN detect_mode TEXT DEFAULT 'text'")
|
||||
except:
|
||||
pass
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
init_db()
|
||||
|
||||
INPAINT_MODE_MAP = {
|
||||
"sttn-auto": lambda: InpaintMode.STTN_AUTO,
|
||||
"sttn-det": lambda: InpaintMode.STTN_DET,
|
||||
"lama": lambda: InpaintMode.LAMA,
|
||||
"propainter": lambda: InpaintMode.PROPAINTER,
|
||||
"opencv": lambda: InpaintMode.OPENCV,
|
||||
}
|
||||
|
||||
# ── state ────────────────────────────────────────────────────────────────────
|
||||
tasks: dict[str, dict] = {}
|
||||
task_lock = threading.Lock()
|
||||
task_queue: q_module.Queue = q_module.Queue()
|
||||
ws_connections: list[WebSocket] = []
|
||||
_loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def url_to_cache_path(url: str) -> Path:
|
||||
h = hashlib.md5(url.encode()).hexdigest()[:12]
|
||||
raw_name = url.rstrip("/").split("/")[-1].split("?")[0][:80]
|
||||
name = raw_name if raw_name and "." in raw_name else f"{h}.mp4"
|
||||
return CACHE_DIR / f"{h}_{name}"
|
||||
|
||||
|
||||
def _task_snapshot(task: dict) -> dict:
|
||||
return {k: v for k, v in task.items() if k != "logs"}
|
||||
|
||||
|
||||
def _push(coro):
|
||||
if _loop and not _loop.is_closed():
|
||||
asyncio.run_coroutine_threadsafe(coro, _loop)
|
||||
|
||||
|
||||
async def _broadcast(msg: dict):
|
||||
dead = []
|
||||
text = json.dumps(msg)
|
||||
for ws in ws_connections:
|
||||
try:
|
||||
await ws.send_text(text)
|
||||
except Exception:
|
||||
dead.append(ws)
|
||||
for ws in dead:
|
||||
if ws in ws_connections:
|
||||
ws_connections.remove(ws)
|
||||
|
||||
|
||||
def update_task(task_id: str, **kwargs):
|
||||
with task_lock:
|
||||
if task_id not in tasks:
|
||||
return
|
||||
tasks[task_id].update(kwargs)
|
||||
snap = _task_snapshot(tasks[task_id])
|
||||
_save_task_to_db(tasks[task_id])
|
||||
_push(_broadcast({"type": "task_update", "task": snap}))
|
||||
|
||||
def _save_task_to_db(task: dict):
|
||||
"""Save task to database"""
|
||||
try:
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.execute("""
|
||||
INSERT OR REPLACE INTO tasks
|
||||
(id, url, is_cache_file, cache_path, inpaint_mode, subtitle_areas,
|
||||
preview_mode, detect_mode, status, progress, download_progress, cached, output_path, error)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
task["id"],
|
||||
task["url"],
|
||||
task.get("is_cache_file", False),
|
||||
task.get("cache_path", ""),
|
||||
task["inpaint_mode"],
|
||||
json.dumps(task["subtitle_areas"]),
|
||||
task.get("preview_mode", False),
|
||||
task.get("detect_mode", "text"),
|
||||
task["status"],
|
||||
task["progress"],
|
||||
task["download_progress"],
|
||||
task["cached"],
|
||||
task["output_path"],
|
||||
task["error"],
|
||||
))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
print(f"Failed to save task to DB: {e}")
|
||||
|
||||
def _load_tasks_from_db():
|
||||
"""Load tasks from database on startup"""
|
||||
try:
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.row_factory = sqlite3.Row # Use column names
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM tasks ORDER BY created_at DESC LIMIT 100"
|
||||
)
|
||||
for row in cursor.fetchall():
|
||||
task = {
|
||||
"id": row["id"],
|
||||
"url": row["url"],
|
||||
"is_cache_file": row["is_cache_file"],
|
||||
"cache_path": row["cache_path"],
|
||||
"inpaint_mode": row["inpaint_mode"],
|
||||
"subtitle_areas": json.loads(row["subtitle_areas"]),
|
||||
"preview_mode": row["preview_mode"] if "preview_mode" in row.keys() else False,
|
||||
"detect_mode": row["detect_mode"] if "detect_mode" in row.keys() else "text",
|
||||
"status": row["status"],
|
||||
"progress": row["progress"],
|
||||
"download_progress": row["download_progress"],
|
||||
"cached": row["cached"],
|
||||
"output_path": row["output_path"],
|
||||
"error": row["error"],
|
||||
"phase": "pending",
|
||||
"detect_pct": 0,
|
||||
"process_pct": 0,
|
||||
"merge_pct": 0,
|
||||
"logs": [],
|
||||
}
|
||||
with task_lock:
|
||||
tasks[task["id"]] = task
|
||||
# Re-queue pending/downloading/processing tasks (processing might have crashed)
|
||||
if task["status"] in ("pending", "downloading", "processing"):
|
||||
task_queue.put(task["id"])
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
print(f"Failed to load tasks from DB: {e}")
|
||||
|
||||
|
||||
def push_log(task_id: str, message: str):
|
||||
with task_lock:
|
||||
if task_id in tasks:
|
||||
tasks[task_id]["logs"].append(message)
|
||||
_push(_broadcast({"type": "log", "task_id": task_id, "message": message}))
|
||||
|
||||
|
||||
# ── worker ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def _download(task_id: str, url: str, dest: Path) -> bool:
|
||||
update_task(task_id, status="downloading", download_progress=0)
|
||||
push_log(task_id, f"下载中: {url}")
|
||||
last_pct = -1
|
||||
try:
|
||||
with httpx.stream("GET", url, follow_redirects=True, timeout=7200) as r:
|
||||
r.raise_for_status()
|
||||
total = int(r.headers.get("content-length", 0))
|
||||
done = 0
|
||||
with open(dest, "wb") as f:
|
||||
for chunk in r.iter_bytes(chunk_size=1024 * 512):
|
||||
f.write(chunk)
|
||||
done += len(chunk)
|
||||
if total:
|
||||
pct = int(done / total * 100)
|
||||
if pct >= last_pct + 2:
|
||||
last_pct = pct
|
||||
mb_done = done // 1048576
|
||||
mb_total = total // 1048576
|
||||
update_task(task_id, download_progress=pct)
|
||||
push_log(task_id, f"下载 {pct}%({mb_done} MB / {mb_total} MB)")
|
||||
push_log(task_id, f"下载完成 → {dest.name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
push_log(task_id, f"下载失败: {e}")
|
||||
dest.unlink(missing_ok=True)
|
||||
return False
|
||||
|
||||
|
||||
def _process(task_id: str, video_path: str, inpaint_mode_str: str, subtitle_areas: list, preview_mode: bool = False, detect_mode: str = "text"):
|
||||
if not VSR_AVAILABLE:
|
||||
push_log(task_id, f"错误:无法加载后端模块。{_import_err_msg}")
|
||||
update_task(task_id, status="error", error="VSR backend not available")
|
||||
return
|
||||
|
||||
modeLabels = {'sttn-auto':'STTN 自动','sttn-det':'STTN 检测','lama':'LaMa','propainter':'ProPainter','opencv':'OpenCV'}
|
||||
push_log(task_id, f"启动 — 算法:{modeLabels.get(inpaint_mode_str, inpaint_mode_str)}")
|
||||
update_task(task_id, status="processing", progress=0, started_at=time.time())
|
||||
|
||||
class _WebRemover(SubtitleRemover):
|
||||
_static_boxes = [] # SAM2检测到的静态区域,格式 (xmin, xmax, ymin, ymax) 像素
|
||||
|
||||
def append_output(self_, *args):
|
||||
msg = " ".join(str(a) for a in args)
|
||||
print(msg, flush=True)
|
||||
push_log(task_id, msg)
|
||||
|
||||
def update_preview_with_comp(self_, frame_ori, frame_comp):
|
||||
pass # no-op: no GUI
|
||||
|
||||
def run(self_):
|
||||
if self_._static_boxes:
|
||||
from backend.tools.subtitle_detect import SubtitleDetect as _SD
|
||||
_orig = _SD.find_subtitle_frame_no
|
||||
_boxes = list(self_._static_boxes)
|
||||
def _patched(sd_self, sub_remover=None):
|
||||
result = _orig(sd_self, sub_remover=sub_remover)
|
||||
ocr_frames = len(result)
|
||||
if ocr_frames > 0:
|
||||
push_log(task_id, f"OCR 在 {ocr_frames} 帧中检测到文字")
|
||||
cap = cv2.VideoCapture(sd_self.video_path)
|
||||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
cap.release()
|
||||
for frame_no in range(1, total + 1):
|
||||
if frame_no not in result:
|
||||
result[frame_no] = []
|
||||
result[frame_no] = list(result[frame_no]) + _boxes
|
||||
push_log(task_id, f"SAM2 静态区域已注入全部 {total} 帧")
|
||||
return result
|
||||
_SD.find_subtitle_frame_no = _patched
|
||||
try:
|
||||
super().run()
|
||||
finally:
|
||||
_SD.find_subtitle_frame_no = _orig
|
||||
else:
|
||||
super().run()
|
||||
|
||||
def merge_audio_to_video(self_):
|
||||
push_log(task_id, "正在合并音频...")
|
||||
update_task(task_id, detect_pct=100, process_pct=100, merge_pct=50, phase="merging")
|
||||
super().merge_audio_to_video()
|
||||
update_task(task_id, detect_pct=100, process_pct=100, merge_pct=100, phase="merging")
|
||||
push_log(task_id, "音频合并完成")
|
||||
|
||||
# Preview mode: cut first 10 seconds using FFmpeg BEFORE processing
|
||||
preview_clip_path = None
|
||||
if preview_mode:
|
||||
try:
|
||||
from backend.tools.ffmpeg_cli import FFmpegCLI
|
||||
preview_clip_path = str(Path(video_path).parent / f"_preview_{task_id}.mp4")
|
||||
cut_cmd = [
|
||||
FFmpegCLI.instance().ffmpeg_path,
|
||||
"-y", "-i", video_path,
|
||||
"-t", "10",
|
||||
"-c", "copy",
|
||||
"-loglevel", "error",
|
||||
preview_clip_path
|
||||
]
|
||||
push_log(task_id, "预览模式:截取前10秒...")
|
||||
import subprocess
|
||||
subprocess.run(cut_cmd, check=True, timeout=60)
|
||||
push_log(task_id, "截取完成,开始处理预览片段")
|
||||
video_path = preview_clip_path # process the short clip
|
||||
except Exception as e:
|
||||
push_log(task_id, f"截取失败,改为处理完整视频: {e}")
|
||||
preview_clip_path = None
|
||||
|
||||
try:
|
||||
sr = _WebRemover(video_path, gui_mode=False)
|
||||
|
||||
if preview_mode:
|
||||
push_log(task_id, f"预览片段:{sr.frame_count} 帧 / {sr.fps:.1f}fps")
|
||||
|
||||
mode_factory = INPAINT_MODE_MAP.get(inpaint_mode_str, INPAINT_MODE_MAP["sttn-auto"])
|
||||
config.set(config.inpaintMode, mode_factory())
|
||||
|
||||
# Auto-detect regions if not provided
|
||||
if not subtitle_areas:
|
||||
combined_areas = []
|
||||
import concurrent.futures
|
||||
|
||||
def _run_ocr():
|
||||
from backend.tools.subtitle_detect import SubtitleDetect
|
||||
detector = SubtitleDetect(video_path, [])
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
ret, frame = cap.read()
|
||||
cap.release()
|
||||
if ret:
|
||||
return detector.detect_subtitle(frame)
|
||||
return []
|
||||
|
||||
def _run_sam2():
|
||||
print(f"[DEBUG] _run_sam2 called, detect_mode={detect_mode}", flush=True)
|
||||
if detect_mode != "smart":
|
||||
print(f"[DEBUG] detect_mode != smart, skipping SAM2", flush=True)
|
||||
return []
|
||||
from backend.tools.smart_detect import smart_detect_watermark
|
||||
return smart_detect_watermark(video_path, sample_frames=30)
|
||||
|
||||
push_log(task_id, "并发检测中(OCR + SAM2)..." if detect_mode == "smart" else "文字检测(OCR)中...")
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
|
||||
f_ocr = ex.submit(_run_ocr)
|
||||
f_sam2 = ex.submit(_run_sam2)
|
||||
ocr_result = f_ocr.result()
|
||||
sam2_result = f_sam2.result()
|
||||
|
||||
if ocr_result:
|
||||
push_log(task_id, f"OCR 检测到 {len(ocr_result)} 个文字区域")
|
||||
combined_areas += [
|
||||
(
|
||||
int(y1 / sr.frame_height * 100),
|
||||
int(y2 / sr.frame_height * 100),
|
||||
int(x1 / sr.frame_width * 100),
|
||||
int(x2 / sr.frame_width * 100),
|
||||
)
|
||||
for x1, x2, y1, y2 in ocr_result
|
||||
]
|
||||
else:
|
||||
push_log(task_id, "OCR 未检测到文字区域")
|
||||
|
||||
if sam2_result:
|
||||
push_log(task_id, f"SAM2 检测到 {len(sam2_result)} 个区域,将强制每帧都处理")
|
||||
for i, (x1, y1, x2, y2) in enumerate(sam2_result):
|
||||
push_log(task_id, f" 区域{i+1}: 左上({x1},{y1}) 右下({x2},{y2}) 大小{x2-x1}x{y2-y1}px")
|
||||
sr._static_boxes = [(x1, x2, y1, y2) for x1, y1, x2, y2 in sam2_result]
|
||||
elif detect_mode == "smart":
|
||||
push_log(task_id, "SAM2 未检测到静态水印区域")
|
||||
|
||||
# 单帧初探仅用于日志,不作为全扫描的空间约束
|
||||
# 若把 21 个 15px 小框设为 sub_areas,全扫描几乎找不到任何匹配帧 → 报错
|
||||
# SubtitleRemover.find_subtitle_frame_no 会自行全帧扫描
|
||||
if not combined_areas and not sr._static_boxes:
|
||||
push_log(task_id, "初探未检测到区域,全扫描将覆盖全屏")
|
||||
elif combined_areas:
|
||||
push_log(task_id, f"初探检测到 {len(combined_areas)} 个文字区域(全扫描不受此约束)")
|
||||
if sr._static_boxes:
|
||||
push_log(task_id, f"SAM2 静态区域已注入全部帧,无需区域约束")
|
||||
# subtitle_areas 保持为空,sub_areas 不设置 → 全帧无约束扫描
|
||||
|
||||
if subtitle_areas:
|
||||
pixel_areas = [
|
||||
(
|
||||
int(a[0] / 100 * sr.frame_height),
|
||||
int(a[1] / 100 * sr.frame_height),
|
||||
int(a[2] / 100 * sr.frame_width),
|
||||
int(a[3] / 100 * sr.frame_width),
|
||||
)
|
||||
for a in subtitle_areas
|
||||
]
|
||||
if detect_mode == "smart":
|
||||
# 水印模式:用户手动指定区域 + 智能检测 → 强制每帧处理该区域
|
||||
# 适合位置变化的动态水印(OCR 无法识别装饰性字体)
|
||||
force_boxes = [
|
||||
(int(a[2]/100*sr.frame_width), int(a[3]/100*sr.frame_width),
|
||||
int(a[0]/100*sr.frame_height), int(a[1]/100*sr.frame_height))
|
||||
for a in subtitle_areas
|
||||
]
|
||||
if sr._static_boxes:
|
||||
sr._static_boxes = sr._static_boxes + force_boxes
|
||||
else:
|
||||
sr._static_boxes = force_boxes
|
||||
push_log(task_id, f"手动指定 {len(force_boxes)} 个水印区域 → 强制每帧去除")
|
||||
print(f"[DEBUG] force_boxes={force_boxes}", flush=True)
|
||||
else:
|
||||
# 字幕模式:限制 OCR 搜索范围
|
||||
sr.sub_areas = pixel_areas
|
||||
|
||||
out_stem = Path(video_path).stem
|
||||
sr.video_out_path = str(OUTPUT_DIR / f"{out_stem}_{task_id[:8]}_no_sub.mp4")
|
||||
|
||||
# 检测阶段:progress_total = (100*frame/total)//2 → 0-50
|
||||
# 处理阶段:update_progress → 0-100(从0重新开始)
|
||||
# 判断切换:raw>50 OR 前值≥40后值<15(跌落 = 阶段重置)
|
||||
_state = {"phase": "detecting", "processing": False, "prev_raw": 0}
|
||||
|
||||
def _on_progress(pct, finished):
|
||||
try:
|
||||
raw = int(pct)
|
||||
prev = _state["prev_raw"]
|
||||
_state["prev_raw"] = raw
|
||||
# 跌落检测:从检测末尾(≥40)掉回处理起始(<15)
|
||||
if not _state["processing"] and prev >= 40 and raw < 15:
|
||||
_state["processing"] = True
|
||||
if _state["processing"] or raw > 50:
|
||||
_state["processing"] = True
|
||||
_state["phase"] = "processing"
|
||||
process_pct = min(raw, 98)
|
||||
update_task(task_id, detect_pct=100, process_pct=process_pct, merge_pct=0, phase="processing")
|
||||
else:
|
||||
detect_pct = min(raw * 2, 99) # 0-50 → 0-99
|
||||
update_task(task_id, detect_pct=detect_pct, process_pct=0, merge_pct=0, phase="detecting")
|
||||
print(f"[PROGRESS] {task_id}: phase={_state['phase']} raw={raw}%", flush=True)
|
||||
except Exception as e:
|
||||
print(f"[ERROR] Progress callback failed: {e}", flush=True)
|
||||
|
||||
push_log(task_id, f"检测模式:{'智能(SAM2)' if detect_mode == 'smart' else '文字(OCR)'}")
|
||||
|
||||
print(f"[DEBUG] Adding progress listener for {task_id}", flush=True)
|
||||
sr.add_progress_listener(_on_progress)
|
||||
print(f"[DEBUG] Progress listeners count: {len(sr.progress_listeners)}", flush=True)
|
||||
print(f"[DEBUG] Starting sr.run()", flush=True)
|
||||
sr.run()
|
||||
print(f"[DEBUG] sr.run() completed", flush=True)
|
||||
|
||||
# Preview mode: skip audio merge, directly use temp file
|
||||
if preview_mode and hasattr(sr, 'video_temp_file') and sr.video_temp_file:
|
||||
push_log(task_id, "预览模式:跳过音频合并,直接输出")
|
||||
import shutil
|
||||
try:
|
||||
shutil.copy2(sr.video_temp_file.name, sr.video_out_path)
|
||||
push_log(task_id, f"预览完成 → {sr.video_out_path}")
|
||||
except Exception as e:
|
||||
push_log(task_id, f"复制文件失败: {e}")
|
||||
|
||||
update_task(task_id, status="done", progress=100, output_path=sr.video_out_path, completed_at=time.time())
|
||||
push_log(task_id, f"完成 → {sr.video_out_path}")
|
||||
|
||||
except Exception:
|
||||
err = traceback.format_exc()
|
||||
push_log(task_id, f"错误:\n{err}")
|
||||
update_task(task_id, status="error", error=err.splitlines()[-1])
|
||||
finally:
|
||||
# Clean up preview clip
|
||||
if preview_clip_path and Path(preview_clip_path).exists():
|
||||
try:
|
||||
Path(preview_clip_path).unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _run_task(task_id: str):
|
||||
with task_lock:
|
||||
task = tasks.get(task_id)
|
||||
if not task:
|
||||
return
|
||||
|
||||
# Handle cache file directly
|
||||
if task.get("is_cache_file"):
|
||||
cache_path = Path(task["cache_path"])
|
||||
if cache_path.exists():
|
||||
push_log(task_id, f"使用本地缓存: {cache_path.name}")
|
||||
update_task(task_id, cached=True)
|
||||
else:
|
||||
update_task(task_id, status="error", error="缓存文件不存在")
|
||||
return
|
||||
else:
|
||||
# Handle URL download
|
||||
url = task["url"]
|
||||
cache_path = url_to_cache_path(url)
|
||||
|
||||
if cache_path.exists():
|
||||
push_log(task_id, f"缓存命中: {cache_path.name}")
|
||||
update_task(task_id, cached=True)
|
||||
else:
|
||||
ok = _download(task_id, url, cache_path)
|
||||
if not ok:
|
||||
update_task(task_id, status="error", error="下载失败")
|
||||
return
|
||||
|
||||
_process(task_id, str(cache_path), task["inpaint_mode"], task["subtitle_areas"],
|
||||
task.get("preview_mode", False), task.get("detect_mode", "text"))
|
||||
|
||||
|
||||
def _worker():
|
||||
while True:
|
||||
task_id = task_queue.get()
|
||||
if task_id is None:
|
||||
break
|
||||
try:
|
||||
_run_task(task_id)
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
task_queue.task_done()
|
||||
|
||||
|
||||
# ── FastAPI ───────────────────────────────────────────────────────────────────
|
||||
|
||||
app = FastAPI(title="VSR Web")
|
||||
|
||||
_worker_thread = threading.Thread(target=_worker, daemon=True, name="vsr-worker")
|
||||
_worker_thread.start()
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def _startup():
|
||||
global _loop
|
||||
_loop = asyncio.get_event_loop()
|
||||
_load_tasks_from_db()
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def index():
|
||||
html = (BASE_DIR / "web" / "index.html").read_text(encoding="utf-8")
|
||||
return HTMLResponse(html)
|
||||
|
||||
|
||||
class TaskCreate(BaseModel):
|
||||
url: str
|
||||
inpaint_mode: str = "sttn-auto"
|
||||
subtitle_areas: list = []
|
||||
preview_mode: bool = False
|
||||
detect_mode: str = "text"
|
||||
|
||||
|
||||
@app.post("/api/tasks")
|
||||
async def create_task(body: TaskCreate):
|
||||
task_id = uuid.uuid4().hex[:8]
|
||||
|
||||
# Check if it's a cache file name (from the recent list)
|
||||
is_cache_file = body.url.startswith("cache_")
|
||||
if is_cache_file:
|
||||
# Remove "cache_" prefix
|
||||
filename = body.url[6:] # strip "cache_"
|
||||
cache_path = CACHE_DIR / filename
|
||||
cached = cache_path.exists()
|
||||
else:
|
||||
cache_path = url_to_cache_path(body.url)
|
||||
cached = cache_path.exists()
|
||||
|
||||
with task_lock:
|
||||
now = time.time()
|
||||
tasks[task_id] = {
|
||||
"id": task_id,
|
||||
"url": body.url,
|
||||
"is_cache_file": is_cache_file,
|
||||
"cache_path": str(cache_path),
|
||||
"inpaint_mode": body.inpaint_mode,
|
||||
"subtitle_areas": body.subtitle_areas,
|
||||
"preview_mode": body.preview_mode,
|
||||
"detect_mode": body.detect_mode,
|
||||
"status": "pending",
|
||||
"progress": 0,
|
||||
"phase": "pending",
|
||||
"detect_pct": 0,
|
||||
"process_pct": 0,
|
||||
"merge_pct": 0,
|
||||
"download_progress": 0,
|
||||
"cached": cached,
|
||||
"output_path": None,
|
||||
"error": None,
|
||||
"logs": [],
|
||||
"created_at": now,
|
||||
"started_at": None,
|
||||
"completed_at": None,
|
||||
}
|
||||
task_queue.put(task_id)
|
||||
return {"task_id": task_id}
|
||||
|
||||
|
||||
@app.get("/api/tasks")
|
||||
async def list_tasks():
|
||||
with task_lock:
|
||||
return [_task_snapshot(t) for t in tasks.values()]
|
||||
|
||||
|
||||
@app.get("/api/tasks/{task_id}")
|
||||
async def get_task(task_id: str):
|
||||
with task_lock:
|
||||
task = tasks.get(task_id)
|
||||
if not task:
|
||||
raise HTTPException(404, "Not found")
|
||||
return task
|
||||
|
||||
|
||||
@app.delete("/api/tasks/{task_id}")
|
||||
async def delete_task(task_id: str):
|
||||
with task_lock:
|
||||
tasks.pop(task_id, None)
|
||||
# Also delete from database
|
||||
try:
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.execute("DELETE FROM tasks WHERE id = ?", (task_id,))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
print(f"Failed to delete task from DB: {e}")
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@app.get("/api/tasks/{task_id}/download")
|
||||
async def download_result(task_id: str):
|
||||
with task_lock:
|
||||
task = tasks.get(task_id)
|
||||
if not task or task["status"] != "done" or not task.get("output_path"):
|
||||
raise HTTPException(404, "Result not ready")
|
||||
out = Path(task["output_path"])
|
||||
if not out.exists():
|
||||
raise HTTPException(404, "Output file missing")
|
||||
return FileResponse(str(out), filename=out.name, media_type="video/mp4")
|
||||
|
||||
|
||||
@app.get("/api/tasks/{task_id}/stream")
|
||||
async def stream_result(task_id: str, request: Request):
|
||||
with task_lock:
|
||||
task = tasks.get(task_id)
|
||||
if not task or task["status"] != "done" or not task.get("output_path"):
|
||||
raise HTTPException(404, "Result not ready")
|
||||
out = Path(task["output_path"])
|
||||
if not out.exists():
|
||||
raise HTTPException(404, "Output file missing")
|
||||
|
||||
file_size = out.stat().st_size
|
||||
range_header = request.headers.get("range")
|
||||
|
||||
if range_header:
|
||||
# Parse Range header: "bytes=start-end"
|
||||
try:
|
||||
range_value = range_header.replace("bytes=", "")
|
||||
start, end = range_value.split("-")
|
||||
start = int(start) if start else 0
|
||||
end = int(end) if end else file_size - 1
|
||||
|
||||
if start >= file_size or end >= file_size or start > end:
|
||||
raise HTTPException(416, "Range Not Satisfiable")
|
||||
|
||||
headers = {
|
||||
"Content-Range": f"bytes {start}-{end}/{file_size}",
|
||||
"Accept-Ranges": "bytes",
|
||||
"Content-Length": str(end - start + 1),
|
||||
}
|
||||
return FileResponse(
|
||||
str(out),
|
||||
media_type="video/mp4",
|
||||
headers=headers,
|
||||
status_code=206,
|
||||
)
|
||||
except (ValueError, IndexError):
|
||||
pass
|
||||
|
||||
return FileResponse(
|
||||
str(out),
|
||||
filename=out.name,
|
||||
media_type="video/mp4",
|
||||
headers={"Accept-Ranges": "bytes"},
|
||||
)
|
||||
|
||||
|
||||
@app.get("/api/cache-check")
|
||||
async def cache_check(url: str):
|
||||
return {"cached": url_to_cache_path(url).exists()}
|
||||
|
||||
|
||||
@app.get("/api/cache-list")
|
||||
async def cache_list():
|
||||
files = []
|
||||
if CACHE_DIR.exists():
|
||||
for f in sorted(CACHE_DIR.iterdir(), key=lambda x: x.stat().st_mtime, reverse=True):
|
||||
if f.is_file():
|
||||
files.append({
|
||||
"name": f.name,
|
||||
"size": f.stat().st_size,
|
||||
"mtime": f.stat().st_mtime,
|
||||
})
|
||||
return {"files": files}
|
||||
|
||||
|
||||
@app.get("/api/status")
|
||||
async def server_status():
|
||||
return {
|
||||
"vsr_available": VSR_AVAILABLE,
|
||||
"vsr_path": VSR_PATH,
|
||||
"queue_size": task_queue.qsize(),
|
||||
}
|
||||
|
||||
|
||||
@app.websocket("/ws")
|
||||
async def ws_endpoint(ws: WebSocket):
|
||||
await ws.accept()
|
||||
ws_connections.append(ws)
|
||||
try:
|
||||
with task_lock:
|
||||
init_tasks = [_task_snapshot(t) for t in tasks.values()]
|
||||
await ws.send_text(json.dumps({"type": "init", "tasks": init_tasks}))
|
||||
while True:
|
||||
await ws.receive_text()
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
if ws in ws_connections:
|
||||
ws_connections.remove(ws)
|
||||
|
||||
|
||||
# ── entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
host = os.environ.get("HOST", "0.0.0.0")
|
||||
port = int(os.environ.get("PORT", "8001"))
|
||||
print(f"VSR_PATH = {VSR_PATH}")
|
||||
print(f"VSR backend available: {VSR_AVAILABLE}")
|
||||
if not VSR_AVAILABLE:
|
||||
print(f" Import error: {_import_err_msg}")
|
||||
print(f" Set VSR_PATH to the project root and ensure dependencies are installed.")
|
||||
print(f"Starting on http://{host}:{port}")
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
Reference in New Issue
Block a user