update readme

README.md

@@ -36,7 +36,7 @@ We study inferring 3D object-centric scene representations from a single image.
 Please run `conda env create -f environment.yml` to install the dependencies.
 
 ## Datasets
-We gathered three datasets (Room-Texture, Kitchen-Matte, Kitchen-Shiny) for evaluating uOCF. The datasets are available at [link (Google Drive)](https://drive.google.com/drive/folders/1v_lZhiI32rvKjUDQVb5B7KHMpNLgQ2P_?usp=drive_link).
+We gathered four datasets (Room-Texture, Room-Furniture, Kitchen-Matte, Kitchen-Shiny) for evaluating uOCF. The datasets are available at [link (Google Drive)](https://drive.google.com/file/d/1xwa5kHCJQq3l7Q8LG4s9cqqlM_BxhFgB/view?usp=sharing).
 
 ### Dataset organization
 
@@ -47,14 +47,14 @@ When running the scripts, please specify the path to the dataset by modifying th
 Note that we need to generate a depth image for each input image, and the MiDAS depth estimator is used in our experiments. Refer to their repo[https://github.com/isl-org/MiDaS] for more details.
 
 ## Pretrained Models
-We provide the pretrained models for uOCF. The models are available at [link (Google Drive)](https://drive.google.com/drive/folders/1zg-uSyFYCFWui3zyO6XXmkh8wlq5ht0Y?usp=drive_link).
+We provide the pretrained models for uOCF. The models are available at [link (Google Drive)](https://drive.google.com/file/d/167v5JYomSxgrIyFnTNP3YnCpJQjaQHAi/view?usp=sharing).
 
 ## Testing
 To test the pretrained model, please run the following command:
 ```
 bash scripts/DATASET-NAME/test-MODE.sh
 ```
-where DATASET-NAME is one of (Room-Texture, Kitchen-Matte, Kitchen-Shiny), and MODE is one of (plane, noplane).
+where DATASET-NAME is one of (Room-Texture, Room-Furniture, Kitchen-Matte, Kitchen-Shiny), and MODE is one of (plane, noplane).
 
 ## Training
 (Optional, we have provided pre-trained checkpoints) Run the following command to train the model from stage 1
@@ -69,7 +69,7 @@ Then run the following command to train the model from stage 2
 ```
 bash scripts/DATASET-NAME/train-stage2-MODE.sh
 ```
-where DATASET-NAME is one of (Room-Texture, Kitchen-Matte, Kitchen-Shiny), and MODE is one of (plane, noplane).
+where DATASET-NAME is one of (Room-Texture, Room-Furniture, Kitchen-Matte, Kitchen-Shiny), and MODE is one of (plane, noplane).
 
 ## Citation
 If you find this work useful in your research, please consider citing:

scripts/room-furniture/test/test-plane.sh

@@ -0,0 +1,16 @@
+DATAROOT=${1:-'/svl/u/redfairy/datasets/room-real/chairs/test-2-4obj'}
+PORT=${2:-12783}
+python -m visdom.server -p $PORT &>/dev/null &
+python test.py --dataroot $DATAROOT --n_scenes 100 --n_img_each_scene 4  \
+    --checkpoints_dir 'checkpoints' --name 'room-furniture' \
+    --display_port $PORT --display_ncols 4 \
+    --load_size 128 --n_samp 256 --input_size 128 --render_size 32 --frustum_size 128 \
+    --model 'uocf_eval' \
+    --num_slots 5 --attn_iter 6 --nss_scale 7 --fg_object_size 3 \
+    --fixed_locality \
+    --attn_momentum 0.5  \
+    --exp_id 'stage2-multiobj-plane' --epoch 120 \
+    --vis_attn --vis_mask --remove_duplicate \
+    --testset_name 'standard' \
+
+

scripts/room-furniture/train/train-stage2-plane.sh

@@ -0,0 +1,18 @@
+DATAROOT=${1:-'/svl/u/redfairy/datasets/room-real/chairs/train-2-4obj'}
+PORT=${2:-12783}
+python -m visdom.server -p $PORT &>/dev/null &
+python train.py --dataroot $DATAROOT --n_scenes 5000 --n_img_each_scene 2 \
+    --checkpoints_dir 'checkpoints' --name 'room-furniture' \
+    --display_port $PORT --display_ncols 4 --print_freq 50 --display_freq 50 --save_epoch_freq 5 \
+    --load_size 128 --n_samp 64 --input_size 128 --supervision_size 64 --frustum_size 64 \
+    --model 'uocf' --lr 0.0003 \
+    --seed 2024 \
+    --num_slots 5 --attn_iter 6 --nss_scale 7 --fg_object_size 3 \
+    --stratified --fixed_locality \
+    --bg_density_loss --depth_supervision --remove_duplicate \
+    --load_pretrain --load_pretrain_path './checkpoints/room-texture/stage1-1obj-plane' \
+    --load_encoder 'load_train' --load_slotattention 'load_train' --load_decoder 'load_train' --load_epoch 100 --one2four \
+    --coarse_epoch 60 --niter 120 --percept_in 10 --no_locality_epoch 20 --dense_sample_epoch 20 \
+    --scaled_depth --depth_scale_pred \
+    --exp_id 'stage2-multiobj-plane' \
+