Update README.md

README.md

@@ -1,32 +1,144 @@
----
-license: apache-2.0
----
-
 # MVInpainter
 [NeurIPS 2024] MVInpainter: Learning Multi-View Consistent Inpainting to Bridge 2D and 3D Editing
 
 [[arXiv]](https://arxiv.org/pdf/2408.08000) [[Project Page]](https://ewrfcas.github.io/MVInpainter/)
 
-Codes and models will be open-released soon!
+## Preparation
+
+### Setup repository and environment
+```
+git clone https://github.com/ewrfcas/MVInpainter.git
+cd MVInpainter
+
+conda create -n mvinpainter python=3.8
+conda activate mvinpainter
 
-Codes organization is working in progress.
+pip install -r requirements.txt
+mim install mmcv-full
+pip install mmflow
+
+# We need to replace the new decoder py of mmflow for faster flow estimation
+cp ./check_points/mmflow/raft_decoder.py /usr/local/conda/envs/mvinpainter/lib/python3.8/site-packages/mmflow/models/decoders/
+```
 
 ### Dataset preparation (training)
 1. Downloading [Co3dv2](https://github.com/facebookresearch/co3d), [MVImgNet](https://github.com/GAP-LAB-CUHK-SZ/MVImgNet) for MVInpainter-O.
 Downloading [Real10k](https://google.github.io/realestate10k/download.html), [DL3DV](https://github.com/DL3DV-10K/Dataset), [Scannet++](https://kaldir.vc.in.tum.de/scannetpp) for MVInpainter-F.
-2. Downloading information of indices and masking formats from [Link](), and put them to ```./data```. Note that we remove some dirty samples from aforementioned datasets. We additionally provide complete foreground masks for MVImgNet through ```CarveKit```.
+2. Downloading information of indices, masking formats, and captions from [Link](). Put them to `./data`. Note that we remove some dirty samples from aforementioned datasets. Since Co3dv2 data contains object masks but MVImgNet does not, we additionally provide complete [foreground masks]() for MVImgNet through `CarveKit`. Please put the MVImgNet masks to `./data/mvimagenet/masks`.
+
+### Pretrained weights
+1. [RAFT weights]() (put it to `./check_points/mmflow/`).
+2. [SD1.5-inpainting]()  (put it to `./check_points/`).
+3. [AnimateDiff weights](). We revise the key name for easier `peft` usages (put it to `./check_points/`).
+
+## Training
+
+Training with fixed nframe=12:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --mixed_precision="fp16" --num_processes=8 --num_machines 1 --main_process_port 29502 \
+  --config_file configs/deepspeed/acc_zero2.yaml train.py \
+  --config_file="configs/mvinpainter_{o,f}.yaml" \
+  --output_dir="check_points/mvinpainter_{o,f}_256" \
+  --train_log_interval=250 \
+  --val_interval=2000 \
+  --val_cfg=7.5 \
+  --img_size=256
+```
+
+Finetuning with dynamic frames (8~24):
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch --mixed_precision="fp16" --num_processes=8 --num_machines 1 --main_process_port 29502 \
+  --config_file configs/deepspeed/acc_zero2.yaml train.py \
+  --config_file="configs/mvinpainter_{o,f}.yaml" \
+  --output_dir="check_points/mvinpainter_{o,f}_256" \
+  --train_log_interval=250 \
+  --val_interval=2000 \
+  --val_cfg=7.5 \
+  --img_size=256 \
+  --resume_from_checkpoint="latest" \
+  --dynamic_nframe \
+  --low_nframe 8 \
+  --high_nframe 24
+```
+Please use `mvinpainter_{o,f}_512.yaml` to train 512x512 models.
+
+## Inference
 
-### Pretrained weights (training)
-Put these weights to ```./check_points```
-1. [SD1.5-inpainting]()
-2. [AnimateDiff weights](). We revise the key name for easier ```peft``` usages.
+### Model weights
+1. [MVSInpainter-O]() (Novel view synthesis, put it to `./check_points/`).
+2. [MVSInpainter-F]() (Removal, put it to `./check_points/`).
 
+### Pipeline
 
+1. Removing or synthesis foreground of the first view through 2D-inpainting. We recommend using [Fooocus-inpainting](https://github.com/lllyasviel/Fooocus) to accomplish this. Getting tracking masks through [Track-Anything](https://github.com/gaomingqi/Track-Anything). 
+Some examples are provided in `./demo`.
+```
+- <folder>
+    - images      # input images with foregrounds
+    - inpainted   # inpainted result of the first view
+    - masks       # masks for images
+```
+2. (Optional) removing foregrounds from all other views through `MVInpainter-F`:
+```
+CUDA_VISIBLE_DEVICES=0 python test_removal.py \
+  --load_path="check_points/mvinpainter_f_256" \
+  --dataset_root="./demo/removal" \
+  --output_path="demo_removal" \
+  --resume_from_checkpoint="best" \
+  --val_cfg=5.0 \
+  --img_size=256 \
+  --sampling_interval=1.0 \
+  --dataset_names realworld \
+  --reference_path="inpainted" \
+  --nframe=24 \
+  --save_images # (whether to save samples respectively)
+```
+![removal](assets/kitchen_DSCF0676_removal_seq_0.jpg)
+
+3. Achieving 3d bbox of the object generated from 2D-inpainting through `python draw_bbox.py`. Put the image `000x.png` and `000x.json` from `./bbox` to `obj_bbox` of the target folder.
+
+![draw_bbox_demo](assets/draw_bbox.gif)
+
+4. Mask adaption to achieve `warp_masks`. If the basic plane where the foreground placed on enjoys a small percentage of the whole image, please use methods like [Grounded-Segment-Anything](https://github.com/IDEA-Research/Grounded-Segment-Anything) to get `plane_masks`.
+```
+CUDA_VISIBLE_DEVICES=0 python mask_adaption.py --input_path="demo/nvs/kitchen" --edited_index=0
+``` 
+You can also use `--no_irregular_mask` to disable irregular mask for more precise warped masks.
+
+![warp_bbox](assets/0000_bbox.jpg)
+
+Make sure the final folder looks like:
+```
+- <folder>
+    - obj_bbox    # inpainted 2d images with new foreground and bbox json
+    - removal     # images without foregrounds
+    - warp_masks  # masks from adaption for the removal folder
+    - plane_masks # (optional, only for mask_adaption) masks of basic plane where the foreground is placed on
+```
+
+5. Run `MVInpainter-O` for novel view synthesis:
+```
+CUDA_VISIBLE_DEVICES=0 python test_nvs.py \
+  --load_path="check_points/mvinpainter_o_256" \
+  --dataset_root="./demo/nvs" \
+  --output_path="demo_nvs" \
+  --edited_index=0 \
+  --resume_from_checkpoint="best" \
+  --val_cfg=7.5 \
+  --img_height=256 \
+  --img_width=256 \
+  --sampling_interval=1.0 \
+  --nframe=24 \
+  --prompt="a red apple with circle and round shape on the table." \
+  --limit_frame=24
+```
+![nvs_result](assets/kitchen_0000_seq_0.jpg)
 
+6. 3D reconstruction: See [Dust3R](https://github.com/naver/dust3r), [MVSFormer++](https://github.com/maybeLx/MVSFormerPlusPlus), and [3DGS](https://github.com/graphdeco-inria/gaussian-splatting) for more details.
 
 
 ## Cite
-If you found our program helpful, please consider citing:
+If you found our project helpful, please consider citing:
 
 ```
 @article{cao2024mvinpainter,