mirror of
https://github.com/sczhou/ProPainter.git
synced 2026-10-18 20:59:51 -05:00
add features for memory-efficient inference.
This commit is contained in:
sczhou committed
1 parent
1cdf3771f8
commit
7587b2b865
7 files changed
+159
-107
No files matched your search
@@ -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.
|
||||
---
|
||||
For the commercial use of the code, please consult Prof. Chen Change Loy (ccloy@ntu.edu.sg)
|
||||
@@ -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
|
||||
<table>
|
||||
<tr>
|
||||
<td>
|
||||
@@ -92,7 +99,6 @@
|
||||
|
||||
|
||||
|
||||
|
||||
## Overview
|
||||
|
||||
|
||||
@@ -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
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
@@ -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 <dataset_root> --save_path <save_flow_root> --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
|
||||
|
||||
+75
-43
@@ -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:
|
||||
|
||||
+4
-2
@@ -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'):
|
||||
"""
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user