You can deploy this model and start earning money today!

Chan7891/
Sources_for_TruVIL

Views

No views yet

Trusted Video Inpainting Localization via Deep Attentive Noise Learning

An official implementation code for paper "Trusted Video Inpainting Localization via Deep Attentive Noise Learning". This repo provides code and trained weights.

Framework

Dependency

torch 1.7.0
python 3.7

Datasets

DAVIS2016
DAVIS2017
MOSE
VOS2k5-800 (in this paper we use 800 videos from VOS2k5)

The MOSE100 dataset in this paper can be found in this

Video inpainting algorithms

云 GPU 完整操作手册

Python 3.7 + PyTorch 1.7.0 · 从零到推理 · 100% 可复现

1. 环境配置

1.1 创建 Conda 环境

bash

# 有 GPU（推荐）：conda 自动匹配 CUDA 版本
conda create -n truvil python=3.7 pytorch==1.7.0 torchvision==0.8.0 cudatoolkit=10.2 -c pytorch -y

# 无 GPU / CPU only：
conda create -n truvil python=3.7 pytorch==1.7.0 torchvision==0.8.0 cpuonly -c pytorch -y

# 激活环境
conda activate truvil

# 验证
python -c "import torch; print('PyTorch:', torch.__version__, '| CUDA:', torch.cuda.is_available())"
# 期望输出: PyTorch: 1.7.0 | CUDA: True

1.2 安装依赖

bash

# 国内云 GPU 建议用清华镜像加速
pip install timm==0.9.16 tqdm opencv-python Pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

# 国外直接用默认源
pip install timm==0.9.16 tqdm opencv-python Pillow

不需要安装 mmengine——项目已改用原生 PyTorch（nn.Module / nn.ModuleList），零额外依赖。

1.3 上传项目代码

bash

# === 在本地执行 ===
cd TruVIL
tar -czf truvil.tar.gz \
    --exclude='checkpoints' \
    --exclude='__pycache__' \
    --exclude='.git' \
    --exclude='.results' \
    --exclude='inputs' \
    --exclude='TruVILStudio' \
    *.py requirements.txt

scp truvil.tar.gz user@your-gpu-ip:/data/

# === 上传权重文件（397MB，单独传，避免重复打包）===
scp checkpoints/TruVIL_train_VI_OP.pth user@your-gpu-ip:/data/checkpoints/

# === 在云 GPU 上执行 ===
ssh user@your-gpu-ip
cd /data && tar -xzf truvil.tar.gz
ls *.py                 # 应看到约 14 个 .py 文件
ls checkpoints/         # 应看到 TruVIL_train_VI_OP.pth

2. 数据准备

2.1 目录结构

/data/dataset/
├── train/                  # 训练集
│   ├── frame/              # 帧文件夹（每个视频一个子文件夹）
│   │   ├── video_001/
│   │   │   ├── frame_01.png
│   │   │   ├── frame_02.png
│   │   │   ├── frame_03.png
│   │   │   ├── frame_04.png
│   │   │   └── frame_05.png
│   │   ├── video_002/
│   │   └── ...
│   └── groundtruth/        # 标签掩膜（每个视频一张）
│       ├── video_001.png
│       ├── video_002.png
│       └── ...
│
├── val_VI/                 # 验证集 1：VI 算法修复
│   ├── frame/
│   └── groundtruth/
├── val_OP/                 # 验证集 2：OP 算法修复
│   ├── frame/
│   └── groundtruth/
└── val_CP/                 # 验证集 3：CP 算法修复
    ├── frame/
    └── groundtruth/

2.2 数据规格

参数	要求
每视频帧数	≥ 5 帧（取前 5 帧）
帧分辨率	任意（代码自动 resize 到 240×432）
帧格式	PNG / JPG / BMP
帧命名	自然排序（frame_1 < frame_2 < ... < frame_10）
标签格式	单通道 PNG，修复区域=255，原图区域=0

3. 权重下载与加载验证

3.1 下载预训练权重

从 Google Drive 下载并放到 checkpoints/ 目录： TruVIL_train_VI_OP.pth（397 MB）

