├── Curation_Pipeline.md
├── README.md
├── WebVi3D
    ├── examples
    │   ├── 000b0a3dd5d6827172e323ab4fc9232db10bec74.mp4
    │   ├── 0b7b58a8bf4200d835dc459bbe1af25bf3f917cc_3924_4085.mp4
    │   ├── 1af3278b744dfc772f43bb550d8dbada93fae2c6.mp4
    │   ├── 1b2590e9aeecd93c5926eea76c4a8dbddcebe980.mp4
    │   ├── 1b37bac6aae4142fdda0146ce5e1ead300ec1830.mp4
    │   ├── 1b8090b14fa7b90807c261c12e3078379f5419cb.mp4
    │   ├── 2b686e7afc43f5fc571d97601ae11c80b01d2228.mp4
    │   ├── 36d4e80bdfe3f59f5281a3126956291879dcb436.mp4
    │   ├── 4a08ffee34fcd0ac524ef723acead61262b65d32.mp4
    │   ├── 94f90e9aad920e534127d112fa24f9a8ab43bd83.mp4
    │   ├── a539c32a2ce9cf79e893a1d30fb1572721b81b5d.mp4
    │   └── a9cf7f4fda030ae0c13da01d08162412389ce84e.mp4
    ├── requirements.txt
    ├── s1_downsample.py
    ├── s2_generate_flow_earlycheck.py
    ├── s3_generate_mask.py
    ├── s4_filter.py
    ├── utils
    │   ├── RAFT
    │   │   ├── __init__.py
    │   │   ├── corr.py
    │   │   ├── datasets.py
    │   │   ├── demo.py
    │   │   ├── extractor.py
    │   │   ├── raft.py
    │   │   ├── update.py
    │   │   └── utils
    │   │   │   ├── __init__.py
    │   │   │   ├── augmentor.py
    │   │   │   ├── flow_viz.py
    │   │   │   ├── frame_utils.py
    │   │   │   └── utils.py
    │   └── flow_utils.py
    └── weights
    │   ├── raft-small.pth
    │   └── raft-things.pth
├── assets
    ├── logo.jpg
    ├── method.jpg
    └── teaser.jpg
├── inference.py
├── mv_diffusion.py
├── mv_diffusion_SR.py
├── mv_unet.py
├── pipeline_mvd_warp_mix_classifier.py
├── pipeline_mvd_warp_mix_classifier_SR.py
├── requirements.txt
├── single_infer.sh
└── sparse_infer.sh


/Curation_Pipeline.md:
--------------------------------------------------------------------------------
 1 | # Curation Pipeline
 2 | 
 3 | Our WebVi3D curation pipeline contains 4 steps, including:
 4 | 
 5 | - Temporal-Spatial Downsampling
 6 | - Semantic-Based Dynamic Recognition 
 7 | - Dynamic Filtering 
 8 | - Tracking-Based Small Viewpoint Filtering
 9 | 
