parser.add_argument('--video', type=str, required=True, help='input video file') parser.add_argument('--result-type', type=str, default='fg', choices=['fg', 'matte'], help='matte - save the alpha matte; fg - save the foreground') parser.add_argument('--fps', type=int, default=30, help='fps of the result video') print('Get CMD Arguments...') args = parser.parse_args() if not os.path.exists(args.video): print('Cannot find the input video: {0}'.format(args.video)) exit() print('Load pre-trained MODNet...') pretrained_ckpt = './pretrained/modnet_webcam_portrait_matting.ckpt' modnet = MODNet(backbone_pretrained=False) modnet = nn.DataParallel(modnet) GPU = True if torch.cuda.device_count() > 0 else False if GPU: print('Use GPU...') modnet = modnet.cuda() modnet.load_state_dict(torch.load(pretrained_ckpt)) else: print('Use CPU...') modnet.load_state_dict(torch.load(pretrained_ckpt, map_location=torch.device('cpu'))) modnet.eval() result = os.path.splitext(args.video)[0] + '_{0}.mp4'.format(args.result_type) alpha_matte = True if args.result_type == 'matte' else False matting(args.video, result, alpha_matte, args.fps)
default=False, help="use cpu inferece") args = parser.parse_args() torch_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) print('Load pre-trained MODNet...') pretrained_ckpt = args.ckpt_path modnet = MODNet(backbone_pretrained=False) modnet = nn.DataParallel(modnet) modnet.load_state_dict(torch.load(pretrained_ckpt, map_location='cpu')) if not args.cpu: modnet.cuda() modnet.eval() print('Init WebCam...') cap = cv2.VideoCapture(0) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) print('Start matting...') while (True): _, frame_np = cap.read() frame_np = cv2.cvtColor(frame_np, cv2.COLOR_BGR2RGB) frame_np = cv2.resize(frame_np, (910, 512), cv2.INTER_AREA) frame_np = frame_np[:, 120:792, :] frame_np = cv2.flip(frame_np, 1)