# --------------- Main ---------------
# Load model
if args.model_type == 'mattingbase':
model = MattingBase(args.model_backbone)
if args.model_type == 'mattingrefine':
model = MattingRefine(
args.model_backbone,
args.model_backbone_scale,
args.model_refine_mode,
args.model_refine_sample_pixels,
args.model_refine_threshold,
args.model_refine_kernel_size)
model = model.cuda().eval()
model.load_state_dict(torch.load(args.model_checkpoint), strict=False)
# Load video and background
vid = VideoDataset(args.video_src)
bgr = [Image.open(args.video_bgr).convert('RGB')]
nbgr = [Image.open(args.new_bgr).convert('RGB')]
dataset = ZipDataset([vid, bgr, nbgr], transforms=A.PairCompose([
A.PairApply(T.Resize(args.video_resize[::-1]) if args.video_resize else nn.Identity()),
HomographicAlignment() if args.preprocess_alignment else A.PairApply(nn.Identity()),
A.PairApply(T.ToTensor())
]))
# Create output directory
if os.path.exists(args.output_dir):
args.model_refine_threshold,
args.model_refine_kernel_size,
refine_prevent_oversampling=False)
if args.model_checkpoint:
model.load_state_dict(torch.load(args.model_checkpoint), strict=False)
if args.precision == 'float32':
precision = torch.float32
else:
precision = torch.float16
if args.backend == 'torchscript':
model = torch.jit.script(model)
model = model.cuda().eval().to(precision)
# Load data
if not args.image_src:
src = torch.rand((args.batch_size, 3, *args.resolution[::-1]),
device='cuda',
dtype=precision)
bgr = torch.rand((args.batch_size, 3, *args.resolution[::-1]),
device='cuda',
dtype=precision)
else:
src = to_tensor(Image.open(args.image_src)).unsqueeze(0).repeat(
args.batch_size, 1, 1, 1).to(device='cuda', dtype=precision)
bgr = to_tensor(Image.open(args.image_bgr)).unsqueeze(0).repeat(
args.batch_size, 1, 1, 1).to(device='cuda', dtype=precision)