10 | ## Running
11 | 
12 | We provide some unfiltered video cases and the demo code of data curation pipeline for your own video dataset. Notice that you may modify the code into parallel processing for faster processing on larger scale dataset. You may also adjust the parameters in each step for your own applications.
13 | 
14 | First, install extra dependencies:
15 | ```sh
16 | cd ./WebVi3D
17 | pip install -r requirements.txt
18 | ```
19 | 
20 | ### Down Sampling and Dynamic Recognition
21 | 
22 | You can run directly with our provided video cases:
23 | ```sh
24 | python s1_downsample.py --input_dir examples --output_dir examples_vid/step1+2_outputs --txt examples_txt/all_video_list.txt
25 | ```
26 | If you want to process your own video dataset, please run:
27 | ```sh
28 | python s1_downsample.py --input_dir <video folder> --output_dir <output dir for seperated frames> --txt <output txt file for indexing>
29 | ```
30 | 
31 | You will get all video names in the `txt` file, and the processed frames are save in `output_dir`, the filtered videos will not be saved in this step.
32 | 
33 | ### Flow Generation
34 | 
35 | We use RAFT to generate optical flow for dynamic masks generation. Simply run:
36 | ```sh
37 | python s2_generate_flow_earlycheck.py --dataset_path examples_vid/step1+2_outputs --txt examples_txt/all_video_list.txt
38 | ```
39 | For your own dataset, please run:
40 | ```sh
41 | python s2_generate_flow_earlycheck.py --dataset_path <seperated frames from last step> --txt <output txt file for further indexing>
42 | ```
43 | 
44 | The computed forward and backward flow will be saved in `flow` subdirectory for further processing. Notice that we do early checking that filter out the videos with very small optical flow values in the earlier frames.
45 | 
46 | ### Dynamic Masking 
47 | 
48 | We generate dynamic masks to filter out videos with moving things, run:
49 | ````sh
50 | python s3_generate_mask.py --dataset_path examples_vid/step1+2_outputs --txt examples_txt/mask_list.txt 
51 | ````
52 | 
53 | For your own dataset, run:
54 | ````sh
55 | python s3_generate_mask.py --dataset_path <seperated frames and flow from last step> --txt <output txt file for further indexing> 
56 | ````
57 | 
58 | to obtain dynamic motion mask, where the video name will be stored in txt.
59 | 
60 | ### Dynamic and Small Camera Movement Filtering
61 | 
62 | We then filter our dynamic scenes and small camera movements in one time, and save the final results to txt. Similarly, run:
63 | ```sh
64 | python s4_filter.py --dataset_path examples_vid/step1+2_outputs  --mask_done_file examples_txt/mask_list.txt --txt examples_txt/final.txt
65 | ```
66 | For your own dataset, run:
67 | ```sh
68 | python s4_filter.py --dataset_path <seperated frames, flow and mask from the last step>  --mask_done_file <dynamic mask txt file from the last step> --txt <final output video names>
69 | ```
70 | 
71 | The filtered video names will be listed in the final txt file.
72 | 
73 | 
74 | ### Acknowledgements
75 | We appreciate [RAFT](https://github.com/princeton-vl/RAFT), [Co-Tracker](https://github.com/facebookresearch/co-tracker) for their awesome open-sourced projects.
76 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | <div align='center'>
 2 | <!-- <img src="assets/logo.jpg" alt="logo" style="zoom:0.5%;" /> -->
 3 | 
 4 | <h2>You See it, You Got it: Learning 3D Creation on Pose-Free Videos at Scale</h2>
 5 | 
 6 | [Baorui Ma*](https://mabaorui.github.io), [Huachen Gao*](https://gaohchen.github.io/), [Haoge Deng*](https://github.com/Bitterdhg/), [Zhengxiong Luo](https://greatlog.github.io/), [Tiejun Huang](https://scholar.google.com/citations?user=knvEK4AAAAAJ&hl=en), [Lulu Tang†](https://scholar.google.com/citations?authuser=1&user=o2fG4xUAAAAJ), [Xinlong Wang†](https://www.xloong.wang/)
 7 | 
 8 | [Beijing Academy of Artificial Intelligence, BAAI](https://www.baai.ac.cn/english.html)<br>
 9 | <sup>*</sup> Equal Contribution, <sup>†</sup> Corresponding Author
10 | 
11 | ## [Paper](https://arxiv.org/abs/2412.06699/) | [Project Page](https://vision.baai.ac.cn/see3d/) | [Models(Google Drive)](https://drive.google.com/drive/folders/1uz4NtAtYYiV53dYlJ_I78fVzqxUBk-Fb?usp=sharing/) |[🤗HF Models](https://huggingface.co/bruiiii/See3D/)| [Dataset]()
12 | 
13 | </div>
14 | We present <strong>See3D</strong>, a visual-conditional multi-view diffusion model trained on large-scale Internet videos for open-world 3D creation. The model aims to </strong>Get</strong> 3D knowledge by solely <strong>Seeing</strong> the visual contents from the vast and rapidly growing video data --- <em>You See it, You Got it</em>. To achieve this, we first scale up the training data using a proposed data curation pipeline that automatically filters out multi-view inconsistencies and insufficient observations from source videos. This results in a high-quality, richly diverse, large-scale dataset of multi-view images, termed <strong>WebVi3D</strong>, containing 320M frames from 16M video clips. Nevertheless, learning generic 3D priors from videos without explicit 3D geometry or camera pose annotations is nontrivial, and annotating poses for web-scale videos is prohibitively expensive. To eliminate the need for pose conditions, we introduce an innovative <em>visual-condition</em> - a purely 2D-inductive visual signal generated by adding time-dependent noise to the masked video data. Finally, we introduce a novel visual-conditional 3D generation framework by integrating <strong>See3D</strong> into a warping-based pipeline for high-fidelity 3D generation. Our numerical and visual comparisons on single and sparse reconstruction benchmarks show that <strong>See3D</strong>, trained on cost-effective and scalable video data, achieves notable zero-shot and open-world generation capabilities, markedly outperforming models trained on costly and constrained 3D datasets. Additionally, our model naturally supports other image-conditioned 3D creation tasks, such as 3D editing, without further fine-tuning.
15 | 
16 | ![](assets/teaser.jpg)
17 | Benefiting from the proposed web-scale dataset **WebVi3D**, **See3D** enables both object- and scene-level 3D creation, including sparse-view-to-3D, (text-) image-to-3D, and 3D editing. It can also be used for Gaussian Splatting to extract meshes or render images.
18 | 
19 | ### Highlights
20 | 
21 | - We present **See3D**, a scalable visual-conditional MVD model for open-world 3D creation, which can be trained on web-scale video collections without pose annotations.
22 | - We curate **WebVi3D**, a multi-view images dataset containing static scenes with sufficient multi-view observations, and establish an automated pipeline for video data curation to train the MVD model.
23 | - We introduce a novel warping-based 3D generation framework with **See3D**, which supports long-sequence generation with complex camera trajectories.
24 | - We achieve state-of-the-art results in single and sparse views reconstruction, demonstrating remarkable zero-shot and open-world generation capability, offering a novel perspective on scalable 3D generation.
25 | 
26 | ## News
27 | **[04/05/2025]**
28 | :fire: See3D is selected as Highlight presentation!
29 | 
30 | **[03/11/2025]**
31 | :rocket: We have released the WebVi3D curation pipeline. Please refer to [this guide](https://github.com/baaivision/See3D/tree/master/Curation_Pipeline.md).
32 | 
33 | **[02/27/2025]**
34 | :fire: See3D is accepted by CVPR 2025.
35 | 
36 | **[12/13/2024]**
37 | :rocket: We have released the pretrained models and example test data in Huggingface🤗.
38 | 
39 | **[12/10/2024]**
40 | :rocket: We have released the pretrained models and inference code. You can download models and example test data [here](https://drive.google.com/drive/folders/1uz4NtAtYYiV53dYlJ_I78fVzqxUBk-Fb?usp=sharing/)
41 | 
42 | ## Installation
43 | ```sh
44 | git clone https://github.com/baaivision/See3D.git
45 | cd See3D
46 | 
47 | pip install -r requirements.txt
48 | ```
49 | 
50 | ## Inference Code
51 | We provide inference code for multi-view generation based on single-view and sparse-view inputs. Please add or remove the `--super_resolution` parameter according to your needs. The multi-view super-resolution model will upscale the default 512 resolution to a consistent 1024 resolution across multiple views, which requires more inference time and GPU memory. Please download the example test data [here](https://drive.google.com/drive/folders/1uz4NtAtYYiV53dYlJ_I78fVzqxUBk-Fb?usp=sharing/) and put it in the `dataset` folder.
52 | ### Generation Based on Single View Input
53 | ```sh
54 | bash single_infer.sh
55 | ```
56 | ### Generation Based on Sparse Views Input
57 | ```sh
58 | bash sparse_infer.sh
59 | ```
60 | ## Data Curation Pipeline of WebVi3D
61 | We provide the demo code of data curation pipeline given your own video dataset. Please refer to this [dataset README](https://github.com/baaivision/See3D/tree/master/Curation_Pipeline.md) for more details.
62 | 
63 | ## TODO List
64 | - [x] Release pretrained models.
65 | - [x] Release inference code.
66 | - [x] Release data curation pipeline from Internet Video.
67 | - [ ] Release training scripts.
68 | - [ ] Release 3D generation framework utilizing the warping-based pipeline.
69 | - [ ] Release the evaluation code.
70 | 
71 | ## Acknowledgement
72 | 
73 | See3D is built using the awesome open-source projects: [Stable Diffusion](https://github.com/Stability-AI/stablediffusion), [MVDream](https://github.com/bytedance/MVDream), [ViewCrafter](https://github.com/Drexubery/ViewCrafter/tree/main), [FrozenRecon](https://github.com/aim-uofa/FrozenRecon).
74 | 
75 | Thanks to the maintainers of these projects for their contribution to the community!
76 | 
77 | ## Citation
78 | 
79 | If you find See3D helpful, please consider citing:
80 | 
81 | ```
82 | @inproceedings{Ma2025See3D,
83 |     title = {You See it, You Got it: Learning 3D Creation on Pose-Free Videos at Scale},
84 |     author = {Baorui Ma and Huachen Gao and Haoge Deng and Zhengxiong Luo and Tiejun Huang and Lulu Tang and Xinlong Wang},
85 |     booktitle={IEEE/CVF conference on computer vision and pattern recognition},
86 |     year={2025}
87 | }
88 | ```  
89 | 


--------------------------------------------------------------------------------
/WebVi3D/examples/000b0a3dd5d6827172e323ab4fc9232db10bec74.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/000b0a3dd5d6827172e323ab4fc9232db10bec74.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/0b7b58a8bf4200d835dc459bbe1af25bf3f917cc_3924_4085.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/0b7b58a8bf4200d835dc459bbe1af25bf3f917cc_3924_4085.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/1af3278b744dfc772f43bb550d8dbada93fae2c6.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/1af3278b744dfc772f43bb550d8dbada93fae2c6.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/1b2590e9aeecd93c5926eea76c4a8dbddcebe980.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/1b2590e9aeecd93c5926eea76c4a8dbddcebe980.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/1b37bac6aae4142fdda0146ce5e1ead300ec1830.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/1b37bac6aae4142fdda0146ce5e1ead300ec1830.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/1b8090b14fa7b90807c261c12e3078379f5419cb.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/1b8090b14fa7b90807c261c12e3078379f5419cb.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/2b686e7afc43f5fc571d97601ae11c80b01d2228.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/2b686e7afc43f5fc571d97601ae11c80b01d2228.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/36d4e80bdfe3f59f5281a3126956291879dcb436.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/36d4e80bdfe3f59f5281a3126956291879dcb436.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/4a08ffee34fcd0ac524ef723acead61262b65d32.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/4a08ffee34fcd0ac524ef723acead61262b65d32.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/94f90e9aad920e534127d112fa24f9a8ab43bd83.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/94f90e9aad920e534127d112fa24f9a8ab43bd83.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/a539c32a2ce9cf79e893a1d30fb1572721b81b5d.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/a539c32a2ce9cf79e893a1d30fb1572721b81b5d.mp4


--------------------------------------------------------------------------------
/WebVi3D/examples/a9cf7f4fda030ae0c13da01d08162412389ce84e.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/examples/a9cf7f4fda030ae0c13da01d08162412389ce84e.mp4


--------------------------------------------------------------------------------
/WebVi3D/requirements.txt:
--------------------------------------------------------------------------------
1 | scikit-image
2 | filelock
3 | jsonlines
4 | einops
5 | git+https://github.com/cvg/LightGlue.git
6 | 


--------------------------------------------------------------------------------
/WebVi3D/s1_downsample.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import imageio
  3 | from tqdm import tqdm
  4 | import cv2
  5 | import shutil
  6 | import torchvision
  7 | from PIL import Image
  8 | import argparse
  9 | try:
 10 |     import jsonlines
 11 | except:
 12 |     os.system('pip install jsonlines')
 13 |     import jsonlines
 14 | import numpy as np
 15 | from utils.flow_utils import *
 16 | 
 17 | def run_maskrcnn(model, img):
 18 |     threshold = 0.5
 19 |     cnt = 0
 20 |     # o_image = Image.open(img).convert("RGB")
 21 |     width, height = o_image.size
 22 |     if width > height:
 23 |         intHeight = 576
 24 |         intWidth = 1024
 25 |     else:
 26 |         intHeight = 1024
 27 |         intWidth = 576
 28 | 
 29 |     image = o_image.resize((intWidth, intHeight), Image.LANCZOS)
 30 | 
 31 |     image_tensor = torchvision.transforms.functional.to_tensor(image).cuda()
 32 | 
 33 |     objPredictions = model([image_tensor])[0]
 34 | 
 35 |     for intMask in range(objPredictions["masks"].size(0)):
 36 |         if objPredictions["scores"][intMask].item() > threshold:
 37 |             # person, vehicle, accessory, animal, sports
 38 |             if objPredictions["labels"][intMask].item() == 1:  # person
 39 |                 cnt += 1
 40 |     return cnt
 41 | 
 42 | 
 43 | if __name__ == "__main__":
 44 |     args = argparse.ArgumentParser()
 45 |     args.add_argument('--gpu', type=int,default=0)
 46 |     args.add_argument('--input_dir', type=str, default='examples', help="input video folder")
 47 |     args.add_argument('--output_dir', type=str, default='examples_vid/step1+2_outputs', help="output folder to store frames")
 48 |     args.add_argument('--txt', type=str, default="examples_txt/all_video_list.txt", help="txt file to store processed video names,use it for the following steps.")
 49 |     args.add_argument('--human_threshold', type=float, default=10, help="number of frames with human, if more than this threshold, remove the video")
 50 |     
 51 |     args = args.parse_args()
 52 |     os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
 53 | 
 54 |     videoname = 'alls'
 55 |     video_folder = args.input_dir
 56 |     videos_list = microsoft_sorted(os.listdir(video_folder))
 57 | 
 58 |     duration = 150    
 59 |     sep = 2 #  2 frame per second``
 60 |     output_path = args.output_dir
 61 |     txt_file = args.txt
 62 | 
 63 |     netMaskrcnn = (
 64 |         torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
 65 |         .cuda()
 66 |         .eval()
 67 |         )
 68 |     done_list = []
 69 |     if os.path.exists(txt_file):
 70 |         with open (txt_file, 'r') as f:
 71 |             for line in f:
 72 |                 done_list.append(line.strip())
 73 |     os.makedirs(output_path, exist_ok=True)
 74 |     os.makedirs(txt_file.split('/')[0], exist_ok=True)
 75 |     for filename in tqdm(videos_list):
 76 | 
 77 |         if filename.endswith('.mp4'):
 78 |             video_path = os.path.join(video_folder, filename)
 79 |             video_name = filename.split('/')[-1].split('.')[0]
 80 |             if filename in done_list:
 81 |                 print("skipping " + video_name)
 82 |                 continue
 83 |             with open (txt_file, 'a+') as f:
 84 |                 f.write(filename + '\n')
 85 |                 f.close()
 86 | 
 87 |             frame_folder = os.path.join(output_path, video_name)
 88 |             frame_folder = os.path.join(frame_folder, 'images')
 89 |             if os.path.exists(os.path.join(output_path, video_name)) and not os.path.getsize(os.path.join(output_path, video_name)):
 90 |                 print("skipping " + video_name)
 91 |                 continue
 92 |             os.makedirs(frame_folder, exist_ok=True)
 93 |             
 94 |             reader = imageio.get_reader(video_path, format='ffmpeg')
 95 |             
 96 | 
 97 |             cnt = 0
 98 |             thres = 0
 99 |             for frame_count, frame in tqdm(enumerate(reader)):
100 |                 if cnt <= 20:
101 |                     # turn into Image
102 |                     o_image = Image.fromarray(frame)
103 |                     thres = thres + run_maskrcnn(netMaskrcnn, o_image)
104 |                     if thres > args.human_threshold:
105 |                         print("human detected, remove!")
106 |                         shutil.rmtree(os.path.join(output_path, video_name))
107 |                         break
108 |                     cnt += 1
109 |                 current_time = reader.get_meta_data()['duration'] * frame_count / (reader.get_meta_data()['fps'] + 1e-9)
110 |                 if current_time >= duration: break
111 |                 if frame_count % sep == 0: 
112 |                     
113 |                     frame_filename = f"frame_{frame_count}.jpg"
114 |                     frame_path = os.path.join(frame_folder, frame_filename)
115 |                     
116 |                     H,W = frame.shape[:2]
117 |                     if H > 2160 or W > 2160:    # 4k
118 |                         frame = cv2.resize(frame, (W // 8, H // 8))
119 |                     elif H >1080 or W > 1080:   # 
120 |                         frame = cv2.resize(frame, (W // 4, H // 4))
121 |                     elif H > 720 or W > 720:
122 |                         frame = cv2.resize(frame, (W // 3, H // 3))
123 | 
124 |                     imageio.imwrite(frame_path, frame)
125 | 
126 |     print(f"All videos processed done. please see {output_path}")


--------------------------------------------------------------------------------
/WebVi3D/s2_generate_flow_earlycheck.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import os
  3 | import cv2
  4 | import glob
  5 | import numpy as np
  6 | import torch
  7 | from PIL import Image
  8 | 
  9 | from utils.RAFT.raft import RAFT
 10 | from utils.RAFT.utils.utils import InputPadder
 11 | 
 12 | from utils.flow_utils import *
 13 | from tqdm import tqdm
 14 | DEVICE = "cuda"
 15 | import shutil
 16 | try:
 17 |     from filelock import FileLock
 18 | except ImportError:
 19 |     os.system("pip install filelock")
 20 |     from filelock import FileLock
 21 | 
 22 | def create_dir(dir):
 23 |     if not os.path.exists(dir):
 24 |         os.makedirs(dir)
 25 | 
 26 | 
 27 | def load_image(imfile):
 28 |     long_dim = 768
 29 |     img = np.array(Image.open(imfile)).astype(np.uint8)
 30 | 
 31 |     # Portrait Orientation
 32 |     if img.shape[0] > img.shape[1]:
 33 |         input_h = long_dim
 34 |         input_w = int(round(float(input_h) / img.shape[0] * img.shape[1]))
 35 |     # Landscape Orientation
 36 |     else:
 37 |         input_w = long_dim
 38 |         input_h = int(round(float(input_w) / img.shape[1] * img.shape[0]))
 39 | 
 40 |     # print("flow input w %d h %d" % (input_w, input_h))
 41 |     img = cv2.resize(img, (input_w, input_h), cv2.INTER_LINEAR)
 42 |     img = torch.from_numpy(img).permute(2, 0, 1).float()
 43 |     return img[None].to(DEVICE)
 44 | 
 45 | 
 46 | def resize_flow(flow, img_h, img_w):
 47 |     flow_h, flow_w = flow.shape[0], flow.shape[1]
 48 |     flow[:, :, 0] *= float(img_w) / float(flow_w)
 49 |     flow[:, :, 1] *= float(img_h) / float(flow_h)
 50 |     flow = cv2.resize(flow, (img_w, img_h), cv2.INTER_LINEAR)
 51 | 
 52 |     return flow
 53 | 
 54 | 
 55 | def warp_flow(img, flow):
 56 |     h, w = flow.shape[:2]
 57 |     flow_new = flow.copy()
 58 |     flow_new[:, :, 0] += np.arange(w)
 59 |     flow_new[:, :, 1] += np.arange(h)[:, np.newaxis]
 60 | 
 61 |     res = cv2.remap(
 62 |         img, flow_new, None, cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT
 63 |     )
 64 |     return res
 65 | 
 66 | 
 67 | def compute_fwdbwd_mask(fwd_flow, bwd_flow):
 68 |     alpha_1 = 0.5
 69 |     alpha_2 = 0.5
 70 | 
 71 |     bwd2fwd_flow = warp_flow(bwd_flow, fwd_flow)
 72 |     fwd_lr_error = np.linalg.norm(fwd_flow + bwd2fwd_flow, axis=-1)
 73 |     fwd_mask = (
 74 |         fwd_lr_error
 75 |         < alpha_1
 76 |         * (np.linalg.norm(fwd_flow, axis=-1) + np.linalg.norm(bwd2fwd_flow, axis=-1))
 77 |         + alpha_2
 78 |     )
 79 | 
 80 |     fwd2bwd_flow = warp_flow(fwd_flow, bwd_flow)
 81 |     bwd_lr_error = np.linalg.norm(bwd_flow + fwd2bwd_flow, axis=-1)
 82 | 
 83 |     bwd_mask = (
 84 |         bwd_lr_error
 85 |         < alpha_1
 86 |         * (np.linalg.norm(bwd_flow, axis=-1) + np.linalg.norm(fwd2bwd_flow, axis=-1))
 87 |         + alpha_2
 88 |     )
 89 | 
 90 |     return fwd_mask, bwd_mask
 91 | 
 92 | def check_flow(folder_path):
 93 |     flow = sorted(glob.glob(os.path.join(folder_path, 'flow') + '/*bwd.npz'))
 94 |     count = 0
 95 |     for f in flow[:20]:
 96 |         npz = np.load(f)
 97 |         flow = npz['flow']
 98 |         move = abs(flow)[...,0]+ abs(flow)[...,1]
 99 |         if move.mean() <= 0.5:
100 |             count += 1
101 |     return count
102 | 
103 | def extract_number(filename):
104 |     return int(''.join(filter(str.isdigit, filename)))
105 | 
106 | def camera_movement(imglist):
107 |     thres = 0
108 |     prev = None
109 |     for img in imglist[::2]:
110 |         img = np.array(Image.open(img))
111 |         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
112 |         edges = cv2.Canny(gray, 50, 100)
113 |         if prev is not None:
114 |             ssim = cv2.matchTemplate(edges, prev, cv2.TM_CCOEFF_NORMED)[0][0]
115 |             # print(ssim)
116 |             if ssim >= 0.7: thres += 1  # changable, higher means more strict
117 |         prev = edges
118 | 
119 |     return thres
120 | 
121 | def check_flow(flow):
122 |     # flow = sorted(glob.glob(os.path.join(folder_path, 'flow') + '/*bwd.npz'))
123 |     move = abs(flow)[...,0]+ abs(flow)[...,1]
124 |     if move.mean() <= 0.4:
125 |         return True
126 |     return False
127 | 
128 | 
129 | def run(args, input_path, output_path, output_img_path):
130 |     model = torch.nn.DataParallel(RAFT(args))
131 |     model.load_state_dict(torch.load(args.model))
132 | 
133 |     model = model.module
134 |     model.to(DEVICE)
135 |     model.eval()
136 |     small_cnt = 0
137 |     cnt = 0
138 | 
139 |     with torch.no_grad():
140 |         images = glob.glob(os.path.join(input_path, "*.png")) + glob.glob(
141 |             os.path.join(input_path, "*.jpg")
142 |         )
143 | 
144 |         images = sorted(images)[:200]   # max 200 frames
145 |         img_train = cv2.imread(images[0])
146 |         for i in tqdm(range(len(images) - 1)):
147 |             # print(i)
148 |             image1 = load_image(images[i])
149 |             image2 = load_image(images[i + 1])
150 | 
151 |             padder = InputPadder(image1.shape)
152 |             image1, image2 = padder.pad(image1, image2)
153 | 
154 |             _, flow_fwd = model(image1, image2, iters=20, test_mode=True)
155 |             _, flow_bwd = model(image2, image1, iters=20, test_mode=True)
156 | 
157 |             flow_fwd = padder.unpad(flow_fwd[0]).cpu().numpy().transpose(1, 2, 0)
158 |             flow_bwd = padder.unpad(flow_bwd[0]).cpu().numpy().transpose(1, 2, 0)
159 | 
160 |             flow_fwd = resize_flow(flow_fwd, img_train.shape[0], img_train.shape[1])
161 |             flow_bwd = resize_flow(flow_bwd, img_train.shape[0], img_train.shape[1])
162 | 
163 |             mask_fwd, mask_bwd = compute_fwdbwd_mask(flow_fwd, flow_bwd)
164 |             if cnt <= 20 and check_flow(flow_bwd):
165 |                 small_cnt = small_cnt + 1
166 |             cnt = cnt + 1
167 |             if small_cnt >= args.little_movement_threshold:
168 |                 print("too little movements, drop")     # early check, prevent long time processing
169 |                 shutil.rmtree(output_path)
170 |                 shutil.rmtree(output_img_path)
171 |                 return 
172 |             
173 |             # Save flow
174 |             np.savez(
175 |                 os.path.join(output_path, "%05d_fwd.npz" % i),
176 |                 flow=flow_fwd,
177 |                 mask=mask_fwd,
178 |             )
179 |             np.savez(
180 |                 os.path.join(output_path, "%05d_bwd.npz" % (i + 1)),
181 |                 flow=flow_bwd,
182 |                 mask=mask_bwd,
183 |             )
184 |             #! removed saving flow images
185 |             # # Save flow_img
186 |             # Image.fromarray(flow_viz.flow_to_image(flow_fwd)).save(
187 |             #     os.path.join(output_img_path, "%05d_fwd.png" % i)
188 |             # )
189 |             # Image.fromarray(flow_viz.flow_to_image(flow_bwd)).save(
190 |             #     os.path.join(output_img_path, "%05d_bwd.png" % (i + 1))
191 |             # )
192 | 
193 |             # Image.fromarray(mask_fwd).save(
194 |             #     os.path.join(output_img_path, "%05d_fwd_mask.png" % i)
195 |             # )
196 |             # Image.fromarray(mask_bwd).save(
197 |             #     os.path.join(output_img_path, "%05d_bwd_mask.png" % (i + 1))
198 |             # )
199 | 
200 | 
201 | if __name__ == "__main__":
202 |     parser = argparse.ArgumentParser()
203 |     parser.add_argument("--dataset_path", type=str, default="examples_vid/step1+2_outputs", help="Intermediate dataset output path")
204 |     parser.add_argument("--model", help="restore RAFT checkpoint",default="weights/raft-things.pth")
205 |     parser.add_argument("--small", action="store_true", help="use small model")
206 |     parser.add_argument("--mixed_precision", action="store_true", help="use mixed precision")
207 |     parser.add_argument("--txt", type=str, default="examples_txt/flow_list.txt",help="txt file to save the flow list")
208 |     parser.add_argument("--nomove_threshold", type=float, default=8, help="number of frames with potentially no movement, smaller means more strict")
209 |     parser.add_argument("--little_movement_threshold", type=float, default=20, help="number of frames with potentially little movement, smaller means more strict")
210 |     args = parser.parse_args()
211 |     done = []
212 |     nomove=0
213 | 
214 |     flow_done_file = args.txt
215 | 
216 |     lockpath = flow_done_file + ".lock"
217 |     lock = FileLock(lockpath)
218 |     if os.path.exists(flow_done_file):
219 |         with lock:
220 |             with open (flow_done_file, 'r') as f:
221 |                 for line in f:
222 |                     done.append(line.strip().split(' >')[0])
223 |             f.close()
224 |     for dataset in microsoft_sorted(os.listdir(args.dataset_path)):
225 |         input_path = os.path.join(args.dataset_path, dataset, "images")
226 |         output_path = os.path.join(args.dataset_path, dataset, "flow")
227 |         output_img_path = os.path.join(args.dataset_path, dataset, "flow_png")
228 |         raw_files = sorted(glob.glob(input_path + '/*.jpg'), key=extract_number)
229 |         nomove = camera_movement(raw_files[:20])
230 |         #TODO:  problem with camera movement
231 |         if nomove >= args.nomove_threshold:
232 |             # with open("dynamic_outputs/dynamic_{}.txt".format(name), 'a') as f:
233 |             print("camera maybe not moving.")
234 |             with lock:
235 |                 with open(flow_done_file, "a") as f:
236 |                     f.write(dataset+" >>>nomove\n")
237 |                 f.close()
238 |             continue
239 |         create_dir(output_path)
240 |         create_dir(output_img_path)
241 |         run(args, input_path, output_path, output_img_path)
242 |         with lock:
243 |             with open(flow_done_file, "a") as f:
244 |                 f.write(dataset+" >>>done\n")
245 |                 f.close()
246 |     print(f"flow done. please see {flow_done_file}")
247 | 


--------------------------------------------------------------------------------
/WebVi3D/s3_generate_mask.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import glob
  3 | import os
  4 | 
  5 | import cv2
  6 | import numpy as np
  7 | import skimage.morphology
  8 | import torch
  9 | from PIL import Image
 10 | from tqdm import tqdm
 11 | try:
 12 |     from filelock import FileLock
 13 | except ImportError:
 14 |     os.system("pip install filelock")
 15 |     from filelock import FileLock
 16 | DEVICE = "cpu"
 17 | 
 18 | 
 19 | import threading
 20 | def create_dir(dir):
 21 |     if not os.path.exists(dir):
 22 |         os.makedirs(dir)
 23 | 
 24 | 
 25 | def load_image(imfile):
 26 |     img = np.array(Image.open(imfile)).astype(np.uint8)[..., :3]
 27 |     img = torch.from_numpy(img).permute(2, 0, 1).float()
 28 |     return img[None].to(DEVICE)
 29 | 
 30 | 
 31 | def get_uv_grid(H, W, homo=False, align_corners=False, device=None):
 32 |     """
 33 |     Get uv grid renormalized from -1 to 1
 34 |     :returns (H, W, 2) tensor
 35 |     """
 36 |     if device is None:
 37 |         device = torch.device("cpu")
 38 |     yy, xx = torch.meshgrid(
 39 |         torch.arange(H, dtype=torch.float32, device=device),
 40 |         torch.arange(W, dtype=torch.float32, device=device),
 41 |         indexing="ij",
 42 |     )
 43 |     if align_corners:
 44 |         xx = 2 * xx / (W - 1) - 1
 45 |         yy = 2 * yy / (H - 1) - 1
 46 |     else:
 47 |         xx = 2 * (xx + 0.5) / W - 1
 48 |         yy = 2 * (yy + 0.5) / H - 1
 49 |     if homo:
 50 |         return torch.stack([xx, yy, torch.ones_like(xx)], dim=-1)
 51 |     return torch.stack([xx, yy], dim=-1)
 52 | 
 53 | 
 54 | def compute_sampson_error(x1, x2, F):
 55 |     """
 56 |     :param x1 (*, N, 2)
 57 |     :param x2 (*, N, 2)
 58 |     :param F (*, 3, 3)
 59 |     """
 60 |     h1 = torch.cat([x1, torch.ones_like(x1[..., :1])], dim=-1)  # N 3
 61 |     h2 = torch.cat([x2, torch.ones_like(x2[..., :1])], dim=-1)  # N 3
 62 |     d1 = torch.matmul(h1, F.transpose(-1, -2))  # (B, N, 3)
 63 |     d2 = torch.matmul(h2, F)  # (B, N, 3)
 64 |     z = (h2 * d1).sum(dim=-1)  # (B, N)
 65 |     err = z**2 / (
 66 |         d1[..., 0] ** 2 + d1[..., 1] ** 2 + d2[..., 0] ** 2 + d2[..., 1] ** 2
 67 |     )
 68 |     return err
 69 | 
 70 | 
 71 | def warp_flow(img, flow):
 72 |     h, w = flow.shape[:2]
 73 |     flow_new = flow.copy()
 74 |     flow_new[:, :, 0] += np.arange(w)
 75 |     flow_new[:, :, 1] += np.arange(h)[:, np.newaxis]
 76 | 
 77 |     res = cv2.remap(
 78 |         img, flow_new, None, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT
 79 |     )
 80 |     return res
 81 | 
 82 | 
 83 | def get_stats(X, norm=2):
 84 |     """
 85 |     :param X (N, C, H, W)
 86 |     :returns mean (1, C, 1, 1), scale (1)
 87 |     """
 88 |     mean = X.mean(dim=(0, 2, 3), keepdim=True)  # (1, C, 1, 1)
 89 |     if norm == 1:
 90 |         mag = torch.abs(X - mean).sum(dim=1)  # (N, H, W)
 91 |     else:
 92 |         mag = np.sqrt(2) * torch.sqrt(torch.square(X - mean).sum(dim=1))  # (N, H, W)
 93 |     scale = mag.mean() + 1e-6
 94 |     return mean, scale
 95 | 
 96 | def extract_number(filename):
 97 |     return int(''.join(filter(str.isdigit, filename)))
 98 | 
 99 | if __name__ == "__main__":
100 |     parser = argparse.ArgumentParser()
101 |     parser.add_argument("--dataset_path", type=str, default="examples_vid/step1+2_outputs",help="Dataset path")
102 |     parser.add_argument("--txt", type=str, default="examples_txt/mask_list.txt", help="txt file to store processed video names,use it for the following steps.")
103 |     args = parser.parse_args()
104 |     done = []
105 | 
106 |     mask_done_file = args.txt
107 |     lockpath = mask_done_file + ".lock"
108 |     lock = FileLock(lockpath)
109 |     if os.path.exists(mask_done_file):
110 |         with lock:
111 |             done = set()
112 |             with open(mask_done_file, 'r') as f:
113 |                 for line in f:
114 |                     done.add(line.strip().split(' >')[0])
115 |     for data in os.listdir(args.dataset_path):
116 |         if data in done:
117 |             continue
118 |         data_dir = os.path.join(args.dataset_path, data)
119 |         if not os.path.exists(os.path.join(data_dir, 'flow')):
120 |             print("[WARNING] no flow detected, aborted", os.path.join(data_dir, 'flow'))
121 |             with lock:
122 |                 with open(mask_done_file, "a") as f:
123 |                     f.write(data+" >>> noflow\n")
124 |                     f.close()
125 |                     continue
126 |         with lock:
127 |             with open(mask_done_file, "a") as f:
128 |                 f.write(data+" >>>done\n")
129 |                 f.close()
130 | 
131 |         images = sorted(glob.glob(os.path.join(data_dir, "images", "*.jpg")), key=extract_number)[:200]
132 | 
133 |         img = load_image(images[0])
134 |         H = img.shape[2]
135 |         W = img.shape[3]
136 |         uv = get_uv_grid(H, W, align_corners=False)
137 |         x1 = uv.reshape(-1, 2)  # N, 2
138 |         motion_mask_frames = []
139 |         motion_mask_frames2 = []
140 |         flow_for_bilateral = []
141 | 
142 |         with tqdm(images) as pbar:
143 |             for idx, _ in enumerate(images):
144 |                 # print("idx: " + str(idx))
145 |                 pbar.set_description("Processing {}".format(idx))
146 |                 motion_masks = []
147 |                 weights = []
148 |                 err_list = []
149 |                 normalized_flow = []
150 |                 this_flow = 0
151 |                 counter = 0
152 |                 for step in [1]:
153 |                     # print("step: " + str(step))
154 |                     if idx - step >= 0:
155 |                         # backward flow and mask
156 |                         bwd_flow_path = os.path.join(
157 |                             data_dir, "flow", str(idx).zfill(5) + "_bwd.npz"
158 |                         )
159 |                         if not os.path.exists(bwd_flow_path):
160 |                             print("Flow not preprared, skipping "+ data_dir)
161 |                             continue
162 |                         bwd_data = np.load(bwd_flow_path)
163 |                         bwd_flow, bwd_mask = bwd_data["flow"], bwd_data["mask"]
164 |                         this_flow = np.copy(this_flow - bwd_flow)
165 |                         counter += 1
166 |                         bwd_flow = torch.from_numpy(bwd_flow)
167 |                         bwd_mask = np.float32(bwd_mask)
168 |                         bwd_mask = torch.from_numpy(bwd_mask)
169 |                         flow = torch.from_numpy(
170 |                             np.stack(
171 |                                 [
172 |                                     2.0 * bwd_flow[..., 0] / (W - 1),  # normalized to -1,1
173 |                                     2.0 * bwd_flow[..., 1] / (H - 1),
174 |                                 ],
175 |                                 axis=-1,
176 |                             )
177 |                         )
178 |                         normalized_flow.append(flow)
179 |                         x2 = x1 + flow.view(-1, 2)  # (H*W, 2),
180 |                         randid = [np.random.randint(0, x1.shape[0]) for _ in range(500)]
181 |                         F, mask = cv2.findFundamentalMat(x1.numpy()[randid], x2.numpy()[randid], cv2.FM_LMEDS) 
182 |                         F = torch.from_numpy(F.astype(np.float32))  # (3, 3)
183 |                         err = compute_sampson_error(x1, x2, F).reshape(H, W)
184 |                         fac = (H + W) / 2
185 |                         err = err * fac ** 2
186 |                         err_list.append(err)
187 |                         weights.append(bwd_mask.mean())
188 | 
189 |                     if idx + step < len(images):
190 |                         # forward flow and mask
191 |                         fwd_flow_path = os.path.join(
192 |                             data_dir, "flow", str(idx).zfill(5) + "_fwd.npz"
193 |                         )
194 |                         if not os.path.exists(fwd_flow_path):
195 |                             print("Flow not preprared, skipping "+ data_dir)
196 |                             continue
197 |                         fwd_data = np.load(fwd_flow_path)
198 |                         fwd_flow, fwd_mask = fwd_data["flow"], fwd_data["mask"]
199 |                         this_flow = np.copy(this_flow + fwd_flow)
200 |                         counter += 1
201 |                         fwd_flow = torch.from_numpy(fwd_flow)
202 |                         fwd_mask = np.float32(fwd_mask)
203 |                         fwd_mask = torch.from_numpy(fwd_mask)
204 |                         flow = torch.from_numpy(
205 |                             np.stack(
206 |                                 [
207 |                                     2.0 * fwd_flow[..., 0] / (W - 1),
208 |                                     2.0 * fwd_flow[..., 1] / (H - 1),
209 |                                 ],
210 |                                 axis=-1,
211 |                             )
212 |                         )
213 |                         normalized_flow.append(flow)
214 |                         x2 = x1 + flow.view(-1, 2)  # (H*W, 2)
215 |                         randid = [np.random.randint(0, x1.shape[0]) for _ in range(500)]
216 |                         F, mask = cv2.findFundamentalMat(x1.numpy()[randid], x2.numpy()[randid], cv2.FM_LMEDS)
217 |                         F = torch.from_numpy(F.astype(np.float32))  # (3, 3)
218 |                         err = compute_sampson_error(x1, x2, F).reshape(H, W)
219 |                         fac = (H + W) / 2
220 |                         err = err * fac**2
221 |                         err_list.append(err)
222 |                         weights.append(fwd_mask.mean())
223 |                 if len(err_list) > 0:
224 |                     err = torch.amax(torch.stack(err_list, 0), 0)
225 |                     flow_for_bilateral.append(this_flow / counter)
226 | 
227 |                     thresh = torch.quantile(err, 0.8)   # threshold
228 |                     err = torch.where(err <= thresh, torch.zeros_like(err), err)
229 | 
230 |                     mask = torch.from_numpy(
231 |                         skimage.morphology.binary_opening(
232 |                             err.numpy() > (H * W / (8100.0)), skimage.morphology.disk(1)
233 |                         )
234 |                     )  
235 |                     mask = torch.from_numpy(
236 |                         skimage.morphology.dilation(mask.cpu(), skimage.morphology.disk(2))
237 |                     ).float()
238 |                     if not os.path.exists(os.path.join(data_dir, "dynamic_mask")):
239 |                         os.makedirs(os.path.join(data_dir, "dynamic_mask"))
240 |                     Image.fromarray(((mask).numpy() * 255.0).astype(np.uint8)).save(
241 |                         os.path.join(data_dir, "dynamic_mask", str(idx).zfill(5) + ".png")
242 |                     )
243 | 
244 |                     pbar.update()
245 |     print(f"mask done. please see {mask_done_file}")
246 | 


--------------------------------------------------------------------------------
/WebVi3D/s4_filter.py:
--------------------------------------------------------------------------------
  1 | import glob
  2 | from PIL import Image
  3 | from PIL import ImageFile
  4 | from matplotlib import pyplot as plt
  5 | from tqdm import tqdm
  6 | import os
  7 | import numpy as np
  8 | import torchvision
  9 | import torch
 10 | import cv2
 11 | from lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED, DoGHardNet
 12 | 
 13 | from einops import rearrange
 14 | 
 15 | def extract_number(filename):
 16 |     return int(''.join(filter(str.isdigit, filename)))
 17 | 
 18 | def run_maskrcnn(model, img_path):
 19 |     threshold = 0.5
 20 |     cnt = 0
 21 |     length = len(img_path)
 22 |     sample = [0, 1, 2, 3, 4, length//8, length//4, length//4 + 1, length//2, length//2 +1, length//2 + 2, length//4*3, length-1]
 23 |     for index in sample:
 24 |         img = img_path[index]
 25 |         o_image = Image.open(img).convert("RGB")
 26 |         width, height = o_image.size
 27 |         if width > height:
 28 |             intHeight = 576
 29 |             intWidth = 1024
 30 |         else:
 31 |             intHeight = 1024
 32 |             intWidth = 576
 33 | 
 34 |         image = o_image.resize((intWidth, intHeight), Image.LANCZOS)
 35 | 
 36 |         image_tensor = torchvision.transforms.functional.to_tensor(image).cuda()
 37 | 
 38 |         objPredictions = model([image_tensor])[0]
 39 | 
 40 |         for intMask in range(objPredictions["masks"].size(0)):
 41 |             if objPredictions["scores"][intMask].item() > threshold:
 42 |                 # person, vehicle, accessory, animal, sports
 43 |                 if objPredictions["labels"][intMask].item() == 1:  # person
 44 |                     cnt += 1
 45 |     return cnt
 46 | 
 47 | def white_count(img):
 48 |     pixel_count = img.size
 49 |     white_pixel_count = np.count_nonzero(img == 255)
 50 |     proportion = white_pixel_count / pixel_count
 51 |     return proportion
 52 | 
 53 | def central_count(img):
 54 | 
 55 |     image_width, image_height = img.shape[1], img.shape[0]
 56 | 
 57 |     central_width_start = int(image_width * 0.5)
 58 |     central_width_end = int(image_width * 0.75)
 59 |     central_height_start = int(image_height * 0.5)
 60 |     central_height_end = int(image_height * 0.75)
 61 | 
 62 |     central_region = img[central_height_start:central_height_end, central_width_start:central_width_end]
 63 | 
 64 |     white_pixel_count_central = np.count_nonzero(central_region == 255)
 65 | 
 66 |     proportion_central = white_pixel_count_central / central_region.size
 67 |     return proportion_central
 68 | 
 69 | 
 70 | 
 71 | def check_flow(folder_path):
 72 |     flow = sorted(glob.glob(os.path.join(folder_path, 'flow') + '/*bwd.npz'))
 73 |     count = 0
 74 |     for f in flow[:20]:
 75 |         npz = np.load(f)
 76 |         flow = npz['flow']
 77 |         move = abs(flow)[...,0]+ abs(flow)[...,1]
 78 |         if move.mean() <= 0.55:
 79 |             count += 1
 80 |     return count
 81 | 
 82 | def cal_oclidean_distance(coord1, coord2):  # N,2
 83 |     x2s,y2s = coord2[:,0], coord2[:,1]
 84 |     x1s,y1s = coord1[:,0], coord1[:,1]
 85 |     return torch.sqrt((x2s - x1s)**2 + (y2s - y1s)**2)
 86 |     
 87 | 
 88 | def dist(p1, p2):
 89 |     return np.linalg.norm(p1 - p2)
 90 | 
 91 | def circle_from_three_points(p1, p2, p3):
 92 |     A = 2 * (p2[0] - p1[0])
 93 |     B = 2 * (p2[1] - p1[1])
 94 |     C = 2 * (p3[0] - p1[0])
 95 |     D = 2 * (p3[1] - p1[1])
 96 |     E = p2[0]**2 - p1[0]**2 + p2[1]**2 - p1[1]**2
 97 |     F = p3[0]**2 - p1[0]**2 + p3[1]**2 - p1[1]**2
 98 |     G = A * D - B * C
 99 | 
100 |     if G == 0:
101 |         return None
102 | 
103 |     center_x = (D * E - B * F) / G
104 |     center_y = (A * F - C * E) / G
105 |     center = np.array([center_x, center_y])
106 |     radius = dist(center, p1)
107 |     return center, radius
108 | 
109 | def distances_to_center(points, center):
110 |     return np.linalg.norm(points - center, axis=1)
111 | 
112 | def ransac_circle(points, iterations=1000, min_inlier_ratio=0.5):
113 |     best_inliers = []
114 |     best_circle = None
115 |     n_points = points.shape[0]
116 |     min_inliers = int(n_points * min_inlier_ratio)
117 |     mini_radius = 1e9
118 |     
119 |     for _ in range(iterations):
120 |         try:
121 |             sample_points = points[np.random.choice(points.shape[0], 3, replace=False)]
122 |         except ValueError:
123 |             continue
124 |         circle = circle_from_three_points(sample_points[0], sample_points[1], sample_points[2])
125 |         
126 |         if circle is None:
127 |             continue
128 |         
129 |         center, radius = circle
130 |         distances = distances_to_center(points, center)
131 |         inliers = points[distances <= radius]
132 |         
133 |         if len(inliers) >= min_inliers and mini_radius > radius:
134 |             mini_radius = radius
135 |             best_inliers = inliers
136 |             best_circle = (center, radius)
137 |     if best_circle:
138 |         center, radius = best_circle
139 |         return center, radius, best_inliers
140 |     else:
141 |         return None
142 | 
143 | def main_func(points):
144 |     # breakpoint()
145 |     result = ransac_circle(points, iterations=500, min_inlier_ratio=0.6)
146 |     if result:
147 |         center, radius, inliers = result
148 |         return center, radius, inliers
149 |     else:
150 |         # print("No circle found with the required inlier ratio")
151 |         return -1,-1,-1
152 | def track_detect(img_path, extractor, cotracker):
153 |     frames = []
154 |     for img in img_path[::4]:
155 |         frames.append(cv2.cvtColor(cv2.imread(img), cv2.COLOR_BGR2RGB))
156 |     frames = np.array(frames)
157 |     device = 'cuda'
158 |     # breakpoint()
159 |     video = torch.tensor(frames).permute(0, 3, 1, 2)[None].float().to(device)  # B T C H W
160 |     feats0 = extractor.extract(rearrange(torch.tensor(np.array(frames[0])).cuda(), 'H W C -> C H W') / 255.)
161 |     keypoints = feats0['keypoints'].int()
162 |     zeros = torch.zeros_like(keypoints[...,0])
163 |     keypoints = torch.cat([zeros.unsqueeze(-1), keypoints], dim=-1).float()
164 |     pred_tracks, pred_visibility = cotracker(video, queries=keypoints, backward_tracking=True) # B T N 2,  B T N 1
165 |     radius = []
166 |     for i in range(pred_tracks.shape[2]):   # B T N 2
167 |         tall = pred_tracks[0,:,i,:]
168 |         mask = pred_visibility[0,:,i]  # i point, all time steps
169 |         tall = tall[mask == 1]
170 |         center, rad, inliers = main_func(tall.cpu().numpy())
171 |         radius.append(rad)
172 |     radius = sorted(torch.tensor(radius))
173 |     count = 0
174 |     for i in range(len(radius)):
175 |         if radius[i] <= 20:
176 |             count += 1
177 |     try:
178 |         mean = torch.tensor(radius).mean()
179 |     except RuntimeError:
180 |         mean = 0
181 |     return count, mean
182 | 
183 | 
184 | 
185 | def check_dynamic(folder_path, output_txt, extractor=None,cotracker=None):
186 | 
187 | 
188 |     if os.path.exists(os.path.join(folder_path,'epipolar_error_png')):  # initial version
189 |         png_files = glob.glob(folder_path + '/epipolar_error_png/*.png')[:200]
190 |     else:
191 |         try:
192 |             png_files = glob.glob(folder_path + '/dynamic_mask/*.png')[:200]
193 |         except:
194 |             print(f"Error with {folder_path}")
195 |             return
196 |     vid_name = folder_path.split('/')[-1]
197 |     raw_files = sorted(glob.glob(folder_path + '/images/*.jpg'),key=extract_number)[:200]
198 | 
199 |     num_imgs = len(png_files)
200 |     certain = 0     
201 |     has_mask = 0
202 |     # for file in png_files:
203 |     centralsave = 0
204 |     manysave = 0
205 |     for id in sorted(png_files):
206 | 
207 |         many_motion = 0     
208 |         central_motion = 0  
209 | 
210 |         img = Image.open(id)
211 |         try:
212 |             img = np.array(img)
213 |         except:
214 |             continue
215 | 
216 |         many_motion = white_count(img)
217 |         central_motion = central_count(img)
218 |         # print(many_motion, central_motion)
219 |         if many_motion > 0.0001:
220 |             has_mask += 1
221 |         if many_motion >=0.12 and central_motion >= 0.35:
222 |             certain += 2
223 |         elif many_motion >=0.12 and central_motion < 0.2:
224 |             certain += 0.8
225 |         elif many_motion >= 0.05 and central_motion >= 0.2:  
226 |             certain += 1.5
227 |         elif many_motion <0.05 and central_motion >= 0.2:    
228 |             certain += 1
229 |         elif many_motion >= 0.05 and central_motion < 0.2:   
230 |             certain += 0.5
231 |         centralsave += central_motion
232 |         manysave += many_motion
233 | 
234 |     count, mean = track_detect(raw_files, extractor, cotracker)
235 | 
236 |     if has_mask >= 0.6 * num_imgs:  
237 |         print(vid_name+"too much, dynamic")
238 |         # f.write(vid_name + '>>{}/{}_too much \n'.format(has_mask,num_imgs))  
239 |         return
240 |     if count >= 40 and mean <= 10:  # adjustable, count: higher more strict, mean: lower more strict
241 |         print(vid_name+"too small movements")
242 |         return
243 |     else:
244 |         with open(output_txt, 'a') as f:
245 |             f.write(vid_name + ' >\n')
246 |         print(vid_name+"static,saved")
247 | 
248 | if __name__ == "__main__":
249 |     import argparse
250 |     args = argparse.ArgumentParser()
251 |     args.add_argument("--mask_done_file", "-f", type=str, help="Folder path",default="examples_txt/mask_list.txt")
252 |     args.add_argument("--dataset_path" , "-v", type=str,default="examples_vid/step1+2_outputs",help="Processed video clips, not the origin videos.")    # 
253 |     args.add_argument("--txt", default="examples_txt/final.txt", type=str)
254 |     args.add_argument("--gpu", type=str, default="0")
255 |     args = args.parse_args()
256 |     now = []
257 |     output_txt = args.txt
258 |     extractor = SuperPoint(max_num_keypoints=100).eval().cuda()  # load the extractor
259 |     cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker2").to('cuda')
260 | 
261 |     with open (args.mask_done_file, 'r') as f:
262 |         for line in f:
263 |             try:
264 |                 if line.strip().split(' >')[1].endswith('done'):
265 |                     now.append(line.strip().split(' >')[0])
266 |             except:
267 |                 pass
268 |     if os.path.exists(output_txt):
269 |         done = []
270 |         with open(output_txt, 'r') as f:
271 |             for line in f:
272 |                 done.append(line.strip().split(' >')[0])
273 |         now = list(set(now) - set(done))
274 |     for idx in tqdm(list(set(now))):
275 |         vid_path = os.path.join(args.dataset_path, idx)
276 |         check_dynamic(vid_path, output_txt=output_txt, extractor=extractor, cotracker=cotracker)
277 |     print(f"filter done, please see {output_txt}")


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/__init__.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
2 | 
3 | # from .demo import RAFT_infer
4 | from .raft import RAFT
5 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/corr.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import torch
  4 | import torch.nn.functional as F
  5 | from .utils.utils import bilinear_sampler, coords_grid
  6 | 
  7 | try:
  8 |     import alt_cuda_corr
  9 | except:
 10 |     # alt_cuda_corr is not compiled
 11 |     pass
 12 | 
 13 | 
 14 | class CorrBlock:
 15 |     def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
 16 |         self.num_levels = num_levels
 17 |         self.radius = radius
 18 |         self.corr_pyramid = []
 19 | 
 20 |         # all pairs correlation
 21 |         corr = CorrBlock.corr(fmap1, fmap2)
 22 | 
 23 |         batch, h1, w1, dim, h2, w2 = corr.shape
 24 |         corr = corr.reshape(batch*h1*w1, dim, h2, w2)
 25 | 
 26 |         self.corr_pyramid.append(corr)
 27 |         for i in range(self.num_levels-1):
 28 |             corr = F.avg_pool2d(corr, 2, stride=2)
 29 |             self.corr_pyramid.append(corr)
 30 | 
 31 |     def __call__(self, coords):
 32 |         r = self.radius
 33 |         coords = coords.permute(0, 2, 3, 1)
 34 |         batch, h1, w1, _ = coords.shape
 35 | 
 36 |         out_pyramid = []
 37 |         for i in range(self.num_levels):
 38 |             corr = self.corr_pyramid[i]
 39 |             dx = torch.linspace(-r, r, 2*r+1)
 40 |             dy = torch.linspace(-r, r, 2*r+1)
 41 |             delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)
 42 | 
 43 |             centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
 44 |             delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
 45 |             coords_lvl = centroid_lvl + delta_lvl
 46 | 
 47 |             corr = bilinear_sampler(corr, coords_lvl)
 48 |             corr = corr.view(batch, h1, w1, -1)
 49 |             out_pyramid.append(corr)
 50 | 
 51 |         out = torch.cat(out_pyramid, dim=-1)
 52 |         return out.permute(0, 3, 1, 2).contiguous().float()
 53 | 
 54 |     @staticmethod
 55 |     def corr(fmap1, fmap2):
 56 |         batch, dim, ht, wd = fmap1.shape
 57 |         fmap1 = fmap1.view(batch, dim, ht*wd)
 58 |         fmap2 = fmap2.view(batch, dim, ht*wd)
 59 | 
 60 |         corr = torch.matmul(fmap1.transpose(1,2), fmap2)
 61 |         corr = corr.view(batch, ht, wd, 1, ht, wd)
 62 |         return corr  / torch.sqrt(torch.tensor(dim).float())
 63 | 
 64 | 
 65 | class CorrLayer(torch.autograd.Function):
 66 |     @staticmethod
 67 |     def forward(ctx, fmap1, fmap2, coords, r):
 68 |         fmap1 = fmap1.contiguous()
 69 |         fmap2 = fmap2.contiguous()
 70 |         coords = coords.contiguous()
 71 |         ctx.save_for_backward(fmap1, fmap2, coords)
 72 |         ctx.r = r
 73 |         corr, = correlation_cudaz.forward(fmap1, fmap2, coords, ctx.r)
 74 |         return corr
 75 | 
 76 |     @staticmethod
 77 |     def backward(ctx, grad_corr):
 78 |         fmap1, fmap2, coords = ctx.saved_tensors
 79 |         grad_corr = grad_corr.contiguous()
 80 |         fmap1_grad, fmap2_grad, coords_grad = \
 81 |             correlation_cudaz.backward(fmap1, fmap2, coords, grad_corr, ctx.r)
 82 |         return fmap1_grad, fmap2_grad, coords_grad, None
 83 | 
 84 | 
 85 | class AlternateCorrBlock:
 86 |     def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
 87 |         self.num_levels = num_levels
 88 |         self.radius = radius
 89 | 
 90 |         self.pyramid = [(fmap1, fmap2)]
 91 |         for i in range(self.num_levels):
 92 |             fmap1 = F.avg_pool2d(fmap1, 2, stride=2)
 93 |             fmap2 = F.avg_pool2d(fmap2, 2, stride=2)
 94 |             self.pyramid.append((fmap1, fmap2))
 95 | 
 96 |     def __call__(self, coords):
 97 | 
 98 |         coords = coords.permute(0, 2, 3, 1)
 99 |         B, H, W, _ = coords.shape
100 | 
101 |         corr_list = []
102 |         for i in range(self.num_levels):
103 |             r = self.radius
104 |             fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1)
105 |             fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1)
106 | 
107 |             coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()
108 |             corr = alt_cuda_corr(fmap1_i, fmap2_i, coords_i, r)
109 |             corr_list.append(corr.squeeze(1))
110 | 
111 |         corr = torch.stack(corr_list, dim=1)
112 |         corr = corr.reshape(B, -1, H, W)
113 |         return corr / 16.0
114 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/datasets.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | # Data loading based on https://github.com/NVIDIA/flownet2-pytorch
  4 | 
  5 | import numpy as np
  6 | import torch
  7 | import torch.utils.data as data
  8 | import torch.nn.functional as F
  9 | 
 10 | import os
 11 | import math
 12 | import random
 13 | from glob import glob
 14 | import os.path as osp
 15 | 
 16 | from utils import frame_utils
 17 | from utils.augmentor import FlowAugmentor, SparseFlowAugmentor
 18 | 
 19 | 
 20 | class FlowDataset(data.Dataset):
 21 |     def __init__(self, aug_params=None, sparse=False):
 22 |         self.augmentor = None
 23 |         self.sparse = sparse
 24 |         if aug_params is not None:
 25 |             if sparse:
 26 |                 self.augmentor = SparseFlowAugmentor(**aug_params)
 27 |             else:
 28 |                 self.augmentor = FlowAugmentor(**aug_params)
 29 | 
 30 |         self.is_test = False
 31 |         self.init_seed = False
 32 |         self.flow_list = []
 33 |         self.image_list = []
 34 |         self.extra_info = []
 35 | 
 36 |     def __getitem__(self, index):
 37 | 
 38 |         if self.is_test:
 39 |             img1 = frame_utils.read_gen(self.image_list[index][0])
 40 |             img2 = frame_utils.read_gen(self.image_list[index][1])
 41 |             img1 = np.array(img1).astype(np.uint8)[..., :3]
 42 |             img2 = np.array(img2).astype(np.uint8)[..., :3]
 43 |             img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
 44 |             img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
 45 |             return img1, img2, self.extra_info[index]
 46 | 
 47 |         if not self.init_seed:
 48 |             worker_info = torch.utils.data.get_worker_info()
 49 |             if worker_info is not None:
 50 |                 torch.manual_seed(worker_info.id)
 51 |                 np.random.seed(worker_info.id)
 52 |                 random.seed(worker_info.id)
 53 |                 self.init_seed = True
 54 | 
 55 |         index = index % len(self.image_list)
 56 |         valid = None
 57 |         if self.sparse:
 58 |             flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])
 59 |         else:
 60 |             flow = frame_utils.read_gen(self.flow_list[index])
 61 | 
 62 |         img1 = frame_utils.read_gen(self.image_list[index][0])
 63 |         img2 = frame_utils.read_gen(self.image_list[index][1])
 64 | 
 65 |         flow = np.array(flow).astype(np.float32)
 66 |         img1 = np.array(img1).astype(np.uint8)
 67 |         img2 = np.array(img2).astype(np.uint8)
 68 | 
 69 |         # grayscale images
 70 |         if len(img1.shape) == 2:
 71 |             img1 = np.tile(img1[...,None], (1, 1, 3))
 72 |             img2 = np.tile(img2[...,None], (1, 1, 3))
 73 |         else:
 74 |             img1 = img1[..., :3]
 75 |             img2 = img2[..., :3]
 76 | 
 77 |         if self.augmentor is not None:
 78 |             if self.sparse:
 79 |                 img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid)
 80 |             else:
 81 |                 img1, img2, flow = self.augmentor(img1, img2, flow)
 82 | 
 83 |         img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
 84 |         img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
 85 |         flow = torch.from_numpy(flow).permute(2, 0, 1).float()
 86 | 
 87 |         if valid is not None:
 88 |             valid = torch.from_numpy(valid)
 89 |         else:
 90 |             valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000)
 91 | 
 92 |         return img1, img2, flow, valid.float()
 93 | 
 94 | 
 95 |     def __rmul__(self, v):
 96 |         self.flow_list = v * self.flow_list
 97 |         self.image_list = v * self.image_list
 98 |         return self
 99 |         
100 |     def __len__(self):
101 |         return len(self.image_list)
102 |         
103 | 
104 | class MpiSintel(FlowDataset):
105 |     def __init__(self, aug_params=None, split='training', root='datasets/Sintel', dstype='clean'):
106 |         super(MpiSintel, self).__init__(aug_params)
107 |         flow_root = osp.join(root, split, 'flow')
108 |         image_root = osp.join(root, split, dstype)
109 | 
110 |         if split == 'test':
111 |             self.is_test = True
112 | 
113 |         for scene in os.listdir(image_root):
114 |             image_list = sorted(glob(osp.join(image_root, scene, '*.png')))
115 |             for i in range(len(image_list)-1):
116 |                 self.image_list += [ [image_list[i], image_list[i+1]] ]
117 |                 self.extra_info += [ (scene, i) ] # scene and frame_id
118 | 
119 |             if split != 'test':
120 |                 self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))
121 | 
122 | 
123 | class FlyingChairs(FlowDataset):
124 |     def __init__(self, aug_params=None, split='train', root='datasets/FlyingChairs_release/data'):
125 |         super(FlyingChairs, self).__init__(aug_params)
126 | 
127 |         images = sorted(glob(osp.join(root, '*.ppm')))
128 |         flows = sorted(glob(osp.join(root, '*.flo')))
129 |         assert (len(images)//2 == len(flows))
130 | 
131 |         split_list = np.loadtxt('chairs_split.txt', dtype=np.int32)
132 |         for i in range(len(flows)):
133 |             xid = split_list[i]
134 |             if (split=='training' and xid==1) or (split=='validation' and xid==2):
135 |                 self.flow_list += [ flows[i] ]
136 |                 self.image_list += [ [images[2*i], images[2*i+1]] ]
137 | 
138 | 
139 | class FlyingThings3D(FlowDataset):
140 |     def __init__(self, aug_params=None, root='datasets/FlyingThings3D', dstype='frames_cleanpass'):
141 |         super(FlyingThings3D, self).__init__(aug_params)
142 | 
143 |         for cam in ['left']:
144 |             for direction in ['into_future', 'into_past']:
145 |                 image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*')))
146 |                 image_dirs = sorted([osp.join(f, cam) for f in image_dirs])
147 | 
148 |                 flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))
149 |                 flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])
150 | 
151 |                 for idir, fdir in zip(image_dirs, flow_dirs):
152 |                     images = sorted(glob(osp.join(idir, '*.png')) )
153 |                     flows = sorted(glob(osp.join(fdir, '*.pfm')) )
154 |                     for i in range(len(flows)-1):
155 |                         if direction == 'into_future':
156 |                             self.image_list += [ [images[i], images[i+1]] ]
157 |                             self.flow_list += [ flows[i] ]
158 |                         elif direction == 'into_past':
159 |                             self.image_list += [ [images[i+1], images[i]] ]
160 |                             self.flow_list += [ flows[i+1] ]
161 |       
162 | 
163 | class KITTI(FlowDataset):
164 |     def __init__(self, aug_params=None, split='training', root='datasets/KITTI'):
165 |         super(KITTI, self).__init__(aug_params, sparse=True)
166 |         if split == 'testing':
167 |             self.is_test = True
168 | 
169 |         root = osp.join(root, split)
170 |         images1 = sorted(glob(osp.join(root, 'image_2/*_10.png')))
171 |         images2 = sorted(glob(osp.join(root, 'image_2/*_11.png')))
172 | 
173 |         for img1, img2 in zip(images1, images2):
174 |             frame_id = img1.split('/')[-1]
175 |             self.extra_info += [ [frame_id] ]
176 |             self.image_list += [ [img1, img2] ]
177 | 
178 |         if split == 'training':
179 |             self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))
180 | 
181 | 
182 | class HD1K(FlowDataset):
183 |     def __init__(self, aug_params=None, root='datasets/HD1k'):
184 |         super(HD1K, self).__init__(aug_params, sparse=True)
185 | 
186 |         seq_ix = 0
187 |         while 1:
188 |             flows = sorted(glob(os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix)))
189 |             images = sorted(glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix)))
190 | 
191 |             if len(flows) == 0:
192 |                 break
193 | 
194 |             for i in range(len(flows)-1):
195 |                 self.flow_list += [flows[i]]
196 |                 self.image_list += [ [images[i], images[i+1]] ]
197 | 
198 |             seq_ix += 1
199 | 
200 | 
201 | def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):
202 |     """ Create the data loader for the corresponding trainign set """
203 | 
204 |     if args.stage == 'chairs':
205 |         aug_params = {'crop_size': args.image_size, 'min_scale': -0.1, 'max_scale': 1.0, 'do_flip': True}
206 |         train_dataset = FlyingChairs(aug_params, split='training')
207 |     
208 |     elif args.stage == 'things':
209 |         aug_params = {'crop_size': args.image_size, 'min_scale': -0.4, 'max_scale': 0.8, 'do_flip': True}
210 |         clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass')
211 |         final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass')
212 |         train_dataset = clean_dataset + final_dataset
213 | 
214 |     elif args.stage == 'sintel':
215 |         aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.6, 'do_flip': True}
216 |         things = FlyingThings3D(aug_params, dstype='frames_cleanpass')
217 |         sintel_clean = MpiSintel(aug_params, split='training', dstype='clean')
218 |         sintel_final = MpiSintel(aug_params, split='training', dstype='final')        
219 | 
220 |         if TRAIN_DS == 'C+T+K+S+H':
221 |             kitti = KITTI({'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.5, 'do_flip': True})
222 |             hd1k = HD1K({'crop_size': args.image_size, 'min_scale': -0.5, 'max_scale': 0.2, 'do_flip': True})
223 |             train_dataset = 100*sintel_clean + 100*sintel_final + 200*kitti + 5*hd1k + things
224 | 
225 |         elif TRAIN_DS == 'C+T+K/S':
226 |             train_dataset = 100*sintel_clean + 100*sintel_final + things
227 | 
228 |     elif args.stage == 'kitti':
229 |         aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.4, 'do_flip': False}
230 |         train_dataset = KITTI(aug_params, split='training')
231 | 
232 |     train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, 
233 |         pin_memory=False, shuffle=True, num_workers=4, drop_last=True)
234 | 
235 |     print('Training with %d image pairs' % len(train_dataset))
236 |     return train_loader
237 | 
238 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/demo.py:
--------------------------------------------------------------------------------
 1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
 2 | 
 3 | import sys
 4 | import argparse
 5 | import os
 6 | import cv2
 7 | import glob
 8 | import numpy as np
 9 | import torch