3.2 验证权重加载

bash

conda activate truvil

python -c "
from model import TruVIL
from load_checkpoint_fix import load_model_safe
import torch

model = TruVIL()
model, stats = load_model_safe(model, './checkpoints/TruVIL_train_VI_OP.pth')

# 预期输出：
#   Remapped keys   : 96
#   Matched keys    : 1534
#   Missing keys    : 3     ← encoder.srm_layer*.weight（HP3D buffer，已在 __init__ 正确初始化）
#   Unexpected keys : 0
#   Weight params   : 100.0%

# 前向传播测试
model.eval()
x = torch.randn(1, 3, 5, 240, 432)
with torch.no_grad():
    y = model(x)
print(f'Input : {list(x.shape)}')
print(f'Output: {list(y.shape)}')     # [1, 1, 240, 432]
print('[OK] Model loaded and working!')
"

3.3 键名映射说明

预训练权重使用旧版命名，load_checkpoint_fix.py 自动处理映射：

旧名称	新名称	键数
`encoder.tacf.*`	`encoder.CAF.*`	69
`rgd.*`	`AND.*`	27

其余 1438 个键名完全一致，直接匹配。3 个 encoder.srm_layer*.weight（HP3D 高通滤波核）不需要从权重加载——HP3D.__init__ 已用预定义的高通卷积核正确初始化。

4. 推理

4.1 命令行

bash

conda activate truvil

python inference.py \
    --checkpoint ./checkpoints/TruVIL_train_VI_OP.pth \
    --input_dir ./demo \
    --output_dir ./output

输入: ./demo/video_001/ 下有 5 帧 PNG 输出: ./output/video_001.png（二值修复掩膜）

4.2 Python API 调用

python

import torch
from inference import load_model, preprocess_frames, run_inference
from PIL import Image

# 加载模型（自动检测 GPU）
model, device = load_model('./checkpoints/TruVIL_train_VI_OP.pth')

# 预处理 5 帧
frames = [
    '/data/demo/vid_001/frame_01.png',
    '/data/demo/vid_001/frame_02.png',
    '/data/demo/vid_001/frame_03.png',
    '/data/demo/vid_001/frame_04.png',
    '/data/demo/vid_001/frame_05.png',
]
tensor = preprocess_frames(frames, device=device)  # → [1, 3, 5, 240, 432]

# 推理得到二值掩膜
mask = run_inference(model, tensor, threshold=0.5)  # → [240, 432] numpy array

# 保存
Image.fromarray((mask * 255).astype('uint8')).save('output.png')

4.3 推理函数速查

函数	签名	说明
`load_model`	`(checkpoint_path, device=None)` → `(model, device)`	加载模型和权重
`preprocess_frames`	`(frame_paths, resize_size=(240,432), device=None)` → `tensor [1,3,5,H,W]`	读取+resize+堆叠帧
`run_inference`	`(model, input_tensor, threshold=0.5)` → `numpy [H,W]`	推理→sigmoid→二值化
`run_inference_on_tensor`	`(model, input_tensor)` → `tensor [1,1,H,W]`	推理返回原始 logits

5. 测试评估

5.1 命令行

bash

conda activate truvil

python test.py \
    --checkpoint ./checkpoints/TruVIL_train_VI_OP.pth \
    --val_roots /data/dataset/val_VI /data/dataset/val_OP /data/dataset/val_CP \
    --batch_size 16

5.2 参数说明

参数	必填	默认值	说明
`--checkpoint`	✓	-	权重文件路径
`--val_roots`	✓	-	验证集根目录（可多个，空格分隔）
`--batch_size`		16	批次大小
`--num_workers`		2	DataLoader 进程数
`--device`		自动	`cuda` / `cuda:0` / `cpu`

5.3 预期输出

============================================================
Checkpoint keys : 1534
Remapped keys   : 96
Matched keys    : 1534
Missing keys    : 3
Unexpected keys : 0
[OK] Loaded: 1534/1537 keys

Evaluating: /data/dataset/val_VI
Val Set 1: F1: 0.xxxx, IoU: 0.xxxx

