From 693d394f55a00351adf96a740eee9bd360c9787b Mon Sep 17 00:00:00 2001 From: liuchen Date: Sat, 27 Jun 2026 13:27:26 +0800 Subject: [PATCH] Add GPU support and multi-threaded subtitle detection - 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 --- .gitignore | 1 + backend/tools/smart_detect.py | 253 ++++++++ backend/tools/subtitle_detect.py | 118 +++- backend/tools/video_io.py | 106 +-- qfluentwidgets/__init__.py | 155 +++++ web/index.html | 1028 ++++++++++++++++++++++++++++++ web_app.py | 768 ++++++++++++++++++++++ 7 files changed, 2360 insertions(+), 69 deletions(-) create mode 100644 backend/tools/smart_detect.py create mode 100644 qfluentwidgets/__init__.py create mode 100644 web/index.html create mode 100644 web_app.py diff --git a/.gitignore b/.gitignore index fd3441c..b05a3d3 100644 --- a/.gitignore +++ b/.gitignore @@ -375,3 +375,4 @@ test*_no_sub*.mp4 /backend/tools/train/release_model/ inference.onnx /config/config.json +web_cache/ diff --git a/backend/tools/smart_detect.py b/backend/tools/smart_detect.py new file mode 100644 index 0000000..5bdffd3 --- /dev/null +++ b/backend/tools/smart_detect.py @@ -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 diff --git a/backend/tools/subtitle_detect.py b/backend/tools/subtitle_detect.py index 252da22..b478087 100644 --- a/backend/tools/subtitle_detect.py +++ b/backend/tools/subtitle_detect.py @@ -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()) diff --git a/backend/tools/video_io.py b/backend/tools/video_io.py index 7930181..ca6f072 100644 --- a/backend/tools/video_io.py +++ b/backend/tools/video_io.py @@ -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}") diff --git a/qfluentwidgets/__init__.py b/qfluentwidgets/__init__.py new file mode 100644 index 0000000..58033b3 --- /dev/null +++ b/qfluentwidgets/__init__.py @@ -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 diff --git a/web/index.html b/web/index.html new file mode 100644 index 0000000..da7a842 --- /dev/null +++ b/web/index.html @@ -0,0 +1,1028 @@ + + + + + +视频字幕去除器 + + + + + +
+ +

视频字幕去除器

+ +
v1.4.0
+
+ +
+ + +
+ + + +
+

新建任务

+ + + +
+ + + ✓ 已缓存 + +
+ + + + +
+ + +
+ +
+ + + +
+ + +
+ +
+ + + +
+ + + +
+ +
+ + +
+ +
+ + +
+

任务队列

+
+
+ 📋 +

没有任务。输入链接后点击开始处理。

+
+
+
+ + +
+
+ 处理日志 +
+ + + +
+
+
+
+ +
+ + + + diff --git a/web_app.py b/web_app.py new file mode 100644 index 0000000..1dc686e --- /dev/null +++ b/web_app.py @@ -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)