10 | from PIL import Image
11 | 
12 | from .raft import RAFT
13 | from .utils import flow_viz
14 | from .utils.utils import InputPadder
15 | 
16 | 
17 | 
18 | DEVICE = 'cuda'
19 | 
20 | def load_image(imfile):
21 |     img = np.array(Image.open(imfile)).astype(np.uint8)
22 |     img = torch.from_numpy(img).permute(2, 0, 1).float()
23 |     return img
24 | 
25 | 
26 | def load_image_list(image_files):
27 |     images = []
28 |     for imfile in sorted(image_files):
29 |         images.append(load_image(imfile))
30 | 
31 |     images = torch.stack(images, dim=0)
32 |     images = images.to(DEVICE)
33 | 
34 |     padder = InputPadder(images.shape)
35 |     return padder.pad(images)[0]
36 | 
37 | 
38 | def viz(img, flo):
39 |     img = img[0].permute(1,2,0).cpu().numpy()
40 |     flo = flo[0].permute(1,2,0).cpu().numpy()
41 | 
42 |     # map flow to rgb image
43 |     flo = flow_viz.flow_to_image(flo)
44 |     # img_flo = np.concatenate([img, flo], axis=0)
45 |     img_flo = flo
46 | 
47 |     cv2.imwrite('/home/chengao/test/flow.png', img_flo[:, :, [2,1,0]])
48 |     # cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
49 |     # cv2.waitKey()
50 | 
51 | 
52 | def demo(args):
53 |     model = torch.nn.DataParallel(RAFT(args))
54 |     model.load_state_dict(torch.load(args.model))
55 | 
56 |     model = model.module
57 |     model.to(DEVICE)
58 |     model.eval()
59 | 
60 |     with torch.no_grad():
61 |         images = glob.glob(os.path.join(args.path, '*.png')) + \
62 |                  glob.glob(os.path.join(args.path, '*.jpg'))
63 | 
64 |         images = load_image_list(images)
65 |         for i in range(images.shape[0]-1):
66 |             image1 = images[i,None]
67 |             image2 = images[i+1,None]
68 | 
69 |             flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
70 |             viz(image1, flow_up)
71 | 
72 | 
73 | def RAFT_infer(args):
74 |     model = torch.nn.DataParallel(RAFT(args))
75 |     model.load_state_dict(torch.load(args.model))
76 | 
77 |     model = model.module
78 |     model.to(DEVICE)
79 |     model.eval()
80 | 
81 |     return model
82 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/extractor.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import torch
  4 | import torch.nn as nn
  5 | import torch.nn.functional as F
  6 | 
  7 | 
  8 | class ResidualBlock(nn.Module):
  9 |     def __init__(self, in_planes, planes, norm_fn='group', stride=1):
 10 |         super(ResidualBlock, self).__init__()
 11 |   
 12 |         self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
 13 |         self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
 14 |         self.relu = nn.ReLU(inplace=True)
 15 | 
 16 |         num_groups = planes // 8
 17 | 
 18 |         if norm_fn == 'group':
 19 |             self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
 20 |             self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
 21 |             if not stride == 1:
 22 |                 self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
 23 |         
 24 |         elif norm_fn == 'batch':
 25 |             self.norm1 = nn.BatchNorm2d(planes)
 26 |             self.norm2 = nn.BatchNorm2d(planes)
 27 |             if not stride == 1:
 28 |                 self.norm3 = nn.BatchNorm2d(planes)
 29 |         
 30 |         elif norm_fn == 'instance':
 31 |             self.norm1 = nn.InstanceNorm2d(planes)
 32 |             self.norm2 = nn.InstanceNorm2d(planes)
 33 |             if not stride == 1:
 34 |                 self.norm3 = nn.InstanceNorm2d(planes)
 35 | 
 36 |         elif norm_fn == 'none':
 37 |             self.norm1 = nn.Sequential()
 38 |             self.norm2 = nn.Sequential()
 39 |             if not stride == 1:
 40 |                 self.norm3 = nn.Sequential()
 41 | 
 42 |         if stride == 1:
 43 |             self.downsample = None
 44 |         
 45 |         else:    
 46 |             self.downsample = nn.Sequential(
 47 |                 nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
 48 | 
 49 | 
 50 |     def forward(self, x):
 51 |         y = x
 52 |         y = self.relu(self.norm1(self.conv1(y)))
 53 |         y = self.relu(self.norm2(self.conv2(y)))
 54 | 
 55 |         if self.downsample is not None:
 56 |             x = self.downsample(x)
 57 | 
 58 |         return self.relu(x+y)
 59 | 
 60 | 
 61 | 
 62 | class BottleneckBlock(nn.Module):
 63 |     def __init__(self, in_planes, planes, norm_fn='group', stride=1):
 64 |         super(BottleneckBlock, self).__init__()
 65 |   
 66 |         self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
 67 |         self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
 68 |         self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
 69 |         self.relu = nn.ReLU(inplace=True)
 70 | 
 71 |         num_groups = planes // 8
 72 | 
 73 |         if norm_fn == 'group':
 74 |             self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
 75 |             self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
 76 |             self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
 77 |             if not stride == 1:
 78 |                 self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
 79 |         
 80 |         elif norm_fn == 'batch':
 81 |             self.norm1 = nn.BatchNorm2d(planes//4)
 82 |             self.norm2 = nn.BatchNorm2d(planes//4)
 83 |             self.norm3 = nn.BatchNorm2d(planes)
 84 |             if not stride == 1:
 85 |                 self.norm4 = nn.BatchNorm2d(planes)
 86 |         
 87 |         elif norm_fn == 'instance':
 88 |             self.norm1 = nn.InstanceNorm2d(planes//4)
 89 |             self.norm2 = nn.InstanceNorm2d(planes//4)
 90 |             self.norm3 = nn.InstanceNorm2d(planes)
 91 |             if not stride == 1:
 92 |                 self.norm4 = nn.InstanceNorm2d(planes)
 93 | 
 94 |         elif norm_fn == 'none':
 95 |             self.norm1 = nn.Sequential()
 96 |             self.norm2 = nn.Sequential()
 97 |             self.norm3 = nn.Sequential()
 98 |             if not stride == 1:
 99 |                 self.norm4 = nn.Sequential()
100 | 
101 |         if stride == 1:
102 |             self.downsample = None
103 |         
104 |         else:    
105 |             self.downsample = nn.Sequential(
106 |                 nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
107 | 
108 | 
109 |     def forward(self, x):
110 |         y = x
111 |         y = self.relu(self.norm1(self.conv1(y)))
112 |         y = self.relu(self.norm2(self.conv2(y)))
113 |         y = self.relu(self.norm3(self.conv3(y)))
114 | 
115 |         if self.downsample is not None:
116 |             x = self.downsample(x)
117 | 
118 |         return self.relu(x+y)
119 | 
120 | class BasicEncoder(nn.Module):
121 |     def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
122 |         super(BasicEncoder, self).__init__()
123 |         self.norm_fn = norm_fn
124 | 
125 |         if self.norm_fn == 'group':
126 |             self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
127 |             
128 |         elif self.norm_fn == 'batch':
129 |             self.norm1 = nn.BatchNorm2d(64)
130 | 
131 |         elif self.norm_fn == 'instance':
132 |             self.norm1 = nn.InstanceNorm2d(64)
133 | 
134 |         elif self.norm_fn == 'none':
135 |             self.norm1 = nn.Sequential()
136 | 
137 |         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
138 |         self.relu1 = nn.ReLU(inplace=True)
139 | 
140 |         self.in_planes = 64
141 |         self.layer1 = self._make_layer(64,  stride=1)
142 |         self.layer2 = self._make_layer(96, stride=2)
143 |         self.layer3 = self._make_layer(128, stride=2)
144 | 
145 |         # output convolution
146 |         self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
147 | 
148 |         self.dropout = None
149 |         if dropout > 0:
150 |             self.dropout = nn.Dropout2d(p=dropout)
151 | 
152 |         for m in self.modules():
153 |             if isinstance(m, nn.Conv2d):
154 |                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
155 |             elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
156 |                 if m.weight is not None:
157 |                     nn.init.constant_(m.weight, 1)
158 |                 if m.bias is not None:
159 |                     nn.init.constant_(m.bias, 0)
160 | 
161 |     def _make_layer(self, dim, stride=1):
162 |         layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
163 |         layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
164 |         layers = (layer1, layer2)
165 |         
166 |         self.in_planes = dim
167 |         return nn.Sequential(*layers)
168 | 
169 | 
170 |     def forward(self, x):
171 | 
172 |         # if input is list, combine batch dimension
173 |         is_list = isinstance(x, tuple) or isinstance(x, list)
174 |         if is_list:
175 |             batch_dim = x[0].shape[0]
176 |             x = torch.cat(x, dim=0)
177 | 
178 |         x = self.conv1(x)
179 |         x = self.norm1(x)
180 |         x = self.relu1(x)
181 | 
182 |         x = self.layer1(x)
183 |         x = self.layer2(x)
184 |         x = self.layer3(x)
185 | 
186 |         x = self.conv2(x)
187 | 
188 |         if self.training and self.dropout is not None:
189 |             x = self.dropout(x)
190 | 
191 |         if is_list:
192 |             x = torch.split(x, [batch_dim, batch_dim], dim=0)
193 | 
194 |         return x
195 | 
196 | 
197 | class SmallEncoder(nn.Module):
198 |     def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
199 |         super(SmallEncoder, self).__init__()
200 |         self.norm_fn = norm_fn
201 | 
202 |         if self.norm_fn == 'group':
203 |             self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)
204 |             
205 |         elif self.norm_fn == 'batch':
206 |             self.norm1 = nn.BatchNorm2d(32)
207 | 
208 |         elif self.norm_fn == 'instance':
209 |             self.norm1 = nn.InstanceNorm2d(32)
210 | 
211 |         elif self.norm_fn == 'none':
212 |             self.norm1 = nn.Sequential()
213 | 
214 |         self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)
215 |         self.relu1 = nn.ReLU(inplace=True)
216 | 
217 |         self.in_planes = 32
218 |         self.layer1 = self._make_layer(32,  stride=1)
219 |         self.layer2 = self._make_layer(64, stride=2)
220 |         self.layer3 = self._make_layer(96, stride=2)
221 | 
222 |         self.dropout = None
223 |         if dropout > 0:
224 |             self.dropout = nn.Dropout2d(p=dropout)
225 |         
226 |         self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)
227 | 
228 |         for m in self.modules():
229 |             if isinstance(m, nn.Conv2d):
230 |                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
231 |             elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
232 |                 if m.weight is not None:
233 |                     nn.init.constant_(m.weight, 1)
234 |                 if m.bias is not None:
235 |                     nn.init.constant_(m.bias, 0)
236 | 
237 |     def _make_layer(self, dim, stride=1):
238 |         layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)
239 |         layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)
240 |         layers = (layer1, layer2)
241 |     
242 |         self.in_planes = dim
243 |         return nn.Sequential(*layers)
244 | 
245 | 
246 |     def forward(self, x):
247 | 
248 |         # if input is list, combine batch dimension
249 |         is_list = isinstance(x, tuple) or isinstance(x, list)
250 |         if is_list:
251 |             batch_dim = x[0].shape[0]
252 |             x = torch.cat(x, dim=0)
253 | 
254 |         x = self.conv1(x)
255 |         x = self.norm1(x)
256 |         x = self.relu1(x)
257 | 
258 |         x = self.layer1(x)
259 |         x = self.layer2(x)
260 |         x = self.layer3(x)
261 |         x = self.conv2(x)
262 | 
263 |         if self.training and self.dropout is not None:
264 |             x = self.dropout(x)
265 | 
266 |         if is_list:
267 |             x = torch.split(x, [batch_dim, batch_dim], dim=0)
268 | 
269 |         return x
270 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/raft.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import numpy as np
  4 | import torch
  5 | import torch.nn as nn
  6 | import torch.nn.functional as F
  7 | 
  8 | from .update import BasicUpdateBlock, SmallUpdateBlock
  9 | from .extractor import BasicEncoder, SmallEncoder
 10 | from .corr import CorrBlock, AlternateCorrBlock
 11 | from .utils.utils import bilinear_sampler, coords_grid, upflow8
 12 | 
 13 | try:
 14 |     autocast = torch.cuda.amp.autocast
 15 | except:
 16 |     # dummy autocast for PyTorch < 1.6
 17 |     class autocast:
 18 |         def __init__(self, enabled):
 19 |             pass
 20 |         def __enter__(self):
 21 |             pass
 22 |         def __exit__(self, *args):
 23 |             pass
 24 | 
 25 | 
 26 | class RAFT(nn.Module):
 27 |     def __init__(self, args):
 28 |         super(RAFT, self).__init__()
 29 |         self.args = args
 30 | 
 31 |         if args.small:
 32 |             self.hidden_dim = hdim = 96
 33 |             self.context_dim = cdim = 64
 34 |             args.corr_levels = 4
 35 |             args.corr_radius = 3
 36 | 
 37 |         else:
 38 |             self.hidden_dim = hdim = 128
 39 |             self.context_dim = cdim = 128
 40 |             args.corr_levels = 4
 41 |             args.corr_radius = 4
 42 | 
 43 |         if 'dropout' not in args._get_kwargs():
 44 |             args.dropout = 0
 45 | 
 46 |         if 'alternate_corr' not in args._get_kwargs():
 47 |             args.alternate_corr = False
 48 | 
 49 |         # feature network, context network, and update block
 50 |         if args.small:
 51 |             self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout)
 52 |             self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout)
 53 |             self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
 54 | 
 55 |         else:
 56 |             self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout)
 57 |             self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
 58 |             self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
 59 | 
 60 | 
 61 |     def freeze_bn(self):
 62 |         for m in self.modules():
 63 |             if isinstance(m, nn.BatchNorm2d):
 64 |                 m.eval()
 65 | 
 66 |     def initialize_flow(self, img):
 67 |         """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
 68 |         N, C, H, W = img.shape
 69 |         coords0 = coords_grid(N, H//8, W//8).to(img.device)
 70 |         coords1 = coords_grid(N, H//8, W//8).to(img.device)
 71 | 
 72 |         # optical flow computed as difference: flow = coords1 - coords0
 73 |         return coords0, coords1
 74 | 
 75 |     def upsample_flow(self, flow, mask):
 76 |         """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
 77 |         N, _, H, W = flow.shape
 78 |         mask = mask.view(N, 1, 9, 8, 8, H, W)
 79 |         mask = torch.softmax(mask, dim=2)
 80 | 
 81 |         up_flow = F.unfold(8 * flow, [3,3], padding=1)
 82 |         up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
 83 | 
 84 |         up_flow = torch.sum(mask * up_flow, dim=2)
 85 |         up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
 86 |         return up_flow.reshape(N, 2, 8*H, 8*W)
 87 | 
 88 | 
 89 |     def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
 90 |         """ Estimate optical flow between pair of frames """
 91 | 
 92 |         image1 = 2 * (image1 / 255.0) - 1.0
 93 |         image2 = 2 * (image2 / 255.0) - 1.0
 94 | 
 95 |         image1 = image1.contiguous()
 96 |         image2 = image2.contiguous()
 97 | 
 98 |         hdim = self.hidden_dim
 99 |         cdim = self.context_dim
