diff --git a/LICENSE b/LICENSE index 44bf750..bb4d0f5 100644 --- a/LICENSE +++ b/LICENSE @@ -1,35 +1,14 @@ -S-Lab License 1.0 +# S-Lab License 1.0 -Copyright 2022 S-Lab +Copyright 2023 S-Lab -Redistribution and use for non-commercial purpose in source and -binary forms, with or without modification, are permitted provided -that the following conditions are met: +Redistribution and use for non-commercial purpose in source and binary forms, with or without modification, are permitted provided that the following conditions are met: +1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. +3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.\ +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +4. In the event that redistribution and/or use for commercial purpose in source or binary forms, with or without modification is required, please contact the contributor(s) of the work. -1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in - the documentation and/or other materials provided with the - distribution. - -3. Neither the name of the copyright holder nor the names of its - contributors may be used to endorse or promote products derived - from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, -DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY -THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -In the event that redistribution and/or use for commercial purpose in -source or binary forms, with or without modification is required, -please contact the contributor(s) of the work. \ No newline at end of file +--- +For the commercial use of the code, please consult Prof. Chen Change Loy (ccloy@ntu.edu.sg) \ No newline at end of file diff --git a/README.md b/README.md index 12a74da..f7648e8 100644 --- a/README.md +++ b/README.md @@ -49,12 +49,19 @@ ## Update -- **2023.09.07**: Our code and model are publicly available. :whale: +- **2023.09.21**: Add features for memory-efficient inference. Check our [GPU memory](https://github.com/sczhou/ProPainter#memory-efficient-inference) requirements. 🏂 +- **2023.09.07**: Our code and model are publicly available. 🐳 - **2023.09.01**: This repo is created. + +### TODO +- [ ] Make a interactive Gradio demo. +- [ ] Make a Colab demo. +- [x] ~~Update features for memory-efficient inference.~~ + ## Results -#### 🏂 Object Removal +#### 👨🏻‍🎨 Object Removal
@@ -92,7 +99,6 @@ - ## Overview ![overall_structure](assets/ProPainter_pipeline.png) @@ -129,7 +135,7 @@ weights |- README.md ``` -### Quick test +### 🚀 Quick test We provide some examples in the [`inputs`](./inputs) folder. Run the following commands to try it out: ```shell @@ -138,18 +144,38 @@ python inference_propainter.py --video inputs/object_removal/bmx-trees --mask in # The second example (watermark removal) python inference_propainter.py --video inputs/watermark_removal/running_car.mp4 --mask inputs/watermark_removal/mask.png ``` + The results will be saved in the `results` folder. To test your own videos, please prepare the input `mp4 video` (or `split frames`) and `frame-wise mask(s)`. -If you want to specify the video resolution for processing or avoid running out of memory, you can use the `--set_size` flag and set the video size of `--width` and `--height`: +If you want to specify the video resolution for processing or avoid running out of memory, you can set the video size of `--width` and `--height`: ```shell -# process a 576x320 video -python inference_propainter.py --video inputs/watermark_removal/running_car.mp4 --mask inputs/watermark_removal/mask.png --set_size --height 320 --width 576 +# process a 576x320 video; set --fp16 to use fp16 (half precision) during inference. +python inference_propainter.py --video inputs/watermark_removal/running_car.mp4 --mask inputs/watermark_removal/mask.png --height 320 --width 576 --fp16 ``` -Regarding the issue of **running out of memory**, you can mitigate it by reducing the number of local neighbors through decreasing the `--neighbor_length` (default 20) or reducing the number of global references by increasing the `--ref_stride` (default 10). +### 🏂 Memory-efficient inference -### Dataset preparation for training and evaluation +Video inpainting typically requires a significant amount of GPU memory. Here, we offer various features that facilitate memory-efficient inference, effectively avoiding the Out-Of-Memory (OOM) error. You can use the following options to reduce memory usage further: + + - Reduce the number of local neighbors through decreasing the `--neighbor_length` (default 10). + - Reduce the number of global references by increasing the `--ref_stride` (default 10). + - Set the `--resize_ratio` (default 1.0) to resize the processing video. + - Set a smaller video size via specifying the `--width` and `--height`. + - Set `--fp16` to use fp16 (half precision) during inference. + - Reduce the frames of sub-videos `--subvideo_lentgh` (default 80), which effectively decouples GPU memory costs and video length. (coming soon!) + +Blow shows the estimated GPU memory requirements for different sub-video lengths with fp32/fp16 precision: + +| Resolution | 50 frames | 80 frames | +| :--- | :----: | :----: | +| 1280 x 720 | 28G / 19G | OOM / 25G | +| 720 x 480 | 11G / 7G | 13G / 8G | +| 640 x 480 | 10G / 6G | 12G / 7G | +| 320 x 240 | 3G / 2G | 4G / 3G | + + +## Dataset preparation @@ -203,7 +229,26 @@ datasets |- train.json |- test.json ``` -### Evaluation + +## Training +Our training configures are provided in [`train_flowcomp.json`](./configs/train_flowcomp.json) (for Recurrent Flow Completion Network) and [`train_propainter.json`](./configs/train_propainter.json) (for ProPainter). + +Run one of the following commands for training: +```shell + # For training Recurrent Flow Completion Network + python train.py -c configs/train_flowcomp.json + # For training ProPainter + python train.py -c configs/train_propainter.json +``` +You can run the **same command** to **resume** your training. + +To speed up the training process, you can precompute optical flow for the training dataset using the following command: +```shell + # Compute optical flow for training dataset + python scripts/compute_flow.py --root_path --save_path --height 240 --width 432 +``` + +## Evaluation Run one of the following commands for evaluation: ```shell # For evaluating flow completion model @@ -215,17 +260,6 @@ Run one of the following commands for evaluation: The scores and results will also be saved in the `results_eval` folder. Please `--save_results` for further [evaluating temporal warping error](https://github.com/phoenix104104/fast_blind_video_consistency#evaluation). -### Training -Our training configures are provided in [`train_flowcomp.json`](./configs/train_flowcomp.json) (for Recurrent Flow Completion Network) and [`train_propainter.json`](./configs/train_propainter.json) (for ProPainter). - -Run one of the following commands for training: -```shell - # For training Recurrent Flow Completion Network - python train.py -c configs/train_flowcomp.json - # For training ProPainter - python train.py -c configs/train_propainter.json -``` -You can run the **same command** to **resume** your training. ## Citation diff --git a/inference_propainter.py b/inference_propainter.py index df291da..f201c7e 100644 --- a/inference_propainter.py +++ b/inference_propainter.py @@ -62,8 +62,9 @@ def read_frame_from_videos(frame_root): frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) frames.append(frame) fps = None - - return frames, fps, video_name + size = frames[0].size + + return frames, fps, size, video_name def binary_mask(mask, th=0.1): @@ -177,44 +178,55 @@ if __name__ == '__main__': device = get_device() parser = argparse.ArgumentParser() - parser.add_argument('-i', '--video', type=str, default='inputs/object_removal/bmx-trees', - help='Path of the input video or image folder.') - parser.add_argument('-m', '--mask', type=str, default='inputs/object_removal/bmx-trees_mask', - help='Path of the mask(s) or mask folder.') - parser.add_argument('-o', '--output', type=str, default='results', - help='Output folder. Default: results') - parser.add_argument("--set_size", action='store_true', default=False) - parser.add_argument('--height', type=int, default=240, - help='Height of the processing video.') - parser.add_argument('--width', type=int, default=432, - help='Width of the processing video.') - parser.add_argument("--ref_stride", type=int, default=10, - help='stride of global reference frames.') - parser.add_argument("--ref_num", type=int, default=-1, - help='num of reference frames.') - parser.add_argument("--neighbor_length", type=int, default=20, - help='length of local neighboring frames.') - parser.add_argument("--raft_iter", type=int, default=20, - help='iterations for RAFT inference.') - parser.add_argument('--mode', default='video_inpainting', choices=['video_inpainting', 'video_outpainting'], - help="modes: video_inpainting / video_outpainting") - parser.add_argument('--scale_h', type=float, default=1.0, - help='Outpainting scale of height for video_outpainting mode.') - parser.add_argument('--scale_w', type=float, default=1.2, - help='Outpainting scale of width for video_outpainting mode.') - parser.add_argument('--save_fps', type=int, default=24, - help='Frame per second. Default: 24') - parser.add_argument('--save_frames', action='store_true', - help='Save output frames. Default: False') - + parser.add_argument( + '-i', '--video', type=str, default='inputs/object_removal/bmx-trees', help='Path of the input video or image folder.') + parser.add_argument( + '-m', '--mask', type=str, default='inputs/object_removal/bmx-trees_mask', help='Path of the mask(s) or mask folder.') + parser.add_argument( + '-o', '--output', type=str, default='results', help='Output folder. Default: results') + parser.add_argument( + "--resize_ratio", type=float, default=1.0, help='Resize scale for processing video.') + parser.add_argument( + '--height', type=int, default=-1, help='Height of the processing video.') + parser.add_argument( + '--width', type=int, default=-1, help='Width of the processing video.') + parser.add_argument( + '--mask_dilation', type=int, default=4, help='Mask dilation for video and flow masking.') + parser.add_argument( + "--ref_stride", type=int, default=10, help='Stride of global reference frames.') + parser.add_argument( + "--ref_num", type=int, default=-1, help='Number of reference frames.') + parser.add_argument( + "--neighbor_length", type=int, default=10, help='Length of local neighboring frames.') + parser.add_argument( + "--subvideo_lentgh", type=int, default=80, help='Length of sub-video for long video inference.') + parser.add_argument( + "--raft_iter", type=int, default=20, help='Iterations for RAFT inference.') + parser.add_argument( + '--mode', default='video_inpainting', choices=['video_inpainting', 'video_outpainting'], help="Modes: video_inpainting / video_outpainting") + parser.add_argument( + '--scale_h', type=float, default=1.0, help='Outpainting scale of height for video_outpainting mode.') + parser.add_argument( + '--scale_w', type=float, default=1.2, help='Outpainting scale of width for video_outpainting mode.') + parser.add_argument( + '--save_fps', type=int, default=24, help='Frame per second. Default: 24') + parser.add_argument( + '--save_frames', action='store_true', help='Save output frames. Default: False') + parser.add_argument( + '--fp16', action='store_true', help='Use fp16 (half precision) during inference. Default: fp32 (single precision).') + args = parser.parse_args() - if args.set_size: + # Use fp16 precision during inference to reduce running memory cost + use_half = True if args.fp16 else False + + + frames, fps, size, video_name = read_frame_from_videos(args.video) + if not args.width == -1 and not args.height == -1: size = (args.width, args.height) - else: - size = None + if not args.resize_ratio == 1.0: + size = (int(args.resize_ratio * size[0]), int(args.resize_ratio * size[1])) - frames, fps, video_name = read_frame_from_videos(args.video) frames, size, out_size = resize_frames(frames, size) fps = args.save_fps if fps is None else fps @@ -225,7 +237,8 @@ if __name__ == '__main__': if args.mode == 'video_inpainting': frames_len = len(frames) flow_masks, masks_dilated = read_mask(args.mask, frames_len, size, - flow_mask_dilates=4, mask_dilates=4) + flow_mask_dilates=args.mask_dilation, + mask_dilates=args.mask_dilation) w, h = size elif args.mode == 'video_outpainting': assert args.scale_h is not None and args.scale_w is not None, 'Please provide a outpainting scale (s_h, s_w).' @@ -268,6 +281,7 @@ if __name__ == '__main__': fix_flow_complete.to(device) fix_flow_complete.eval() + ############################################## # set up ProPainter model ############################################## @@ -276,20 +290,27 @@ if __name__ == '__main__': model = InpaintGenerator(model_path=ckpt_path).to(device) model.eval() - video_length = frames.size(1) - masked_frames = frames * (1 - masks_dilated) ############################################## # ProPainter inference ############################################## + video_length = frames.size(1) print(f'\nProcessing: {video_name} [{video_length} frames]...') with torch.no_grad(): # ---- compute flow ---- - short_len = 60 - if frames.size(1) > short_len: + if frames.size(-1) <= 640: + short_clip_len = 12 + elif frames.size(-1) <= 720: + short_clip_len = 8 + elif frames.size(-1) <= 1280: + short_clip_len = 4 + else: + short_clip_len = 2 + + if frames.size(1) > short_clip_len: gt_flows_f_list, gt_flows_b_list = [], [] - for f in range(0, video_length, short_len): - end_f = min(video_length, f + short_len) + for f in range(0, video_length, short_clip_len): + end_f = min(video_length, f + short_clip_len) if f == 0: flows_f, flows_b = fix_raft(frames[:,f:end_f], iters=args.raft_iter) else: @@ -304,11 +325,19 @@ if __name__ == '__main__': else: gt_flows_bi = fix_raft(frames, iters=args.raft_iter) + + if use_half: + frames, flow_masks, masks_dilated = frames.half(), flow_masks.half(), masks_dilated.half() + gt_flows_bi = (gt_flows_bi[0].half(), gt_flows_bi[1].half()) + fix_flow_complete = fix_flow_complete.half() + model = model.half() + # ---- complete flow ---- pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, flow_masks) pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, flow_masks) # ---- temporal propagation ---- + masked_frames = frames * (1 - masks_dilated) prop_imgs, updated_local_masks = model.img_propagation(masked_frames, pred_flows_bi, masks_dilated, 'nearest') b, t, _, _, _ = masks_dilated.size() @@ -316,6 +345,7 @@ if __name__ == '__main__': updated_frames = frames * (1-masks_dilated) + prop_imgs.view(b, t, 3, h, w) * masks_dilated # merge del gt_flows_bi, frames, prop_imgs, updated_local_masks + torch.cuda.empty_cache() ori_frames = frames_inp @@ -354,6 +384,8 @@ if __name__ == '__main__': comp_frames[idx] = comp_frames[idx].astype(np.float32) * 0.5 + img.astype(np.float32) * 0.5 comp_frames[idx] = comp_frames[idx].astype(np.uint8) + + torch.cuda.empty_cache() # save each frame if args.save_frames: diff --git a/model/propainter.py b/model/propainter.py index c09a4f1..18b0cb8 100644 --- a/model/propainter.py +++ b/model/propainter.py @@ -25,7 +25,8 @@ def fbConsistencyCheck(flow_fw, flow_bw, alpha1=0.01, alpha2=0.5): mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped) # |wf| + |wb(wf(x))| occ_thresh_fw = alpha1 * mag_sq_fw + alpha2 - fb_valid_fw = (length_sq(flow_diff_fw) < occ_thresh_fw).float() + # fb_valid_fw = (length_sq(flow_diff_fw) < occ_thresh_fw).float() + fb_valid_fw = (length_sq(flow_diff_fw) < occ_thresh_fw).to(flow_fw) return fb_valid_fw @@ -97,7 +98,8 @@ class BidirectionalPropagation(nn.Module): def binary_mask(self, mask, th=0.1): mask[mask>th] = 1 mask[mask<=th] = 0 - return mask.float() + # return mask.float() + return mask.to(mask) def forward(self, x, flows_forward, flows_backward, mask, interpolation='bilinear'): """ diff --git a/scripts/compute_flow.py b/scripts/compute_flow.py index 5700b0f..3ff1c5d 100644 --- a/scripts/compute_flow.py +++ b/scripts/compute_flow.py @@ -40,7 +40,8 @@ if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-i', '--root_path', type=str, default='your_dataset_root/youtube-vos/JPEGImages') parser.add_argument('-o', '--save_path', type=str, default='your_dataset_root/youtube-vos/Flows_flo') - parser.add_argument('-s', '--size', type=tuple, default=(432, 240)) + parser.add_argument('--height', type=int, default=240) + parser.add_argument('--width', type=int, default=432) args = parser.parse_args() @@ -49,7 +50,7 @@ if __name__ == '__main__': root_path = args.root_path save_path = args.save_path - h_up, w_up = args.size + h_new, w_new = (args.height, args.width) file_list = sorted(os.listdir(root_path)) for f in file_list: @@ -69,16 +70,16 @@ if __name__ == '__main__': # upsize to a multiple of 16 # h, w = img1.shape[2:4] - # w_up = w if (w % 16) == 0 else 16 * (w // 16 + 1) - # h_up = h if (h % 16) == 0 else 16 * (h // 16 + 1) + # w_new = w if (w % 16) == 0 else 16 * (w // 16 + 1) + # h_new = h if (h % 16) == 0 else 16 * (h // 16 + 1) img1 = F.interpolate(input=img1, - size=(h_up, w_up), + size=(h_new, w_new), mode='bilinear', align_corners=False) img2 = F.interpolate(input=img2, - size=(h_up, w_up), + size=(h_new, w_new), mode='bilinear', align_corners=False) @@ -93,8 +94,8 @@ if __name__ == '__main__': flow_f = flow_f[0].permute(1,2,0).cpu().numpy() flow_b = flow_b[0].permute(1,2,0).cpu().numpy() - # flow_f = resize_flow(flow_f, w_up, h_up) - # flow_b = resize_flow(flow_b, w_up, h_up) + # flow_f = resize_flow(flow_f, w_new, h_new) + # flow_b = resize_flow(flow_b, w_new, h_new) save_flow_f = os.path.join(save_path, f, f'{m_list[i][:-4]}_{m_list[i+1][:-4]}_f.flo') save_flow_b = os.path.join(save_path, f, f'{m_list[i+1][:-4]}_{m_list[i][:-4]}_b.flo') diff --git a/scripts/evaluate_flow_completion.py b/scripts/evaluate_flow_completion.py index e9e50f3..8623824 100644 --- a/scripts/evaluate_flow_completion.py +++ b/scripts/evaluate_flow_completion.py @@ -51,6 +51,7 @@ def tensor2np(array): def main_worker(args): # set up datasets and data loader + args.size = (args.width, args.height) test_dataset = TestDataset(vars(args)) test_loader = DataLoader(test_dataset, @@ -177,7 +178,8 @@ def main_worker(args): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--size', default=(432, 240), type=tuple) + parser.add_argument('--height', type=int, default=240) + parser.add_argument('--width', type=int, default=432) parser.add_argument('--raft_model_path', default='weights/raft-things.pth', type=str) parser.add_argument('--fc_model_path', default='weights/recurrent_flow_completion.pth', type=str) parser.add_argument('--dataset', choices=['davis', 'youtube-vos'], type=str) diff --git a/scripts/evaluate_propainter.py b/scripts/evaluate_propainter.py index 6fc0e9c..e7d79de 100644 --- a/scripts/evaluate_propainter.py +++ b/scripts/evaluate_propainter.py @@ -32,6 +32,7 @@ def get_ref_index(neighbor_ids, length, ref_stride=10): def main_worker(args): + args.size = (args.width, args.height) w, h = args.size # set up datasets and data loader assert (args.dataset == 'davis') or args.dataset == 'youtube-vos', \ @@ -255,7 +256,8 @@ def main_worker(args): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--size', default=(432, 240), type=tuple) + parser.add_argument('--height', type=int, default=240) + parser.add_argument('--width', type=int, default=432) parser.add_argument("--ref_stride", type=int, default=10) parser.add_argument("--neighbor_length", type=int, default=20) parser.add_argument("--raft_iter", type=int, default=20)