Update val_mm.py to support SUNRGBD infer

VCIP-RGBD · Apr 29, 2024 · 5f20f62 · 5f20f62
1 parent e9e68f3
commit 5f20f62
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/val_mm.py b/utils/val_mm.py
@@ -149,7 +149,7 @@ def evaluate(model, dataloader, config, device, engine, save_dir=None, sliding=F
             if config.dataset_name in ["KITTI-360", "EventScape"]:
                 preds = palette[preds]
                 plt.imsave(save_name, preds)
-            elif config.dataset_name in ["NYUDepthv2"]:
+            elif config.dataset_name in ["NYUDepthv2", "SUNRGBD"]:
                 palette = np.load("./utils/nyucmap.npy")
                 preds = palette[preds]
                 plt.imsave(save_name, preds)
@@ -360,7 +360,7 @@ def evaluate_msf(
             if config.dataset_name in ["KITTI-360", "EventScape"]:
                 preds = palette[preds]
                 plt.imsave(save_name, preds)
-            elif config.dataset_name in ["NYUDepthv2"]:
+            elif config.dataset_name in ["NYUDepthv2", "SUNRGBD"]:
                 palette = np.load("./utils/nyucmap.npy")
                 preds = palette[preds]
                 plt.imsave(save_name, preds)