Update README benchmark numbers

README.md

@@ -312,9 +312,9 @@ python main.py eval  --config configs/config_CRAG_aug.yaml --checkpoint weights/
 | Model                                   | mIoU (%) ↑ | mF1 (%) ↑ | mBIoU (%) ↑ | GFLOPs ↓ | FPS ↑ |
 |----------------------------------------|-----------:|----------:|------------:|---------:|------:|
 | *RGB Baselines (official)*              |            |           |             |          |       |
-| Swin-T + UPerNet                        | 62.01 | 75.27 | – | – | – |
-| Swin-S + UPerNet                        | 61.87 | 75.11 | – | – | – |
-| Swin-B + UPerNet                        | 64.05 | 76.88 | – | – | – |
+| Swin-T + UPerNet                        | 62.01 | 75.27 | – | 237 | 69.2 |
+| Swin-S + UPerNet                        | 61.87 | 75.11 | – | 261 | 41.5 |
+| Swin-B + UPerNet                        | 64.05 | 76.88 | – | 306 | 36.3 |
 | Swin-B + UPerNet (retrained)            | 64.02 | 76.64 | 32.57 | – | – |
 | Swin-L + UPerNet                        | 63.36 | 76.35 | – | 420 | 27.8 |
 | *Dual-branch baselines*                 |            |           |             |          |       |

utils/bench.py

@@ -4,8 +4,6 @@ import torch
 import torch.nn as nn
 from fvcore.nn import FlopCountAnalysis, flop_count_table, parameter_count_table
 
-from CASWiT import CASWiT
-
 try:
     from transformers import UperNetForSemanticSegmentation
 except ImportError:
@@ -145,6 +143,7 @@ def main():
     B, H, W = args.batch, args.h, args.w
 
     if args.model == "caswit":
+        from model.CASWiT_upernet import CASWiT
         base_model = CASWiT(num_head_xa=1, num_classes=15, model_name="openmmlab/upernet-swin-base").to(device).eval()
         model_for_flops = ForwardForFlops(base_model, which=args.which).to(device).eval()