Evaluating: /data/dataset/val_OP
Val Set 2: F1: 0.xxxx, IoU: 0.xxxx

Evaluating: /data/dataset/val_CP
Val Set 3: F1: 0.xxxx, IoU: 0.xxxx

6. 训练

6.1 从头训练

bash

conda activate truvil

python train.py \
    --train_root /data/dataset/train \
    --val_roots /data/dataset/val_VI /data/dataset/val_OP /data/dataset/val_CP \
    --epochs 200 \
    --batch_size 8 \
    --lr 0.0005 \
    --num_workers 4 \
    --save_dir ./weights

6.2 断点续训

bash

# 自动从 ./weights/latest.pth 恢复
python train.py \
    --train_root /data/dataset/train \
    --val_roots /data/dataset/val_VI /data/dataset/val_OP /data/dataset/val_CP \
    --save_dir ./weights

# 或指定具体 checkpoint 恢复
python train.py \
    --train_root /data/dataset/train \
    --val_roots /data/dataset/val_VI /data/dataset/val_OP /data/dataset/val_CP \
    --resume ./weights/epoch_050.pth

6.3 训练参数

参数	必填	默认值	说明
`--train_root`	✓	-	训练集根目录
`--val_roots`	✓	-	验证集根目录（可多个）
`--epochs`		200	训练轮数
`--batch_size`		8	批次大小（显存不足时改 4）
`--lr`		0.0005	AdamW 学习率
`--num_workers`		2	DataLoader 子进程数
`--save_dir`		./weights	权重保存目录
`--resume`		None	指定断点续训的 checkpoint
`--device`		自动	`cuda` / `cuda:0` / `cpu`

6.4 训练输出

weights/
├── epoch_000.pth       # 每 epoch 保存的模型权重
├── epoch_001.pth
├── ...
├── epoch_199.pth
└── latest.pth          # 最新检查点（含 optimizer/scheduler 状态，用于续训）

6.5 多 GPU 训练

python

# 在 train.py main() 中 model = TruVIL().to(device) 之后添加：
if torch.cuda.device_count() > 1:
    print(f"Using {torch.cuda.device_count()} GPUs")
    model = torch.nn.DataParallel(model)

7. 常见问题

Q1: `No module named 'mmengine'`

不需要 mmengine。如果遇到此错误，说明使用的是旧版 uniformer.py。确认 uniformer.py 顶部是：

python

BaseModule = nn.Module
ModuleList = nn.ModuleList

而非：

from mmengine.model import BaseModule, ModuleList # 旧版，不要用

Q2: 权重加载有 3 个 Missing keys

正常现象。Missing 的是 encoder.srm_layer1.weight、encoder.srm_layer2.weight、encoder.srm_layer3.weight——HP3D 高通滤波核。这些值在 HP3D.__init__ 中已用预定义的高通卷积核（拉普拉斯类算子）正确初始化，不需要从 checkpoint 加载。不影响模型精度。

Q3: CUDA Out of Memory

减小 batch size 并增大 num_workers：

python train.py --batch_size 4 --num_workers 2 ...

Q4: 推理输出全黑或全白

检查：

是否传入了 恰好 5 帧连续帧
帧路径是否按时间顺序排列（不是随机）
preprocess_frames 接收的列表顺序 = 帧的时间顺序

Q5: `FileNotFoundError: Frame directory not found`

数据集目录结构错误。确认存在 root_dir/frame/ 和 root_dir/groundtruth/ 两个子目录。

Q6: 如何验证环境配置是否成功

bash

conda activate truvil
python -c "
import torch; import numpy; import timm; import cv2; import tqdm
from PIL import Image
from model import TruVIL
from HP3D import HP3D
from uniformer import Encoder
from segformer_head import Decoder_3d
from AttentionModule import CAF, AND
from base_dataset import FramesDataset
from loss import Focal_IoU_LOSS
from metric import F1, IoU
from load_checkpoint_fix import load_model_safe
print('[OK] All modules imported successfully')
print(f'PyTorch {torch.__version__} | CUDA {torch.cuda.is_available()}')
print(f'NumPy {numpy.__version__}')