100 | 
101 |         # run the feature network
102 |         with autocast(enabled=self.args.mixed_precision):
103 |             fmap1, fmap2 = self.fnet([image1, image2])
104 | 
105 |         fmap1 = fmap1.float()
106 |         fmap2 = fmap2.float()
107 |         if self.args.alternate_corr:
108 |             corr_fn = CorrBlockAlternate(fmap1, fmap2, radius=self.args.corr_radius)
109 |         else:
110 |             corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
111 | 
112 |         # run the context network
113 |         with autocast(enabled=self.args.mixed_precision):
114 |             cnet = self.cnet(image1)
115 |             net, inp = torch.split(cnet, [hdim, cdim], dim=1)
116 |             net = torch.tanh(net)
117 |             inp = torch.relu(inp)
118 | 
119 |         coords0, coords1 = self.initialize_flow(image1)
120 | 
121 |         if flow_init is not None:
122 |             coords1 = coords1 + flow_init
123 | 
124 |         flow_predictions = []
125 |         for itr in range(iters):
126 |             coords1 = coords1.detach()
127 |             corr = corr_fn(coords1) # index correlation volume
128 | 
129 |             flow = coords1 - coords0
130 |             with autocast(enabled=self.args.mixed_precision):
131 |                 net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
132 | 
133 |             # F(t+1) = F(t) + \Delta(t)
134 |             coords1 = coords1 + delta_flow
135 | 
136 |             # upsample predictions
137 |             if up_mask is None:
138 |                 flow_up = upflow8(coords1 - coords0)
139 |             else:
140 |                 flow_up = self.upsample_flow(coords1 - coords0, up_mask)
141 | 
142 |             flow_predictions.append(flow_up)
143 | 
144 |         if test_mode:
145 |             return coords1 - coords0, flow_up
146 | 
147 |         return flow_predictions
148 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/update.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import torch
  4 | import torch.nn as nn
  5 | import torch.nn.functional as F
  6 | 
  7 | 
  8 | class FlowHead(nn.Module):
  9 |     def __init__(self, input_dim=128, hidden_dim=256):
 10 |         super(FlowHead, self).__init__()
 11 |         self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
 12 |         self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
 13 |         self.relu = nn.ReLU(inplace=True)
 14 | 
 15 |     def forward(self, x):
 16 |         return self.conv2(self.relu(self.conv1(x)))
 17 | 
 18 | class ConvGRU(nn.Module):
 19 |     def __init__(self, hidden_dim=128, input_dim=192+128):
 20 |         super(ConvGRU, self).__init__()
 21 |         self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
 22 |         self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
 23 |         self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
 24 | 
 25 |     def forward(self, h, x):
 26 |         hx = torch.cat([h, x], dim=1)
 27 | 
 28 |         z = torch.sigmoid(self.convz(hx))
 29 |         r = torch.sigmoid(self.convr(hx))
 30 |         q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
 31 | 
 32 |         h = (1-z) * h + z * q
 33 |         return h
 34 | 
 35 | class SepConvGRU(nn.Module):
 36 |     def __init__(self, hidden_dim=128, input_dim=192+128):
 37 |         super(SepConvGRU, self).__init__()
 38 |         self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
 39 |         self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
 40 |         self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
 41 | 
 42 |         self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
 43 |         self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
 44 |         self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
 45 | 
 46 | 
 47 |     def forward(self, h, x):
 48 |         # horizontal
 49 |         hx = torch.cat([h, x], dim=1)
 50 |         z = torch.sigmoid(self.convz1(hx))
 51 |         r = torch.sigmoid(self.convr1(hx))
 52 |         q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))        
 53 |         h = (1-z) * h + z * q
 54 | 
 55 |         # vertical
 56 |         hx = torch.cat([h, x], dim=1)
 57 |         z = torch.sigmoid(self.convz2(hx))
 58 |         r = torch.sigmoid(self.convr2(hx))
 59 |         q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))       
 60 |         h = (1-z) * h + z * q
 61 | 
 62 |         return h
 63 | 
 64 | class SmallMotionEncoder(nn.Module):
 65 |     def __init__(self, args):
 66 |         super(SmallMotionEncoder, self).__init__()
 67 |         cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
 68 |         self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0)
 69 |         self.convf1 = nn.Conv2d(2, 64, 7, padding=3)
 70 |         self.convf2 = nn.Conv2d(64, 32, 3, padding=1)
 71 |         self.conv = nn.Conv2d(128, 80, 3, padding=1)
 72 | 
 73 |     def forward(self, flow, corr):
 74 |         cor = F.relu(self.convc1(corr))
 75 |         flo = F.relu(self.convf1(flow))
 76 |         flo = F.relu(self.convf2(flo))
 77 |         cor_flo = torch.cat([cor, flo], dim=1)
 78 |         out = F.relu(self.conv(cor_flo))
 79 |         return torch.cat([out, flow], dim=1)
 80 | 
 81 | class BasicMotionEncoder(nn.Module):
 82 |     def __init__(self, args):
 83 |         super(BasicMotionEncoder, self).__init__()
 84 |         cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
 85 |         self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
 86 |         self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
 87 |         self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
 88 |         self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
 89 |         self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1)
 90 | 
 91 |     def forward(self, flow, corr):
 92 |         cor = F.relu(self.convc1(corr))
 93 |         cor = F.relu(self.convc2(cor))
 94 |         flo = F.relu(self.convf1(flow))
 95 |         flo = F.relu(self.convf2(flo))
 96 | 
 97 |         cor_flo = torch.cat([cor, flo], dim=1)
 98 |         out = F.relu(self.conv(cor_flo))
 99 |         return torch.cat([out, flow], dim=1)
