feat(ml): debug for vid metrics

jolibrain · Oct 18, 2024 · 7c57471 · 7c57471
1 parent eb39ec5
commit 7c57471
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 39 deletions.
diff --git a/models/base_model.py b/models/base_model.py
@@ -774,6 +774,11 @@ def get_current_visuals(self, nb_imgs, phase="train", test_name=""):
             ):  # GANs have more outputs in practice, including semantics
                 if i == nb_imgs - 1:
                     break
+        if phase == "test" and self.opt.G_netG == "unet_vid":
+            visual_ret = visual_ret[
+                : self.opt.test_batch_size * (self.opt.data_temporal_number_frames)
+            ]
+
         return visual_ret
 
     def get_display_param(self):
@@ -1563,22 +1568,34 @@ def compute_metrics_test(
 
             i = 0
             for sub_list in self.visual_names:
-                if i < offset:
+                if self.opt.G_netG == "unet_vid":
+                    for name in sub_list:
+                        if hasattr(self, name):
+                            setattr(
+                                self, name + "_test_" + test_name, getattr(self, name)
+                            )
+
+                else:
+                    if i < offset:
+                        i += 1
+                        continue
+                    for name in sub_list:
+                        if hasattr(self, name):
+                            setattr(
+                                self, name + "_test_" + test_name, getattr(self, name)
+                            )
                     i += 1
-                    continue
-                for name in sub_list:
-                    if hasattr(self, name):
-                        setattr(self, name + "_test_" + test_name, getattr(self, name))
-                i += 1
-                if i - offset == self.opt.test_batch_size:
-                    break
+                    if i - offset == self.opt.test_batch_size:
+                        break
 
             if progress:
                 progress.n = min(len(fake_list), progress.total)
                 progress.refresh()
 
             if len(fake_list) >= self.opt.train_nb_img_max_fid:
                 break
+            if self.opt.G_netG == "unet_vid" and i < self.opt.test_batch_size:
+                break
 
         fake_list = fake_list[: self.opt.train_nb_img_max_fid]
         real_list = real_list[: self.opt.train_nb_img_max_fid]

diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
@@ -15,8 +15,6 @@ echo "Current dir is [$current_dir]"
 export PYTHONDONTWRITEBYTECODE=1
 export NCCL_P2P_DISABLE=1
 
-
-
 ######## doc auto generation
 echo "Running doc auto generation"
 python3 ${current_dir}/../scripts/generate_doc.py --save_to ""
@@ -87,7 +85,8 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-##### test img2img_turbo
+
+###### test img2img_turbo
 echo "Running test img2img_turbo"
 
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_img2img_turbo.py" --dataroot "$TARGET_NOSEM_DIR"
@@ -107,7 +106,6 @@ mkdir $TARGET_MASK_SEM_DIR
 unzip $ZIP_FILE -d $DIR
 rm $ZIP_FILE
 
-
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_semantic_mask.py" --dataroot "$TARGET_MASK_SEM_DIR"
 OUT=$?
 
@@ -124,6 +122,18 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
+###### test cut
+echo "Running cut test"
+python3 "${current_dir}/../test.py" \
+        --save_config \
+        --test_model_dir $DIR/joligen_utest_cut/ \
+        --test_metrics_list FID KID PSNR LPIPS
+OUT=$?
+
+if [ $OUT != 0 ]; then
+    exit 1
+fi
+
 ###### diffusion process test
 echo "Running diffusion process test"
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_diffusion.py" --dataroot "$TARGET_MASK_SEM_DIR"
@@ -160,7 +170,6 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-
 ###### mask cls semantics test with online dataloading
 echo "Running mask online semantics training tests"
 URL=https://joligen.com/datasets/online_mario2sonic_lite2.zip 
@@ -173,15 +182,13 @@ rm $ZIP_FILE
 ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainA $TARGET_MASK_SEM_ONLINE_DIR/testA
 ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainB $TARGET_MASK_SEM_ONLINE_DIR/testB
 
-
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_semantic_mask_online.py" --dataroot "$TARGET_MASK_SEM_ONLINE_DIR"
 OUT=$?
 
 if [ $OUT != 0 ]; then
     exit 1
 fi
 
-
 ###### diffusion process test online
 echo "Running diffusion process test online"
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_diffusion_online.py" --dataroot "$TARGET_MASK_SEM_ONLINE_DIR"
@@ -191,12 +198,12 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-###### test cut
-echo "Running cut test"
+###### test palette
+echo "Running test palette"
 python3 "${current_dir}/../test.py" \
-	--save_config \
-	--test_model_dir $DIR/joligen_utest_cut/ \
-	--test_metrics_list FID KID PSNR LPIPS
+        --save_config \
+        --test_model_dir $DIR/joligen_utest_palette/ \
+        --test_metrics_list FID KID PSNR LPIPS
 OUT=$?
 
 if [ $OUT != 0 ]; then
@@ -215,14 +222,24 @@ rm $ZIP_FILE
 ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainA $TARGET_MASK_SEM_ONLINE_DIR/testA
 ln -s $TARGET_MASK_SEM_ONLINE_DIR/trainB $TARGET_MASK_SEM_ONLINE_DIR/testB
 
-
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_vid_diffusion_online.py" --dataroot "$TARGET_MASK_SEM_ONLINE_DIR"
 OUT=$?
 
 if [ $OUT != 0 ]; then
     exit 1
 fi
 
+###### test vid palette
+echo "Running test vid palette"
+python3 "${current_dir}/../test.py" \
+        --save_config \
+        --test_model_dir $DIR/joligen_utest_vid_palette/ \
+        --test_metrics_list SSIM PSNR LPIPS
+OUT=$?
+
+if [ $OUT != 0 ]; then
+    exit 1
+fi
 
 ###### test cycle_gan
 # echo "Running test cycle_gan"
@@ -235,18 +252,6 @@ fi
 #     exit 1
 # fi
 
-###### test palette
-echo "Running test palette"
-python3 "${current_dir}/../test.py" \
-	--save_config \
-	--test_model_dir $DIR/joligen_utest_palette/ \
-	--test_metrics_list FID KID PSNR LPIPS
-OUT=$?
-
-if [ $OUT != 0 ]; then
-    exit 1
-fi
-
 ####### mask cls semantics test
 echo "Running mask and class semantics training tests"
 URL=https://joligen.com/datasets/daytime2dawn_dusk_lite.zip
@@ -277,7 +282,6 @@ rm $ZIP_FILE
 
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_semantic_cls.py" --dataroot "$TARGET_CLS_SEM_DIR"
 
-
 OUT=$?
 
 if [ $OUT != 0 ]; then
@@ -336,7 +340,6 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-
 python3 -m pytest --rootdir ${current_dir} -p no:cacheprovider -s "${current_dir}/../tests/test_run_pix2pix_diffusion.py" --dataroot "$TARGET_PIX2PIX_DIR"
 OUT=$?
 
@@ -356,7 +359,7 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-# ####### api server tests
+######### api server tests
 echo "Running api server tests"
 
 python3 -m pytest --rootdir ${current_dir} \
@@ -388,7 +391,6 @@ if [ $OUT != 0 ]; then
     exit 1
 fi
 
-
 # End of tests
 # Clean up
 echo "Deleting target dir $DIR"

diff --git a/tests/test_run_vid_diffusion_online.py b/tests/test_run_vid_diffusion_online.py
@@ -12,8 +12,8 @@
 from scripts.gen_vid_diffusion import InferenceDiffusionOptions, inference
 
 json_like_dict = {
-    "name": "joligen_utest",
-    "output_display_env": "joligen_utest",
+    "name": "joligen_utest_vid",
+    "output_display_env": "joligen_utest_vid",
     "output_display_id": 0,
     "gpu_ids": "0",
     "data_load_size": 64,