import torch import cv2 import numpy as np from PIL import Image from torchvision import transforms from model import MattingNetwork # RVM 的模型定义文件 # 加载模型 model = MattingNetwork("resnet50") model.load_state_dict(torch.load(r"D:\project\seg\jacke121-rvm_128_json\model_a\rvm_resnet50.pth", map_location="cpu")) model = model.eval() # 输入图像 # img = cv2.imread("input.jpg")[:, :, ::-1] # BGR -> RGB # img = cv2.imread(r"D:\project_2025\live2d\LayerDiffuse_DiffusersCLI-main\imgs\inputs\mm01.jpg") # BGR -> RGB img = cv2.imread(r"D:\soft\801.jpeg") # BGR -> RGB img = cv2.imread(r"D:\soft\805.png") # BGR -> RGB h, w, _ = img.shape # 转 tensor img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = img.astype(np.float32) / 255.0 # 归一化到 0-1 # 转 tensor: [H,W,3] -> [1,3,H,W] src = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) src = src.cpu() # RVM 需要 recurrent states (init None) rec = [None] * 4 # 推理，RVM 输入为：src, downsample_ratio, *rec with torch.no_grad(): fgr, pha, *rec = model(src, downsample_ratio=0.25, *rec) # fgr: foreground，pha: alpha matte fgr = fgr[0].permute(1, 2, 0).cpu().numpy() pha = pha[0].permute(1, 2, 0).cpu().numpy() # 合成为 rgba rgba = np.concatenate([fgr, pha], axis=2) # [H,W,4] rgba = (rgba * 255).astype(np.uint8) Image.fromarray(rgba).save("output_rvm.png") print("Done: output_rvm.png")