100 | 
101 | class SmallUpdateBlock(nn.Module):
102 |     def __init__(self, args, hidden_dim=96):
103 |         super(SmallUpdateBlock, self).__init__()
104 |         self.encoder = SmallMotionEncoder(args)
105 |         self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82+64)
106 |         self.flow_head = FlowHead(hidden_dim, hidden_dim=128)
107 | 
108 |     def forward(self, net, inp, corr, flow):
109 |         motion_features = self.encoder(flow, corr)
110 |         inp = torch.cat([inp, motion_features], dim=1)
111 |         net = self.gru(net, inp)
112 |         delta_flow = self.flow_head(net)
113 | 
114 |         return net, None, delta_flow
115 | 
116 | class BasicUpdateBlock(nn.Module):
117 |     def __init__(self, args, hidden_dim=128, input_dim=128):
118 |         super(BasicUpdateBlock, self).__init__()
119 |         self.args = args
120 |         self.encoder = BasicMotionEncoder(args)
121 |         self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
122 |         self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
123 | 
124 |         self.mask = nn.Sequential(
125 |             nn.Conv2d(128, 256, 3, padding=1),
126 |             nn.ReLU(inplace=True),
127 |             nn.Conv2d(256, 64*9, 1, padding=0))
128 | 
129 |     def forward(self, net, inp, corr, flow, upsample=True):
130 |         motion_features = self.encoder(flow, corr)
131 |         inp = torch.cat([inp, motion_features], dim=1)
132 | 
133 |         net = self.gru(net, inp)
134 |         delta_flow = self.flow_head(net)
135 | 
136 |         # scale mask to balence gradients
137 |         mask = .25 * self.mask(net)
138 |         return net, mask, delta_flow
139 | 
140 | 
141 | 
142 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/utils/__init__.py:
--------------------------------------------------------------------------------
1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
2 | 
3 | from .flow_viz import flow_to_image
4 | from .frame_utils import writeFlow
5 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/utils/augmentor.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import numpy as np
  4 | import random
  5 | import math
  6 | from PIL import Image
  7 | 
  8 | import cv2
  9 | cv2.setNumThreads(0)
 10 | cv2.ocl.setUseOpenCL(False)
 11 | 
 12 | import torch
 13 | from torchvision.transforms import ColorJitter
 14 | import torch.nn.functional as F
 15 | 
 16 | 
 17 | class FlowAugmentor:
 18 |     def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True):
 19 |         
 20 |         # spatial augmentation params
 21 |         self.crop_size = crop_size
 22 |         self.min_scale = min_scale
 23 |         self.max_scale = max_scale
 24 |         self.spatial_aug_prob = 0.8
 25 |         self.stretch_prob = 0.8
 26 |         self.max_stretch = 0.2
 27 | 
 28 |         # flip augmentation params
 29 |         self.do_flip = do_flip
 30 |         self.h_flip_prob = 0.5
 31 |         self.v_flip_prob = 0.1
 32 | 
 33 |         # photometric augmentation params
 34 |         self.photo_aug = ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5/3.14)
 35 |         self.asymmetric_color_aug_prob = 0.2
 36 |         self.eraser_aug_prob = 0.5
 37 | 
 38 |     def color_transform(self, img1, img2):
 39 |         """ Photometric augmentation """
 40 | 
 41 |         # asymmetric
 42 |         if np.random.rand() < self.asymmetric_color_aug_prob:
 43 |             img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
 44 |             img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)
 45 | 
 46 |         # symmetric
 47 |         else:
 48 |             image_stack = np.concatenate([img1, img2], axis=0)
 49 |             image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
 50 |             img1, img2 = np.split(image_stack, 2, axis=0)
 51 | 
 52 |         return img1, img2
 53 | 
 54 |     def eraser_transform(self, img1, img2, bounds=[50, 100]):
 55 |         """ Occlusion augmentation """
 56 | 
 57 |         ht, wd = img1.shape[:2]
 58 |         if np.random.rand() < self.eraser_aug_prob:
 59 |             mean_color = np.mean(img2.reshape(-1, 3), axis=0)
 60 |             for _ in range(np.random.randint(1, 3)):
 61 |                 x0 = np.random.randint(0, wd)
 62 |                 y0 = np.random.randint(0, ht)
 63 |                 dx = np.random.randint(bounds[0], bounds[1])
 64 |                 dy = np.random.randint(bounds[0], bounds[1])
 65 |                 img2[y0:y0+dy, x0:x0+dx, :] = mean_color
 66 | 
 67 |         return img1, img2
 68 | 
 69 |     def spatial_transform(self, img1, img2, flow):
 70 |         # randomly sample scale
 71 |         ht, wd = img1.shape[:2]
 72 |         min_scale = np.maximum(
 73 |             (self.crop_size[0] + 8) / float(ht), 
 74 |             (self.crop_size[1] + 8) / float(wd))
 75 | 
 76 |         scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
 77 |         scale_x = scale
 78 |         scale_y = scale
 79 |         if np.random.rand() < self.stretch_prob:
 80 |             scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
 81 |             scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
 82 |         
 83 |         scale_x = np.clip(scale_x, min_scale, None)
 84 |         scale_y = np.clip(scale_y, min_scale, None)
 85 | 
 86 |         if np.random.rand() < self.spatial_aug_prob:
 87 |             # rescale the images
 88 |             img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
 89 |             img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
 90 |             flow = cv2.resize(flow, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
 91 |             flow = flow * [scale_x, scale_y]
 92 | 
 93 |         if self.do_flip:
 94 |             if np.random.rand() < self.h_flip_prob: # h-flip
 95 |                 img1 = img1[:, ::-1]
 96 |                 img2 = img2[:, ::-1]
 97 |                 flow = flow[:, ::-1] * [-1.0, 1.0]
 98 | 
 99 |             if np.random.rand() < self.v_flip_prob: # v-flip
100 |                 img1 = img1[::-1, :]
101 |                 img2 = img2[::-1, :]
102 |                 flow = flow[::-1, :] * [1.0, -1.0]
103 | 
104 |         y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
105 |         x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
106 |         
107 |         img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
108 |         img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
109 |         flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
110 | 
111 |         return img1, img2, flow
112 | 
113 |     def __call__(self, img1, img2, flow):
114 |         img1, img2 = self.color_transform(img1, img2)
115 |         img1, img2 = self.eraser_transform(img1, img2)
116 |         img1, img2, flow = self.spatial_transform(img1, img2, flow)
117 | 
118 |         img1 = np.ascontiguousarray(img1)
119 |         img2 = np.ascontiguousarray(img2)
120 |         flow = np.ascontiguousarray(flow)
121 | 
122 |         return img1, img2, flow
123 | 
124 | class SparseFlowAugmentor:
125 |     def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False):
126 |         # spatial augmentation params
127 |         self.crop_size = crop_size
128 |         self.min_scale = min_scale
129 |         self.max_scale = max_scale
130 |         self.spatial_aug_prob = 0.8
131 |         self.stretch_prob = 0.8
132 |         self.max_stretch = 0.2
133 | 
134 |         # flip augmentation params
135 |         self.do_flip = do_flip
136 |         self.h_flip_prob = 0.5
137 |         self.v_flip_prob = 0.1
138 | 
139 |         # photometric augmentation params
140 |         self.photo_aug = ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3/3.14)
141 |         self.asymmetric_color_aug_prob = 0.2
142 |         self.eraser_aug_prob = 0.5
143 |         
144 |     def color_transform(self, img1, img2):
145 |         image_stack = np.concatenate([img1, img2], axis=0)
146 |         image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
147 |         img1, img2 = np.split(image_stack, 2, axis=0)
148 |         return img1, img2
149 | 
150 |     def eraser_transform(self, img1, img2):
151 |         ht, wd = img1.shape[:2]
152 |         if np.random.rand() < self.eraser_aug_prob:
153 |             mean_color = np.mean(img2.reshape(-1, 3), axis=0)
154 |             for _ in range(np.random.randint(1, 3)):
155 |                 x0 = np.random.randint(0, wd)
156 |                 y0 = np.random.randint(0, ht)
157 |                 dx = np.random.randint(50, 100)
158 |                 dy = np.random.randint(50, 100)
159 |                 img2[y0:y0+dy, x0:x0+dx, :] = mean_color
160 | 
161 |         return img1, img2
162 | 
163 |     def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0):
164 |         ht, wd = flow.shape[:2]
165 |         coords = np.meshgrid(np.arange(wd), np.arange(ht))
166 |         coords = np.stack(coords, axis=-1)
167 | 
168 |         coords = coords.reshape(-1, 2).astype(np.float32)
169 |         flow = flow.reshape(-1, 2).astype(np.float32)
170 |         valid = valid.reshape(-1).astype(np.float32)
171 | 
172 |         coords0 = coords[valid>=1]
173 |         flow0 = flow[valid>=1]
174 | 
175 |         ht1 = int(round(ht * fy))
176 |         wd1 = int(round(wd * fx))
177 | 
178 |         coords1 = coords0 * [fx, fy]
179 |         flow1 = flow0 * [fx, fy]
180 | 
181 |         xx = np.round(coords1[:,0]).astype(np.int32)
182 |         yy = np.round(coords1[:,1]).astype(np.int32)
183 | 
184 |         v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
185 |         xx = xx[v]
186 |         yy = yy[v]
187 |         flow1 = flow1[v]
188 | 
189 |         flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32)
190 |         valid_img = np.zeros([ht1, wd1], dtype=np.int32)
191 | 
192 |         flow_img[yy, xx] = flow1
193 |         valid_img[yy, xx] = 1
194 | 
195 |         return flow_img, valid_img
196 | 
197 |     def spatial_transform(self, img1, img2, flow, valid):
198 |         # randomly sample scale
199 | 
200 |         ht, wd = img1.shape[:2]
201 |         min_scale = np.maximum(
202 |             (self.crop_size[0] + 1) / float(ht), 
203 |             (self.crop_size[1] + 1) / float(wd))
204 | 
205 |         scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
206 |         scale_x = np.clip(scale, min_scale, None)
207 |         scale_y = np.clip(scale, min_scale, None)
208 | 
209 |         if np.random.rand() < self.spatial_aug_prob:
210 |             # rescale the images
211 |             img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
212 |             img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
213 |             flow, valid = self.resize_sparse_flow_map(flow, valid, fx=scale_x, fy=scale_y)
214 | 
215 |         if self.do_flip:
216 |             if np.random.rand() < 0.5: # h-flip
217 |                 img1 = img1[:, ::-1]
218 |                 img2 = img2[:, ::-1]
219 |                 flow = flow[:, ::-1] * [-1.0, 1.0]
220 |                 valid = valid[:, ::-1]
221 | 
222 |         margin_y = 20
223 |         margin_x = 50
224 | 
225 |         y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0] + margin_y)
226 |         x0 = np.random.randint(-margin_x, img1.shape[1] - self.crop_size[1] + margin_x)
227 | 
228 |         y0 = np.clip(y0, 0, img1.shape[0] - self.crop_size[0])
229 |         x0 = np.clip(x0, 0, img1.shape[1] - self.crop_size[1])
230 | 
231 |         img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
232 |         img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
233 |         flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
234 |         valid = valid[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
235 |         return img1, img2, flow, valid
236 | 
237 | 
238 |     def __call__(self, img1, img2, flow, valid):
239 |         img1, img2 = self.color_transform(img1, img2)
240 |         img1, img2 = self.eraser_transform(img1, img2)
241 |         img1, img2, flow, valid = self.spatial_transform(img1, img2, flow, valid)
242 | 
243 |         img1 = np.ascontiguousarray(img1)
244 |         img2 = np.ascontiguousarray(img2)
245 |         flow = np.ascontiguousarray(flow)
246 |         valid = np.ascontiguousarray(valid)
247 | 
248 |         return img1, img2, flow, valid
249 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/utils/flow_viz.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | # Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization
  4 | 
  5 | 
  6 | # MIT License
  7 | #
  8 | # Copyright (c) 2018 Tom Runia
  9 | #
 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy
 11 | # of this software and associated documentation files (the "Software"), to deal
 12 | # in the Software without restriction, including without limitation the rights
 13 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 14 | # copies of the Software, and to permit persons to whom the Software is
 15 | # furnished to do so, subject to conditions.
 16 | #
 17 | # Author: Tom Runia
 18 | # Date Created: 2018-08-03
 19 | 
 20 | import numpy as np
 21 | 
 22 | def make_colorwheel():
 23 |     """
 24 |     Generates a color wheel for optical flow visualization as presented in:
 25 |         Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
 26 |         URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
 27 | 
 28 |     Code follows the original C++ source code of Daniel Scharstein.
 29 |     Code follows the the Matlab source code of Deqing Sun.
 30 | 
 31 |     Returns:
 32 |         np.ndarray: Color wheel
 33 |     """
 34 | 
 35 |     RY = 15
 36 |     YG = 6
 37 |     GC = 4
 38 |     CB = 11
 39 |     BM = 13
 40 |     MR = 6
 41 | 
 42 |     ncols = RY + YG + GC + CB + BM + MR
 43 |     colorwheel = np.zeros((ncols, 3))
 44 |     col = 0
 45 | 
 46 |     # RY
 47 |     colorwheel[0:RY, 0] = 255
 48 |     colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
 49 |     col = col+RY
 50 |     # YG
 51 |     colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
 52 |     colorwheel[col:col+YG, 1] = 255
 53 |     col = col+YG
 54 |     # GC
 55 |     colorwheel[col:col+GC, 1] = 255
 56 |     colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
 57 |     col = col+GC
 58 |     # CB
 59 |     colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
 60 |     colorwheel[col:col+CB, 2] = 255
 61 |     col = col+CB
 62 |     # BM
 63 |     colorwheel[col:col+BM, 2] = 255
 64 |     colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
 65 |     col = col+BM
 66 |     # MR
 67 |     colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
 68 |     colorwheel[col:col+MR, 0] = 255
 69 |     return colorwheel
 70 | 
 71 | 
 72 | def flow_uv_to_colors(u, v, convert_to_bgr=False):
 73 |     """
 74 |     Applies the flow color wheel to (possibly clipped) flow components u and v.
 75 | 
 76 |     According to the C++ source code of Daniel Scharstein
 77 |     According to the Matlab source code of Deqing Sun
 78 | 
 79 |     Args:
 80 |         u (np.ndarray): Input horizontal flow of shape [H,W]
 81 |         v (np.ndarray): Input vertical flow of shape [H,W]
 82 |         convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
 83 | 
 84 |     Returns:
 85 |         np.ndarray: Flow visualization image of shape [H,W,3]
 86 |     """
 87 |     flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
 88 |     colorwheel = make_colorwheel()  # shape [55x3]
 89 |     ncols = colorwheel.shape[0]
 90 |     rad = np.sqrt(np.square(u) + np.square(v))
 91 |     a = np.arctan2(-v, -u)/np.pi
 92 |     fk = (a+1) / 2*(ncols-1)
 93 |     k0 = np.floor(fk).astype(np.int32)
 94 |     k1 = k0 + 1
 95 |     k1[k1 == ncols] = 0
 96 |     f = fk - k0
 97 |     for i in range(colorwheel.shape[1]):
 98 |         tmp = colorwheel[:,i]
 99 |         col0 = tmp[k0] / 255.0
100 |         col1 = tmp[k1] / 255.0
101 |         col = (1-f)*col0 + f*col1
102 |         idx = (rad <= 1)
103 |         col[idx]  = 1 - rad[idx] * (1-col[idx])
104 |         col[~idx] = col[~idx] * 0.75   # out of range
105 |         # Note the 2-i => BGR instead of RGB
106 |         ch_idx = 2-i if convert_to_bgr else i
107 |         flow_image[:,:,ch_idx] = np.floor(255 * col)
108 |     return flow_image
109 | 
110 | 
111 | def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
112 |     """
113 |     Expects a two dimensional flow image of shape.
114 | 
115 |     Args:
116 |         flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
117 |         clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
118 |         convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
119 | 
120 |     Returns:
121 |         np.ndarray: Flow visualization image of shape [H,W,3]
122 |     """
123 |     assert flow_uv.ndim == 3, 'input flow must have three dimensions'
124 |     assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
125 |     if clip_flow is not None:
126 |         flow_uv = np.clip(flow_uv, 0, clip_flow)
127 |     u = flow_uv[:,:,0]
128 |     v = flow_uv[:,:,1]
129 |     rad = np.sqrt(np.square(u) + np.square(v))
130 |     rad_max = np.max(rad)
131 |     epsilon = 1e-5
132 |     u = u / (rad_max + epsilon)
133 |     v = v / (rad_max + epsilon)
134 |     return flow_uv_to_colors(u, v, convert_to_bgr)


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/utils/frame_utils.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import numpy as np
  4 | from PIL import Image
  5 | from os.path import *
  6 | import re
  7 | 
  8 | import cv2
  9 | cv2.setNumThreads(0)
 10 | cv2.ocl.setUseOpenCL(False)
 11 | 
 12 | TAG_CHAR = np.array([202021.25], np.float32)
 13 | 
 14 | def readFlow(fn):
 15 |     """ Read .flo file in Middlebury format"""
 16 |     # Code adapted from:
 17 |     # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
 18 | 
 19 |     # WARNING: this will work on little-endian architectures (eg Intel x86) only!
 20 |     # print 'fn = %s'%(fn)
 21 |     with open(fn, 'rb') as f:
 22 |         magic = np.fromfile(f, np.float32, count=1)
 23 |         if 202021.25 != magic:
 24 |             print('Magic number incorrect. Invalid .flo file')
 25 |             return None
 26 |         else:
 27 |             w = np.fromfile(f, np.int32, count=1)
 28 |             h = np.fromfile(f, np.int32, count=1)
 29 |             # print 'Reading %d x %d flo file\n' % (w, h)
 30 |             data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
 31 |             # Reshape data into 3D array (columns, rows, bands)
 32 |             # The reshape here is for visualization, the original code is (w,h,2)
 33 |             return np.resize(data, (int(h), int(w), 2))
 34 | 
 35 | def readPFM(file):
 36 |     file = open(file, 'rb')
 37 | 
 38 |     color = None
 39 |     width = None
 40 |     height = None
 41 |     scale = None
 42 |     endian = None
 43 | 
 44 |     header = file.readline().rstrip()
 45 |     if header == b'PF':
 46 |         color = True
 47 |     elif header == b'Pf':
 48 |         color = False
 49 |     else:
 50 |         raise Exception('Not a PFM file.')
 51 | 
 52 |     dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
 53 |     if dim_match:
 54 |         width, height = map(int, dim_match.groups())
 55 |     else:
 56 |         raise Exception('Malformed PFM header.')
 57 | 
 58 |     scale = float(file.readline().rstrip())
 59 |     if scale < 0: # little-endian
 60 |         endian = '<'
 61 |         scale = -scale
 62 |     else:
 63 |         endian = '>' # big-endian
 64 | 
 65 |     data = np.fromfile(file, endian + 'f')
 66 |     shape = (height, width, 3) if color else (height, width)
 67 | 
 68 |     data = np.reshape(data, shape)
 69 |     data = np.flipud(data)
 70 |     return data
 71 | 
 72 | def writeFlow(filename,uv,v=None):
 73 |     """ Write optical flow to file.
 74 |     
 75 |     If v is None, uv is assumed to contain both u and v channels,
 76 |     stacked in depth.
 77 |     Original code by Deqing Sun, adapted from Daniel Scharstein.
 78 |     """
 79 |     nBands = 2
 80 | 
 81 |     if v is None:
 82 |         assert(uv.ndim == 3)
 83 |         assert(uv.shape[2] == 2)
 84 |         u = uv[:,:,0]
 85 |         v = uv[:,:,1]
 86 |     else:
 87 |         u = uv
 88 | 
 89 |     assert(u.shape == v.shape)
 90 |     height,width = u.shape
 91 |     f = open(filename,'wb')
 92 |     # write the header
 93 |     f.write(TAG_CHAR)
 94 |     np.array(width).astype(np.int32).tofile(f)
 95 |     np.array(height).astype(np.int32).tofile(f)
 96 |     # arrange into matrix form
 97 |     tmp = np.zeros((height, width*nBands))
 98 |     tmp[:,np.arange(width)*2] = u
 99 |     tmp[:,np.arange(width)*2 + 1] = v
100 |     tmp.astype(np.float32).tofile(f)
101 |     f.close()
102 | 
103 | 
104 | def readFlowKITTI(filename):
105 |     flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
106 |     flow = flow[:,:,::-1].astype(np.float32)
107 |     flow, valid = flow[:, :, :2], flow[:, :, 2]
108 |     flow = (flow - 2**15) / 64.0
109 |     return flow, valid
110 | 
111 | def readDispKITTI(filename):
112 |     disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
113 |     valid = disp > 0.0
114 |     flow = np.stack([-disp, np.zeros_like(disp)], -1)
115 |     return flow, valid
116 | 
117 | 
118 | def writeFlowKITTI(filename, uv):
119 |     uv = 64.0 * uv + 2**15
120 |     valid = np.ones([uv.shape[0], uv.shape[1], 1])
121 |     uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
122 |     cv2.imwrite(filename, uv[..., ::-1])
123 |     
124 | 
125 | def read_gen(file_name, pil=False):
126 |     ext = splitext(file_name)[-1]
127 |     if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
128 |         return Image.open(file_name)
129 |     elif ext == '.bin' or ext == '.raw':
130 |         return np.load(file_name)
131 |     elif ext == '.flo':
132 |         return readFlow(file_name).astype(np.float32)
133 |     elif ext == '.pfm':
134 |         flow = readPFM(file_name).astype(np.float32)
135 |         if len(flow.shape) == 2:
136 |             return flow
137 |         else:
138 |             return flow[:, :, :-1]
139 |     return []


--------------------------------------------------------------------------------
/WebVi3D/utils/RAFT/utils/utils.py:
--------------------------------------------------------------------------------
 1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
 2 | 
 3 | import torch
 4 | import torch.nn.functional as F
 5 | import numpy as np
 6 | from scipy import interpolate
 7 | 
 8 | 
 9 | class InputPadder:
10 |     """ Pads images such that dimensions are divisible by 8 """
11 |     def __init__(self, dims, mode='sintel'):
12 |         self.ht, self.wd = dims[-2:]
13 |         pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8
14 |         pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8
15 |         if mode == 'sintel':
16 |             self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
17 |         else:
18 |             self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
19 | 
20 |     def pad(self, *inputs):
21 |         return [F.pad(x, self._pad, mode='replicate') for x in inputs]
22 | 
23 |     def unpad(self,x):
24 |         ht, wd = x.shape[-2:]
25 |         c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
26 |         return x[..., c[0]:c[1], c[2]:c[3]]
27 | 
28 | def forward_interpolate(flow):
29 |     flow = flow.detach().cpu().numpy()
30 |     dx, dy = flow[0], flow[1]
31 | 
32 |     ht, wd = dx.shape
33 |     x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
34 | 
35 |     x1 = x0 + dx
36 |     y1 = y0 + dy
37 |     
38 |     x1 = x1.reshape(-1)
39 |     y1 = y1.reshape(-1)
40 |     dx = dx.reshape(-1)
41 |     dy = dy.reshape(-1)
42 | 
43 |     valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
44 |     x1 = x1[valid]
45 |     y1 = y1[valid]
46 |     dx = dx[valid]
47 |     dy = dy[valid]
48 | 
49 |     flow_x = interpolate.griddata(
50 |         (x1, y1), dx, (x0, y0), method='nearest', fill_value=0)
51 | 
52 |     flow_y = interpolate.griddata(
53 |         (x1, y1), dy, (x0, y0), method='nearest', fill_value=0)
54 | 
55 |     flow = np.stack([flow_x, flow_y], axis=0)
56 |     return torch.from_numpy(flow).float()
57 | 
58 | 
59 | def bilinear_sampler(img, coords, mode='bilinear', mask=False):
60 |     """ Wrapper for grid_sample, uses pixel coordinates """
61 |     H, W = img.shape[-2:]
62 |     xgrid, ygrid = coords.split([1,1], dim=-1)
63 |     xgrid = 2*xgrid/(W-1) - 1
64 |     ygrid = 2*ygrid/(H-1) - 1
65 | 
66 |     grid = torch.cat([xgrid, ygrid], dim=-1)
67 |     img = F.grid_sample(img, grid, align_corners=True)
68 | 
69 |     if mask:
70 |         mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
71 |         return img, mask.float()
72 | 
73 |     return img
74 | 
75 | 
76 | def coords_grid(batch, ht, wd):
77 |     coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
78 |     coords = torch.stack(coords[::-1], dim=0).float()
79 |     return coords[None].repeat(batch, 1, 1, 1)
80 | 
81 | 
82 | def upflow8(flow, mode='bilinear'):
83 |     new_size = (8 * flow.shape[2], 8 * flow.shape[3])
84 |     return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
85 | 


--------------------------------------------------------------------------------
/WebVi3D/utils/flow_utils.py:
--------------------------------------------------------------------------------
  1 | # Copyright (c) Meta Platforms, Inc. and affiliates.
  2 | 
  3 | import os
  4 | import cv2
  5 | import numpy as np
  6 | from PIL import Image
  7 | from os.path import *
  8 | import re
  9 | def microsoft_sorted(input_list):
 10 |     return sorted(input_list, key=lambda s: [int(s) if s.isdigit() else s for s in sum(re.findall(r'(\D+)(\d+)', 'a'+s+'0'), ())])
 11 |     
 12 | UNKNOWN_FLOW_THRESH = 1e7
 13 | 
 14 | 
 15 | def flow_to_image(flow, global_max=None):
 16 |     """
 17 |     Convert flow into middlebury color code image
 18 |     :param flow: optical flow map
 19 |     :return: optical flow image in middlebury color
 20 |     """
 21 |     u = flow[:, :, 0]
 22 |     v = flow[:, :, 1]
 23 | 
 24 |     maxu = -999.0
 25 |     maxv = -999.0
 26 |     minu = 999.0
 27 |     minv = 999.0
 28 | 
 29 |     idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
 30 |     u[idxUnknow] = 0
 31 |     v[idxUnknow] = 0
 32 | 
 33 |     maxu = max(maxu, np.max(u))
 34 |     minu = min(minu, np.min(u))
 35 | 
 36 |     maxv = max(maxv, np.max(v))
 37 |     minv = min(minv, np.min(v))
 38 | 
 39 |     rad = np.sqrt(u**2 + v**2)
 40 | 
 41 |     if global_max == None:
 42 |         maxrad = max(-1, np.max(rad))
 43 |     else:
 44 |         maxrad = global_max
 45 | 
 46 |     u = u / (maxrad + np.finfo(float).eps)
 47 |     v = v / (maxrad + np.finfo(float).eps)
 48 | 
 49 |     img = compute_color(u, v)
 50 | 
 51 |     idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
 52 |     img[idx] = 0
 53 | 
 54 |     return np.uint8(img)
 55 | 
 56 | 
 57 | def compute_color(u, v):
 58 |     """
 59 |     compute optical flow color map
 60 |     :param u: optical flow horizontal map
 61 |     :param v: optical flow vertical map
 62 |     :return: optical flow in color code
 63 |     """
 64 |     [h, w] = u.shape
 65 |     img = np.zeros([h, w, 3])
 66 |     nanIdx = np.isnan(u) | np.isnan(v)
 67 |     u[nanIdx] = 0
 68 |     v[nanIdx] = 0
 69 | 
 70 |     colorwheel = make_color_wheel()
 71 |     ncols = np.size(colorwheel, 0)
 72 | 
 73 |     rad = np.sqrt(u**2 + v**2)
 74 | 
 75 |     a = np.arctan2(-v, -u) / np.pi
 76 | 
 77 |     fk = (a + 1) / 2 * (ncols - 1) + 1
 78 | 
 79 |     k0 = np.floor(fk).astype(int)
 80 | 
 81 |     k1 = k0 + 1
 82 |     k1[k1 == ncols + 1] = 1
 83 |     f = fk - k0
 84 | 
 85 |     for i in range(0, np.size(colorwheel, 1)):
 86 |         tmp = colorwheel[:, i]
 87 |         col0 = tmp[k0 - 1] / 255
 88 |         col1 = tmp[k1 - 1] / 255
 89 |         col = (1 - f) * col0 + f * col1
 90 | 
 91 |         idx = rad <= 1
 92 |         col[idx] = 1 - rad[idx] * (1 - col[idx])
 93 |         notidx = np.logical_not(idx)
 94 | 
 95 |         col[notidx] *= 0.75
 96 |         img[:, :, i] = np.uint8(np.floor(255 * col * (1 - nanIdx)))
 97 | 
 98 |     return img
 99 | 
100 | 
101 | def make_color_wheel():
102 |     """
103 |     Generate color wheel according Middlebury color code
104 |     :return: Color wheel
105 |     """
106 |     RY = 15
107 |     YG = 6
108 |     GC = 4
109 |     CB = 11
110 |     BM = 13
111 |     MR = 6
112 | 
113 |     ncols = RY + YG + GC + CB + BM + MR
114 | 
115 |     colorwheel = np.zeros([ncols, 3])
116 | 
117 |     col = 0
118 | 
119 |     # RY
120 |     colorwheel[0:RY, 0] = 255
121 |     colorwheel[0:RY, 1] = np.transpose(np.floor(255 * np.arange(0, RY) / RY))
122 |     col += RY
123 | 
124 |     # YG
125 |     colorwheel[col : col + YG, 0] = 255 - np.transpose(
126 |         np.floor(255 * np.arange(0, YG) / YG)
127 |     )
128 |     colorwheel[col : col + YG, 1] = 255
129 |     col += YG
130 | 
131 |     # GC
132 |     colorwheel[col : col + GC, 1] = 255
133 |     colorwheel[col : col + GC, 2] = np.transpose(np.floor(255 * np.arange(0, GC) / GC))
134 |     col += GC
135 | 
136 |     # CB
137 |     colorwheel[col : col + CB, 1] = 255 - np.transpose(
138 |         np.floor(255 * np.arange(0, CB) / CB)
139 |     )
140 |     colorwheel[col : col + CB, 2] = 255
141 |     col += CB
142 | 
143 |     # BM
144 |     colorwheel[col : col + BM, 2] = 255
145 |     colorwheel[col : col + BM, 0] = np.transpose(np.floor(255 * np.arange(0, BM) / BM))
146 |     col += +BM
147 | 
148 |     # MR
149 |     colorwheel[col : col + MR, 2] = 255 - np.transpose(
150 |         np.floor(255 * np.arange(0, MR) / MR)
151 |     )
152 |     colorwheel[col : col + MR, 0] = 255
153 | 
154 |     return colorwheel
155 | 
156 | 
157 | def resize_flow(flow, H_new, W_new):
158 |     H_old, W_old = flow.shape[0:2]
159 |     flow_resized = cv2.resize(flow, (W_new, H_new), interpolation=cv2.INTER_LINEAR)
160 |     flow_resized[:, :, 0] *= H_new / H_old
161 |     flow_resized[:, :, 1] *= W_new / W_old
162 |     return flow_resized
163 | 
164 | 
165 | def warp_flow(img, flow):
166 |     h, w = flow.shape[:2]
167 |     flow_new = flow.copy()
168 |     flow_new[:, :, 0] += np.arange(w)
169 |     flow_new[:, :, 1] += np.arange(h)[:, np.newaxis]
170 | 
171 |     res = cv2.remap(
172 |         img, flow_new, None, cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT
173 |     )
174 |     return res
175 | 
176 | 
177 | def consistCheck(flowB, flowF):
178 |     # |--------------------|  |--------------------|
179 |     # |       y            |  |       v            |
180 |     # |   x   *            |  |   u   *            |
181 |     # |                    |  |                    |
182 |     # |--------------------|  |--------------------|
183 | 
184 |     # sub: numPix * [y x t]
185 | 
186 |     imgH, imgW, _ = flowF.shape
187 | 
188 |     (fy, fx) = np.mgrid[0:imgH, 0:imgW].astype(np.float32)
189 |     fxx = fx + flowB[:, :, 0]  # horizontal
190 |     fyy = fy + flowB[:, :, 1]  # vertical
191 | 
192 |     u = fxx + cv2.remap(flowF[:, :, 0], fxx, fyy, cv2.INTER_LINEAR) - fx
193 |     v = fyy + cv2.remap(flowF[:, :, 1], fxx, fyy, cv2.INTER_LINEAR) - fy
194 |     BFdiff = (u**2 + v**2) ** 0.5
195 | 
196 |     return BFdiff, np.stack((u, v), axis=2)
197 | 
198 | 
199 | def read_optical_flow(basedir, img_i_name, read_fwd):
200 |     flow_dir = os.path.join(basedir, "flow")
201 | 
202 |     fwd_flow_path = os.path.join(flow_dir, "%s_fwd.npz" % img_i_name[:-4])
203 |     bwd_flow_path = os.path.join(flow_dir, "%s_bwd.npz" % img_i_name[:-4])
204 | 
205 |     if read_fwd:
206 |         fwd_data = np.load(fwd_flow_path)
207 |         fwd_flow, fwd_mask = fwd_data["flow"], fwd_data["mask"]
208 |         return fwd_flow, fwd_mask
209 |     else:
210 |         bwd_data = np.load(bwd_flow_path)
211 |         bwd_flow, bwd_mask = bwd_data["flow"], bwd_data["mask"]
212 |         return bwd_flow, bwd_mask
213 | 
214 | 
215 | def compute_epipolar_distance(T_21, K, p_1, p_2):
216 |     R_21 = T_21[:3, :3]
217 |     t_21 = T_21[:3, 3]
218 | 
219 |     E_mat = np.dot(skew(t_21), R_21)
220 |     # compute bearing vector
221 |     inv_K = np.linalg.inv(K)
222 | 
223 |     F_mat = np.dot(np.dot(inv_K.T, E_mat), inv_K)
224 | 
225 |     l_2 = np.dot(F_mat, p_1)
226 |     algebric_e_distance = np.sum(p_2 * l_2, axis=0)
227 |     n_term = np.sqrt(l_2[0, :] ** 2 + l_2[1, :] ** 2) + 1e-8
228 |     geometric_e_distance = algebric_e_distance / n_term
229 |     geometric_e_distance = np.abs(geometric_e_distance)
230 | 
231 |     return geometric_e_distance
232 | 
233 | 
234 | def skew(x):
235 |     return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
236 | 


--------------------------------------------------------------------------------
/WebVi3D/weights/raft-small.pth:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/weights/raft-small.pth


--------------------------------------------------------------------------------
/WebVi3D/weights/raft-things.pth:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/WebVi3D/weights/raft-things.pth


--------------------------------------------------------------------------------
/assets/logo.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/assets/logo.jpg


--------------------------------------------------------------------------------
/assets/method.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/assets/method.jpg


--------------------------------------------------------------------------------
/assets/teaser.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/baaivision/See3D/52594ef34b02ee75fd8dc8d19411c60902cce87b/assets/teaser.jpg


--------------------------------------------------------------------------------
/inference.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | import glob
  3 | 
  4 | import warnings
  5 | 
  6 | from random import randint
  7 | from argparse import ArgumentParser
  8 | 
  9 | import torchvision.transforms.functional as tf
 10 | 
 11 | warnings.filterwarnings(action='ignore')
 12 | 
 13 | 
 14 | import numpy as np
 15 | 
 16 | from PIL import Image
 17 | from tqdm import tqdm
 18 | from scipy.interpolate import griddata as interp_grid
 19 | from scipy.ndimage import minimum_filter, maximum_filter
 20 | 
 21 | import torch
 22 | import torch.nn.functional as F
 23 | from mv_diffusion import mvdream_diffusion_model
 24 | from mv_diffusion_SR import mvdream_diffusion_model as mvdream_diffusion_model_SR
 25 | from transformers import CLIPTextModel, CLIPTokenizer
 26 | import cv2
 27 | from sklearn.linear_model import RANSACRegressor
 28 | from sklearn.linear_model import LinearRegression
 29 | import re
 30 | from torch.autograd import Variable
 31 | from math import exp
 32 | import argparse
 33 | import gc
 34 | 
 35 | parser = argparse.ArgumentParser(description="Process some integers.")
 36 | parser.add_argument('--single_view', action='store_true', help='Use single view mode')
 37 | parser.add_argument('--super_resolution', action='store_true', help='Enable super resolution')
 38 | parser.add_argument('--base_model_path', type=str, help='Base model directory path')
 39 | parser.add_argument('--source_imgs_dir', type=str, help='Source images directory path')
 40 | parser.add_argument('--warp_root_dir', type=str, help='Warp images directory path')
 41 | parser.add_argument('--output_dir', type=str, help='Output directory path')
 42 | 
 43 | # Parse arguments
 44 | args = parser.parse_args()
 45 | 
 46 | def PIL2tensor(height,width,num_frames,masks,warps,logicalNot=False):
 47 |     channels = 3
 48 |     pixel_values = torch.empty((num_frames, channels, height, width))
 49 |     condition_pixel_values = torch.empty((num_frames, channels, height, width))
 50 |     masks_pixel_values = torch.ones((num_frames, 1, height, width))
 51 |     
 52 |     # input_ids
 53 |     prompt = ''
 54 | 
 55 |     for i, img in enumerate(masks):
 56 |         img = masks[i]
 57 |         img = img.convert('L') # make sure channel 1
 58 |         img_resized = img.resize((width, height)) # hard code here
 59 |         img_tensor = torch.from_numpy(np.array(img_resized)).float()
 60 | 
 61 |         # Normalize the image by scaling pixel values to [0, 1]
 62 |         img_normalized = img_tensor / 255
 63 |         mask_condition = (img_normalized > 0.9).float()
 64 |         
 65 |         masks_pixel_values[i] = mask_condition
 66 |     
 67 |     for i, img in enumerate(warps):
 68 |         # Resize the image and convert it to a tensor
 69 |         img_resized = img.resize((width, height)) # hard code here
 70 |         img_tensor = torch.from_numpy(np.array(img_resized)).float()
 71 | 
 72 |         # Normalize the image by scaling pixel values to [-1, 1]
 73 |         img_normalized = img_tensor / 127.5 - 1
 74 | 
 75 |         img_normalized = img_normalized.permute(2, 0, 1)  # For RGB images
 76 | 
 77 |         if(logicalNot):
 78 |             img_normalized = torch.logical_not(masks_pixel_values[i])*(-1) + masks_pixel_values[i]*img_normalized
 79 |         condition_pixel_values[i] = img_normalized
 80 |         
 81 |     return [prompt], {
 82 |             'conditioning_pixel_values': condition_pixel_values, # [-1,1]
 83 |             'masks': masks_pixel_values# [0,1]
 84 |             }
 85 |     
 86 | def get_image_files(folder_path):
 87 |     image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.gif', '*.bmp', '*.tiff', '*.webp']
 88 |     
 89 |     image_files = []
 90 |     for ext in image_extensions:
 91 |         image_files.extend(glob.glob(os.path.join(folder_path, ext)))
 92 | 
 93 |     image_names = [os.path.basename(file) for file in image_files]
 94 |     
 95 |     return image_names
 96 | 
 97 | single_view = args.single_view
 98 | super_resolution = args.super_resolution
 99 | base_model_path = args.base_model_path
100 | if(single_view):
101 |     mv_unet_path = base_model_path + "/unet/single/ema-checkpoint"
102 |     print(mv_unet_path)
103 |     tokenizer = CLIPTokenizer.from_pretrained(base_model_path, subfolder="tokenizer")
104 | else:
105 |     mv_unet_path = base_model_path + "/unet/sparse/ema-checkpoint"
106 |     print(mv_unet_path)
107 |     tokenizer = CLIPTokenizer.from_pretrained(base_model_path, subfolder="tokenizer")
108 | rgb_model = mvdream_diffusion_model(base_model_path,mv_unet_path,tokenizer,seed=12345)
109 | mv_unet_path = base_model_path + "/unet/SR/ema-checkpoint"
110 | rgb_model_SR = mvdream_diffusion_model_SR(base_model_path,mv_unet_path,tokenizer,quantization=False,seed=12345)
111 |         
112 | 
113 | source_imgs_dir = args.source_imgs_dir
114 | warp_root_dir = args.warp_root_dir
115 | output_root_dir = args.output_dir
116 | os.makedirs(output_root_dir, exist_ok=True)
117 | 
118 | 
119 | height_mvd = 512
120 | width_mvd = 512
121 | masks_infer = []
122 | warps_infer = []
123 | input_names = []
124 | 
125 | gt_num_b = 0
126 | mask2 = np.ones((height_mvd,width_mvd), dtype=np.float32)
127 | 
128 | image_names_ref = get_image_files(source_imgs_dir)
129 | fimage = Image.open(os.path.join(source_imgs_dir + image_names_ref[0]))
130 | (width, height)= fimage.size
131 | for imn in image_names_ref:
132 |     masks_infer.append(Image.fromarray(np.repeat(np.expand_dims(np.round(mask2*255.).astype(np.uint8),axis=2),3,axis=2)).resize((width_mvd, height_mvd)))
133 |     warps_infer.append(Image.open(os.path.join(source_imgs_dir + imn)))
134 |     input_names.append(imn)
135 |     gt_num_b = gt_num_b + 1
136 | 
137 | 
138 | 
139 | image_files = glob.glob(os.path.join(warp_root_dir, "warp_*"))
140 | image_names = [os.path.basename(image) for image in image_files]
141 | 
142 | image_names.sort()
143 | print(image_names)
144 | for ins in image_names:
145 |     warps_infer.append(Image.open(os.path.join(warp_root_dir, ins))) 
146 |     masks_infer.append(Image.open(os.path.join(warp_root_dir, ins.replace('warp','mask'))))
147 |     input_names.append(ins)
148 | 
149 | 
150 | 
151 | print('sequence length:', len(warps_infer))
152 | images_predict = []
153 | images_mask_p = []
154 | images_predict_names = []
155 | 
156 | grounp_size = len(masks_infer)
157 | for i in range(0, len(masks_infer[gt_num_b:]), grounp_size):
158 |     if(len(images_predict)!=0):
159 |         masks_infer_batch = masks_infer[:gt_num_b] + [masks_infer_batch[-1]] + masks_infer[(gt_num_b+i):(i+gt_num_b+grounp_size)]
160 |         warp_infer_batch = warps_infer[:gt_num_b] + [images_predict[-1]] + warps_infer[(gt_num_b+i):(i+gt_num_b+grounp_size)]
161 |         input_names_batch = input_names[:gt_num_b] + [input_names_batch[len(masks_infer_batch)//2]] + [input_names_batch[-1]] + input_names[(gt_num_b+i):(i+gt_num_b+grounp_size)]
162 |     else:
163 |         masks_infer_batch = masks_infer[:gt_num_b] + masks_infer[(gt_num_b+i):(i+gt_num_b+grounp_size)]
164 |         warp_infer_batch = warps_infer[:gt_num_b] + warps_infer[(gt_num_b+i):(i+gt_num_b+grounp_size)]
165 |         input_names_batch = input_names[:gt_num_b] + input_names[(gt_num_b+i):(i+gt_num_b+grounp_size)]
166 | 
167 |     prompt, batch = PIL2tensor(height_mvd,width_mvd,len(masks_infer_batch),masks_infer_batch,warp_infer_batch,logicalNot=False)
168 |     if(len(images_predict)!=0):
169 |         images_predict_batch = rgb_model.inference_next_frame(prompt,batch,len(masks_infer_batch),height_mvd,width_mvd,gt_num_frames=gt_num_b,output_type='pil')
170 |         for jj in range(gt_num_b+1,len(images_predict_batch)):
171 |             images_predict.append(images_predict_batch[jj])
172 |             images_mask_p.append(batch['masks'][0][jj][0].cpu().numpy())
173 |             images_predict_names.append(input_names_batch[jj])
174 |     else:
175 |         images_predict_batch = rgb_model.inference_next_frame(prompt,batch,len(masks_infer_batch),height_mvd,width_mvd,gt_num_frames=gt_num_b,output_type='pil')
176 |         for jj in range(gt_num_b,len(images_predict_batch)):
177 |             images_predict.append(images_predict_batch[jj])
178 |             images_mask_p.append(batch['masks'][0][jj][0].cpu().numpy())
179 |             images_predict_names.append(input_names_batch[jj])
180 |         
181 | for jj in range(len(images_predict)):
182 |     images_predict[jj].resize((width, height)).save(os.path.join(output_root_dir,"predict_{}.jpg".format(images_predict_names[jj])))
183 | 
184 | if(super_resolution):
185 |     del mvdream_diffusion_model
186 |     gc.collect()
187 |     torch.cuda.empty_cache()
188 | 
189 |     masks_infer_SR = []
190 |     warps_infer_SR = []
191 |     mask2 = np.ones((height_mvd*2,width_mvd*2), dtype=np.float32)
192 | 
193 |     for imn in image_names_ref:
194 |         masks_infer_SR.append(Image.fromarray(np.repeat(np.expand_dims(np.round(mask2*255.).astype(np.uint8),axis=2),3,axis=2)).resize((width_mvd, height_mvd)))
195 |         warps_infer_SR.append(Image.open(os.path.join(source_imgs_dir + imn)))
196 | 
197 |     for i in range(len(images_predict)):
198 |         masks_infer_SR.append(masks_infer[i])
199 |         warps_infer_SR.append(images_predict[i])
200 | 
201 |     images_predict = []
202 |     images_predict_names = []
203 |     # grounp_size = min((len(masks_infer_SR) + 5)//2,50)
204 |     grounp_size = (len(masks_infer_SR) + 3) // 2
205 |     # grounp_size = (len(masks_infer_SR) + 3)
206 |     print('grounp_size:',grounp_size)
207 |     for i in range(0, len(masks_infer_SR[gt_num_b:]), grounp_size):
208 |         if(len(images_predict)!=0):
209 |             masks_infer_batch = masks_infer_SR[:gt_num_b] + [masks_infer_batch[len(masks_infer_batch)//2]] + [masks_infer_batch[-1]] + masks_infer_SR[(gt_num_b+i):(i+gt_num_b+grounp_size)]
210 |             warp_infer_batch = warps_infer_SR[:gt_num_b] + [images_predict[len(images_predict)//2]] + [images_predict[-1]] + warps_infer_SR[(gt_num_b+i):(i+gt_num_b+grounp_size)]
211 |             input_names_batch = input_names[:gt_num_b] + [input_names_batch[len(masks_infer_batch)//2]] + [input_names_batch[-1]] + input_names[(gt_num_b+i):(i+gt_num_b+grounp_size)]
212 |         else:
213 |             masks_infer_batch = masks_infer_SR[:gt_num_b] + masks_infer_SR[(gt_num_b+i):(i+gt_num_b+grounp_size)]
214 |             warp_infer_batch = warps_infer_SR[:gt_num_b] + warps_infer_SR[(gt_num_b+i):(i+gt_num_b+grounp_size)]
215 |             input_names_batch = input_names[:gt_num_b] + input_names[(gt_num_b+i):(i+gt_num_b+grounp_size)]
216 | 
217 |         
218 |         prompt, batch = PIL2tensor(height_mvd*2,width_mvd*2,len(masks_infer_batch),masks_infer_batch,warp_infer_batch)
219 |         if(len(images_predict)!=0):
220 |             images_predict_batch = rgb_model_SR.inference_next_frame(prompt,batch,len(masks_infer_batch),height_mvd*2,width_mvd*2,gt_num_frames=gt_num_b,output_type='pil')
221 |             for jj in range(gt_num_b+2,len(images_predict_batch)):
222 |                 images_predict.append(images_predict_batch[jj])
223 |                 images_predict_names.append(input_names_batch[jj])
224 |         else:
225 |             images_predict_batch = rgb_model_SR.inference_next_frame(prompt,batch,len(masks_infer_batch),height_mvd*2,width_mvd*2,gt_num_frames=gt_num_b,output_type='pil')
226 |             for jj in range(gt_num_b,len(images_predict_batch)):
227 |                 images_predict.append(images_predict_batch[jj])
228 |                 images_predict_names.append(input_names_batch[jj])
229 |         gc.collect()
230 |         torch.cuda.empty_cache()
231 | 
232 | 
233 |     for jj in range(len(images_predict)):
234 |         images_predict[jj].resize((width, height)).save(os.path.join(output_root_dir,"SR_predict_{}.jpg".format(images_predict_names[jj])))
235 |     
236 | 


--------------------------------------------------------------------------------
/mv_diffusion.py:
--------------------------------------------------------------------------------
 1 | from diffusers import (
 2 |     UniPCMultistepScheduler,
 3 |     DDIMScheduler
 4 | )
 5 | 
 6 | from pipeline_mvd_warp_mix_classifier import MVDreamPipeline
 7 | import torch
 8 | import numpy as np
 9 | from PIL import Image
10 | from mv_unet import MultiViewUNetModel
11 | from accelerate.utils import ProjectConfiguration, set_seed
12 | import argparse
13 | import os
14 | import sys
15 | from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection
16 | 
17 | 
18 | 
19 | class mvdream_diffusion_model:
20 |     def __init__(self, base_model_path,mv_unet_path,tokenizer,seed=12345):
21 | 
22 |         generator = torch.Generator("cuda").manual_seed(seed)
23 |         
24 | 
25 |         unet = MultiViewUNetModel.from_pretrained(
26 |             mv_unet_path,
27 |             torch_dtype=torch.float16
28 |         )
29 |         feature_extractor: CLIPImageProcessor = CLIPImageProcessor.from_pretrained(base_model_path + "/CLIP-ViT-H-14-laion2B-s32B-b79K")
30 |         image_encoder: CLIPVisionModelWithProjection = CLIPVisionModelWithProjection.from_pretrained(base_model_path + "/CLIP-ViT-H-14-laion2B-s32B-b79K")
31 |         self.pipe = MVDreamPipeline.from_pretrained(
32 |             base_model_path, 
33 |             unet=unet,
34 |             torch_dtype=torch.float16,
35 |             feature_extractor=feature_extractor,
36 |             image_encoder=image_encoder,
37 |         )
38 | 
39 |         
40 |         self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config, timestep_scaling="trailing", rescale_betas_zero_snr=True)
41 | 
42 |         self.pipe.enable_model_cpu_offload()
43 |         self.pipe.vae.enable_slicing()
44 |         self.pipe.vae.enable_tiling()
45 | 
46 | 
47 | 
48 | 
49 |     def inference_next_frame(self, prompt, batch, num_frames, height, width, gt_num_frames=1, output_type='pil'):
50 |         batch['conditioning_pixel_values'] = torch.unsqueeze(batch['conditioning_pixel_values'], 0)
51 |         batch['masks'] = torch.unsqueeze(batch['masks'], 0)
52 | 
53 |         image, image_warp = self.pipe(
54 |             prompt=prompt,
55 |             image=batch['conditioning_pixel_values'],
56 |             masks=batch['masks'],
57 |             height=height,
58 |             width=width,
59 |             guidance_scale=2.0,
60 |             # guidance_scale=1.0,
61 |             num_inference_steps=50, 
62 |             guidance_rescale=0.0,
63 |             # guidance_rescale=0.5,
64 |             elevation=0,
65 |             num_frames=batch['conditioning_pixel_values'].shape[1],
66 |             condition_num_frames=gt_num_frames,
67 |             gt_frame=None,  
68 |             output_type=output_type
69 |         )
70 |         
71 |         return image
72 | 


--------------------------------------------------------------------------------
/mv_diffusion_SR.py:
--------------------------------------------------------------------------------
 1 | import sys
 2 | from diffusers import (
 3 |     UniPCMultistepScheduler,
 4 |     DDIMScheduler
 5 | )
 6 | from pipeline_mvd_warp_mix_classifier_SR import MVDreamPipeline
 7 | import torch
 8 | import cv2
 9 | import numpy as np
10 | from PIL import Image
11 | from mv_unet import MultiViewUNetModel
12 | from accelerate.utils import ProjectConfiguration, set_seed
13 | import argparse
14 | import os
15 | import sys
16 | from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor, CLIPVisionModelWithProjection
17 | 
18 | 
19 | 
20 | class mvdream_diffusion_model:
21 |     def __init__(self, base_model_path,mv_unet_path,tokenizer,quantization=False,seed=12345):
22 | 
23 |         generator = torch.Generator("cuda").manual_seed(seed)
24 | 
25 |         unet = MultiViewUNetModel.from_pretrained(
26 |         mv_unet_path,
27 |         torch_dtype=torch.float16
28 |         )
29 |         feature_extractor: CLIPImageProcessor = CLIPImageProcessor.from_pretrained(base_model_path + "/CLIP-ViT-H-14-laion2B-s32B-b79K")
30 |         image_encoder: CLIPVisionModelWithProjection = CLIPVisionModelWithProjection.from_pretrained(base_model_path + "/CLIP-ViT-H-14-laion2B-s32B-b79K")
31 |         self.pipe = MVDreamPipeline.from_pretrained(
32 |             base_model_path, # remote weights
33 |             unet=unet,
34 |             torch_dtype=torch.float16,
35 |             feature_extractor=feature_extractor,
36 |             image_encoder=image_encoder,
37 |         )
38 |         
39 |         
40 |         self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config, timestep_scaling="trailing", rescale_betas_zero_snr=True)
41 |         
42 |         self.pipe.enable_model_cpu_offload()
43 |         self.pipe.vae.enable_slicing()
44 |         self.pipe.vae.enable_tiling()
45 | 
46 | 
47 | 
48 | 
49 |     def inference_next_frame(self, prompt, batch, num_frames, height, width, gt_num_frames=1, output_type='pil'):
50 |         batch['conditioning_pixel_values'] = torch.unsqueeze(batch['conditioning_pixel_values'], 0)
51 |         batch['masks'] = torch.unsqueeze(batch['masks'], 0)
52 |         
53 |         image, image_warp = self.pipe(
54 |             prompt=prompt,
55 |             image=batch['conditioning_pixel_values'],
56 |             masks=batch['masks'],
57 |             height=height,
58 |             width=width,
59 |             guidance_scale=1.0,
60 |             num_inference_steps=30, 
61 |             guidance_rescale=0.0, 
62 |             elevation=0,
63 |             num_frames=batch['conditioning_pixel_values'].shape[1],
64 |             condition_num_frames=gt_num_frames,
65 |             gt_frame=None,  
66 |             output_type=output_type
67 |         )
68 |         
69 |         return image
70 | 
71 | 


--------------------------------------------------------------------------------
/pipeline_mvd_warp_mix_classifier.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn.functional as F
  3 | import inspect
  4 | import numpy as np
  5 | from typing import Callable, List, Optional, Union
  6 | from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor,CLIPVisionModelWithProjection
  7 | from diffusers import AutoencoderKL, DiffusionPipeline
  8 | from diffusers.utils import (
  9 |     deprecate,
 10 |     is_accelerate_available,
 11 |     is_accelerate_version,
 12 |     logging,
 13 | )
 14 | from diffusers import DDPMScheduler
 15 | from diffusers.configuration_utils import FrozenDict
 16 | from diffusers.schedulers import DDIMScheduler
 17 | from diffusers.utils.torch_utils import randn_tensor
 18 | 
 19 | from einops import rearrange, repeat
 20 | 
 21 | from mv_unet import MultiViewUNetModel, get_camera
 22 | 
 23 | 
 24 | logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 25 | 
 26 | 
 27 | def custom_decay_function_weight(t):
 28 |     # t = torch.tensor(t, dtype=torch.float32)
 29 |     
 30 |     t_peak = 200
 31 |     # t_peak = 333
 32 |     t_slow_decay_end = 60
 33 |     # t_slow_decay_end = 40
 34 |     value_at_slow_decay_end = 0.8
 35 |     
 36 |     # Initialize output tensor
 37 |     values = torch.zeros_like(t)
 38 |     
 39 |     # Slow decay mask (t >= t_slow_decay_end)
 40 |     slow_decay_mask = t >= t_slow_decay_end
 41 |     values[slow_decay_mask] = 1.0 - (1.0 - value_at_slow_decay_end) * (t_peak - t[slow_decay_mask]) / (t_peak - t_slow_decay_end)
 42 |     
 43 |     # Fast decay mask (t < t_slow_decay_end)
 44 |     fast_decay_mask = t < t_slow_decay_end
 45 |     values[fast_decay_mask] = value_at_slow_decay_end * torch.exp(-0.075 * (t_slow_decay_end - t[fast_decay_mask]))
 46 |     
 47 |     # Ensure values are within [0, 1]
 48 |     values = torch.clamp(values, 0, 1)
 49 |     
 50 |     return values
 51 | 
 52 | 
 53 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
 54 |     """
 55 |     Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
 56 |     Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
 57 |     """
 58 |     std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
 59 |     std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
 60 |     # rescale the results from guidance (fixes overexposure)
 61 |     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
 62 |     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
 63 |     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
 64 |     return noise_cfg
 65 | 
 66 | def compute_weights(timesteps, max_time=333.0, min_time=0.0):
 67 |     weights = (timesteps - min_time) / (max_time - min_time)
 68 |     weights = torch.clamp(weights, 0.0, 1.0)  # 限制权重在0和1之间
 69 |     return weights
 70 | 
 71 | class MVDreamPipeline(DiffusionPipeline):
 72 | 
 73 |     # _optional_components = ["feature_extractor", "image_encoder"]
 74 |     # _optional_components = ["image_encoder"]
 75 | 
 76 |     def __init__(
 77 |         self,
 78 |         vae: AutoencoderKL,
 79 |         unet: MultiViewUNetModel,
 80 |         tokenizer: CLIPTokenizer,
 81 |         text_encoder: CLIPTextModel,
 82 |         scheduler: DDIMScheduler,
 83 |         # imagedream variant
 84 |         feature_extractor: CLIPImageProcessor,
 85 |         image_encoder: CLIPVisionModelWithProjection,
 86 |         requires_safety_checker: bool = False,
 87 |     ):
 88 |         super().__init__()
 89 | 
 90 |         if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:  # type: ignore
 91 |             deprecation_message = (
 92 |                 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
 93 |                 f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "  # type: ignore
 94 |                 "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
 95 |                 " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
 96 |                 " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
 97 |                 " file"
 98 |             )
 99 |             deprecate(
100 |                 "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
101 |             )
102 |             new_config = dict(scheduler.config)
103 |             new_config["steps_offset"] = 1
104 |             scheduler._internal_dict = FrozenDict(new_config)
105 | 
106 |         if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:  # type: ignore
107 |             deprecation_message = (
108 |                 f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
109 |                 " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
110 |                 " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
111 |                 " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
112 |                 " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
113 |             )
114 |             deprecate(
115 |                 "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
116 |             )
117 |             new_config = dict(scheduler.config)
118 |             new_config["clip_sample"] = False
119 |             scheduler._internal_dict = FrozenDict(new_config)
120 | 
121 |         self.register_modules(
122 |             vae=vae,
123 |             unet=unet,
124 |             scheduler=scheduler,
125 |             tokenizer=tokenizer,
126 |             text_encoder=text_encoder,
127 |             feature_extractor=feature_extractor,
128 |             image_encoder=image_encoder,
129 |         )
130 |         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
131 |         self.register_to_config(requires_safety_checker=requires_safety_checker)
132 | 
133 |     def enable_vae_slicing(self):
134 |         r"""
135 |         Enable sliced VAE decoding.
136 | 
137 |         When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
138 |         steps. This is useful to save some memory and allow larger batch sizes.
139 |         """
140 |         self.vae.enable_slicing()
141 | 
142 |     def disable_vae_slicing(self):
143 |         r"""
144 |         Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
145 |         computing decoding in one step.
146 |         """
147 |         self.vae.disable_slicing()
148 | 
149 |     def enable_vae_tiling(self):
150 |         r"""
151 |         Enable tiled VAE decoding.
152 | 
153 |         When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
154 |         several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
155 |         """
156 |         self.vae.enable_tiling()
157 | 
158 |     def disable_vae_tiling(self):
159 |         r"""
160 |         Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
161 |         computing decoding in one step.
162 |         """
163 |         self.vae.disable_tiling()
164 | 
165 |     def enable_sequential_cpu_offload(self, gpu_id=0):
166 |         r"""
167 |         Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
168 |         text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
169 |         `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
170 |         Note that offloading happens on a submodule basis. Memory savings are higher than with
171 |         `enable_model_cpu_offload`, but performance is lower.
172 |         """
173 |         if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
174 |             from accelerate import cpu_offload
175 |         else:
176 |             raise ImportError(
177 |                 "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
178 |             )
179 | 
180 |         device = torch.device(f"cuda:{gpu_id}")
181 | 
182 |         if self.device.type != "cpu":
183 |             self.to("cpu", silence_dtype_warnings=True)
184 |             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
185 | 
186 |         for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
187 |             cpu_offload(cpu_offloaded_model, device)
188 | 
189 |     def enable_model_cpu_offload(self, gpu_id=0):
190 |         r"""
191 |         Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
192 |         to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
193 |         method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
194 |         `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
195 |         """
196 |         if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
197 |             from accelerate import cpu_offload_with_hook
198 |         else:
199 |             raise ImportError(
200 |                 "`enable_model_offload` requires `accelerate v0.17.0` or higher."
201 |             )
202 | 
203 |         device = torch.device(f"cuda:{gpu_id}")
204 | 
205 |         if self.device.type != "cpu":
206 |             self.to("cpu", silence_dtype_warnings=True)
207 |             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
208 | 
209 |         hook = None
210 |         for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
211 |             _, hook = cpu_offload_with_hook(
212 |                 cpu_offloaded_model, device, prev_module_hook=hook
213 |             )
214 | 
215 |         # We'll offload the last model manually.
216 |         self.final_offload_hook = hook
217 | 
218 |     @property
219 |     def _execution_device(self):
220 |         r"""
221 |         Returns the device on which the pipeline's models will be executed. After calling
222 |         `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
223 |         hooks.
224 |         """
225 |         if not hasattr(self.unet, "_hf_hook"):
226 |             return self.device
227 |         for module in self.unet.modules():
228 |             if (
229 |                 hasattr(module, "_hf_hook")
230 |                 and hasattr(module._hf_hook, "execution_device")
231 |                 and module._hf_hook.execution_device is not None
232 |             ):
233 |                 return torch.device(module._hf_hook.execution_device)
234 |         return self.device
235 | 
236 |     def _encode_prompt(
237 |         self,
238 |         prompt,
239 |         device,
240 |         num_images_per_prompt,
241 |         do_classifier_free_guidance: bool,
242 |         negative_prompt=None,
243 |     ):
244 |         r"""
245 |         Encodes the prompt into text encoder hidden states.
246 | 
247 |         Args:
248 |              prompt (`str` or `List[str]`, *optional*):
249 |                 prompt to be encoded
250 |             device: (`torch.device`):
251 |                 torch device
252 |             num_images_per_prompt (`int`):
253 |                 number of images that should be generated per prompt
254 |             do_classifier_free_guidance (`bool`):
255 |                 whether to use classifier free guidance or not
256 |             negative_prompt (`str` or `List[str]`, *optional*):
257 |                 The prompt or prompts not to guide the image generation. If not defined, one has to pass
258 |                 `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
259 |                 Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
260 |             prompt_embeds (`torch.FloatTensor`, *optional*):
261 |                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
262 |                 provided, text embeddings will be generated from `prompt` input argument.
263 |             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
264 |                 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
265 |                 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
266 |                 argument.
267 |         """
268 |         if prompt is not None and isinstance(prompt, str):
269 |             batch_size = 1
270 |         elif prompt is not None and isinstance(prompt, list):
271 |             batch_size = len(prompt)
272 |         else:
273 |             raise ValueError(
274 |                 f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
275 |             )
276 | 
277 |         text_inputs = self.tokenizer(
278 |             prompt,
279 |             padding="max_length",
280 |             max_length=self.tokenizer.model_max_length,
281 |             truncation=True,
282 |             return_tensors="pt",
283 |         )
284 |         text_input_ids = text_inputs.input_ids
285 |         untruncated_ids = self.tokenizer(
286 |             prompt, padding="longest", return_tensors="pt"
287 |         ).input_ids
288 | 
289 |         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
290 |             text_input_ids, untruncated_ids
291 |         ):
292 |             removed_text = self.tokenizer.batch_decode(
293 |                 untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
294 |             )
295 |             logger.warning(
296 |                 "The following part of your input was truncated because CLIP can only handle sequences up to"
297 |                 f" {self.tokenizer.model_max_length} tokens: {removed_text}"
298 |             )
299 | 
300 |         if (
301 |             hasattr(self.text_encoder.config, "use_attention_mask")
302 |             and self.text_encoder.config.use_attention_mask
303 |         ):
304 |             attention_mask = text_inputs.attention_mask.to(device)
305 |         else:
306 |             attention_mask = None
307 | 
308 |         prompt_embeds = self.text_encoder(
309 |             text_input_ids.to(device),
310 |             attention_mask=attention_mask,
311 |         )
312 |         prompt_embeds = prompt_embeds[0]
313 | 
314 |         prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
315 | 
316 |         bs_embed, seq_len, _ = prompt_embeds.shape
317 |         # duplicate text embeddings for each generation per prompt, using mps friendly method
318 |         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
319 |         prompt_embeds = prompt_embeds.view(
320 |             bs_embed * num_images_per_prompt, seq_len, -1
321 |         )
322 | 
323 |         # get unconditional embeddings for classifier free guidance
324 |         if do_classifier_free_guidance:
325 |             uncond_tokens: List[str]
326 |             if negative_prompt is None:
327 |                 uncond_tokens = [""] * batch_size
328 |             elif type(prompt) is not type(negative_prompt):
329 |                 raise TypeError(
330 |                     f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
331 |                     f" {type(prompt)}."
332 |                 )
333 |             elif isinstance(negative_prompt, str):
334 |                 uncond_tokens = [negative_prompt]
335 |             elif batch_size != len(negative_prompt):
336 |                 raise ValueError(
337 |                     f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
338 |                     f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
339 |                     " the batch size of `prompt`."
340 |                 )
341 |             else:
342 |                 uncond_tokens = negative_prompt
343 | 
344 |             max_length = prompt_embeds.shape[1]
345 |             uncond_input = self.tokenizer(
346 |                 uncond_tokens,
347 |                 padding="max_length",
348 |                 max_length=max_length,
349 |                 truncation=True,
350 |                 return_tensors="pt",
351 |             )
352 |             if (
353 |                 hasattr(self.text_encoder.config, "use_attention_mask")
354 |                 and self.text_encoder.config.use_attention_mask
355 |             ):
356 |                 attention_mask = uncond_input.attention_mask.to(device)
357 |             else:
358 |                 attention_mask = None
359 | 
360 |             negative_prompt_embeds = self.text_encoder(
361 |                 uncond_input.input_ids.to(device),
362 |                 attention_mask=attention_mask,
363 |             )
364 |             negative_prompt_embeds = negative_prompt_embeds[0]
365 | 
366 |             # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
367 |             seq_len = negative_prompt_embeds.shape[1]
368 | 
369 |             negative_prompt_embeds = negative_prompt_embeds.to(
370 |                 dtype=self.text_encoder.dtype, device=device
371 |             )
372 | 
373 |             negative_prompt_embeds = negative_prompt_embeds.repeat(
374 |                 1, num_images_per_prompt, 1
375 |             )
376 |             negative_prompt_embeds = negative_prompt_embeds.view(
377 |                 batch_size * num_images_per_prompt, seq_len, -1
378 |             )
379 | 
380 |             # For classifier free guidance, we need to do two forward passes.
381 |             # Here we concatenate the unconditional and text embeddings into a single batch
382 |             # to avoid doing two forward passes
383 |             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
384 | 
385 |         return prompt_embeds
386 | 
387 |     def decode_latents(self, latents):
388 |         latents = 1 / self.vae.config.scaling_factor * latents
389 |         image = self.vae.decode(latents).sample
390 |         image = (image / 2 + 0.5).clamp(0, 1)
391 |         # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
392 |         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
393 |         return image
394 | 
395 |     def prepare_extra_step_kwargs(self, generator, eta):
396 |         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
397 |         # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
398 |         # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
399 |         # and should be between [0, 1]
400 | 
401 |         accepts_eta = "eta" in set(
402 |             inspect.signature(self.scheduler.step).parameters.keys()
403 |         )
404 |         extra_step_kwargs = {}
405 |         if accepts_eta:
406 |             extra_step_kwargs["eta"] = eta
407 | 
408 |         # check if the scheduler accepts generator
409 |         accepts_generator = "generator" in set(
410 |             inspect.signature(self.scheduler.step).parameters.keys()
411 |         )
412 |         if accepts_generator:
413 |             extra_step_kwargs["generator"] = generator
414 |         return extra_step_kwargs
415 | 
416 |     def prepare_latents(
417 |         self,
418 |         batch_size,
419 |         num_channels_latents,
420 |         height,
421 |         width,
422 |         dtype,
423 |         device,
424 |         generator,
425 |         latents=None,
426 |     ):
427 |         shape = (
428 |             batch_size,
429 |             num_channels_latents,
430 |             height // self.vae_scale_factor,
431 |             width // self.vae_scale_factor,
432 |         )
433 |         if isinstance(generator, list) and len(generator) != batch_size:
434 |             raise ValueError(
435 |                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
436 |                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
437 |             )
438 | 
439 |         if latents is None:
440 |             latents = randn_tensor(
441 |                 shape, generator=generator, device=device, dtype=dtype
442 |             )
443 |         else:
444 |             latents = latents.to(device)
445 | 
446 |         # scale the initial noise by the standard deviation required by the scheduler
447 |         latents = latents * self.scheduler.init_noise_sigma
448 |         return latents
449 | 
450 |     def encode_image(self, image, device, num_images_per_prompt):
451 |         # dtype = next(self.feature_extractor.parameters()).dtype
452 |         dtype = self.text_encoder.dtype
453 |         # if image.dtype == np.float32:
454 |         #     image = (image * 255).astype(np.uint8)
455 |         
456 |         # image_clip = (((image + 1.0)/2.0)* 255)
457 |         image_clip = (((image + 1.0)/2.0)* 1.0)
458 |         image_clip = self.feature_extractor(image_clip, return_tensors="pt",do_rescale=False).pixel_values
459 |         image_clip = image_clip.to(device=device, dtype=dtype)
460 |         
461 |         encoded_embeds = self.image_encoder(image_clip, output_hidden_states=False)
462 |         image_embeds = encoded_embeds.image_embeds
463 |         image_embeds = 0.2 * image_embeds.unsqueeze(-2).repeat(1, 77, 1).to(dtype)
464 | 
465 |         return torch.zeros_like(image_embeds), image_embeds
466 | 
467 |     def encode_image_latents(self, image, device, num_images_per_prompt):
468 |         
469 |         dtype = next(self.image_encoder.parameters()).dtype
470 | 
471 |         image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
472 |         image = 2 * image - 1
473 |         image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
474 |         image = image.to(dtype=dtype)
475 | 
476 |         posterior = self.vae.encode(image).latent_dist
477 |         latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
478 |         latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
479 | 
480 |         return torch.zeros_like(latents), latents
481 | 
482 |     def decode_latents_to_pil_save(self, latents, path):
483 |         image_list = self.decode_latents(latents)
484 |         image_list = self.numpy_to_pil(image_list)
485 |         for i, image in enumerate(image_list):
486 |             image.save(path+f'_{i}.jpg')
487 | 
488 |     
489 |     def save_latents(self, noisy_latents_warp,latents,mix,name,t):
490 |         image_warp = self.decode_latents(noisy_latents_warp)
491 |         image_warp = self.numpy_to_pil(image_warp)
492 |         image = self.decode_latents(latents)
493 |         image = self.numpy_to_pil(image)
494 |         image_mix = self.decode_latents(mix)
495 |         image_mix = self.numpy_to_pil(image_mix)
496 |         grid = np.concatenate(
497 |             [
498 |                 np.concatenate([image_warp[0], image_warp[1], image_warp[2]], axis=0),
499 |                 np.concatenate([image_warp[3], image_warp[4], image_warp[5]], axis=0),
500 |                 np.concatenate([image[0], image[1], image[2]], axis=0),
501 |                 np.concatenate([image[3], image[4], image[5]], axis=0),
502 |                 np.concatenate([image_mix[0], image_mix[1], image_mix[2]], axis=0),
503 |                 np.concatenate([image_mix[3], image_mix[4], image_mix[5]], axis=0),
504 |             ],
505 |             axis=1,
506 |         )
507 |         kiui.write_image(f'/share/project/mabaorui/lucidDream/LucidDreamer-main/outputs/tt/warp/{t}_{name}.jpg', grid)
508 |     
509 |     @torch.no_grad()
510 |     def __call__(
511 |         self,
512 |         prompt: str = "",# 'a cute owl'
513 |         image: Optional[np.ndarray] = None,# NOne
514 |         masks: Optional[np.ndarray] = None,# NOne
515 |         height: int = 256,# 256
516 |         width: int = 256,# 256
517 |         elevation: float = 0,# 0
518 |         num_inference_steps: int = 50,# 30
519 |         guidance_scale: float = 7.0,# 5
520 |         negative_prompt: str = None,# ''
521 |         num_images_per_prompt: int = 1,# 1
522 |         eta: float = 0.0,# 0.0
523 |         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,# None
524 |         output_type: Optional[str] = "numpy", # pil, numpy, latents | numpy
525 |         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
526 |         callback_steps: int = 1,
527 |         num_frames: int = 4, # 4
528 |         device=torch.device("cuda:0"),
529 |         guidance_rescale: float = 0.0,
530 |         # gt_num: int=3,
531 |         condition_num_frames: int=3,
532 |         gt_frame: Optional[np.ndarray] = None,
533 |     ):
534 |         gt_num = condition_num_frames
535 |         
536 |         self.unet = self.unet.to(device=device)
537 |         self.vae = self.vae.to(device=device)
538 |         self.text_encoder = self.text_encoder.to(device=device)
539 |         weight_dtype = next(self.vae.parameters()).dtype
540 |         self.guidance_rescale = guidance_rescale
541 | 
542 |         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
543 |         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
544 |         # corresponds to doing no classifier free guidance.
545 |         # do_classifier_free_guidance = guidance_scale > 1.0
546 |         do_classifier_free_guidance = (guidance_scale != 0.0)
547 |         
548 |         # Prepare timesteps
549 |         self.scheduler.set_timesteps(num_inference_steps, device=device)# DDIMScheduler
550 |         timesteps = self.scheduler.timesteps
551 |         # import pdb;pdb.set_trace()
552 |         if self.scheduler.timestep_scaling == "trailing":
553 |             step_ratio = 999 // num_inference_steps
554 |             timesteps = torch.from_numpy(np.round(np.arange(999, 0, -step_ratio)).astype(np.int64).copy()).to(timesteps.device)
555 | 
556 |         # imagedream variant
557 |         if image is not None:
558 |             image = rearrange(image, "b f c h w -> (b f) c h w", f=num_frames)
559 |             img_latents = self.vae.encode(image.to(weight_dtype).to(device)).latent_dist.sample()
560 |             img_latents = img_latents * self.vae.config.scaling_factor # [b*f, c, h, w] shape=[6, 4, 32, 32]
561 |             # import pdb;pdb.set_trace()
562 |             # img_latents[:gt_num] = img_latents[:gt_num] * self.vae.config.scaling_factor
563 |         if gt_frame is not None:    
564 |             image_gt = rearrange(gt_frame, "b f c h w -> (b f) c h w", f=num_frames)
565 |             img_latents_gt = self.vae.encode(image_gt.to(weight_dtype).to(device)).latent_dist.sample()
566 |             img_latents_gt = img_latents_gt * self.vae.config.scaling_factor # [b*f, c, h, w] shape=[6, 4, 32, 32]
567 |         
568 |         self.image_encoder = self.image_encoder.to(device=device)
569 |         image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt) 
570 |         image_embeds_pos = image_embeds_pos[0:1]
571 |         
572 |         # import pdb;pdb.set_trace()
573 |         if masks is not None:
574 |             masks = rearrange(masks, "b f c h w -> (b f) c h w", f=num_frames)# [b*f, c, h, w]
575 |             mask_latents = torch.nn.functional.interpolate(
576 |                 masks, 
577 |                 size=(
578 |                     height // 8, 
579 |                     width // 8
580 |                 )
581 |             ).to(weight_dtype).to(device)# [b*f, c, h, w]
582 |         
583 |         if do_classifier_free_guidance:
584 |             uncond_img_latents = torch.zeros_like(img_latents)
585 |         
586 |         prompt_embeds = self._encode_prompt(
587 |             prompt=prompt,
588 |             device=device,
589 |             num_images_per_prompt=num_images_per_prompt,
590 |             do_classifier_free_guidance=do_classifier_free_guidance,
591 |             negative_prompt=negative_prompt,
592 |         )  # type: ignore torch.Size([1, 77, 1024])
593 |         # prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
594 | 
595 |         # Prepare latent variables
596 |         actual_num_frames = num_frames
597 |         noisy_latents = self.prepare_latents(
598 |             actual_num_frames * num_images_per_prompt,
599 |             4,
600 |             height,
601 |             width,
602 |             prompt_embeds.dtype,
603 |             device,
604 |             generator,
605 |             latents=None,
606 |         )# shape=[3, 4, 32, 32]
607 |         
608 |         
609 |         
610 |         noise_scheduler = self.scheduler
611 |         if gt_frame is not None:
612 |             noise_gt = torch.randn_like(img_latents)
613 |             # t = timesteps[0]
614 |             t = 999
615 |             tt = torch.tensor([t] * actual_num_frames, dtype=img_latents.dtype, device=device)
616 |             noisy_latents = noise_scheduler.add_noise(img_latents_gt, noise_gt, tt.long())
617 |         
618 | 
619 |         latents = torch.cat([img_latents[:gt_num], noisy_latents[gt_num:]], dim=0)
620 |         
621 |         camera = None
622 | 
623 |         # Prepare extra step kwargs.
624 |         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
625 | 
626 |         # Denoising loop
627 |         assert image.shape[0] == actual_num_frames, "Currently only supports batchsize==1"
628 |         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
629 | 
630 |         time_ratio = 5
631 |         timestep_warp = (timesteps[0]//time_ratio).long()
632 |         print(timestep_warp)
633 |         noise_warp = torch.randn_like(noisy_latents)
634 |         tt = torch.tensor([timestep_warp] * actual_num_frames, dtype=noisy_latents.dtype, device=device)
635 |         noisy_latents_warp = noise_scheduler.add_noise(img_latents, noise_warp, tt.long())
636 |         
637 |    
638 |         
639 |         mix_latents_warp = torch.cat([img_latents[:gt_num], noisy_latents_warp[gt_num:]], dim=0)
640 | 
641 |         with self.progress_bar(total=num_inference_steps) as progress_bar:
642 |             for i, t in enumerate(timesteps):
643 |                 # expand the latents if we are doing classifier free guidance
644 |                 multiplier = 2 if do_classifier_free_guidance else 1
645 |                 
646 |                 latents = torch.cat([img_latents[:gt_num], latents[gt_num:]], dim=0)# [6, 4, 32, 32]
647 |                 
648 |                 timestep_warp = (t//time_ratio).long()
649 |                 noise_warp = torch.randn_like(noisy_latents)
650 |                 tt = torch.tensor([timestep_warp] * actual_num_frames, dtype=noisy_latents.dtype, device=device)
651 |                 noisy_latents_warp = noise_scheduler.add_noise(img_latents, noise_warp, tt.long())
652 |                 weights = custom_decay_function_weight(tt.float())
653 |                 weights = weights.view(tt.shape[0], 1, 1, 1).to(noisy_latents_warp.dtype)
654 |                 noisy_latents_warp_weight = weights*noisy_latents_warp + (1. - weights) * latents
655 |                 mix_latents_warp = torch.cat([img_latents[:gt_num], noisy_latents_warp_weight[gt_num:]], dim=0)
656 |                 
657 |                 
658 |                 if self.unet.input_blocks[0].__dict__['_modules']['0'].__dict__['in_channels'] == 9:
659 |                     latents = torch.cat([latents, mix_latents_warp, mask_latents], dim=1)# [b*f, c, h, w] shape=[6, 9, 32, 32]
660 |                 elif self.unet.input_blocks[0].__dict__['_modules']['0'].__dict__['in_channels'] == 5:
661 |                     latents = torch.cat([latents, mask_latents], dim=1)# [b*f, c, h, w] shape=[6, 9, 32, 32]
662 |                 
663 |                 if do_classifier_free_guidance:
664 |                     
665 |                     unc_mix_latents_warp = torch.cat([img_latents[:gt_num], uncond_img_latents[gt_num:]], dim=0)
666 |                     
667 |                     
668 |                     unc_mask_latents = torch.cat([mask_latents[:gt_num], torch.zeros_like(mask_latents[gt_num:])], dim=0)
669 |                     uncond_latents = torch.cat([latents[:,:4], unc_mix_latents_warp, unc_mask_latents], dim=1)
670 |                     latent_model_input = torch.cat([latents, uncond_latents], dim=0)# torch.Size([12, 9, 64, 64])
671 |                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
672 |                 else:
673 |                     latent_model_input = torch.cat([latents] * multiplier)
674 |                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
675 | 
676 |                 image_embeds_pos_input = image_embeds_pos.repeat(multiplier,1,1)
677 |                 unet_inputs = {
678 |                     'x': latent_model_input,# torch.Size([num_frames, 5, 32, 32])
679 |                     'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),# torch.Size([num_frames])
680 |                     'context': (prompt_embeds + image_embeds_pos_input).repeat(actual_num_frames,1,1),
681 |                     'num_frames': actual_num_frames,# 4
682 |                     'camera': None,# 
683 |                 }
684 | 
685 |                 
686 |                 noise_pred = self.unet.forward(**unet_inputs)# [6, 4, 32, 32]
687 | 
688 |                 # perform guidance
689 |                 if do_classifier_free_guidance:
690 |                     noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
691 |                     noise_pred = (1 + guidance_scale) * noise_pred_cond - guidance_scale * noise_pred_uncond
692 |                     
693 |                     
694 |                 if do_classifier_free_guidance and self.guidance_rescale > 0.0:
695 |                     
696 |                     noise_pred[:,:4] = rescale_noise_cfg(noise_pred[:,:4], noise_pred_cond[:,:4], guidance_rescale=self.guidance_rescale)
697 | 
698 |   
699 |                 latents: torch.Tensor = self.scheduler.step(
700 |                     noise_pred[:,:4], t, latents[:, :4], **extra_step_kwargs, return_dict=False
701 |                 )[0]
702 |                 
703 |                 # call the callback, if provided
704 |                 if i == len(timesteps) - 1 or (
705 |                     (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
706 |                 ):
707 |                     progress_bar.update()
708 |                     if callback is not None and i % callback_steps == 0:
709 |                         callback(i, t, latents)  # type: ignore
710 | 
711 |         # Post-processing
712 |         if output_type == "latent":
713 |             image = latents
714 |             noisy_latents_warp = noisy_latents_warp
715 |         elif output_type == "pil":
716 |             image = self.decode_latents(latents)
717 |             image = self.numpy_to_pil(image)
718 |             image_warp = self.decode_latents(noisy_latents_warp)
719 |             image_warp = self.numpy_to_pil(image_warp)
720 |         else: # numpy
721 |             image = self.decode_latents(latents)
722 |             image_warp = self.decode_latents(noisy_latents_warp)
723 | 
724 |         # Offload last model to CPU
725 |         if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
726 |             self.final_offload_hook.offload()
727 | 
728 |         return image, image_warp


--------------------------------------------------------------------------------
/pipeline_mvd_warp_mix_classifier_SR.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn.functional as F
  3 | import inspect
  4 | import numpy as np
  5 | from typing import Callable, List, Optional, Union
  6 | from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor,CLIPVisionModelWithProjection
  7 | from diffusers import AutoencoderKL, DiffusionPipeline
  8 | from diffusers.utils import (
  9 |     deprecate,
 10 |     is_accelerate_available,
 11 |     is_accelerate_version,
 12 |     logging,
 13 | )
 14 | from diffusers import DDPMScheduler
 15 | from diffusers.configuration_utils import FrozenDict
 16 | from diffusers.schedulers import DDIMScheduler
 17 | from diffusers.utils.torch_utils import randn_tensor
 18 | 
 19 | from einops import rearrange, repeat
 20 | 
 21 | from mv_unet import MultiViewUNetModel, get_camera
 22 | 
 23 | 
 24 | logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 25 | 
 26 | 
 27 | def custom_decay_function_weight(t):
 28 |     # t = torch.tensor(t, dtype=torch.float32)
 29 |     
 30 |     t_peak = 200
 31 |     # t_peak = 333
 32 |     t_slow_decay_end = 60
 33 |     value_at_slow_decay_end = 0.8
 34 |     
 35 |     # Initialize output tensor
 36 |     values = torch.zeros_like(t)
 37 |     
 38 |     # Slow decay mask (t >= t_slow_decay_end)
 39 |     slow_decay_mask = t >= t_slow_decay_end
 40 |     values[slow_decay_mask] = 1.0 - (1.0 - value_at_slow_decay_end) * (t_peak - t[slow_decay_mask]) / (t_peak - t_slow_decay_end)
 41 |     
 42 |     # Fast decay mask (t < t_slow_decay_end)
 43 |     fast_decay_mask = t < t_slow_decay_end
 44 |     values[fast_decay_mask] = value_at_slow_decay_end * torch.exp(-0.075 * (t_slow_decay_end - t[fast_decay_mask]))
 45 |     
 46 |     # Ensure values are within [0, 1]
 47 |     values = torch.clamp(values, 0, 1)
 48 |     
 49 |     return values
 50 | 
 51 | 
 52 | def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
 53 |     """
 54 |     Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
 55 |     Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
 56 |     """
 57 |     std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
 58 |     std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
 59 |     # rescale the results from guidance (fixes overexposure)
 60 |     noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
 61 |     # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
 62 |     noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
 63 |     return noise_cfg
 64 | 
 65 | def compute_weights(timesteps, max_time=333.0, min_time=0.0):
 66 |     weights = (timesteps - min_time) / (max_time - min_time)
 67 |     weights = torch.clamp(weights, 0.0, 1.0)  # 限制权重在0和1之间
 68 |     return weights
 69 | 
 70 | class MVDreamPipeline(DiffusionPipeline):
 71 | 
 72 |     # _optional_components = ["feature_extractor", "image_encoder"]
 73 |     # _optional_components = ["image_encoder"]
 74 | 
 75 |     def __init__(
 76 |         self,
 77 |         vae: AutoencoderKL,
 78 |         unet: MultiViewUNetModel,
 79 |         tokenizer: CLIPTokenizer,
 80 |         text_encoder: CLIPTextModel,
 81 |         scheduler: DDIMScheduler,
 82 |         # imagedream variant
 83 |         feature_extractor: CLIPImageProcessor,
 84 |         image_encoder: CLIPVisionModelWithProjection,
 85 |         requires_safety_checker: bool = False,
 86 |     ):
 87 |         super().__init__()
 88 | 
 89 |         if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:  # type: ignore
 90 |             deprecation_message = (
 91 |                 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
 92 |                 f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "  # type: ignore
 93 |                 "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
 94 |                 " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
 95 |                 " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
 96 |                 " file"
 97 |             )
 98 |             deprecate(
 99 |                 "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
100 |             )
101 |             new_config = dict(scheduler.config)
102 |             new_config["steps_offset"] = 1
103 |             scheduler._internal_dict = FrozenDict(new_config)
104 | 
105 |         if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:  # type: ignore
106 |             deprecation_message = (
107 |                 f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
108 |                 " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
109 |                 " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
110 |                 " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
111 |                 " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
112 |             )
113 |             deprecate(
114 |                 "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
115 |             )
116 |             new_config = dict(scheduler.config)
117 |             new_config["clip_sample"] = False
118 |             scheduler._internal_dict = FrozenDict(new_config)
119 | 
120 |         self.register_modules(
121 |             vae=vae,
122 |             unet=unet,
123 |             scheduler=scheduler,
124 |             tokenizer=tokenizer,
125 |             text_encoder=text_encoder,
126 |             feature_extractor=feature_extractor,
127 |             image_encoder=image_encoder,
128 |         )
129 |         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
130 |         self.register_to_config(requires_safety_checker=requires_safety_checker)
131 | 
132 |     def enable_vae_slicing(self):
133 |         r"""
134 |         Enable sliced VAE decoding.
135 | 
136 |         When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
137 |         steps. This is useful to save some memory and allow larger batch sizes.
138 |         """
139 |         self.vae.enable_slicing()
140 | 
141 |     def disable_vae_slicing(self):
142 |         r"""
143 |         Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
144 |         computing decoding in one step.
145 |         """
146 |         self.vae.disable_slicing()
147 | 
148 |     def enable_vae_tiling(self):
149 |         r"""
150 |         Enable tiled VAE decoding.
151 | 
152 |         When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
153 |         several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
154 |         """
155 |         self.vae.enable_tiling()
156 | 
157 |     def disable_vae_tiling(self):
158 |         r"""
159 |         Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
160 |         computing decoding in one step.
161 |         """
162 |         self.vae.disable_tiling()
163 | 
164 |     def enable_sequential_cpu_offload(self, gpu_id=0):
165 |         r"""
166 |         Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
167 |         text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
168 |         `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
169 |         Note that offloading happens on a submodule basis. Memory savings are higher than with
170 |         `enable_model_cpu_offload`, but performance is lower.
171 |         """
172 |         if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
173 |             from accelerate import cpu_offload
174 |         else:
175 |             raise ImportError(
176 |                 "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
177 |             )
178 | 
179 |         device = torch.device(f"cuda:{gpu_id}")
180 | 
181 |         if self.device.type != "cpu":
182 |             self.to("cpu", silence_dtype_warnings=True)
183 |             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
184 | 
185 |         for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
186 |             cpu_offload(cpu_offloaded_model, device)
187 | 
188 |     def enable_model_cpu_offload(self, gpu_id=0):
189 |         r"""
190 |         Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
191 |         to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
192 |         method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
193 |         `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
194 |         """
195 |         if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
196 |             from accelerate import cpu_offload_with_hook
197 |         else:
198 |             raise ImportError(
199 |                 "`enable_model_offload` requires `accelerate v0.17.0` or higher."
200 |             )
201 | 
202 |         device = torch.device(f"cuda:{gpu_id}")
203 | 
204 |         if self.device.type != "cpu":
205 |             self.to("cpu", silence_dtype_warnings=True)
206 |             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
207 | 
208 |         hook = None
209 |         for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
210 |             _, hook = cpu_offload_with_hook(
211 |                 cpu_offloaded_model, device, prev_module_hook=hook
212 |             )
213 | 
214 |         # We'll offload the last model manually.
215 |         self.final_offload_hook = hook
216 | 
217 |     @property
218 |     def _execution_device(self):
219 |         r"""
220 |         Returns the device on which the pipeline's models will be executed. After calling
221 |         `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
222 |         hooks.
223 |         """
224 |         if not hasattr(self.unet, "_hf_hook"):
225 |             return self.device
226 |         for module in self.unet.modules():
227 |             if (
228 |                 hasattr(module, "_hf_hook")
229 |                 and hasattr(module._hf_hook, "execution_device")
230 |                 and module._hf_hook.execution_device is not None
231 |             ):
232 |                 return torch.device(module._hf_hook.execution_device)
233 |         return self.device
234 | 
235 |     def _encode_prompt(
236 |         self,
237 |         prompt,
238 |         device,
239 |         num_images_per_prompt,
240 |         do_classifier_free_guidance: bool,
241 |         negative_prompt=None,
242 |     ):
243 |         r"""
244 |         Encodes the prompt into text encoder hidden states.
245 | 
246 |         Args:
247 |              prompt (`str` or `List[str]`, *optional*):
248 |                 prompt to be encoded
249 |             device: (`torch.device`):
250 |                 torch device
251 |             num_images_per_prompt (`int`):
252 |                 number of images that should be generated per prompt
253 |             do_classifier_free_guidance (`bool`):
254 |                 whether to use classifier free guidance or not
255 |             negative_prompt (`str` or `List[str]`, *optional*):
256 |                 The prompt or prompts not to guide the image generation. If not defined, one has to pass
257 |                 `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
258 |                 Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
259 |             prompt_embeds (`torch.FloatTensor`, *optional*):
260 |                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
261 |                 provided, text embeddings will be generated from `prompt` input argument.
262 |             negative_prompt_embeds (`torch.FloatTensor`, *optional*):
263 |                 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
264 |                 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
265 |                 argument.
266 |         """
267 |         if prompt is not None and isinstance(prompt, str):
268 |             batch_size = 1
269 |         elif prompt is not None and isinstance(prompt, list):
270 |             batch_size = len(prompt)
271 |         else:
272 |             raise ValueError(
273 |                 f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
274 |             )
275 | 
276 |         text_inputs = self.tokenizer(
277 |             prompt,
278 |             padding="max_length",
279 |             max_length=self.tokenizer.model_max_length,
280 |             truncation=True,
281 |             return_tensors="pt",
282 |         )
283 |         text_input_ids = text_inputs.input_ids
284 |         untruncated_ids = self.tokenizer(
285 |             prompt, padding="longest", return_tensors="pt"
286 |         ).input_ids
287 | 
288 |         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
289 |             text_input_ids, untruncated_ids
290 |         ):
291 |             removed_text = self.tokenizer.batch_decode(
292 |                 untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
293 |             )
294 |             logger.warning(
295 |                 "The following part of your input was truncated because CLIP can only handle sequences up to"
296 |                 f" {self.tokenizer.model_max_length} tokens: {removed_text}"
297 |             )
298 | 
299 |         if (
300 |             hasattr(self.text_encoder.config, "use_attention_mask")
301 |             and self.text_encoder.config.use_attention_mask
302 |         ):
303 |             attention_mask = text_inputs.attention_mask.to(device)
304 |         else:
305 |             attention_mask = None
306 | 
307 |         prompt_embeds = self.text_encoder(
308 |             text_input_ids.to(device),
309 |             attention_mask=attention_mask,
310 |         )
311 |         prompt_embeds = prompt_embeds[0]
312 | 
313 |         prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
314 | 
315 |         bs_embed, seq_len, _ = prompt_embeds.shape
316 |         # duplicate text embeddings for each generation per prompt, using mps friendly method
317 |         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
318 |         prompt_embeds = prompt_embeds.view(
319 |             bs_embed * num_images_per_prompt, seq_len, -1
320 |         )
321 | 
322 |         # get unconditional embeddings for classifier free guidance
323 |         if do_classifier_free_guidance:
324 |             uncond_tokens: List[str]
325 |             if negative_prompt is None:
326 |                 uncond_tokens = [""] * batch_size
327 |             elif type(prompt) is not type(negative_prompt):
328 |                 raise TypeError(
329 |                     f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
330 |                     f" {type(prompt)}."
331 |                 )
332 |             elif isinstance(negative_prompt, str):
333 |                 uncond_tokens = [negative_prompt]
334 |             elif batch_size != len(negative_prompt):
335 |                 raise ValueError(
336 |                     f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
337 |                     f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
338 |                     " the batch size of `prompt`."
339 |                 )
340 |             else:
341 |                 uncond_tokens = negative_prompt
342 | 
343 |             max_length = prompt_embeds.shape[1]
344 |             uncond_input = self.tokenizer(
345 |                 uncond_tokens,
346 |                 padding="max_length",
347 |                 max_length=max_length,
348 |                 truncation=True,
349 |                 return_tensors="pt",
350 |             )
351 |             if (
352 |                 hasattr(self.text_encoder.config, "use_attention_mask")
353 |                 and self.text_encoder.config.use_attention_mask
354 |             ):
355 |                 attention_mask = uncond_input.attention_mask.to(device)
356 |             else:
357 |                 attention_mask = None
358 | 
359 |             negative_prompt_embeds = self.text_encoder(
360 |                 uncond_input.input_ids.to(device),
361 |                 attention_mask=attention_mask,
362 |             )
363 |             negative_prompt_embeds = negative_prompt_embeds[0]
364 | 
365 |             # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
366 |             seq_len = negative_prompt_embeds.shape[1]
367 | 
368 |             negative_prompt_embeds = negative_prompt_embeds.to(
369 |                 dtype=self.text_encoder.dtype, device=device
370 |             )
371 | 
372 |             negative_prompt_embeds = negative_prompt_embeds.repeat(
373 |                 1, num_images_per_prompt, 1
374 |             )
375 |             negative_prompt_embeds = negative_prompt_embeds.view(
376 |                 batch_size * num_images_per_prompt, seq_len, -1
377 |             )
378 | 
379 |             # For classifier free guidance, we need to do two forward passes.
380 |             # Here we concatenate the unconditional and text embeddings into a single batch
381 |             # to avoid doing two forward passes
382 |             prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
383 | 
384 |         return prompt_embeds
385 | 
386 |     def decode_latents(self, latents):
387 |         latents = 1 / self.vae.config.scaling_factor * latents
388 |         image = self.vae.decode(latents).sample
389 |         image = (image / 2 + 0.5).clamp(0, 1)
390 |         # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
391 |         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
392 |         return image
393 | 
394 |     def prepare_extra_step_kwargs(self, generator, eta):
395 |         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
396 |         # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
397 |         # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
398 |         # and should be between [0, 1]
399 | 
400 |         accepts_eta = "eta" in set(
401 |             inspect.signature(self.scheduler.step).parameters.keys()
402 |         )
403 |         extra_step_kwargs = {}
404 |         if accepts_eta:
405 |             extra_step_kwargs["eta"] = eta
406 | 
407 |         # check if the scheduler accepts generator
408 |         accepts_generator = "generator" in set(
409 |             inspect.signature(self.scheduler.step).parameters.keys()
410 |         )
411 |         if accepts_generator:
412 |             extra_step_kwargs["generator"] = generator
413 |         return extra_step_kwargs
414 | 
415 |     def prepare_latents(
416 |         self,
417 |         batch_size,
418 |         num_channels_latents,
419 |         height,
420 |         width,
421 |         dtype,
422 |         device,
423 |         generator,
424 |         latents=None,
425 |     ):
426 |         shape = (
427 |             batch_size,
428 |             num_channels_latents,
429 |             height // self.vae_scale_factor,
430 |             width // self.vae_scale_factor,
431 |         )
432 |         if isinstance(generator, list) and len(generator) != batch_size:
433 |             raise ValueError(
434 |                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
435 |                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
436 |             )
437 | 
438 |         if latents is None:
439 |             latents = randn_tensor(
440 |                 shape, generator=generator, device=device, dtype=dtype
441 |             )
442 |         else:
443 |             latents = latents.to(device)
444 | 
445 |         # scale the initial noise by the standard deviation required by the scheduler
446 |         latents = latents * self.scheduler.init_noise_sigma
447 |         return latents
448 | 
449 |     def encode_image(self, image, device, num_images_per_prompt):
450 |         # dtype = next(self.feature_extractor.parameters()).dtype
451 |         dtype = self.text_encoder.dtype
452 |         # if image.dtype == np.float32:
453 |         #     image = (image * 255).astype(np.uint8)
454 |         
455 |         # image_clip = (((image + 1.0)/2.0)* 255)
456 |         image_clip = (((image + 1.0)/2.0)* 1.0)
457 |         image_clip = self.feature_extractor(image_clip, return_tensors="pt",do_rescale=False).pixel_values
458 |         image_clip = image_clip.to(device=device, dtype=dtype)
459 |         
460 |         # tt = self.image_encoder(image_clip, output_hidden_states=True)
461 |         # import pdb;pdb.set_trace() 
462 |         # image_embeds = self.image_encoder(image_clip, output_hidden_states=True).hidden_states[-2]
463 |         # image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
464 |         encoded_embeds = self.image_encoder(image_clip, output_hidden_states=False)
465 |         image_embeds = encoded_embeds.image_embeds
466 |         image_embeds = 0.2 * image_embeds.unsqueeze(-2).repeat(1, 77, 1).to(dtype)
467 |         # import pdb;pdb.set_trace()
468 | 
469 |         return torch.zeros_like(image_embeds), image_embeds
470 | 
471 |     def encode_image_latents(self, image, device, num_images_per_prompt):
472 |         
473 |         dtype = next(self.image_encoder.parameters()).dtype
474 | 
475 |         image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
476 |         image = 2 * image - 1
477 |         image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
478 |         image = image.to(dtype=dtype)
479 | 
480 |         posterior = self.vae.encode(image).latent_dist
481 |         latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
482 |         latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
483 | 
484 |         return torch.zeros_like(latents), latents
485 | 
486 |     def decode_latents_to_pil_save(self, latents, path):
487 |         image_list = self.decode_latents(latents)
488 |         image_list = self.numpy_to_pil(image_list)
489 |         for i, image in enumerate(image_list):
490 |             image.save(path+f'_{i}.jpg')
491 | 
492 |     
493 |     
494 |     
495 |     @torch.no_grad()
496 |     def __call__(
497 |         self,
498 |         prompt: str = "",# 'a cute owl'
499 |         image: Optional[np.ndarray] = None,# NOne
500 |         masks: Optional[np.ndarray] = None,# NOne
501 |         height: int = 256,# 256
502 |         width: int = 256,# 256
503 |         elevation: float = 0,# 0
504 |         num_inference_steps: int = 50,# 30
505 |         guidance_scale: float = 7.0,# 5
506 |         negative_prompt: str = None,# ''
507 |         num_images_per_prompt: int = 1,# 1
508 |         eta: float = 0.0,# 0.0
509 |         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,# None
510 |         output_type: Optional[str] = "numpy", # pil, numpy, latents | numpy
511 |         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
512 |         callback_steps: int = 1,
513 |         num_frames: int = 4, # 4
514 |         device=torch.device("cuda:0"),
515 |         guidance_rescale: float = 0.0,
516 |         # gt_num: int=3,
517 |         condition_num_frames: int=3,
518 |         gt_frame: Optional[np.ndarray] = None,
519 |     ):
520 |         gt_num = condition_num_frames
521 |         
522 |         self.unet = self.unet.to(device=device)
523 |         self.vae = self.vae.to(device=device)
524 |         self.text_encoder = self.text_encoder.to(device=device)
525 |         weight_dtype = next(self.vae.parameters()).dtype
526 |         self.guidance_rescale = guidance_rescale
527 | 
528 |         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
529 |         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
530 |         # corresponds to doing no classifier free guidance.
531 |         # do_classifier_free_guidance = guidance_scale > 1.0
532 |         do_classifier_free_guidance = (guidance_scale != 0.0)
533 |         
534 |         # Prepare timesteps
535 |         self.scheduler.set_timesteps(num_inference_steps, device=device)# DDIMScheduler
536 |         timesteps = self.scheduler.timesteps
537 |         # import pdb;pdb.set_trace()
538 |         if self.scheduler.timestep_scaling == "trailing":
539 |             step_ratio = 999 // num_inference_steps
540 |             timesteps = torch.from_numpy(np.round(np.arange(999, 0, -step_ratio)).astype(np.int64).copy()).to(timesteps.device)
541 | 
542 |         # imagedream variant
543 |         if image is not None:
544 |             image = rearrange(image, "b f c h w -> (b f) c h w", f=num_frames)
545 |             img_latents = self.vae.encode(image.to(weight_dtype).to(device)).latent_dist.sample()
546 |             img_latents = img_latents * self.vae.config.scaling_factor # [b*f, c, h, w] shape=[6, 4, 32, 32]
547 |         if gt_frame is not None:    
548 |             image_gt = rearrange(gt_frame, "b f c h w -> (b f) c h w", f=num_frames)
549 |             img_latents_gt = self.vae.encode(image_gt.to(weight_dtype).to(device)).latent_dist.sample()
550 |             img_latents_gt = img_latents_gt * self.vae.config.scaling_factor # [b*f, c, h, w] shape=[6, 4, 32, 32]
551 |         
552 |         self.image_encoder = self.image_encoder.to(device=device)
553 |         image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt) 
554 |         image_embeds_pos = image_embeds_pos[0:1]
555 |         
556 |         if masks is not None:
557 |             masks = rearrange(masks, "b f c h w -> (b f) c h w", f=num_frames)# [b*f, c, h, w]
558 |             mask_latents = torch.nn.functional.interpolate(
559 |                 masks, 
560 |                 size=(
561 |                     height // 8, 
562 |                     width // 8
563 |                 )
564 |             ).to(weight_dtype).to(device)# [b*f, c, h, w]
565 |         
566 |         if do_classifier_free_guidance:
567 |             uncond_img_latents = torch.zeros_like(img_latents)
568 |         
569 |         prompt_embeds = self._encode_prompt(
570 |             prompt=prompt,
571 |             device=device,
572 |             num_images_per_prompt=num_images_per_prompt,
573 |             do_classifier_free_guidance=do_classifier_free_guidance,
574 |             negative_prompt=negative_prompt,
575 |         )  # type: ignore torch.Size([1, 77, 1024])
576 |         # prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
577 | 
578 |         # Prepare latent variables
579 |         actual_num_frames = num_frames
580 |         noisy_latents = self.prepare_latents(
581 |             actual_num_frames * num_images_per_prompt,
582 |             4,
583 |             height,
584 |             width,
585 |             prompt_embeds.dtype,
586 |             device,
587 |             generator,
588 |             latents=None,
589 |         )# shape=[3, 4, 32, 32]
590 |         
591 |         
592 |         
593 |         noise_scheduler = self.scheduler
594 |         if gt_frame is not None:
595 |             noise_gt = torch.randn_like(img_latents)
596 |             # t = timesteps[0]
597 |             t = 999
598 |             # print('init time:',t)
599 |             # print(torch.min(noise_gt),torch.max(noise_gt),torch.mean(noise_gt))
600 |             tt = torch.tensor([t] * actual_num_frames, dtype=img_latents.dtype, device=device)
601 |             noisy_latents = noise_scheduler.add_noise(img_latents_gt, noise_gt, tt.long())
602 |             # noisy_latents = noise_scheduler.add_noise(noisy_latents, noise_gt, tt.long())
603 |             print(torch.min(noisy_latents),torch.max(noisy_latents),torch.mean(noisy_latents))
604 |         
605 | 
606 |         latents = torch.cat([img_latents[:gt_num], noisy_latents[gt_num:]], dim=0)
607 |         
608 |         camera = None
609 | 
610 |         # Prepare extra step kwargs.
611 |         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
612 | 
613 |         # Denoising loop
614 |         # import pdb;pdb.set_trace()
615 |         assert image.shape[0] == actual_num_frames, "Currently only supports batchsize==1"
616 |         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
617 |     
618 |         time_ratio = 5
619 |         timestep_warp = (timesteps[0]//time_ratio).long()
620 |         print(timestep_warp)
621 |         noise_warp = torch.randn_like(noisy_latents)
622 |         tt = torch.tensor([timestep_warp] * actual_num_frames, dtype=noisy_latents.dtype, device=device)
623 |         noisy_latents_warp = noise_scheduler.add_noise(img_latents, noise_warp, tt.long())
624 |         
625 | 
626 |         
627 |         mix_latents_warp = torch.cat([img_latents[:gt_num], noisy_latents_warp[gt_num:]], dim=0)
628 |         with self.progress_bar(total=num_inference_steps) as progress_bar:
629 |             for i, t in enumerate(timesteps):
630 |                 # expand the latents if we are doing classifier free guidance
631 |                 multiplier = 2 if do_classifier_free_guidance else 1
632 |                 
633 |                 latents = torch.cat([img_latents[:gt_num], latents[gt_num:]], dim=0)# [6, 4, 32, 32]
634 |                 
635 |                 timestep_warp = (t//time_ratio).long()
636 |                 noise_warp = torch.randn_like(noisy_latents)
637 |                 tt = torch.tensor([timestep_warp] * actual_num_frames, dtype=noisy_latents.dtype, device=device)
638 |                 noisy_latents_warp = noise_scheduler.add_noise(img_latents, noise_warp, tt.long())
639 |                 weights = custom_decay_function_weight(tt.float())
640 |                 weights = weights.view(tt.shape[0], 1, 1, 1).to(noisy_latents_warp.dtype)
641 |                 # noisy_latents_warp_weight = weights*noisy_latents_warp + (1. - weights) * latents
642 |                 noisy_latents_warp_weight = img_latents
643 |                 mix_latents_warp = torch.cat([img_latents[:gt_num], noisy_latents_warp_weight[gt_num:]], dim=0)
644 |                 
645 |                 
646 |                 
647 |                 if self.unet.input_blocks[0].__dict__['_modules']['0'].__dict__['in_channels'] == 9:
648 |                     latents = torch.cat([latents, mix_latents_warp, mask_latents], dim=1)# [b*f, c, h, w] shape=[6, 9, 32, 32]
649 |                 elif self.unet.input_blocks[0].__dict__['_modules']['0'].__dict__['in_channels'] == 5:
650 |                     latents = torch.cat([latents, mask_latents], dim=1)# [b*f, c, h, w] shape=[6, 9, 32, 32]
651 |                 
652 |                 if do_classifier_free_guidance:
653 |                     
654 |                     unc_mix_latents_warp = torch.cat([img_latents[:gt_num], uncond_img_latents[gt_num:]], dim=0)
655 |                     
656 |                     
657 |                     unc_mask_latents = torch.cat([mask_latents[:gt_num], torch.zeros_like(mask_latents[gt_num:])], dim=0)
658 |                     uncond_latents = torch.cat([latents[:,:4], unc_mix_latents_warp, unc_mask_latents], dim=1)
659 |                     latent_model_input = torch.cat([latents, uncond_latents], dim=0)# torch.Size([12, 9, 64, 64])
660 |                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
661 |                 else:
662 |                     latent_model_input = torch.cat([latents] * multiplier)
663 |                     latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
664 | 
665 |                 image_embeds_pos_input = image_embeds_pos.repeat(multiplier,1,1)
666 |                 unet_inputs = {
667 |                     'x': latent_model_input,# torch.Size([num_frames, 5, 32, 32])
668 |                     'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),# torch.Size([num_frames])
669 |                     # 'context': prompt_embeds.repeat(actual_num_frames,1,1),# torch.Size([num_frames, 77, 1024])
670 |                     'context': (prompt_embeds + image_embeds_pos_input).repeat(actual_num_frames,1,1),
671 |                     'num_frames': actual_num_frames,# 4
672 |                     'camera': None,# 
673 |                 }
674 | 
675 |                 
676 |                 noise_pred = self.unet.forward(**unet_inputs)# [6, 4, 32, 32]
677 | 
678 |                 # perform guidance
679 |                 if do_classifier_free_guidance:
680 |                     noise_pred_cond, noise_pred_uncond = noise_pred.chunk(2)
681 |                     noise_pred = (1 + guidance_scale) * noise_pred_cond - guidance_scale * noise_pred_uncond
682 |                     
683 |                     
684 |                 if do_classifier_free_guidance and self.guidance_rescale > 0.0:
685 |                     # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
686 |                     
687 |                     noise_pred[:,:4] = rescale_noise_cfg(noise_pred[:,:4], noise_pred_cond[:,:4], guidance_rescale=self.guidance_rescale)
688 | 
689 |                 
690 |                 latents: torch.Tensor = self.scheduler.step(
691 |                     noise_pred[:,:4], t, latents[:, :4], **extra_step_kwargs, return_dict=False
692 |                 )[0]
693 |                 
694 |                 # call the callback, if provided
695 |                 if i == len(timesteps) - 1 or (
696 |                     (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
697 |                 ):
698 |                     progress_bar.update()
699 |                     if callback is not None and i % callback_steps == 0:
700 |                         callback(i, t, latents)  # type: ignore
701 |         # Post-processing
702 |         if output_type == "latent":
703 |             image = latents
704 |             noisy_latents_warp = noisy_latents_warp
705 |         elif output_type == "pil":
706 |             image = self.decode_latents(latents)
707 |             image = self.numpy_to_pil(image)
708 |             image_warp = self.decode_latents(noisy_latents_warp)
709 |             image_warp = self.numpy_to_pil(image_warp)
710 |         else: # numpy
711 |             image = self.decode_latents(latents)
712 |             image_warp = self.decode_latents(noisy_latents_warp)
713 | 
714 |         # Offload last model to CPU
715 |         if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
716 |             self.final_offload_hook.offload()
717 | 
718 |         return image, image_warp


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | datasets==2.19.1
 2 | decorator==5.1.1
 3 | deepspeed==0.14.2
 4 | -e git+https://github.com/graphdeco-inria/diff-gaussian-rasterization@59f5f77e3ddbac3ed9db93ec2cfe99ed6c5d121d#egg=diff_gaussian_rasterization&subdirectory=../../../submodules/diff-gaussian-rasterization
 5 | diffusers==0.27.2
 6 | einops==0.8.0
 7 | Flask==3.0.3
 8 | huggingface-hub==0.23.0
 9 | hydra-core==1.3.2
10 | imageio==2.34.1
11 | imageio-ffmpeg==0.4.9
12 | Jinja2==3.1.4
13 | jsonpatch==1.33
14 | kiwisolver==1.4.5
15 | -e git+https://github.com/princeton-vl/lietorch.git@0fa9ce8ffca86d985eca9e189a99690d6f3d4df6#egg=lietorch
16 | # Editable install with no version control (lightglue==0.0)
17 | -e /share/project/mabaorui/vggSFM/vggsfm-main/dependency/LightGlue
18 | lightning-utilities==0.11.5
19 | lpips==0.1.4
20 | matplotlib==3.9.0
21 | matplotlib-inline==0.1.7
22 | ninja==1.11.1.1
23 | numba==0.59.1
24 | numpy==1.26.4
25 | nvidia-cublas-cu12==12.1.3.1
26 | nvidia-cuda-cupti-cu12==12.1.105
27 | nvidia-cuda-nvrtc-cu12==12.1.105
28 | nvidia-cuda-runtime-cu12==12.1.105
29 | nvidia-cudnn-cu12==8.9.2.26
30 | nvidia-cufft-cu12==11.0.2.54
31 | nvidia-curand-cu12==10.3.2.106
32 | nvidia-cusolver-cu12==11.4.5.107
33 | nvidia-cusparse-cu12==12.1.0.106
34 | nvidia-nccl-cu12==2.18.1
35 | nvidia-nvjitlink-cu12==12.4.127
36 | nvidia-nvtx-cu12==12.1.105
37 | omegaconf==2.3.0
38 | opencv-python==4.9.0.80
39 | packaging==24.0
40 | pandas==2.2.2
41 | pillow==10.3.0
42 | plyfile==1.0.3
43 | pyceres==2.3
44 | pyparsing==3.1.2
45 | pyquaternion==0.9.9
46 | pytorch-cuda==12.1
47 | PyYAML==6.0.1
48 | requests==2.31.0
49 | safetensors==0.4.3
50 | scikit-image==0.24.0
51 | scikit-learn==1.5.0
52 | scipy==1.13.1
53 | simple_knn @ file:///share/project/gaohuachen/PosefreeGS/gaussian-splatting/InstantSplat/submodules/simple-knn
54 | tensorboard==2.16.2
55 | tensorboard-data-server==0.7.2
56 | timm==1.0.3
57 | tokenizers==0.19.1
58 | torch==2.1.0
59 | torchvision==0.16.0
60 | tqdm==4.66.4
61 | transformers==4.41.1
62 | trimesh==4.4.0
63 | xformers==0.0.22.post7
64 | 


--------------------------------------------------------------------------------
/single_infer.sh:
--------------------------------------------------------------------------------
1 | CUDA_VISIBLE_DEVICES=0 python inference.py \
2 | --super_resolution \
3 | --single_view \
4 | --base_model_path "./checkpoint/MVD_weights/" \
5 | --source_imgs_dir "./dataset/co3d_1views_save/12_104_640/reference_images/" \
6 | --warp_root_dir "./dataset/co3d_1views_save/12_104_640/warp_images/" \
7 | --output_dir "./output/12_104_640/"


--------------------------------------------------------------------------------
/sparse_infer.sh:
--------------------------------------------------------------------------------
1 | CUDA_VISIBLE_DEVICES=0 python inference.py \
2 | --super_resolution \
3 | --base_model_path "./checkpoint/MVD_weights/" \
4 | --source_imgs_dir ".dataset/DTU_3views_save_warp-2/scan21_colmap/reference_images/" \
5 | --warp_root_dir ".dataset/DTU_3views_save_warp-2/scan21_colmap/warp_images/" \
6 | --output_dir "./output/scan21/" 


--------------------------------------------------------------------------------