# 模型前向传播
model = TruVIL()
model, _ = load_model_safe(model, './checkpoints/TruVIL_train_VI_OP.pth')
model.eval()
x = torch.randn(1, 3, 5, 240, 432)
with torch.no_grad():
    y = model(x)
print(f'Input {list(x.shape)} -> Output {list(y.shape)}')
print('[OK] Full pipeline verified!')
"

附录 A: 完整文件清单

文件	用途
`model.py`	TruVIL 顶层模型
`uniformer.py`	双流 UniFormer 编码器
`segformer_head.py`	3D SegFormer 解码器
`AttentionModule.py`	注意力模块（CAF / AND / PAM / CAM / TAM）
`HP3D.py`	高通 3D 滤波器（buffer 注册，性能优化）
`base_model.py`	模型抽象基类
`base_dataset.py`	数据集加载器（自然排序，错误检查）
`loss.py`	Focal + IoU 组合损失
`metric.py`	F1 / IoU 评估指标
`wrappers.py`	上采样工具
`train.py`	训练脚本（argparse CLI）
`test.py`	测试脚本（argparse CLI）
`inference.py`	推理脚本 + 可复用 Python API
`load_checkpoint_fix.py`	权重键名自动映射
`requirements.txt`	Python 依赖清单

附录 B: 模型架构速览

输入 [B, 3, 5, H, W]
    │
    ├─ HP3D 高通滤波 ──→ Res 残差流
    │                    │
    ├─ RGB 流            │
    │                    │
    Stage 1: CBlock×5    CBlock×5  ──→ + HP3D ──→ outs[0]
    Stage 2: CBlock×8    CBlock×8  ──→ + HP3D ──→ outs[1]
    Stage 3: SABlock×20  SABlock×20 ──→ CAF交叉融合 ──→ outs[2]
    Stage 4: SABlock×7   SABlock×7  ──→ outs[3], outs[4]
    │                                        │
    └──── outs[0~4] ──→ Decoder_3d ──→ AND门控解码 ──→ 输出 [B, 1, H, W]

总参数: 103,753,616 | 输入: 5 帧 240×432 | 输出: 单通道修复掩膜

Citation

If you use this code for your research, please cite our paper

@article{lou2025trusted,
  title={Trusted Video Inpainting Localization via Deep Attentive Noise Learning},
  author={Lou, Zijie and Cao, Gang and Lin, Man and Yu, Lifang and Weng, Shaowei},
  journal={IEEE Transactions on Dependable and Secure Computing},
  year={2025},
  publisher={IEEE}
}

License

Licensed under a Creative Commons Attribution-NonCommercial 4.0 International for Non-commercial use only. Any commercial use should get formal permission first.

Chan7891/Sources_for_TruVIL

Trusted Video Inpainting Localization via Deep Attentive Noise Learning

Framework

Dependency

Datasets

Video inpainting algorithms

云 GPU 完整操作手册

目录

1. 环境配置

1.1 创建 Conda 环境

1.2 安装依赖

1.3 上传项目代码

2. 数据准备

2.1 目录结构

2.2 数据规格

3. 权重下载与加载验证

3.1 下载预训练权重

3.2 验证权重加载

3.3 键名映射说明

4. 推理

4.1 命令行

4.2 Python API 调用

4.3 推理函数速查

5. 测试评估

5.1 命令行

5.2 参数说明

5.3 预期输出

6. 训练

6.1 从头训练

6.2 断点续训

6.3 训练参数

6.4 训练输出

6.5 多 GPU 训练

7. 常见问题

Q1: No module named 'mmengine'

Q2: 权重加载有 3 个 Missing keys

Q3: CUDA Out of Memory

Q4: 推理输出全黑或全白

Q5: FileNotFoundError: Frame directory not found

Q6: 如何验证环境配置是否成功

附录 A: 完整文件清单

附录 B: 模型架构速览

Citation

License

Chan7891/
Sources_for_TruVIL

Q1: `No module named 'mmengine'`

Q5: `FileNotFoundError: Frame directory not found`